下载
中文
注册

简介

本节内容为使用基础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 矩阵算子实现流程
  • 算子分析:分析算子的数学表达式、输入、输出以及计算逻辑的实现,明确需要调用的Ascend C接口。
  • 核函数定义:定义Ascend C算子入口函数。
  • 根据矩阵编程范式实现算子类:完成核函数的内部实现,调用私有成员函数CopyIn、SplitA、SplitB、Compute、Aggregate、CopyOut完成矩阵算子的五级流水操作。

下文将以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();
}