简介

本节内容为使用基础API进行矩阵乘法的编程指导。使用基础API进行矩阵编程的功能支持的产品型号为：
  • Atlas 推理系列产品
  • Atlas 训练系列产品

编程范式

Cube编程范式把算子的实现流程分为5个基本任务：CopyIn，Split，Compute，Aggregate，CopyOut。CopyIn负责搬入操作，Split负责数据切分操作，Compute负责矩阵指令计算操作，Aggregate负责数据汇聚操作，CopyOut负责搬出操作。

图1 矩阵编程基本任务设计

具体任务之间的交互流程和流程图如下。

  1. Stage1：CopyIn任务。

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

  2. Stage2：Split任务。

    1. 使用DeQue从A1/B1中取出LocalTensor。
    2. 使用Ascend C接口将LocalTensor从A1/B1中搬运到A2/B2。
    3. 使用EnQue将计算结果LocalTensor放入到A2/B2的Queue中。

  3. Stage3：Compute任务。

    1. 使用DeQue从A2/B2中取出LocalTensor。
    2. 使用Ascend C接口完成矩阵计算。
    3. 使用EnQue将计算结果LocalTensor放入到CO1的Queue中。

  4. Stage4：Aggregate任务。

    1. 使用DeQue从CO1中取出LocalTensor。
    2. 使用Ascend C接口拷贝结果矩阵到CO2。
    3. 使用EnQue将计算结果LocalTensor放入到CO2的Queue中。

  5. Stage5：CopyOut任务。

    1. 使用DeQue接口从CO2的Queue中取出LocalTensor。
    2. 使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。

图2 矩阵编程Queue队列

开发流程

基于Ascend C方式实现矩阵算子的流程如下图所示。

图3 矩阵算子实现流程

下文将以Matmul算子为例对上述步骤进行详细介绍，Matmul算子的代码框架如下，完整代码请参见实现样例。

#include "kernel_operator.h"

// 根据编程范式实现算子类
class KernelMatmul {
public:
    __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR c)
    {
        // ...
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        SplitA();
        AscendC::LocalTensor<half> b1Local = inQueueB1.DeQue<half>();
        AscendC::LocalTensor<half> a2Local = inQueueA2.DeQue<half>();
        AscendC::LocalTensor<float> c2Local = outQueueCO2.AllocTensor<float>();
        // split matrix b into 2 parts, [32, 16] and [32, 16]
        for (int i = 0; i < 2; ++i) {
            SplitB(b1Local, i);
            Compute(a2Local);
            Aggregate(c2Local, i);
        }
        inQueueB1.FreeTensor(b1Local);
        inQueueA2.FreeTensor(a2Local);
        outQueueCO2.EnQue<float>(c2Local);
        CopyOut();
    }
private:
    __aicore__ inline void CopyIn()
    {
        // ...
    }
    __aicore__ inline void SplitA()
    {
        // ...
    }
    __aicore__ inline void SplitB(const LocalTensor<half>& b1Local, const int bSplitIdx)
    {
        // ...
    }
    __aicore__ inline void Compute(const LocalTensor<half>& a2Local)
    {
        // ...
    }
    __aicore__ inline void Aggregate(const LocalTensor<float>& c2Local, const int bSplitIdx)
    {
        // ...
    }
    __aicore__ inline void CopyOut()
    {
        // ...
    }
private:
    // ...

};

//核函数定义
extern "C" __global__ __aicore__ void matmul_custom(GM_ADDR a, GM_ADDR b, GM_ADDR c)
{
    KernelMatmul op;
    op.Init(a, b, c);
    op.Process();
}
父主题： 矩阵编程（基础API）