算子类实现

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

根据矢量编程范式对Add算子的实现流程进行设计的思路如下,矢量编程范式请参考流水编程范式中的矢量编程部分,设计完成后得到的Add算子实现流程图参见图1

图1 Add算子实现流程

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

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:
    TPipe pipe;  //Pipe内存管理对象
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;  //输入数据Queue队列管理对象,QuePosition为VECIN
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;  //输出数据Queue队列管理对象,QuePosition为VECOUT
    GlobalTensor<half> xGm, yGm, zGm;  //管理输入输出Global Memory内存地址的对象,其中xGm, yGm为输入,zGm为输出
};

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; // seperate to 2 parts, due to double buffer
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    // get start index for current core, core parallel
    xGm.SetGlobalBuffer((__gm__ half*)x + BLOCK_LENGTH * GetBlockIdx(), BLOCK_LENGTH);
    yGm.SetGlobalBuffer((__gm__ half*)y + BLOCK_LENGTH * GetBlockIdx(), BLOCK_LENGTH);
    zGm.SetGlobalBuffer((__gm__ half*)z + BLOCK_LENGTH * GetBlockIdx(), 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));
}

Process函数实现

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

    __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);
        }
    }

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

  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);
        }

  2. Stage2:Compute函数实现。

    1. 使用DeQue从VecIn中取出LocalTensor。
    2. 使用Ascend C接口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);
    }

  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);
    }