本样例的目的是通过一个简单的算子实现展示SIMD与SIMT混合编程方式,不是该算子功能的最佳实践。
基于SIMD与SIMT混合编程方式实现矢量算子核函数的流程如下图所示。
以下内容以从长度为10万的一维向量中提取指定索引的8192个数据,并对提取的数据分别执行加1运算的gather & adds算子为例,对上述步骤进行详细说明。本样例中介绍的算子完整代码请参见SIMD与SIMT混合编程实现gather&adds算子样例。
算子分析具体步骤如下:
gather & adds算子输出output中第i个数据为:
output[i] = input[index[i]] + 1
计算逻辑如下:
simd_adds中加1运算实际可以在simt_gather函数中快速实现,本样例的目的是通过一个简单的算子实现展示SIMD与SIMT混合编程方式,不是该算子功能的最佳实践。
本例中算子输入index的形状为8192,可设置核数为8,每个核处理数据量为1024。
对于SIMT实现,可设置线程数为1024,每个线程处理1个数据,单个核只需调用1次simt_gather函数即可完成gather运算。
对于SIMD Reg矢量计算实现,单核处理数据量为1024,Reg矢量计算API单次处理的数据长度one_repeat_size为GetVecLen/sizeof(float),API的循环调用次数repeat_times为1024/one_repeat_size。
通过以上分析,得到Ascend C gather & adds算子的设计规格如下:
|
name |
shape |
data type |
format |
|---|---|---|---|
|
input(输入) |
100000 |
float |
ND |
|
index(输入) |
8192 |
uint32_t |
ND |
|
output(输出) |
8192 |
float |
ND |
根据核函数中介绍的规则进行核函数的定义。
本样例中,函数名为gather_and_adds_kernel(核函数名称可自定义),根据上述分析,函数原型定义如下:
1 2 3 |
__global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length) { } |
1
|
uint32_t index_total_length_per_block = index_total_length / AscendC::GetBlockNum(); |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
constexpr uint32_t THREAD_COUNT = 1024; __global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length) { // 设置kernel type为AIV_ONLY KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 定义UB内存分配对象 AscendC::LocalMemAllocator<AscendC::Hardware::UB> ub_allocator; // 计算单核应处理的数据量 uint32_t index_total_length_per_block = index_total_length / AscendC::GetBlockNum(); // 申请UB内存作为simt_gather的输出 AscendC::LocalTensor<float> gather_output = ub_allocator.Alloc<float>(index_total_length_per_block); // 1. 调用simt函数获取指定索引的1024个数据 asc_vf_call<simt_gather>(dim3(THREAD_COUNT), input, index, (__ubuf__ float *)gather_output.GetPhyAddr(), input_total_length, index_total_length, index_total_length_per_block); // 2. 调用SIMD函数完成加1操作 ... // 3. 将数据搬运到GM ... } |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
__global__ __aicore__ void gather_and_adds_kernel(__gm__ float *input, __gm__ uint32_t *index, __gm__ float *output, uint32_t input_total_length, uint32_t index_total_length) { // 1. 调用simt函数获取指定索引的1024个数据 ... // 申请UB作为simd_adds的输出 AscendC::LocalTensor<float> adds_output = ub_allocator.Alloc<float>(index_total_length_per_block); // 计算Reg矢量计算API单次调用处理的数据量 constexpr uint32_t one_repeat_size = AscendC::GetVecLen() / sizeof(float); // 计算Reg矢量计算API循环调用次数 uint16_t repeat_times = (index_total_length_per_block + one_repeat_size - 1) / one_repeat_size; // 2. 调用SIMD函数完成加1操作 asc_vf_call<simd_adds>((__ubuf__ float *)adds_output.GetPhyAddr(), (__ubuf__ float *)gather_output.GetPhyAddr(), index_total_length_per_block, one_repeat_size, repeat_times); // 依赖PIPE_V和PIPE_MTE3流水同步 AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0); AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0); // 3. 将数据搬运到GM ... } |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
__global__ __aicore__ void gather_and_adds_kernel(__gm__ float* input, __gm__ uint32_t* index, __gm__ float* output, uint32_t input_total_length, uint32_t index_total_length) { // 1. 调用simt函数获取指定索引的1024个数据 ... // 2. 调用SIMD函数完成加1操作 ... // 3. 将数据搬运到GM // 定义GlobalTensor对象,用于数据搬运 AscendC::GlobalTensor<float> output_global_tensor; // 根据核偏移初始化GlobalTensor地址 output_global_tensor.SetGlobalBuffer(output + index_total_length_per_block * AscendC::GetBlockIdx()); // 调用数据搬运接口将数据搬运到GM AscendC::DataCopy(output_global_tensor, adds_output, index_total_length_per_block); } |
根据上述对SIMT VF函数的参数分析,定义SIMT VF函数原型。使用__simt_vf__函数类型限定符标识SIMT VF核函数入口,使其可以被asc_vf_call调用。
在SIMT 编程中,__launch_bounds__(thread_num)是可选配置,用于在编译期指定核函数启动的最大线程数(如果不配置,thread_num默认为1024),使用时请注意:thread_num >= x * y * z (即:asc_vf_call的第一个参数:dim3{x, y, z}), 线程数thread_num的取值范围为1到2048。最大线程数决定了每个线程可分配的寄存器数量,具体对应关系请见表5,寄存器用于存储线程中的局部变量,若局部变量的个数超出寄存器个数,容易出现栈溢出等问题。
1 2 3 4 5 6 7 8 9 10 11 |
constexpr uint32_t THREAD_COUNT = 1024; __simt_vf__ __launch_bounds__(THREAD_COUNT) inline void simt_gather( __gm__ float* input, __gm__ uint32_t* index, __ubuf__ float* gather_output, uint32_t input_total_length, uint32_t index_total_length, uint32_t output_total_length) { } |
simt_gather函数实现从输入input(Global Memory)中获取指定索引的数据。基于上述数据切分策略,首先计算线程应处理数据的索引,然后通过赋值操作将数据存储到Unified Buffer上。
1
|
int idx = blockIdx.x * blockDim.x + threadIdx.x; |
blockIdx用于获取当前核id。blockDim用于获取线程三维层次结构{x, y, z},本例中为{1024, 1, 1},其中第2,3维度均为1,使用一维层次结构,因此线程数可写作blockDim.x。threadIdx用于获取三维线程索引{x, y, z},本例中仅使用第1维x,可通过threadIdx.x获取当前线程的id。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
__simt_vf__ __launch_bounds__(THREAD_COUNT) inline void simt_gather( __gm__ float* input, __gm__ uint32_t* index, __ubuf__ float* gather_output, uint32_t input_total_length, uint32_t index_total_length, uint32_t output_total_length) { // 异常判断,防止越界 if (threadIdx.x >= output_total_length) { return; } // 计算线程应处理的数据在输入index的索引 int idx = blockIdx.x * blockDim.x + threadIdx.x; // 异常判断,防止越界 if (idx >= index_total_length) { return; } // 读取线程应获取的数据在输入input中的索引 uint32_t gather_idx = index[idx]; // 异常判断,防止越界 if (gather_idx >= input_total_length) { return; } // 将input中索引为gather_idx的数据存到UB上 gather_output[threadIdx.x] = input[gather_idx]; } |
1 2 3 4 |
__simd_vf__ inline void simd_adds(__ubuf__ float *output, __ubuf__ float *input, uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times) { } |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
constexpr float ADDS_ADDEND = 1.0f; __simd_vf__ inline void simd_adds(__ubuf__ float *output, __ubuf__ float *input, uint32_t count, uint32_t one_repeat_size, uint16_t repeat_times) { // 初始化Reg矢量计算单元,源操作数 AscendC::Reg::RegTensor<float> src_reg0; // 初始化Reg矢量计算单元,目的操作数 AscendC::Reg::RegTensor<float> dst_reg0; // 初始化Reg矢量计算标记寄存器 AscendC::Reg::MaskReg mask; for (uint16_t i = 0; i < repeat_times; i++) { // 从UB搬运数据到Reg矢量计算基本单元 mask = AscendC::Reg::UpdateMask<float>(count); AscendC::Reg::LoadAlign(src_reg0, input + i * one_repeat_size); // 调用Adds接口完成加1运算 AscendC::Reg::Adds(dst_reg0, src_reg0, ADDS_ADDEND, mask_reg); // 从Reg矢量计算基本单元搬运数据到UB AscendC::Reg::StoreAlign(output + i * one_repeat_size, dst_reg0, mask_reg); } } |