CANN/asc-devkit分离模式矩阵乘法实现

分离模式【免费下载链接】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。为了方便这里取mkn32。Matmul算子的输入输出、计算逻辑以及需要调用的Ascend C接口如下表所示。表1Ascend C Matmul算子设计规格项目内容算子类型OpTypeMatmul算子输入ashape(m, k) (32, 32)data typehalfformatND算子输入bshape(k, n) (32, 32)data typehalfformatND算子输出cshape(m, n) (32, 32)data typehalfformatND核函数名称matmul_custom使用的主要接口DataCopyGlobal Memory到L1 Buffer数据搬运 ND转NZ格式转换接口LoadData2D矩阵搬运L1 Buffer到L0 Buffer数据搬运 NZ转ZZ/ZN格式转换接口Mmad矩阵乘计算接口FixpipeL0C Buffer到Global Memory数据搬出 NZ转ND格式转换 精度转换接口核函数定义根据核函数中介绍的规则进行核函数的定义。核函数名为matmul_custom有3个参数abc其中ab都为输入内存c为输出内存。使用函数类型限定符__global__来标识它是一个核函数可以被调用使用函数类型限定符__cube__来标识该核函数在设备端aicore上的Cube核执行。核函数中算子类的Init函数完成内存初始化相关工作Process函数完成算子实现的核心逻辑。核函数在开始时调用AscendC::InitSocState()初始化硬件状态结束时调用AscendC::PipeBarrierPIPE_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(); KernelMatmulM, K, N op; op.Init(a, b, c); op.Process(); AscendC::PipeBarrierPIPE_ALL(); }[!NOTE]说明 核函数使用模板参数传入矩阵的维度信息这样可以在编译期确定循环次数有利于编译器优化。其中M/K/N为总矩阵维度。算子类实现在分离模式中一个Matmul矩阵乘算子的实现分为以下四步流水。CopyIn任务Globale Memory -- L1 Buffer使用DataCopy接口将Global Memory 中的数据搬运到L1 Buffer将ND格式转换为NZ格式。Split任务L1 Buffer-- L0 Buffer使用LoadData2D矩阵搬运接口将数据从L1 Buffer中搬运到L0A Buffer/L0B Buffer同时完成NZ到ZZ/ZN的格式转换。Compute任务矩阵乘计算使用Mmad接口在完成矩阵乘计算结果存放在L0C Buffer。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::LocalTensorhalf a1Local(AscendC::TPosition::A1, a1Addr, M * K); AscendC::LocalTensorhalf a2Local(AscendC::TPosition::A2, a2Addr, M * K); AscendC::LocalTensorhalf b1Local(AscendC::TPosition::B1, b1Addr, K * N); AscendC::LocalTensorhalf b2Local(AscendC::TPosition::B2, b2Addr, K * N); AscendC::LocalTensorfloat cLocal(AscendC::TPosition::CO1, cAddr, M * N); CopyInA(a1Local); CopyInB(b1Local); AscendC::PipeBarrierPIPE_ALL(); DataLoadA(a1Local, a2Local); DataLoadB(b1Local, b2Local); AscendC::PipeBarrierPIPE_ALL(); Compute(cLocal, a2Local, b2Local); AscendC::PipeBarrierPIPE_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考