公司动态

CANN/asc-devkit耦合模式矩阵编程指南

📅 2026/7/18 10:25:27
CANN/asc-devkit耦合模式矩阵编程指南
耦合模式【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit[!NOTE]说明 本节内容为针对耦合模式使用基础API进行矩阵乘法的编程指导。编程范式Cube编程范式把算子的实现流程分为5个基本任务CopyInSplitComputeAggregateCopyOut。CopyIn负责搬入操作Split负责数据切分操作Compute负责矩阵指令计算操作Aggregate负责数据汇聚操作CopyOut负责搬出操作。图1矩阵编程基本任务设计![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/矩阵编程基本任务设计.png 矩阵编程基本任务设计?utm_sourcegitcode_repo_files)具体任务之间的交互流程和流程图如下。Stage1CopyIn任务。使用DataCopy接口将GlobalTensor数据拷贝到LocalTensor。使用EnQue将LocalTensor放入L1 BufferA1/B1的Queue中。Stage2Split任务。使用DeQue从L1 BufferA1/B1中取出LocalTensor。使用LoadData接口将LocalTensor从L1 BufferA1/B1搬运到L0A/L0B BufferA2/B2。使用EnQue将计算结果LocalTensor放入到L0A/L0B BufferA2/B2的Queue中。Stage3Compute任务。使用DeQue从L0A/L0B BufferA2/B2中取出LocalTensor。使用Mmad接口完成矩阵计算。使用EnQue将计算结果LocalTensor放入到L0C BufferCO1的Queue中。Stage4Aggregate任务。使用DeQue从L0C BufferCO1中取出LocalTensor。使用Ascend C接口拷贝结果矩阵到逻辑位置CO2。CO2的物理存储映射与产品型号相关具体请参考逻辑位置和物理存储的映射关系。使用EnQue将计算结果LocalTensor放入到CO2对应的Queue中。Stage5CopyOut任务。使用DeQue接口从CO2对应的Queue中取出LocalTensor。使用DataCopy接口将LocalTensor拷贝到GlobalTensor上。图2矩阵编程Queue队列![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/矩阵编程Queue队列.png 矩阵编程Queue队列?utm_sourcegitcode_repo_files)开发流程基于Ascend C方式实现矩阵算子的流程如下图所示。图3矩阵算子实现流程![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/矩阵算子实现流程.png 矩阵算子实现流程?utm_sourcegitcode_repo_files)算子分析分析算子的数学表达式、输入、输出以及计算逻辑的实现明确需要调用的Ascend C接口。核函数定义定义Ascend C算子入口函数。根据矩阵编程范式实现算子类完成核函数的内部实现调用私有成员函数CopyIn、SplitA、SplitB、Compute、Aggregate、CopyOut完成矩阵算子的五级流水操作。下文将以Matmul算子为例对上述步骤进行详细介绍Matmul算子的代码框架如下。#include kernel_operator.h // 根据编程范式实现算子类 class KernelMatmul { public: __aicore__ inline void Init(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c) { // ... } __aicore__ inline void Process() { CopyIn(); SplitA(); AscendC::LocalTensorhalf b1Local inQueueB1.DeQuehalf(); AscendC::LocalTensorhalf a2Local inQueueA2.DeQuehalf(); AscendC::LocalTensorfloat c2Local outQueueCO2.AllocTensorfloat(); // 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.EnQuefloat(c2Local); CopyOut(); } private: __aicore__ inline void CopyIn() { // ... } __aicore__ inline void SplitA() { // ... } __aicore__ inline void SplitB(const LocalTensorhalf b1Local, const int bSplitIdx) { // ... } __aicore__ inline void Compute(const LocalTensorhalf a2Local) { // ... } __aicore__ inline void Aggregate(const LocalTensorfloat c2Local, const int bSplitIdx) { // ... } __aicore__ inline void CopyOut() { // ... } private: // ... }; //核函数定义 extern C __global__ __aicore__ void matmul_custom(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c) { KernelMatmul op; op.Init(a, b, c); op.Process(); }算子分析在开发算子代码之前需要分析算子的数学表达式、输入、输出以及计算逻辑的实现明确需要调用的Ascend C接口。明确算子的数学表达式及计算逻辑。Matmul算子完成矩阵乘操作其数学表达式如下形状为[m, k]的矩阵a和形状为[k, n]的矩阵b相乘得到形状为[m, n]的矩阵c。为了方便令mkn32。c a * b注意需要处理的数据过大时需要对数据进行切分并分块搬运到L0A BufferA2、L0B BufferB2分别计算后再进行汇聚。下文的计算逻辑为了展示Split和Aggregate阶段的样例请您根据实际需要处理的数据大小决定是否需要切分和汇聚。计算逻辑如下将输入数据矩阵a、b分别搬运至L1 BufferA1、L1 BufferB1。将a矩阵从L1 BufferA1搬运至L0A BufferA2。为实现部分并行将b矩阵切分为part1和part2形状均为[k, n / 2]切分后再分块搬运至L0B BufferB2。a矩阵和b矩阵part1、part2分别做矩阵乘运算获得矩阵c的part1和part2形状均为[m, n / 2]。计算结果在L0C BufferCO1存储。将矩阵c的part1和part2分别拷贝到逻辑位置CO2进行合并。将合并后的输出数据从逻辑位置CO2搬出。明确输入和输出。Matmul算子有两个输入a与b输出为c。本样例中算子输入支持的数据类型为halffloat16算子输出的数据类型为float32。矩阵a、b、c的形状均为[32, 32]。算子输入输出支持的数据格式为ND。确定核函数名称和参数。您可以自定义核函数名称本样例中核函数命名为matmul_custom。根据对算子输入输出的分析确定核函数有3个参数abcab为输入在Global Memory上的内存地址c为输出在Global Memory上的内存地址。约束分析。由于硬件架构对矩阵乘计算的输入输出有格式约束需要在算子实现中增加格式转换的流程。将矩阵a、b分别搬运至L1 BufferA1、L1 BufferB1时将ND格式的矩阵a、b转换为NZ格式。从L1 BufferA1搬运矩阵a至L0A BufferA2时将NZ格式的a矩阵转换为ZZ格式从L1 BufferB1搬运矩阵b到L0B BufferB2时将NZ格式的b矩阵转换为ZN格式。将计算结果从逻辑位置CO2搬出时将NZ格式的c矩阵转换为ND格式。数据排布格式的相关介绍详见数据排布格式。确定算子实现所需接口。实现外部存储和内部存储间的数据搬运查看Ascend C API参考中的数据搬移接口具体参考DataCopy。实现矩阵数据格式转换查看Ascend C API参考中的数据转换接口具体参考LoadData。矩阵计算过程涉及矩阵乘法查看Ascend C API参考中的矩阵计算接口具体参考Mmad。计算中使用到的Tensor数据结构使用Queue队列进行管理会使用到EnQue、DeQue等接口。通过以上分析得到Ascend C Matmul算子的计算流程图和设计规格如下图4Matmul算子的计算流程图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/Matmul算子的计算流程图.png Matmul算子的计算流程图?utm_sourcegitcode_repo_files)表1Ascend C Matmul算子设计规格算子类型OpTypeMatmul算子输入nameshapedata typeformata(m, k) (32, 32)halfNDb(k, n) (32, 32)halfND算子输出c(m, n) (32, 32)float32ND核函数名称matmul_custom使用的主要接口DataCopy数据搬移接口LoadData矩阵数据格式转换接口Mmad矩阵乘计算接口EnQue、DeQue等接口Queue队列管理接口核函数定义根据核函数中介绍的规则进行核函数的定义。函数原型定义。本样例中函数名为matmul_custom核函数名称可自定义根据算子分析中对算子输入输出的分析确定有3个参数abc其中ab都为输入内存c为输出内存。根据核函数中核函数的规则介绍函数原型定义如下所示使用__global__函数类型限定符来标识它是一个核函数可以被调用使用__aicore__函数类型限定符来标识该核函数在设备端aicore上执行。extern C __global__ __aicore__ void matmul_custom(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c) { }调用算子类的Init和Process函数。算子类的Init函数完成内存初始化相关工作Process函数完成算子实现的核心逻辑具体介绍参见算子类实现。extern C __global__ __aicore__ void matmul_custom(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c) { KernelMatmul op; op.Init(a, b, c); op.Process(); }对核函数进行封装得到matmul_custom_do函数便于主程序调用。#ifndef ASCENDC_CPU_DEBUG表示该封装函数仅在编译运行NPU侧的算子时会用到编译运行CPU侧的算子时可以直接调用matmul_custom函数。根据核函数定义和调用章节调用核函数时除了需要传入参数abc还需要传入numBlocks核函数执行的核数、动态UB大小无动态UB需求时设置为0、stream应用程序中维护异步操作执行顺序的stream来规定核函数的执行配置。#ifndef ASCENDC_CPU_DEBUG // call of kernel function void matmul_custom_do(uint32_t numBlocks, uint32_t dynUBufSize, void* stream, uint8_t* a, uint8_t* b, uint8_t* c) { matmul_customnumBlocks, dynUBufSize, stream(a, b, c); } #endif算子类实现根据上一章节介绍核函数中会调用算子类的Init和Process函数本章具体讲解基于编程范式实现算子类。矩阵编程范式请参考编程范式。算子类中主要包含对外开放的初始化Init函数和核心处理函数Process以及一些实现中会用到的私有成员。KernelMatmul算子类的定义如下class KernelMatmul { public: __aicore__ inline KernelMatmul(){} // 初始化函数完成内存初始化相关操作 __aicore__ inline void Init(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c){} // 核心处理函数实现算子逻辑 // 调用私有成员函数CopyIn、SplitA、SplitB、Compute、Aggregate、CopyOut完成矩阵算子的五级流水操作 __aicore__ inline void Process(){} private: __aicore__ inline void CopyND2NZ(const LocalTensorhalf dst, const GlobalTensorhalf src, const uint16_t height, const uint16_t width){} // 搬进函数完成编程范式中的CopyIn阶段的处理由Process函数调用 __aicore__ inline void CopyIn(){} // 搬进函数完成编程范式中的Split阶段的处理由Process函数调用 __aicore__ inline void SplitA(){} // 搬进函数完成编程范式中的Split阶段的处理由Process函数循环调用两次分别搬运b矩阵的两个part __aicore__ inline void SplitB(const LocalTensorhalf b1Local, const int bSplitIdx){} // 计算函数完成编程范式中的Compute阶段的处理由Process函数循环调用两次分别计算出矩阵c的两个part __aicore__ inline void Compute(const LocalTensorhalf a2Local){} // 搬出函数完成编程范式中的Aggregate阶段的处理由Process函数循环调用两次分别搬出矩阵c的两个part __aicore__ inline void Aggregate(const LocalTensorfloat c2Local, const int bSplitIdx){} // 搬出函数完成编程范式中的CopyOut阶段的处理由Process函数调用 __aicore__ inline void CopyOut(){} private: AscendC::TPipe pipe; // Pipe内存管理对象管理Queue队列的内存 AscendC::TQueAscendC::TPosition::A1, 1 inQueueA1; // 输入数据的队列TPosition为A1 AscendC::TQueAscendC::TPosition::A2, 1 inQueueA2; // 输入数据的队列TPosition为A2 AscendC::TQueAscendC::TPosition::B1, 1 inQueueB1; // 输入数据的队列TPosition为B1 AscendC::TQueAscendC::TPosition::B2, 2 inQueueB2; // 输入数据的队列TPosition为B2 AscendC::TQueAscendC::TPosition::CO1, 2 outQueueCO1; // 输出数据的队列TPosition为CO1 AscendC::TQueAscendC::TPosition::CO2, 1 outQueueCO2; // 输出数据的队列TPosition为CO2 // 管理输入输出Global Memory内存地址的对象其中aGMbGM为输入cGM为输出 AscendC::GlobalTensorhalf aGM, bGM; AscendC::GlobalTensorfloat cGM; uint16_t m 32; uint16_t n 32; uint16_t k 32; uint16_t aSize, bSize, cSize, mBlocks, nBlocks, kBlocks; };KernelMatmul构造函数实现构造函数中对私有成员变量进行初始化具体代码如下__aicore__ inline KernelMatmul() { aSize m * k; bSize k * n; cSize m * n; mBlocks m / 16; nBlocks n / 16; kBlocks k / 16; }矩阵a的形状为[m, k]矩阵b的形状为[k, n]矩阵c的形状为[m,n]此样例中m、n、k均设置为32。aSize、bSize、cSize分别为矩阵a、b、c的数值个数。mBlocks、 nBlocks、 kBlocks为m、n、k所占分形数量half类型一个分形Shape为16 * 16blocks计算公式为mBlocks m / 16nBlocks n / 16kBlocks k / 16分形具体介绍可参考数据排布格式。Init函数实现Init函数主要完成以下内容设置输入输出Global Tensor的Global Memory内存地址。以设置输入a在Global Memory上的内存偏移地址为例aGM.SetGlobalBuffer((__gm__ half*)a);注意因为本样例中Init函数的入参统一设置为uint8_t*这里需要强转成具体的数据类型(__gm__ half*)再进行偏移。通过Pipe内存管理对象为输入输出Queue分配内存。比如为输入数据队列inQueueB2分配内存可以通过如下代码段实现pipe.InitBuffer(inQueueB2, 2, bSize * sizeof(half) / 2);此样例中将b矩阵切分为两个part为inQueueB2分配内存时需要申请两块内存空间每一块的大小为b矩阵大小的一半outQueueCO1的内存初始化同理。具体的初始化函数代码如下__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__ float*)c); pipe.InitBuffer(inQueueA1, 1, aSize * sizeof(half)); pipe.InitBuffer(inQueueA2, 1, aSize * sizeof(half)); pipe.InitBuffer(inQueueB1, 1, bSize * sizeof(half)); pipe.InitBuffer(inQueueB2, 2, bSize * sizeof(half) / 2); pipe.InitBuffer(outQueueCO1, 2, cSize * sizeof(float) / 2); pipe.InitBuffer(outQueueCO2, 1, cSize * sizeof(float)); }Process函数实现基于矩阵编程范式将核函数的实现分为5个基本阶段CopyInSplitComputeAggregateCopyOut。SplitComputeAggregate阶段需要区分a、b矩阵。Process函数中通过如下方式调用这几个函数。__aicore__ inline void Process() { CopyIn(); SplitA(); AscendC::LocalTensorhalf b1Local inQueueB1.DeQuehalf(); AscendC::LocalTensorhalf a2Local inQueueA2.DeQuehalf(); AscendC::LocalTensorfloat c2Local outQueueCO2.AllocTensorfloat(); // 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.EnQuefloat(c2Local); CopyOut(); }两次循环内SplitB需要从inQueueB1中分别搬运两个part的b矩阵Compute需要分别计算a矩阵和两个part b矩阵的乘法Aggregate要分别搬运两个part的c矩阵具体五个阶段数据流通示意图如下图5数据流通示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/数据流通示意图.png 数据流通示意图?utm_sourcegitcode_repo_files)切分b矩阵可以实现一部分的并行本样例的流水并行示意图如下图6并行示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/并行示意图.png 并行示意图?utm_sourcegitcode_repo_files)Stage1CopyIn函数实现。使用AllocTensor分别从L1 BufferA1、L1 BufferB1的Queue中申请a1Local和b1Local。使用DataCopy接口将矩阵a、b搬运到Local Memory同时将其数据格式从ND转换为NZ。一次DataCopy指令搬运height*16个数循环执行width/16次。DataCopy的参数设置如下blockCount设置为height共搬运height次。blockLen设置为1一次搬运16个类型为half的数。srcStride设置为width/16 - 1源矩阵每搬运一个block需要跳跃一行。dstStride设置为0目的矩阵每个block在内存中连续存储。每次循环迭代源矩阵首地址移动16个数目的矩阵首地址移动16*height个数。格式转换示意图如下第一次循环搬运蓝色部分第二次循环搬运绿色部分图中width为32占两个分形height为32占两个分形一共搬运4个16*16分形。图7ND to NZ转换示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/nd_to_nz.png ND-to-NZ转换示意图?utm_sourcegitcode_repo_files)注意上述ND到NZ的格式转换仅作为举例说明开发者可根据实际情况选择合适的转换方式。使用EnQue将a1Local、b1Local分别放入L1 BufferA1、L1 BufferB1的Queue中。具体代码如下__aicore__ inline void CopyND2NZ(const LocalTensorhalf dst, const GlobalTensorhalf src, const uint16_t height, const uint16_t width) { for (int i 0; i width / 16; i) { int srcOffset i * 16; int dstOffset i * 16 * height; AscendC::DataCopy(dst[dstOffset], src[srcOffset], { height, 1, uint16_t(width / 16 - 1), 0 }); } } __aicore__ inline void CopyIn() { AscendC::LocalTensorhalf a1Local inQueueA1.AllocTensorhalf(); AscendC::LocalTensorhalf b1Local inQueueB1.AllocTensorhalf(); CopyND2NZ(a1Local, aGM, m, k); CopyND2NZ(b1Local, bGM, k, n); inQueueA1.EnQue(a1Local); inQueueB1.EnQue(b1Local); }Stage2SplitA函数实现。使用DeQue从L1 BufferA1的Queue中取出a1Local。使用AllocTensor从L0A BufferA2的Queue中申请a2Local。使用LoadData将矩阵a搬运到L0A BufferA2同时将a矩阵从NZ格式转换为ZZ格式。搬运及格式转换示意图如下图中k为32占kBlocksk/162个分形m为32占mBlocksm/162个分形一共搬运4个16*16分形。本示例中调用一次LoadData接口完成两个16*16分形的搬运循环调用两次LoadData。第一次循环搬运蓝色部分两个分形第二次循环搬运绿色部分两个分形。单次循环中LoadData本样例中要完成2个分形的搬运蓝色部分或者绿色部分的参数设置如下repeatTimes表示数据处理的迭代次数因为LoadData每个迭代处理一个分形所以也可以理解为待搬运分形的个数。本样例中即为k轴方向的分形个数设置为kBlocks表示搬运kBlocks个分形。srcStride表示相邻迭代间源操作数分形首地址之间的间隔以搬运蓝色部分分形为例下图中左侧源操作数矩阵第一个蓝色分形和第二个蓝色分形起始地址之间的间隔为mBlocks个分形此处设置为mBlocks。dstGap使用默认值目的矩阵两个分形连续存储。ifTranspose设置为false每块分形搬运前搬运后都为Z格式不开启转置功能。每次循环迭代源矩阵首地址偏移16*16目的矩阵首地址偏移16*k。图8NZ to ZZ格式转换示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/nz_to_zz.png NZ-to-ZZ格式转换示意图?utm_sourcegitcode_repo_files)使用EnQue将计算结果a2Local放入到L0A BufferA2的Queue中。具体代码如下__aicore__ inline void SplitA() { int srcOffset 0; int dstOffset 0; AscendC::LocalTensorhalf a1Local inQueueA1.DeQuehalf(); AscendC::LocalTensorhalf a2Local inQueueA2.AllocTensorhalf(); // transform nz to zz for (int i 0; i mBlocks; i) { AscendC::LoadData2DParams loadDataParams; loadDataParams.repeatTimes kBlocks; loadDataParams.srcStride mBlocks; loadDataParams.ifTranspose false; AscendC::LoadData(a2Local[dstOffset], a1Local[srcOffset], loadDataParams); srcOffset 16 * 16; dstOffset k * 16; } inQueueA2.EnQuehalf(a2Local); inQueueA1.FreeTensor(a1Local); }Stage2SplitB函数实现。SplitB需要传入两个参数使用DeQue从L1 BufferB1的Queue中取出的b1Local和循环迭代变量index。使用AllocTensor从L0B BufferB2的Queue中申请b2Local。使用LoadData将b矩阵搬运到L0B BufferB2同时从NZ格式转换为ZN格式。搬运及格式转换示意图如下图中k为32占kBlocksk/162个分形n为32占nBlocksn/162个分形一共搬运4个16*16分形。本示例中调用一次LoadData接口完成两个16*16分形的搬运循环调用两次LoadData。第一次循环搬运蓝色部分两个分形第二次循环搬运绿色部分两个分形。单次循环中LoadData本样例中要完成2个分形的搬运蓝色部分或者绿色部分的参数设置如下repeatTimes表示数据处理的迭代次数因为LoadData每个迭代处理一个分形所以也可以理解为待搬运分形的个数。本样例中即为k轴方向的分形个数设置为kBlocks表示搬运kBlocks个分形。srcStride相邻迭代间源操作数分形首地址之间的间隔以搬运蓝色部分分形为例下图中左侧源操作数矩阵第一个蓝色分形和第二个蓝色分形起始地址之间的间隔为1个分形此处设置为1源矩阵两个分形连续存储。dstGap使用默认值0目的矩阵两个分形连续存储。ifTranspose设置为true每块分形搬运前为Z格式搬运后需要为N格式需要开启转置功能。每次循环迭代源矩阵首地址需要偏移k*n/2。图9NZ to ZN格式转换示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/nz_to_zn.png NZ-to-ZN格式转换示意图?utm_sourcegitcode_repo_files)使用EnQue将计算结果b2Local放入到L0B BufferB2的Queue中。具体代码如下__aicore__ inline void SplitB(const AscendC::LocalTensorhalf b1Local, const int bSplitIdx) { AscendC::LocalTensorhalf b2Local inQueueB2.AllocTensorhalf(); // transform nz to zn AscendC::LoadData2DParams loadDataParams; loadDataParams.repeatTimes kBlocks; loadDataParams.srcStride 1; loadDataParams.ifTranspose true; AscendC::LoadData(b2Local, b1Local[bSplitIdx * bSize / 2], loadDataParams); inQueueB2.EnQuehalf(b2Local); }Stage3Compute函数实现完成核心的矩阵计算功能。Compute函数需要传入参数a2Locala2Local从L0A BufferA2的Queue中使用DeQue取出。使用AllocTensor从L0C BufferCO1的Queue中申请c1Local。使用DeQue从L0B BufferB2中取出b2Local。使用Mmad完成矩阵乘计算。使用EnQue将计算结果c1Local放入到L0C BufferCO1的Queue中。具体代码如下__aicore__ inline void Compute(const AscendC::LocalTensorhalf a2Local) { AscendC::LocalTensorhalf b2Local inQueueB2.DeQuehalf(); AscendC::LocalTensorfloat c1Local outQueueCO1.AllocTensorfloat(); AscendC::MmadParams mmadParams; mmadParams.m m; mmadParams.n n / 2; mmadParams.k k; AscendC::Mmad(c1Local, a2Local, b2Local, mmadParams); outQueueCO1.EnQuefloat(c1Local); inQueueB2.FreeTensor(b2Local); }Stage4Aggregate函数实现完成数据汇聚操作。Aggregate需要传入两个参数使用AllocTensor从CO2对应的Queue中申请的c2Local和循环迭代变量index。使用DeQue从L0C BufferCO1中取出c1Local。使用DataCopy将结果矩阵从L0C BufferCO1搬运到逻辑位置CO2。DataCopy参数设置如下blockCount设置为1blockLen设置为2连续搬运两个分形无需格式转换。blockMode设置为BlockMode::BLOCK_MODE_MATRIX表示需要按分形搬运。c2Local首地址偏移量设置为index * cSize / 2。具体代码如下__aicore__ inline void Aggregate(const AscendC::LocalTensorfloat c2Local, const int bSplitIdx) { AscendC::LocalTensorfloat c1Local outQueueCO1.DeQuefloat(); AscendC::DataCopyParams dataCopyParams; dataCopyParams.blockCount 1; dataCopyParams.blockLen 2; AscendC::DataCopyEnhancedParams enhancedParams; enhancedParams.blockMode AscendC::BlockMode::BLOCK_MODE_MATRIX; AscendC::DataCopy(c2Local[bSplitIdx * cSize / 2], c1Local, dataCopyParams, enhancedParams); outQueueCO1.FreeTensor(c1Local); }Stage5CopyOut函数实现。使用DeQue从CO2中取出c2Local。使用DataCopy将结果矩阵从逻辑位置CO2搬运到Global Memory同时需要将格式从NZ转换为ND。每次循环移动一个分形搬运m*16个数。DataCopy参数说明如下blockCount设置为m共搬运m次。blockLen设置为2DataCopy指令一次搬运2个block每个block16个数。srcStride设置为0每两次搬运间没有间隙。dstStride设置为(nBlocks - 1) * 2每两次搬运间隔2个block。每次循环迭代目的矩阵偏移16源矩阵偏移m*16。格式转换示意图如下第一次循环搬运蓝色部分数据第二次循环搬运绿色部分数据。图10NZ to ND格式转换示意图![](https://raw.gitcode.com/cann/asc-devkit/raw/258f0725972f41f1a23d1fcda692fd94f2d443ce/docs/zh/guide/figures/nz_to_nd.png NZ-to-ND格式转换示意图?utm_sourcegitcode_repo_files)具体代码如下__aicore__ inline void CopyOut() { AscendC::LocalTensorfloat c2Local outQueueCO2.DeQuefloat(); // transform nz to nd for (int i 0; i nBlocks; i) { AscendC::DataCopy(cGM[i * 16], c2Local[i * m * 16], { m, 2, 0, uint16_t((nBlocks - 1) * 2) }); } outQueueCO2.FreeTensor(c2Local); }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考