昇腾Ascend TIK自定义算子开发教程（概念版）_算子内部约束

一、参考资料

【2023 · CANN训练营第一季】Ascend C算子开发入门（中）

二、重要说明

• TIK2编程范式把算子核内的处理程序，分成多个流水任务，任务之间通过队列（Queue）进行通信和同步，并通过统一的内存管理模块（Pipe）管理任务间通信内存。
• TIK2分别针对Vector、Cube编程设计了不同的流水任务。开发者只需要完成基本任务的代码实现即可，底层的指令同步和并行调度由TIK2框架实现，开发者无需关注。
• 由于开发高性能Cube算子难度较大，当前仅支持用户开发Vector算子。
• 当前TIK2支持的AI处理器型号为昇腾310P AI处理器、昇腾910 AI处理器，其他型号暂不支持。
• 当前支持用户使用g++等C/C++编译器编译在cpu侧执行的TIK2算子，并使用gdb单步调试；支持用户使用CCEC编译器编译在npu侧执行的TIK2算子，实现加速计算，暂不支持加载至网络模型中进行整网验证。
• 算子输出的数据类型与输入数据类型相同。
• 输出shape与输入shape相同。

三、相关介绍

1. CANN算子

CANN算子有两种类型，TBE算子与AI CPU算子。

• AI Core是昇腾AI处理器的计算核心，负责执行矩阵、向量、标量计算密集的算子任务，在AI Core上执行的算子称为TBE（Tensor Boost Engine）算子。
• AI CPU负责执行不适合跑在AI Core上的算子，是AI Core算子的补充，主要承担非矩阵类、逻辑比较复杂的分支密集型计算，在AI CPU上执行的算子称为AI CPU算子。

1.1 TBE算子

TBE（Tensor Boost Engine，张量加速引擎）提供了基于TVM（Tensor Virtual Machine，张量虚拟机）框架的自定义算子开发能力，提供了用户开发自定义算子所需工具。TBE框架给用户提供了两种算子开发方式 ：DSL与TIK。开发者可以根据需求自由选择，两种开发方式的区别如下：

• DSL（ Domain-Specific Language ，基于特性域语言）

  DSL接口已高度封装，用户仅需要使用DSL接口完成计算过程的表达，后续的算子调度、算子优化及编译都可通过已有的接口一键式完成，适合初级开发用户。

• TIK（ Tensor Iterator Kernel， 张量嵌套内核）

  开发者可以通过调用TIK提供的API基于Python语言编写自定义算子，然后TIK编译器会将其编译为适配昇腾AI处理器应用程序的二进制文件。TIK需要用户手工控制数据搬运和计算流程，入门较高，但开发方式比较灵活，能够充分挖掘硬件能力，在性能上有一定的优势。

1.2 AI CPU算子

以下几种场景下，可使用AI CPU方式实现自定义算子：

• 不适合跑在AI Core上的算子，例如非矩阵类的复杂计算，逻辑比较复杂的分支密集型算子等。

  例如，Dump、profiling等控制算子，Queue、Stack等资源状态类算子，TopK、Where等检索类算子。

• AI Core不支持的算子，算子需要某些数据类型，但AI Core不支持，例如Complex32、Complex64。

• 某些场景下，为了快速打通网络在昇腾AI处理器的执行流程，在TBE实现自定义算子较为困难的情况下，可通过自定义AI CPU算子进行功能调测，提升调测效率。功能调通之后，后续性能调测过程中再将AI CPU自定义算子转换为TBE算子实现。

2. TIK

TIK（Tensor Iterator Kernel）是一种基于Python语言的动态编程框架，呈现为一个Python模块。 开发者可以通过调用TIK提供的API基于Python语言编写自定义算子，然后TIK编译器会编译为适配昇腾AI处理器应用程序的二进制文件。

TIK编程模型

使用TIK进行编程的过程，如下图所示，用户调用TIK API编写算子对应的Python程序后，TIK会将其转化为TIK DSL(TIK DSL是一种DSL语言，它可以在比CCE更高的抽象层次上定义CCEC程序的行为)，经过编译器编译后生成CCEC文件（CCEC代码目前对于TIK编程人员无法感知），再经过CCE编译器编译后生成可运行在昇腾AI处理器上的应用程序。

3. TIK2

TIK2是一种使用C/C++作为前端语言的编程框架，开发者可以使用TIK2提供的API编写自定义算子，并通过CCEC编译器对自定义算子进行编译，生成可运行在昇腾AI处理器上的应用程序。

TIK与TIK2开发方式对比

四、AI Core架构

AI Core是昇腾AI处理器的计算核心，可以看成是一个相对简化的现代微处理器的基本架构，负责执行矩阵、向量、标量计算密集的算子任务。它包括了三种基础计算资源：矩阵计算单元（Cube Unit）、向量计算单元（Vector Unit）和标量计算单元（Scalar Unit）。这三种计算单元各司其职，形成了三条独立的执行流水线，在系统软件的统一调度下互相配合达到优化的计算效率。AI Core中包含计算单元、存储单元、控制单元、搬运单元。

1. 计算单元

计算单元是AI Core中提供强大算力的核心单元，相当于AI Core的主力军，主要包括：Cube Unit（矩阵计算单元）、Vector Unit（向量计算单元）和Scalar Unit（标量计算单元），完成AI Core中不同类型的数据计算。

2. 存储单元

AI Core需要把外部存储中的数据加载到内部存储中，才能完成相应的计算。

2.1 外部存储

通常，AI Core的外部存储包括L2、HBM、DDR等。

2.2 内部存储

AI Core的内部存储，统称为Local Memory，主要包括：L1 Buffer（L1缓冲区），L0 Buffer（L0缓冲区），Unified Buffer（统一缓冲区）和Scalar Buffer（标量缓冲区）。

2.3 存储单元分类

2.4 存储单元大小

不同类型的昇腾AI处理器，存储单元大小不同，用户可通过get_soc_spec接口获取。

2.4.1 函数原型

def get_soc_spec(key)

2.4.2 参数说明

2.4.3 返回值

根据输入的key返回对应的值：

• SOC_VERSION：返回标识SOC类型的字符串。
• AICORE_TYPE：返回Core的类型，有AiCore或VectorCore两种返回值。
• CORE_NUM：返回核数，int类型。
• UB_SIZE：返回UB大小，int类型，单位Byte。
• L2_SIZE：返回L2大小，int类型，单位Byte。
• L1_SIZE：返回L1大小，int类型，单位Byte。
• CUBE_SIZE：返回CUBE大小，tuple类型，如（16,16,16），单位为Byte。
• L0A_SIZE：返回L0A大小，int类型，单位为Byte。
• L0B_SIZE：返回L0B大小，int类型，单位为Byte。
• L0C_SIZE：返回L0C大小，int类型，单位为Byte。
• SMASK_SIZE：返回Smask buffer大小，int类型，单位为Byte。

2.4.4 示例代码

实际调用时，将变量soc_version的值修改为实际的昇腾AI处理器型号。

import tbe
soc_version="xxx"
tbe.common.platform.set_current_compile_soc_info(soc_version)
tbe.common.platform.get_soc_spec("CORE_NUM")
1
2
3
4

2.5 指令与存储访问关系

上图的存储单元是软件层面概念，其中：

• Scalar Buffer对应硬件存储单元Scalar Buffer。
• Unified Buffer对应硬件存储单元Unified Buffer。
• L1 Buffer对应硬件存储单元L1 Buffer。
• L1Out Buffer为从L0C上抽象出来的存储Cube计算输出数据的存储单元

2.6 QuePosition与硬件存储单元映射关系

2.7 硬件存储单元对齐

不同scope的对齐要求，如下表所示：

3. 控制单元

控制单元为整个计算过程提供了指令控制，相当于AI Core的司令部，负责整个AI Core的运行。系统控制模块（System Control）负责指挥和协调AI Core的整体运行模式，配置参数和实现功耗控制等。当指令通过指令发射模块（Instruction Dispatch）顺次发射出去后，根据指令的不同类型，将会分别被发送到矩阵运算队列（Cube Queue）、向量运算队列（Vector Queue）和存储转换队列（MTE Queue）。指令执行过程中，可以提前预取后续指令，并一次读入多条指令进入缓存，提升指令执行效率。多条指令从系统内存通过总线接口（BIU）进入到AI Core的指令缓存模块（Instruction Cache）中等待，后续硬件快速自动解码或运算。指令被解码后便会被导入标量指令处理队列（Scalar PSQ）中，实现地址解码与运算控制。

3.1 控制单元分类

AI Core包含的控制单元，如下表所示。

3.2 指令队列分类

根据调度分类的不同，可以把指令分类，加上被译码过程直接解释的Scalar指令（缩写为S），可以有6种指令分类：S、V、M、MTE1、MTE2、MTE3。

除S队列之外，不同队列的指令能够乱序执行，但是队列内部指令为顺序执行，即在满足数据依赖的前提下，指令的物理执行顺序不一定与代码的书写顺序一致。

硬件按照下发顺序，将不同队列的指令分发到相应的队列上执行，昇腾AI处理器提供Barrier、set_flag/wait_flag两种指令，保证队列内部以及队列之间按照逻辑关系执行。

• Barrier本身是一条指令，用于在队列内部约束执行顺序。其作用是，保证前序队列中所有数据的读写工作全部完成，后序指令才能执行。
• set_flag/wait_flag为两条指令，在set_flag/wait_flag的指令中，可以指定一对指令队列的关系，表示两个队列之间完成一组“锁”机制，其作用方式为：
  • set_flag：当前序指令的所有读写操作都完成之后，当前指令开始执行，并将硬件中的对应标志位设置为1。
  • wait_flag：当执行到该指令时，如果发现对应标志位为0，该队列的后续指令将一直被阻塞；如果发现对应标志位为1，则将对应标志位设置为0，同时后续指令开始执行。

注意：TBE封装了这种依赖关系，所以应用开发人员不必对Barrier或者Flag进行编程。但应用开发人员仍需要理解这个基本原理，才能通过合适的代码调度，实现更好的同步关系。基于DSL方式进行算子开发无需关注代码调度，DSL提供了自动调度（auto_schedule）机制。

3.3 AI Core指令调度方式

AI Core采用顺序取指令、并行执行指令的调度方式，流水线执行过程如下图所示：

指令序列被顺序译码。根据指令的类型，有两种可能：

• 如果指令是Scalar指令，指令会被直接执行。
• 其他指令，指令会被调度到5个独立的指令队列，然后再分配到某个空间的执行部件执行。

4. 搬运单元

DMA搬运单元，负责在Global Memory和Local Memory之间搬运数据，具体来说，把数据搬运到Local Memory，Vector/Cube计算单元完成数据计算，并把计算结果写回Local Memory，DMA搬出单元把处理好的数据搬运回Global Memory。DMA搬运单元包括：MTE2（Memory Transfer Engine，数据搬入单元），MTE3（数据搬出单元）。

五、核函数

核函数是直接在Device设备端执行的代码。在核函数中，需要为在一个核上执行的代码规定要进行的数据访问和计算操作，当核函数被调用时，多个核将并行执行同一个计算任务。

extern "C" __global__ __aicore__ void add_tik2(__gm__ uint8_t* x, __gm__ uint8_t* y, __gm__ uint8_t* z)
{
}
1
2
3

1. 函数类型限定符

编写核函数

核函数的函数类型限定符，包括 __global__ 和 __aicore__，其中__global__ 标识核函数，__aicore__ 表示核函数在设备端aicore上执行。

2. 变量类型限定符

指针入参变量统一的类型定义为 __gm__ uint8_t*，Init()函数的入参统一设置为uint8_t*类型的指针，在后续的使用中需要将其转化为实际的指针类型；用户亦可直接传入实际的指针类型。

3. 核函数调用符

#ifndef __CCE_KT_TEST__ 表示核函数在NPU侧运行，核函数通过核函数调用符 <<<...>>> 调用。<<<...>>> 仅在NPU侧调用，在CPU侧直接调用核函数即可。

#ifndef __CCE_KT_TEST__
// 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);
}
#endif
1
2
3
4
5
6
7

4. Init()函数实现

constexpr 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)
{
    //获取核函数的输入输出在Global Memory上的内存偏移地址
    // 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内存管理对象为输入输出Queue分配内存
    // 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));
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

解释说明：

1. 数据整体长度TOTAL_LENGTH为8* 2048，平均分配到8个核上运行，每个核上处理的数据大小BLOCK_LENGTH为2048。
2. block_idx为核的逻辑ID，(__gm__ half*)x + block_idx * BLOCK_LENGTH 即为单核处理程序中x在Global Memory上的内存偏移地址。注意，因为Init函数的入参统一设置为 uint8_t*，这里需要强转成具体的数据类型 (__gm__ half*)，再进行偏移。

4.1 BLOCK

// 数据整体长度
// total length of data
constexpr int32_t TOTAL_LENGTH = 8 * 2048;

// 使用多核
// num of core used
constexpr int32_t USE_CORE_NUM = 8;

//每个核处理数据的大小
// length computed of each core
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM; 
1
2
3
4
5
6
7
8
9
10
11

block_num

block_num默认取值为1，即不分核；而采用分核并行时，其取值上限为65535，用户需要保证block_num的值不超过此阈值。

在for_range的原型定义里，用户通过设置参数block_num来实现分核并行，简单代码示例如下：

with tik_instance.for_range( 0, 10, block_num=10) as i:
1

for_range循环中的表达式会被作用在10个执行实例上，最终10个执行实例会被分配到10个核上并行运行，每个核拿到一个执行实例和一个不同的Block ID。如果当前可用的核的数量小于10，则执行实例会在当前可用的核上分批调度执行；如果当前可用的核的数量大于等于10，则会根据实际情况调度执行，实际运行的核数可能小于等于10。

一个算子中只能调用一次for_range实现分核，即设置block_num >=2，不允许多次开启多核。

CORE_NUM

用户可以通过get_soc_spec接口获取AI Core的个数。

# 请根据实际昇腾AI处理器型号进行设置
soc_version="xxx"
# 设置昇腾AI处理器的型号及目标核的类型
tbe.common.platform.set_current_compile_soc_info(soc_version) 
tbe.common.platform.get_soc_spec("CORE_NUM") # 使用该接口前需要先设置芯片类型
1
2
3
4
5

为保证负载均衡，block_num一般尽量设置为实际核数量的倍数。假设芯片内含32个AI Core，假如一个张量的形状为（16, 2, 32, 32, 32），如果以张量的第一维度（最外层）进行分核，则只能绑定16个核。此时，可通过将张量的第一维度和第二维度合并，使得最外层的长度变成32，以此将任务均摊到32个AI Core上，使用尽可能多的核并行处理。需要注意的是，顾及后端内存自动分配机制限制，用户实施分核并行时必须从最外层开始做维度合并。

4.2 Tiling

对于单核上的处理数据，可以进行数据切块（Tiling）。

// split data into 8 tiles for each core
constexpr int32_t TILE_NUM = 8;

// tensor num for each queue
constexpr int32_t BUFFER_NUM = 2;

// each tile length is seperated to 2 part, due to double buffer
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; 
1
2
3
4
5
6
7
8

5. Process()函数实现

基于矢量编程范式，将核函数的实现分为3个基本任务：CopyIn，Compute，CopyOut。

__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);
    }
}
1
2
3
4
5
6
7
8
9
10
11

核函数内通过数据切块（Tiling），实现流水线之间的并行。举例来说，将单核处理数据分成n份，使用progress0_{processn-1表示处理第1}n个数据切片。progress0经过CopyIn Stage之后进入Compute Stage，CopyIn即可以处理progress1，做到了流水线间并行。根据编程范式上面的算法分析，将整个计算拆分成三个Stage，用户单独编写每个Stage的代码，三阶段流程示意图如下：

5.1 Stage1：CopyIn函数实现。

1. 使用DataCopy接口将GlobalTensor数据拷贝到LocalTensor。
2. 使用EnQue将LocalTensor放入VecIn的Queue中。

__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);
}
1
2
3
4
5
6
7
8
9
10
11
12

5.2 Stage2：Compute函数实现。

1. 使用DeQue从VecIn中取出LocalTensor。
2. 使用TIK2接口Add完成矢量计算。
3. 使用EnQue将计算结果LocalTensor放入到VecOut的Queue中。
4. 使用FreeTensor将不再使用的LocalTensor进行回收。

__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);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14

5.3 Stage3：CopyOut函数实现。

1. 使用DeQue接口从VecOut的Queue中取出LocalTensor。
2. 使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。
3. 使用FreeTensor将不再使用的LocalTensor进行回收。

 __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);
}
1
2
3
4
5
6
7
8
9

六、Queue通信和同步

任务间通信和同步

不同的流水任务之间存在数据依赖，需要进行数据传递。TIK2中使用Queue队列完成任务之间的数据通信和同步，提供EnQue、DeQue等基础API。

1. QuePosition逻辑位置

Queue队列管理NPU上不同层级的物理内存时，用一种抽象的逻辑位置 (QuePosition) 来表达各个级别的存储(Storage Scope)，代替了片上物理存储的概念，开发者无需感知硬件架构，达到隐藏芯片架构的目的。Queue类型包括：VECIN、VECOUT、A1、A2、B1、B2、CO1、CO2，其中VECIN、VECOUT主要用于矢量编程，具体说明参见[矢量编程](javascript: