分离模式

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

[!NOTE]说明 本节内容为针对分离模式，使用基础API进行矩阵乘法的编程指导。

针对分离模式，由于硬件架构差异，在使用基础API进行矩阵乘法算子的实现上存在差异：

• 分离模式中，Cube计算单元与Vector计算单元分离部署，每个核有自己的Scalar单元，能独立加载自己的代码段。在核函数编写时，可通过函数类型限定符__cube__标识该核函数在设备端aicore上的Cube计算单元执行。
• 分离模式中支持Fixpipe硬件加速，支持从L0C Buffer直接搬出数据到L1 Buffer/Global Memory/Unified Buffer的同时，完成量化、反量化、数据排布格式转换等功能。

遵循算子分析、核函数定义、算子类实现的开发流程，下面以Matmul算子为例，给出分离模式下Matmul算子的代码框架。

算子分析

Matmul算子完成形状为[m, k]的矩阵a和形状为[k, n]的矩阵b的矩阵乘，得到形状为[m, n]的矩阵c。为了方便这里取m=k=n=32。Matmul算子的输入输出、计算逻辑以及需要调用的Ascend C接口如下表所示。

表1Ascend C Matmul算子设计规格

核函数定义

根据核函数中介绍的规则进行核函数的定义。核函数名为matmul_custom，有3个参数a，b，c，其中a，b都为输入内存，c为输出内存。使用函数类型限定符__global__来标识它是一个核函数，可以被<<<>>>调用；使用函数类型限定符__cube__来标识该核函数在设备端aicore上的Cube核执行。

核函数中，算子类的Init函数，完成内存初始化相关工作，Process函数完成算子实现的核心逻辑。核函数在开始时调用AscendC::InitSocState()初始化硬件状态，结束时调用AscendC::PipeBarrier<PIPE_ALL>()等待所有指令完成。

template <uint32_t M, uint32_t K, uint32_t N> __global__ __cube__ void mmad_custom(__gm__ uint8* a, __gm__ uint8* b, __gm__ uint8* c) { AscendC::InitSocState(); KernelMatmul<M, K, N> op; op.Init(a, b, c); op.Process(); AscendC::PipeBarrier<PIPE_ALL>(); }

[!NOTE]说明 核函数使用模板参数传入矩阵的维度信息，这样可以在编译期确定循环次数，有利于编译器优化。其中M/K/N为总矩阵维度。

算子类实现

在分离模式中，一个Matmul矩阵乘算子的实现分为以下四步流水。

1. CopyIn任务（Globale Memory --> L1 Buffer）
   使用DataCopy接口将Global Memory 中的数据搬运到L1 Buffer，将ND格式转换为NZ格式。

2. Split任务（L1 Buffer--> L0 Buffer）
   使用LoadData（2D矩阵搬运）接口将数据从L1 Buffer中搬运到L0A Buffer/L0B Buffer，同时完成NZ到ZZ/ZN的格式转换。

3. Compute任务（矩阵乘计算）
   使用Mmad接口在完成矩阵乘计算，结果存放在L0C Buffer。

4. CopyOut任务（L0C Buffer--> Globale Memory）
   使用Fixpipe接口将结果从L0C Buffer直接搬运到Global Memory，同时利用Fixpipe的随路格式转换功能将NZ格式转回ND格式。

一个典型的Matmul算子类的代码框架如下所示。

template <uint32_t M, uint32_t K, uint32_t N> class KernelMatmul { public: __aicore__ inline KernelMatmul() {} __aicore__ inline void Init(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c) { aGM.SetGlobalBuffer((__gm__ half*)a); bGM.SetGlobalBuffer((__gm__ half*)b); cGM.SetGlobalBuffer((__gm__ half*)c); } __aicore__ inline void Process() { AscendC::LocalTensor<half> a1Local(AscendC::TPosition::A1, a1Addr, M * K); AscendC::LocalTensor<half> a2Local(AscendC::TPosition::A2, a2Addr, M * K); AscendC::LocalTensor<half> b1Local(AscendC::TPosition::B1, b1Addr, K * N); AscendC::LocalTensor<half> b2Local(AscendC::TPosition::B2, b2Addr, K * N); AscendC::LocalTensor<float> cLocal(AscendC::TPosition::CO1, cAddr, M * N); CopyInA(a1Local); CopyInB(b1Local); AscendC::PipeBarrier<PIPE_ALL>(); DataLoadA(a1Local, a2Local); DataLoadB(b1Local, b2Local); AscendC::PipeBarrier<PIPE_ALL>(); Compute(cLocal, a2Local, b2Local); AscendC::PipeBarrier<PIPE_ALL>(); CopyOut(cLocal); } private: // 搬入函数，完成矩阵A从GM到L1的CopyIn任务 __aicore__ inline void CopyInA(...) () { // ... } // 搬入函数，完成矩阵B从GM到L1的CopyIn任务 __aicore__ inline void CopyInB(..) () { // ... } // 搬入函数，完成矩阵A从L1到L0A的Split任务 __aicore__ inline void DataLoadA(...) () { // ... } // 搬入函数，完成矩阵B从L1到L0B的Split任务 __aicore__ inline void DataLoadB(...) () { // ... } // 计算函数，完成L0上的计算任务 __aicore__ inline void Compute(...) () { // ... } // 搬出函数，完成矩阵C从L0C到GM的CopyOut任务 __aicore__ inline void CopyOut(...) () { // ... } private: // 私有成员变量 // ... };

[!NOTE]说明 CopyIn、DataLoad、Compute和CopyOut四个阶段的操作分别在不同的硬件流水线上执行。为保证数据依赖正确——即每个阶段读取的数据必须由前一阶段写入完毕，需要在各阶段之间设置同步屏障。这里使用全局流水同步PipeBarrier.md)。

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit