基础矢量算子

基于Ascend C方式实现基础矢量算子核函数的流程如下图所示。

图1 矢量算子核函数实现流程

• 算子分析：分析算子的数学表达式、输入、输出以及计算逻辑的实现，明确需要调用的Ascend C接口。
• 核函数定义：定义Ascend C算子入口函数。
• 根据矢量编程范式实现算子类：完成核函数的内部实现，包括3个基本任务：CopyIn，Compute，CopyOut。

下文以输入为half数据类型且shape的最后一维为32Bytes对齐、在单核上运行的、一次完成计算的Add算子为例，对上述步骤进行详细介绍。本样例中介绍的算子完整代码请参见基础Add算子样例。

算子类实现

根据上一节介绍，核函数中会调用算子类的Init和Process函数，本节具体讲解如何基于编程范式实现算子类。

根据矢量编程范式对Add算子的实现流程进行设计的思路如下，矢量编程范式请参考Vector编程范式，设计完成后得到的Add算子实现流程图参见图3 Add算子实现流程：

• 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内存管理对象进行管理。

图3 Add算子实现流程

算子类中主要实现上述流程，包含对外开放的初始化Init函数和核心处理函数Process，Process函数中会对上图中的三个基本任务进行调用；同时包括一些算子实现中会用到的私有成员，比如上图中的GlobalTensor（xGm、yGm、zGm）和VECIN、VECOUT队列等。KernelAdd算子类具体成员如下：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

class KernelAdd { public: __aicore__ inline KernelAdd() {} // 初始化函数，完成内存初始化相关操作 __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z){} // 核心处理函数，实现算子逻辑，调用私有成员函数CopyIn、Compute、CopyOut完成矢量算子的三级流水操作 __aicore__ inline void Process(){} private: // 搬入函数，完成CopyIn阶段的处理，被核心Process函数调用 __aicore__ inline void CopyIn(int32_t progress){} // 计算函数，完成Compute阶段的处理，被核心Process函数调用 __aicore__ inline void Compute(int32_t progress){} // 搬出函数，完成CopyOut阶段的处理，被核心Process函数调用 __aicore__ inline void CopyOut(int32_t progress){} private: AscendC::TPipe pipe; // Pipe内存管理对象 AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY; // 输入数据Queue队列管理对象，TPosition为VECIN AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ; // 输出数据Queue队列管理对象，TPosition为VECOUT AscendC::GlobalTensor<half> xGm; // 管理输入输出Global Memory内存地址的对象，其中xGm, yGm为输入，zGm为输出 AscendC::GlobalTensor<half> yGm; AscendC::GlobalTensor<half> zGm; };

初始化函数主要完成以下内容：

• 设置输入输出Global Tensor的Global Memory内存地址。

  本样例中的分配方案是：数据整体长度TOTAL_LENGTH为1 * 2048，使用GlobalTensor类的SetGlobalBuffer接口设定该核上Global Memory的起始地址以及长度。

  1	xGm.SetGlobalBuffer((__gm__ half *)x, TOTAL_LENGTH);

• 通过Pipe内存管理对象为输入输出Queue分配内存。

  比如，为输入x的Queue分配内存，可以通过如下代码段实现：

  1	pipe.InitBuffer(inQueueX, 1, TOTAL_LENGTH * sizeof(half))

具体的初始化函数代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

constexpr int32_t TOTAL_LENGTH = 1 * 2048;  // 数据总长
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    // 设置Global Memory的起始地址以及长度
    xGm.SetGlobalBuffer((__gm__ half *)x, TOTAL_LENGTH);
    yGm.SetGlobalBuffer((__gm__ half *)y, TOTAL_LENGTH);
    zGm.SetGlobalBuffer((__gm__ half *)z, TOTAL_LENGTH);

    // 通过Pipe内存管理对象为输入输出Queue分配内存
    pipe.InitBuffer(inQueueX, BUFFER_NUM, TOTAL_LENGTH * sizeof(half));
    pipe.InitBuffer(inQueueY, BUFFER_NUM, TOTAL_LENGTH * sizeof(half));
    pipe.InitBuffer(outQueueZ, BUFFER_NUM, TOTAL_LENGTH * sizeof(half));
}

基于矢量编程范式，将核函数的实现分为3个基本任务：CopyIn，Compute，CopyOut。Process函数中通过如下方式调用这三个函数。

1
2
3
4
5
6

__aicore__ inline void Process()
{
    CopyIn();
    Compute();
    CopyOut();
}

根据编程范式上面的算法分析，将整个计算拆分成三个Stage，用户单独编写每个Stage的代码，三阶段流程示意图参见图3，具体流程如下：

1. Stage1：CopyIn函数实现。
   1. 使用DataCopy接口将GlobalTensor数据拷贝到LocalTensor。
   2. 使用EnQue将LocalTensor放入VECIN的Queue中。

   1 2 3 4 5 6 7 8 9 10 11 12

   __aicore__ inline void CopyIn(int32_t progress) { // 从Que中为LocalTensor分配内存 AscendC::LocalTensor<half> xLocal = inQueueX.AllocTensor<half>(); AscendC::LocalTensor<half> yLocal = inQueueY.AllocTensor<half>(); // 将GlobalTensor数据拷贝到LocalTensor AscendC::DataCopy(xLocal, xGm, TOTAL_LENGTH); AscendC::DataCopy(yLocal, yGm, TOTAL_LENGTH); // LocalTensor放入VECIN的Queue中 inQueueX.EnQue(xLocal); inQueueY.EnQue(yLocal); }

2. Stage2：Compute函数实现。
   1. 使用DeQue从VECIN中取出LocalTensor。
   2. 使用Ascend C接口Add完成矢量计算。
   3. 使用EnQue将计算结果LocalTensor放入到VECOUT的Queue中。
   4. 使用FreeTensor释放不再使用的LocalTensor。

   1 2 3 4 5 6 7 8 9 10 11 12 13 14

   __aicore__ inline void Compute(int32_t progress) { // 将Input从VECIN的Queue中取出 AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>(); AscendC::LocalTensor<half> yLocal = inQueueY.DeQue<half>(); AscendC::LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>(); // 调用Add算子进行计算 AscendC::Add(zLocal, xLocal, yLocal, TOTAL_LENGTH); // 将计算结果LocalTensor放入到VECOUT的Queue中 outQueueZ.EnQue<half>(zLocal); // 释放LocalTensor inQueueX.FreeTensor(xLocal); inQueueY.FreeTensor(yLocal); }

3. Stage3：CopyOut函数实现。
   1. 使用DeQue接口从VECOUT的Queue中取出LocalTensor。
   2. 使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。
   3. 使用FreeTensor将不再使用的LocalTensor进行回收。

   1 2 3 4 5 6 7 8 9

   __aicore__ inline void CopyOut(int32_t progress) { // 将计算结果从到VECOUT的Queue中取出 AscendC::LocalTensor<half> zLocal = outQueueZ.DeQue<half>(); // 将计算结果从LocalTensor数据拷贝到GlobalTensor AscendC::DataCopy(zGm, zLocal, TOTAL_LENGTH); // 释放LocalTensor outQueueZ.FreeTensor(zLocal); }