【2023 · CANN训练营第一季】TIK C++算子编程范式及算子实现
编程范式
编程范式描述了算子实现的固定流程,基于编程范式进行编程,可以快速搭建算子实现的代码框架。
TIK C++编程范式把算子核内的处理程序,分成多个流水任务,以张量为数据载体,任务之间通过队列(Queue)进行通信和同步,并通过统一的内存管理模块(Pipe)管理任务间通信内存。
流水任务
流水任务指的是单核处理程序中主程序调度的并行任务。在核函数内部,可以通过流水任务实现数据的并行处理,进一步提升性能。以下面的流水任务示意图为例,单核处理程序的功能被拆分成3个流水任务:Stage1、Stage2、Stage3,每个任务专注于完成单一功能;需要处理的数据被切分成n片,使用Progress1~n表示,每个任务需要依次完成n个数据切片的处理。Stage间的箭头表达数据间的依赖关系,比如Stage1处理完Progress1之后,Stage2才能对Progress1进行处理。
若n=3,即待处理的数据被切分成3片,则上图中的流水任务运行起来的示意图如下,从运行图中可以看出,对于同一片数据,Stage1、Stage2、Stage3之间的处理具有依赖关系,需要串行处理;不同的数据切片,同一时间点,可以有多个任务在并行处理,由此达到任务并行、提升性能的目的。
矢量算子编程范式把算子的实现流程分为3个基本任务:CopyIn,Compute,CopyOut。CopyIn负责搬入操作,Compute负责矢量指令计算操作,CopyOut负责搬出操作。
矩阵算子编程范式把算子的实现流程分为5个基本任务:CopyIn,Split,Compute,Aggregate,CopyOut。CopyIn负责搬入操作,Split负责数据切分操作,Compute负责矩阵指令计算操作,Aggregate负责数据汇聚操作,CopyOut负责搬出操作。
矢量编程 | 矩阵编程 |
---|---|
任务间通信和同步
不同的流水任务之间存在数据依赖,需要进行数据传递。
TIK C++中使用Queue队列完成任务之间的数据通信和同步,提供EnQue、DeQue等基础API。
Queue队列管理不同层级的物理内存时,用一种抽象的逻辑位置(QuePosition)来表达各级别的存储,代替了片上物理存储的概念,达到隐藏芯片架构的目的。Queue类型包括:VECIN、VECOUT、A1、A2、B1、B2、CO1、CO2,其中VECIN、VECOUT主要用于矢量编程,A1、A2、B1、B2、CO1、CO2用于矩阵编程。
矢量编程
矢量编程中使用到的逻辑位置(QuePosition)定义如下:
- 搬入数据的存放位置:VECIN;
- 搬出数据的存放位置:VECOUT。
矢量编程主要分为CopyIn、Compute、CopyOut三个任务:
- CopyIn任务中将输入数据从Global内存搬运至Local内存后,需要使用EnQue将LocalTensor放入VECIN的Queue中;
- Compute任务等待VECIN的Queue中LocalTensor出队之后才可以完成矢量计算,计算完成后使用EnQue将计算结果LocalTensor放入到VECOUT的Queue中;
- CopyOut任务等待VECOUT的Queue中LocalTensor出队,再将其拷贝到Global内存。
- Stage1:CopyIn任务。
- 使用DataCopy接口将GlobalTensor数据拷贝到LocalTensor。
- 使用EnQue将LocalTensor放入VECIN的Queue中。
- Stage2:Compute任务。
- 使用DeQue从VECIN中取出LocalTensor。
- 使用TIK C++接口完成矢量计算。
- 使用EnQue将计算结果LocalTensor放入到VECOUT的Queue中。
- Stage3:CopyOut任务。
- 使用DeQue接口从VECOUT的Queue中去除LocalTensor。
- 使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。
内存管理
任务间数据传递使用到的内存统一由内存管理模块Pipe进行管理。
如下图所示,Pipe作为片上内存管理者,通过InitBuffer接口对外提供Queue内存初始化功能,开发者可以通过该接口为指定的Queue分配内存。
Queue队列内存初始化完成后,需要使用内存时,通过调用AllocTensor来为LocalTensor分配内存,当创建的LocalTensor完成相关计算无需再使用时,再调用FreeTensor来回收LocalTensor的内存。
1 | TPipe pipe; |
编程过程中使用到的临时变量内存同样通过Pipe进行管理。临时变量可以使用TBuf数据结构来申请指定QuePosition上的存储空间。使用TBuf申请的内存空间只能参与计算,无法执行Queue队列的入队出队操作。
1 | // 为TBuf初始化分配内存,分配内存长度为1024字节 |
算子开发流程
TIK C++算子开发分为快速流程和标准流程,快速开发流程为
- 完成算子核函数的开发;
- 基于内核调用符方式进行算子运行验证。
标注流程为:
- 完成算子核函数的开发;
- 完成单算子网络应用程序的开发;
- 基于ACL单算子调用方式进行算子运行验证。
基于TIK C++方式实现矢量算子的流程如下图所示:
- 算子分析:分析算子的数学表达式、输入、输出以及计算逻辑的实现,明确需要调用的TIK C++接口。
- 核函数定义:定义TIK C++算子入口函数。
- 根据矢量编程范式实现算子类:完成核函数的内部实现。
下面以张量相加为例:
算子分析
在开发算子代码之前需要分析算子的数学表达式、输入、输出以及计算逻辑的实现。
- 明确算子的数学表达式及计算逻辑。
数学表达式为:z = x + y
,其中x
和y
为(8, 2048)
的张量,数据类型为half,输入排列方式为ND。
计算逻辑:TIK C++提供的矢量计算接口的操作元素都为LocalTensor,输入数据需要先搬运进片上存储,然后使用计算接口完成两个输入参数相加,得到最终结果,再搬出到外部存储上。 - 明确输入和输出。
- Add算子有两个输入:x与y,输出为z。
- 算子输入的数据类型为half(float16),算子输出的数据类型与输入数据类型相同。
- 算子输入支持形状为(8,2048),输出形状与输入形状相同。
- 算子输入支持的数据排列为:ND。
- 确定核函数名称和参数。
- 自定义核函数名称,核函数命名为add_tik2。
- 根据对算子输入输出的分析,确定核函数有3个参数x,y,z;
- x,y为输入在Global Memory上的内存地址,z为输出在Global Memory上的内存地址。
- 确定算子实现所需接口。
- 涉及外部存储和内部存储间的数据搬运,需要使用DataCopy来实;
- 矢量计算的加法操作,使用双目指令Add接口Add实现x+y;
- 计算中使用到的Tensor数据结构为LocalTensor,使用Queue队列进行管理,会使用到EnQue、DeQue等接口。
核函数定义
-
函数原型定义
函数名为add_tik2;根据算子分析中对算子输入输出的分析,确定有3个参数x,y,z,参数类型统一设置成uint8_t*,其中x,y都为输入内存,z为输出内存;根据编写核函数核函数的规则介绍,返回值为void,并增加extern "C"标识。除了需要按照C/C++函数声明的方式定义核函数之外,还要为核函数加上额外的函数类型限定符和变量类型限定符。
由此,可以得到函数原型定义为:1
2
3extern "C" __global__ __aicore__ void add_tik2(__gm__ uint8_t*x, __gm__ uint8_t* y, __gm__ uint8_t* z)
{
} -
调用算子类的Init和Process函数。
算子类的Init函数,完成内存初始化相关工作,Process函数完成算子实现的核心逻辑.1
2
3
4
5
6extern "C" __global__ __aicore__ void add_tik2(__gm__ uint8_t*x, __gm__ uint8_t* y, __gm__ uint8_t* z)
{
KernelAdd op;
op.Init(x, y, z);
op.Process();
} -
对核函数的调用进行封装,得到add_tik2_do函数,便于主程序调用。
#ifndef __CCE_KT_TEST__表示该封装函数仅在编译运行NPU侧的算子时会用到,编译运行CPU侧的算子时,可以直接调用add_tik2函数。
根据调用核函数章节,调用核函数时,除了需要传入参数x,y,z,还需要传入blockDim(核函数执行的核数), l2ctrl(保留参数,设置为nullptr), stream(应用程序中维护异步操作执行顺序的stream)来规定核函数的执行配置。1
2
3
4
5
6
7
8
9
// call of kernel function
void add_tik2_do(uint32_t blockDim, void*l2ctrl, void* stream, uint8_t*x, uint8_t* y, uint8_t* z)
{
add_tik2<<<blockDim, l2ctrl, stream>>>(x, y, z);
}
算子类实现
-
流水任务
Add算子的实现流程分为3个基本任务:CopyIn,Compute,CopyOut。
- CopyIn任务负责将Global Memory上的输入Tensor xGm和yGm搬运至Local Memory,分别存储在xLocal, yLocal;
- Compute任务负责对xLocal, yLocal执行加法操作,计算结果存储在zLocal中;
- CopyOut任务负责将输出数据从zLocal搬运至Global Memory上的输出Tensor zGm中。
-
数据通信
- CopyIn,Compute任务间通过VECIN队列inQueueX,inQueueY进行通信和同步,
- Compute,CopyOut任务间通过VECOUT队列outQueueZ进行通信和同步。
-
内存管理
任务间交互使用到的内存、临时变量使用到的内存统一使用pipe内存管理对象进行管理。
算子类中主要实现上述流程,包含对外开放的初始化Init函数和核心处理函数Process,Process函数中会对上图中的三个基本任务进行调用;同时包括一些算子实现中会用到的私有成员,比如上图中的Global Tensor和VECIN、VECOUT队列等。KernelAdd算子类具体成员如下:
1 | class KernelAdd { |
-
Init函数实现
- 获取该核函数需要处理的输入输出在Global Memory上的内存偏移地址。
对于多核并行计算,需要把数据进行分片,分配到多个核上进行处理。TIK C++核函数是在一个核上的处理函数,所以只处理部分数据,需要在初始化函数中获取该核函数需要处理的输入输出在Global Memory上的内存偏移地址。
例如:数据整体长度TOTAL_LENGTH为(8, 2048),平均分配到8个核上运行,每个核上处理的数据大小BLOCK_LENGTH为2048。block_idx为核的逻辑ID,(__gm__ half*)x + block_idx *BLOCK_LENGTH
即为单核处理程序中x在Global Memory上的内存偏移地址。
对于单核上的处理数据,可以进行数据切块(Tiling),将数据切分成8块(并不意味着8块就是性能最优)。切分后的每个数据块再次切分成2块,即可开启double buffer,实现流水线之间的并行。这样单核上的数据(2048个数)被切分成16块,每块TILE_LENGTH(128)个数据。上文代码表示Pipe为inQueueX分配了两块大小为TILE_LENGTH * sizeof(half)个字节的内存块,每个内存块能容纳TILE_LENGTH(128)个half类型数据。 - 通过Pipe内存管理对象为输入输出Queue分配内存。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18constexpr int32_t TOTAL_LENGTH = 8 * 2048; // total length of data
constexpr int32_t USE_CORE_NUM = 8; // num of core used
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM; // length computed of each core
constexpr int32_t TILE_NUM = 8; // split data into 8 tiles for each core
constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // each tile length is seperated to 2 part, due to double buffer
__aicore__ inline void Init(__gm__ uint8_t*x, __gm__ uint8_t* y, __gm__ uint8_t*z)
{
// get start index for current core, core parallel
xGm.SetGlobalBuffer((__gm__ half*)x + block_idx *BLOCK_LENGTH);
yGm.SetGlobalBuffer((__gm__ half*)y + block_idx *BLOCK_LENGTH);
zGm.SetGlobalBuffer((__gm__ half*)z + block_idx *BLOCK_LENGTH);
// pipe alloc memory to queue, the unit is Bytes
pipe.InitBuffer(inQueueX, BUFFER_NUM, TILE_LENGTH* sizeof(half));
pipe.InitBuffer(inQueueY, BUFFER_NUM, TILE_LENGTH *sizeof(half));
pipe.InitBuffer(outQueueZ, BUFFER_NUM, TILE_LENGTH* sizeof(half));
} - 获取该核函数需要处理的输入输出在Global Memory上的内存偏移地址。
-
Process函数实现
基于矢量编程范式,将核函数的实现分为3个基本任务:CopyIn,Compute,CopyOut。1
2
3
4
5
6
7
8
9
10
11__aicore__ inline void Process()
{
// loop count need to be doubled, due to double buffer
constexpr int32_t loopCount = TILE_NUM * BUFFER_NUM;
// tiling strategy, pipeline parallel
for (int32_t i = 0; i < loopCount; i++) {
CopyIn(i);
Compute(i);
CopyOut(i);
}
}-
CopyIn阶段
- 使用DataCopy接口将GlobalTensor数据拷贝到LocalTensor。
- 使用EnQue将LocalTensor放入VecIn的Queue中。
1
2
3
4
5
6
7
8
9
10
11
12__aicore__ inline void CopyIn(int32_t progress)
{
// alloc tensor from queue memory
LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
LocalTensor<half> yLocal = inQueueY.AllocTensor<half>();
// copy progress_th tile from global tensor to local tensor
DataCopy(xLocal, xGm[progress * TILE_LENGTH], TILE_LENGTH);
DataCopy(yLocal, yGm[progress * TILE_LENGTH], TILE_LENGTH);
// enque input tensors to VECIN queue
inQueueX.EnQue(xLocal);
inQueueY.EnQue(yLocal);
} -
Compute阶段
- 使用DeQue从VecIn中取出LocalTensor。
- 使用TIK C++接口Add完成矢量计算。
- 使用EnQue将计算结果LocalTensor放入到VecOut的Queue中。
- 使用FreeTensor将释放不再使用的LocalTensor。
1
2
3
4
5
6
7
8
9
10
11
12
13
14__aicore__ inline void Compute(int32_t progress)
{
// deque input tensors from VECIN queue
LocalTensor<half> xLocal = inQueueX.DeQue<half>();
LocalTensor<half> yLocal = inQueueY.DeQue<half>();
LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
// call Add instr for computation
Add(zLocal, xLocal, yLocal, TILE_LENGTH);
// enque the output tensor to VECOUT queue
outQueueZ.EnQue<half>(zLocal);
// free input tensors for reuse
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
} -
CopyOut阶段
- 使用DeQue接口从VecOut的Queue中取出LocalTensor。
- 使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。
- 使用FreeTensor将不再使用的LocalTensor进行回收。
1
2
3
4
5
6
7
8
9__aicore__ inline void CopyOut(int32_t progress)
{
// deque output tensor from VECOUT queue
LocalTensor<half> zLocal = outQueueZ.DeQue<half>();
// copy progress_th tile from local tensor to global tensor
DataCopy(zGm[progress * TILE_LENGTH], zLocal, TILE_LENGTH);
// free output tensor for reuse
outQueueZ.FreeTensor(zLocal);
}
-
-
double buffer
double buffer通过将数据搬运与矢量计算并行执行以隐藏数据搬运时间并降低矢量指令的等待时间,最终提高矢量计算单元的利用效率。为减少Vector等待时间,double buffer机制将待处理的数据一分为二,比如Tensor1、Tensor2。
- 当Vector对Tensor1中数据进行Compute时,Tensor2可以执行CopyIn的过程;
- 而当Vector切换到计算Tensor2时,Tensor1可以执行CopyOut的过程。