cann/asc-devkit批量矩阵乘法 📅 2026/7/15 15:35:24 Batch Matmul样例概述本样例介绍在输入为float数据类型并且左、右矩阵均不转置的场景下带batch的矩阵乘法其中从GM--L1采用了DataCopy ND2NZL0C--GM、L0C--L1采用了Fixpipe批量搬运数据从L1--L0A/L0B搬运数据以及Mmad执行矩阵乘这两个步骤则是循环batch次每次循环内只处理一对左、右矩阵。支持的产品Ascend 950PR/Ascend 950DTAtlas A3 训练系列产品/Atlas A3 推理系列产品Atlas A2 训练系列产品/Atlas A2 推理系列产品目录结构介绍├── batch_matmul │ ├── scripts │ │ ├── gen_data.py // 输入数据和真值数据生成脚本 │ │ └── verify_result.py // 验证输出数据和真值数据是否一致的验证脚本 │ ├── CMakeLists.txt // 编译工程文件 │ ├── data_utils.h // 数据读入写出函数 │ └── batch_matmul.asc // Ascend C样例实现 调用样例算子描述batch mmad定义批量矩阵乘法batch mmad是普通矩阵乘法在批次维度Batch Dimension上的扩展核心逻辑是对一个包含多个矩阵的批次数据逐一对批次内的矩阵执行标准矩阵乘法最终输出同批次数量的结果矩阵。简单来说若有两个批量矩阵 A 和 B它们的形状分别为 [B, M, K] 和 [B, K, N]其中 B 是批次大小M/K/N 是矩阵维度批量矩阵乘法会为每一个批次索引 ii ∈ [0, B-1]取 A[i]形状 [M, K]和 B[i]形状 [K, N]执行普通矩阵乘法最终得到形状为 [B, M, N] 的批量结果矩阵 C。对任意批次 i0 ≤ i BC 的第 i 个矩阵满足 C[i]A[i]×B[i]。需要注意的是不同批次的矩阵之间不会互相计算。样例规格本样例中输入输出矩阵的规格如下表1所示 表1输入输出的规格输入输出数据类型Shape是否转置输入矩阵Afloat[4, 30, 40]false输入矩阵Bfloat[4, 40, 70]false输出矩阵Cfloat[4, 30, 70]-矩阵批量搬入GM-L1根据batch mmad的定义可知共计B对A、B矩阵进行矩阵乘法。数据从GM--L1通路时如下所示调用随路转换ND2NZ搬运接口时通过配置nd2nzA1Params.ndNum B实现一次性搬入B对A、B矩阵。// GM--L1搬运A矩阵 AscendC::Nd2NzParams nd2nzA1Params; // 传输ND矩阵的数目 nd2nzA1Params.ndNum B; // ND矩阵的行数 nd2nzA1Params.nValue m; // ND矩阵的列数 nd2nzA1Params.dValue k; // 源操作数相邻ND矩阵起始地址间的偏移单位是元素 nd2nzA1Params.srcNdMatrixStride m * k; // 源操作数同一ND矩阵的相邻行起始地址间的偏移单元是元素 nd2nzA1Params.srcDValue k; // ND转换到NZ格式后源操作数中的一行会转换为目的操作数的多行。 // 该参数表示目的NZ矩阵中来自源操作数同一行的多行数据相邻行起始地址间的偏移单位C0_SIZE32B。 // 数据搬运到L1时会进行对齐 nd2nzA1Params.dstNzC0Stride CeilAlign(m, cubeShape[0]); // 目的NZ矩阵中Z型矩阵相邻行起始地址之间的偏移。单位C0_SIZE32B。 nd2nzA1Params.dstNzNStride 1; // 目的NZ矩阵中相邻NZ矩阵起始地址间的偏移单位是元素 nd2nzA1Params.dstNzMatrixStride aSizeAlignL0;L1-L0A/L0B搬运和矩阵乘Mmad循环执行B次for循环B次每次从L1-L0A/L0B搬运每个batch的A、B矩阵mmad指令每次计算一对A、B矩阵矩阵乘的结果for (int32_t batchIndex 0; batchIndex B; batchIndex) { SplitA(a1Local[batchIndex * aSizeAlignL0]); SplitBTranspose(b1Local[batchIndex * bSizeAlignL0]); AscendC::SetFlagAscendC::HardEvent::MTE1_M(EVENT_ID0); AscendC::WaitFlagAscendC::HardEvent::MTE1_M(EVENT_ID0); Compute(batchIndex, c1Local); AscendC::SetFlagAscendC::HardEvent::M_MTE1(EVENT_ID0); AscendC::WaitFlagAscendC::HardEvent::M_MTE1(EVENT_ID0); }for循环计算得到每个batch的矩阵计算结果存储到L0C的相应位置上待循环结束后将L0C上的完整结果搬出到GM。AscendC::Mmad(c1Local[batchIndex * CeilAlign(m, cubeShape[0]) * CeilAlign(n, cubeShape[0])], a2Local, b2Local, mmadParams);矩阵批量搬出数据从L0C--GM通路时如下所示调用fixpipe搬运接口时通过配置fixpipeParams.ndNum B实现一次性搬出B对C矩阵。注意的是L0C中的C矩阵是对齐后的而搬出到GM的C矩阵是原始非对齐的shape。对于Atlas A2 训练系列产品/Atlas A2 推理系列产品以及Atlas A3 训练系列产品/Atlas A3 推理系列产品样例还额外演示了L0C-L1搬出路径Fixpipe将L0C上的float结果量化为half后写入L1并通过DumpTensor打印L1上的NZ结果。// 源NZ矩阵在N方向上的大小。 fixpipeParams.nSize n; // 源NZ矩阵在M方向上的大小。 fixpipeParams.mSize m; // 源NZ矩阵中相邻Z排布的起始地址偏移单位C0_Size(16*sizeof(T)T为src的数据类型) fixpipeParams.srcStride CeilAlign(m, cubeShape[0]); // 使能NZ2ND功能时代表目的ND矩阵每一行中的元素个数取值不为0单位元素 fixpipeParams.dstStride n; // 源NZ矩阵的数目也就是传输ND矩阵的数目 fixpipeParams.ndNum B; // 不同NZ矩阵起始地址之间的间隔单位1024B fixpipeParams.srcNdStride (CeilAlign(m, cubeShape[0]) * CeilAlign(n, cubeShape[0])) / (cubeShape[0] * cubeShape[0]); // 目的相邻ND矩阵起始地址之间的偏移单位element避免数据占用总内存超过存储空间限制用户应该保证batch mmad整个过程中的数据所占总的内存不超过存储空间限制。 用户可以通过PlatformAscendC类成员函数GetCoreMemSize获取获取硬件平台中L1、L0A、L0B、L0C存储空间的内存大小。编译运行在本样例根目录下执行如下步骤编译并执行算子。配置环境变量请根据当前环境上CANN开发套件包的安装方式选择对应配置环境变量的命令。默认路径root用户安装CANN软件包source /usr/local/Ascend/cann/set_env.sh默认路径非root用户安装CANN软件包source $HOME/Ascend/cann/set_env.sh指定路径install_path安装CANN软件包source ${install_path}/cann/set_env.sh样例执行mkdir -p build cd build; # 创建并进入build目录 cmake -DCMAKE_ASC_ARCHITECTURESdav-2201 -DCMAKE_ASC_RUN_MODEnpu ..;make -j; # 编译工程 python3 ../scripts/gen_data.py # 生成测试输入数据 ./demo # 执行编译生成的可执行程序执行样例 python3 ../scripts/verify_result.py output/output.bin output/golden.bin # 验证输出结果是否正确确认算法逻辑正确使用 CPU调试 或 NPU仿真 模式时添加-DCMAKE_ASC_RUN_MODEcpu或-DCMAKE_ASC_RUN_MODEsim参数即可。示例如下cmake -DCMAKE_ASC_RUN_MODEcpu -DCMAKE_ASC_ARCHITECTURESdav-2201 ..;make -j; # cpu调试模式 cmake -DCMAKE_ASC_RUN_MODEsim -DCMAKE_ASC_ARCHITECTURESdav-2201 ..;make -j; # NPU仿真模式注意切换编译模式前需清理 cmake 缓存可在 build 目录下执行rm CMakeCache.txt后重新 cmake。编译选项说明选项可选值说明CMAKE_ASC_RUN_MODEnpu默认、cpu、sim运行模式NPU运行、CPU调试、NPU仿真CMAKE_ASC_ARCHITECTURESdav-2201默认、dav-3510NPU架构dav-2201 对应 Atlas A2 训练系列产品/Atlas A2 推理系列产品/Atlas A3 训练系列产品/Atlas A3 推理系列产品dav-3510 对应 Ascend 950PR/Ascend 950DT执行结果如下说明精度对比成功。test pass!创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考