Device侧代码与说明

Device侧代码示例如下，并参考本节内容了解主要操作。

// Device侧文件：kernel_aic.cce
#ifdef __CCE_KT_TEST__
#define __aicore__ 
#else
#define __aicore__ [aicore]
#endif
extern "C" __global__ __aicore__ void mat_mul_kernel(__gm__ float* __restrict__ tensor_a, __gm__ float* __restrict__ tensor_b, __gm__ float* __restrict__ tensor_c_gm) {
    set_l1_3d_size(0);  // 设置L1循环缓冲区大小为整个L1
    set_padding(0);     // 设置pad大小为0
    set_mask_norm();    // 设置mask模式为norm，使能mask寄存器
    set_atomic_none();  // 设置不使用Atomic模式
    if (block_idx < 2) { // tiling相关分核策略，这里只使用两个核
        __cbuf__ float* tensor_a_fractal = (__cbuf__ float*)get_imm(0);       // 在L1上分配矩阵A的起始地址
        __cbuf__ float* tensor_b_fractal = (__cbuf__ float*)get_imm(2048);    // 在L1上分配矩阵B的起始地址
        __cc__ float* tensor_mmad = (__cc__ float*)get_imm(0);                // 在L0C上分配结果矩阵的起始地址
        __ca__ float* tensor_a_zz = (__ca__ float*)get_imm(0);                // 在L0A上分配矩阵A的起始地址
        __cb__ float* tensor_b_zn = (__cb__ float*)get_imm(0);                // 在L0B上分配矩阵B的起始地址
        copy_gm_to_cbuf_multi_nd2nz_b32s(tensor_a_fractal, tensor_a + block_idx * 256, 0, 1, 16, 16, 1, 16, 16, 1, 1);
        set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
        copy_gm_to_cbuf_multi_nd2nz_b32s(tensor_b_fractal, tensor_b, 0, 1, 16, 32, 1, 32, 16, 1, 1);
        set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
        wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
        load_cbuf_to_ca(tensor_a_zz, tensor_a_fractal, 0, 2, 1, 0, 0, 0, inc);
        wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
        load_cbuf_to_cb(tensor_b_zn, tensor_b_fractal, 0, 4, 1, 0, 0, 0, inc);
        set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
        wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
        mad(tensor_mmad, tensor_a_zz, tensor_b_zn, 16, 16, 32, 0, 0, 0, 1);
        set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
        set_nd_para(4295032833);
        wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
        copy_matrix_cc_to_gm(tensor_c_gm + block_idx * 512, tensor_mmad, 0, 32, 16, 32, 16, 0, NoQuant, 0, 0, 1);
    }
    pipe_barrier(PIPE_ALL);
}

1. GM搬入L1

   在配置完各种所需寄存器，并分配好计算的输入输出地址空间后，首先用copy_gm_to_cbuf_multi_nd2nz_b32s分别将A、B矩阵从GM拷贝到L1上，其接口原型为：

   void copy_gm_to_cbuf_multi_nd2nz_b32s(__cbuf__ float *dst, __gm__ float *src, uint8_t sid, uint16_t ndNum, uint16_t nValue, uint16_t dValue, uint16_t srcNdMatrixStride, uint16_t srcDValue, uint16_t dstNzC0Stride, uint16_t dstNzNStride, uint16_t dstNzMatrixStride)

   • dst：L1上的起始地址，对应__cbuf__标识符。
   • src：GM上的起始地址，对应__gm__标识符；由于A矩阵切分为两块，一个核上处理的数据为16 * 16 = 256，所以第二个核处理的数据的起始地址偏移了block_idx * 256。
   • sid：预留参数，此处默认为0即可。
   • ndNum：一次计算一个nd块，故填1。
   • nValue：A矩阵竖直大小为16，B矩阵竖直大小为16。
   • dValue：A矩阵水平大小为16，B矩阵水平大小为32。
   • srcNDMatrixStride：只有一个nd块。
   • srcDValue：实际水平方向无无效数据，所以与dValue相同。
   • dstNzC0Stride：目的数据块竖直大小不变，与nValue相同。
   • dstNzNStride：只有一个nd块。
   • dstNzMatrixStride：只有一个nd块。
2. 设置同步等待搬运完成

   以上两条搬运指令结束后，都紧跟了一条同步指令set_flag，是由于以上从GM搬运到L1的数据，还需要搬运到L0A或L0B。所以在前置指令完成前，不能进行后续操作，故在copy_gm_to_cbuf_multi_nd2nz_b32s插入set_flag，并在load_cbuf_to_ca/load_cbuf_to_cb前插入wait_flag，以确保其顺序执行。其接口原型为：

   void set_flag(pipe_t pipe, pipe_t tpipe, event_t pipeID)
   void wait_flag(pipe_t pipe, pipe_t tpipe, event_t pipeID)

   • pipe：前置操作流水线类型，对应copy_gm_to_cbuf_multi_nd2nz_b32s的流水为PIPE_MTE2。
   • tpipe：依赖操作的流水线类型，对应load_cbuf_to_ca/load_cbuf_to_cb的流水线为PIPE_MTE1。
   • pipeID：用两个ID区分A、B矩阵的搬运操作。
3. L1搬入计算缓存（L0A/L0B）

   等待前置指令结束，数据被完整搬运到L1上之后，同步指令等待结束，接下来使用load_cbuf_to_ca/load_cbuf_to_cb将A、B矩阵分别搬运到L0A和L0B上，其接口原型为：

   void load_cbuf_to_ca(__ca__ float *dst, __cbuf__ float *src, uint16_t baseIdx, uint8_t repeat, uint16_t srcStride, uint16_t dstStride, uint8_t sid, bool transpose, addr_cal_mode_t addr_cal_mode)
   void load_cbuf_to_cb(__cb__ float *dst, __cbuf__ float *src, uint16_t baseIdx, uint8_t repeat, uint16_t srcStride, uint16_t dstStride, uint8_t sid, bool transpose, addr_cal_mode_t addr_cal_mode)

   • dst：L0A/L0B上的起始地址，分别对应__ca__和__cb__标识符。
   • src：L1上的起始地址。
   • baseIdx：A，B矩阵都从头开始所以都为0。
   • repeat：对于A矩阵大小为16*16*sizeof（float）=1024B，每个分形大小为512B，所以重复两次；对于B矩阵大小为16*32*sizeof（float）=2048B，所以重复四次。
   • srcStride：每次偏移一个分型。
   • dstStride：连续存储设为0。
   • sid：预留参数，此处默认为0即可。
   • transpose：无需转置。
   • addr_cal_mode：inc表示src地址是增加，表示数据连续往后顺序读取。
4. 设置同步等待搬运完成。

   由于后续mad计算依赖数据完整地搬运到L0A和L0B，所以在mad之前也需要插入同步指令，需要完成的流水类型改为PIPE_MTE1，需要等待后执行的流水类型为PIPE_M。

5. 矩阵乘法计算。

   同步等待结束后，即可使用mad接口进行矩阵乘法，其接口原型为：

   void mad(__cc__ float *c, __ca__ float *a, __cb__ float *b, uint16_t m, uint16_t k, uint16_t n, uint8_t unitFlag, bool kDirectionAlign, bool cmatrixSource, bool cmatrixInitVal)

   • c：L0C起始地址，对应__cc__标识符。
   • a：L0A起始地址。
   • b：L0B起始地址。
   • m：A矩阵竖直方向大小为16。
   • k：A矩阵水平方向大小为16，也就是B矩阵竖直方向大小为16。
   • n：B矩阵水平方向大小为32。
   • unitFlag：设置为0即可。
   • kDirectionAlign：无需对齐设置为0。
   • cmatrixcSource：无bais，初始c矩阵在L0C上，设置为0。
   • cmatrixInitVal：设置为1，使得C矩阵初始值为0。
6. 计算结果搬出。

   之后依旧需要一对同步指令，使得矩阵乘的结果全部计算完之后，再搬出。这里为了展示NZ2ND模式，使用set_nd_para对nd块进行设置，但由于只有一块数据所以在第0位、第16位、第32位设置1即可，转换为十进制对应4295032833。搬出时使用的接口为copy_matrix_cc_to_gm，接口原型为：

   void copy_matrix_cc_to_gm(__gm__ float *dst, __cc__ float *src, uint8_t sid, uint16_t NSize, uint16_t MSize, uint32_t dstStride_dst_D, uint16_t srcStride, uint8_t UnitFlagMode, uint64_t QuantPRE, uint8_t ReLUPRE, bool channelSplit, bool NZ2ND_EN)

   • dst：GM上的起始地址。由于分两个核计算，一个核上处理的数据为16 * 32 = 512，所以每个核处理的数据的起始地址偏移了block_idx * 512。
   • src：L0C上的起始地址。
   • sid：预留参数，此处默认为0即可。
   • NSize：C矩阵水平方向大小32。
   • MSize：C矩阵竖直方向大小为16。
   • dstStride_dst_D：使能了NZ2ND，ND矩阵每行中元素个数为NSize即32。
   • srcStride：只有一个数据块，要求为16的倍数。
   • UnitFlagMode：填0即可。
   • QuantPRE：不适用量化，设为NoQuant。
   • ReLUPRE：不使用，设为0。
   • channelSpilt：不进行通道切分，设为0。
   • NZ2ND_EN：使能NZ2ND类型转换，设为1。
7. 添加以下指令，确保此kernel全部执行完再执行其他kernel。

   pipe_barrier(PIPE_ALL)