Kernel直调算子开发

开发流程

算子分析：分析算子的数学表达式，输入、输出以及计算逻辑的实现，明确需要调用的Ascend C接口
核函数定义：定义Ascend C算子入口函数
根据编程范式实现算子类：完成核函数的内部实现
编写算子的应用程序：完成调用核函数main.cpp代码
其他脚本：数据生成脚本，数据比对脚本

算子设计

以ElemWise(Add)算子为例，数学公式

$$\overrightarrow{z}= \overrightarrow{x}+\overrightarrow{y}$$

为简单起见，设定输入张量x,y和输出张量z为固定shape(8,2048)，数据类型dtype为half类型，数据排布类型format为ND，核函数名称为add_custom

算子分析

明确算子的数学表达式及计算逻辑
Add算子的数学表达式为：

$$\overrightarrow{z}= \overrightarrow{x}+\overrightarrow{y}$$

计算逻辑：输入数据需要先搬入到片上存储，然后使用计算接口完成两个加法运算，得到最终结果，再搬出到外部存储

明确输入和输出

$$Add算子有两个输入：\overrightarrow{x}和\overrightarrow{y}，输出为\overrightarrow{z}$$

输入数据类型为half，输出数据类型与输入数据类型相同，输入支持固定shape(8,2048)，输出shape与输入shape相同，输入数据排布类型为ND

确定核函数名和参数
自定义核函数名，如add_custom，根据输入输出，确定核函数有3个入参x,y,z
x,y为输入在Global Memory上的内存地址，z为输出在Global Memory上的内存地址

确定算子实现所需接口
涉及内外部存储间的数据搬运，使用数据搬移接口：DataCopy实现
涉及矢量计算的加法操作，使用适量双目命令：Add实现
使用到LocalTensor，使用Queue队列管理，会使用到EnQue，DeQue等接口。

算子类实现

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内存管理对象对任务间交互使用到的内存，临时变量使用到的内存统一进行管理。

算子类实现

算子类名：KernelAdd
初始化函数Init()
核心处理函数Process
流水任务：
CopyIn()
Compute()
CopyOut()

Init函数实现

Init()函数实现

使用多核并行计算，需要将数据切片，获取到每个核实际需要处理在Global Memory上的内存偏移地址
数据整体长度TOTAL_LENGTH为8*2048，平均分配到8个核上运行，每个核上处理的数据大小BLOCK_LENGTH为2048。block_idx为核的逻辑ID，(gm half*)x+block_idx*BLOCK_LENGTH即索引为block_idx的核的输入数据在Global Memory上的内存偏移地址

对于单核处理数据，可以进行数据切块(Tilling)，将数据切分成8块，切分后的每个数据块再次切分成BUFFER_NUM=2块，可开启double buffer，实现流水线之间的并行

单核需要处理的2048个数被切分成16块，每块TILE_LENGTH=128个数据，Pipe为inQueue分配了BUFFER_NUM块大小为TILE_LENGTH*sizeof(half)个字节的内存块，每个内存块能容量TILE_LENGTH=128个half类型数据

Process()的实现

完整的代码

#include "kernel_operator.h"

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; // separate to 2 parts, due to double buffer

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}
    // 初始化函数，完成内存初始化相关操作
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z)
    // 核心处理函数，实现算子逻辑，调用私有成员函数CopyIn，Compute，CopyOute完成算子逻辑
    {
        xGm.SetGlobalBuffer((__gm__ half *)x + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        yGm.SetGlobalBuffer((__gm__ half *)y + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        zGm.SetGlobalBuffer((__gm__ half *)z + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        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));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = TILE_NUM * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    // 调入函数，完成CopyIn阶段的处理，被Process函数调用
    {
        AscendC::LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        AscendC::LocalTensor<half> yLocal = inQueueY.AllocTensor<half>();
        AscendC::DataCopy(xLocal, xGm[progress * TILE_LENGTH], TILE_LENGTH);
        AscendC::DataCopy(yLocal, yGm[progress * TILE_LENGTH], TILE_LENGTH);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    // 计算函数，完成Compute阶段的处理，被Process函数调用
    {
        AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        AscendC::LocalTensor<half> yLocal = inQueueY.DeQue<half>();
        AscendC::LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
        AscendC::Add(zLocal, xLocal, yLocal, TILE_LENGTH);
        outQueueZ.EnQue<half>(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    // 调出函数，完成CopyOut阶段的处理，被Process函数调用
    {
        AscendC::LocalTensor<half> zLocal = outQueueZ.DeQue<half>();
        AscendC::DataCopy(zGm[progress * TILE_LENGTH], zLocal, TILE_LENGTH);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    // Tpipe内存管理对象
    AscendC::TPipe pipe;
    // 输入数据Queue队列管理对象，QuePositon为VECIN
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    // 输出数据Queue队列管理对象，QuePosition为VECOUT
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    // 管理输入输出Global Memory内存地址的对象，其中xGM，yGM为输入，zGm为输出
    AscendC::GlobalTensor<half> xGm;
    AscendC::GlobalTensor<half> yGm;
    AscendC::GlobalTensor<half> zGm;
};

extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    KernelAdd op;
    op.Init(x, y, z);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void add_custom_do(uint32_t blockDim, void *stream, uint8_t *x, uint8_t *y, uint8_t *z)
{
    add_custom<<<blockDim, nullptr, stream>>>(x, y, z);
}
#endif