基于静态Tensor编程实现Matmul计算概述本样例基于静态Tensor编程范式实现多核矩阵乘计算。支持的产品Ascend 950PR/Ascend 950DTAtlas A3 训练系列产品/Atlas A3 推理系列产品Atlas A2 训练系列产品/Atlas A2 推理系列产品目录结构介绍├── matmul_basic_api │ ├── scripts │ │ ├── gen_data.py // 输入数据和真值数据生成脚本文件 │ │ └── verify_result.py // 真值对比文件 │ ├── CMakeLists.txt // 编译工程文件 │ ├── data_utils.h // 数据读入写出函数 │ └── matmul_basic_api.asc // Ascend C样例实现 调用样例样例描述样例功能本样例使用Ascend C基础API实现一个最基础的矩阵乘法Matmul核函数。矩阵乘法的计算公式如下 $$ C A * B $$ 其中A矩阵的形状为[M, K]B矩阵的形状为[K, N]输出C矩阵的形状为[M, N]。对输出矩阵C中的每一个元素C[m, n]都会累加A矩阵第m行和B矩阵第n列在K轴上的乘积。样例规格本样例参数M 256, N 256, K 64输入输出均为half类型、ND格式。样例启动2个核完成计算每个核负责输出矩阵C在M轴方向的128行、N轴方向的全部256列第0个核计算C矩阵的第0~127行。第1个核计算C矩阵的第128~255行。输入输出规格如下表所示样例类型(OpType)Matmul样例输入nameshapedata typeformatA[M, K]halfNDB[K, N]halfND样例输出C[M, N]halfND核函数名mmad_custom样例实现Kernel侧整体思路mmad_custom是一个__global__ __cube__核函数表示该函数运行在AI Core的Cube计算单元上主要用于矩阵计算。样例使用静态Tensor编程方式通过LocalMemAllocator创建LocalTensor。CUBE_BLOCK 16表示half数据类型分形为16 x 16代码中按16 x 16的分形为单位进行LoadData搬运。Kernel侧详细流程创建GlobalTensorhalf对象aGM、bGM、cGM分别表示GMGlobal Memory全局内存中的A、B、C矩阵。通过AscendC::GetBlockIdx()获取当前核号并计算mIterIdx。本样例只沿M轴切分任务因此每个核只需要处理A矩阵和C矩阵中属于自己的M轴分片。设置GM地址偏移aGM偏移mIterIdx * singleCoreM * K使当前核读取自己负责的A矩阵行块。bGM不偏移因为每个核都需要读取完整B矩阵。cGM偏移mIterIdx * singleCoreM * N使当前核把结果写回C矩阵中自己负责的行块。通过LocalMemAllocator创建静态LocalTensora1LocalA矩阵在L1中的临时存储。a2LocalA矩阵在L0A中的临时存储供Mmad读取。b1LocalB矩阵在L1中的临时存储。a1Local和b1Local使用同一个L1 allocator按申请顺序分配避免手动维护L1地址偏移。b2LocalB矩阵在L0B中的临时存储供Mmad读取。cLocal矩阵乘结果在L0C中的临时存储。调用DataCopy将A、B矩阵从GM搬运到L1。这里使用Nd2NzParams参数在搬运过程中将输入的ND格式数据转换为Cube计算需要的Nz格式。调用SetFlagHardEvent::MTE2_MTE1和WaitFlagHardEvent::MTE2_MTE1进行同步。DataCopy属于MTE2流水后续LoadData属于MTE1流水MTE1必须等待MTE2完成避免读取到尚未搬运完成的L1数据。调用LoadData将A矩阵从L1搬运到L0A将B矩阵从L1搬运到L0B。L0A和L0B是Cube矩阵计算单元直接读取的输入缓存。调用SetFlagHardEvent::MTE1_M和WaitFlagHardEvent::MTE1_M进行同步。LoadData属于MTE1流水后续Mmad属于M流水M流水必须等待MTE1完成避免读取到尚未搬运完成的L0A/L0B数据。调用Mmad(cLocal, a2Local, b2Local, {baseM, baseN, baseK, 0, false, true})执行矩阵乘。这里baseM 128、baseN 256、baseK 64对应单个核一次计算的矩阵块大小。调用SetFlagHardEvent::M_FIX和WaitFlagHardEvent::M_FIX进行同步。Mmad属于M流水后续Fixpipe属于FIX流水FIX流水必须等待M流水完成避免读取到尚未计算完成的L0C结果。调用Fixpipe将L0C中的float累加结果转换为half并搬运回GM中的C矩阵输出位置。最后调用PipeBarrierPIPE_ALL()确保当前核内相关流水任务完成。调用实现使用内核调用符调用核函数。调用时模板参数传入矩阵规格、单核计算量和基础Tile大小运行时参数传入Device侧A、B、C矩阵地址。接口参数说明以下结构体均以花括号{}方式传参各字段含义如下字段顺序与API文档保持一致实际struct声明中部分字段顺序可能不同AscendC::Nd2NzParams—DataCopy接口使用描述ND→Nz格式转换参数struct Nd2NzParams { int32_t ndNum; // 传输ND矩阵的数目[0, 4095] uint16_t nValue; // ND矩阵的行数[0, 16384] int32_t dValue; // ND矩阵的列数[0, 65535] int32_t srcNdMatrixStride; // 相邻ND矩阵起始地址偏移单位元素[0, 65535] int32_t srcDValue; // 同一ND矩阵相邻行偏移单位元素[1, 65535] uint16_t dstNzC0Stride; // 目的Nz中同源行转换后多行相邻偏移单位C0_SIZE(32B)[1, 16384] uint16_t dstNzNStride; // 目的Nz中Z型矩阵相邻行偏移单位C0_SIZE(32B)[1, 16384] int32_t dstNzMatrixStride; // 目的Nz中相邻Nz矩阵起始地址偏移单位元素[1, 65535] };例如搬运A矩阵时{1, baseM, baseK, 0, K, baseM, 1, 0}将baseM×baseK的ND数据转为Nz格式。AscendC::LoadData2DParams—LoadData接口使用描述Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品中A矩阵L1到L0A和B矩阵L1到L0B的数据搬运参数struct LoadData2DParams { int32_t startIndex; // 分形矩阵ID0为第1个单位512B[0, 65535] int32_t repeatTimes; // 迭代次数每个迭代处理512B[1, 255] int32_t srcStride; // 相邻迭代源分形起始地址间隔单位512B[0, 65535] int32_t sid; // 预留配置为0 int32_t dstGap; // 目的端相邻迭代分形间隔单位512B[0, 65535] bool ifTranspose; // 是否转置每个分形默认false bool addrMode; // 地址更新方式false递增true递减默认false };例如Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品中L0A上的排布格式为Zz搬运A矩阵时{0, baseK / CUBE_BLOCK, baseM / CUBE_BLOCK, 0, 0, false, 0}搬运B矩阵时ifTransposetrue完成Nz到Zn的转置搬运。AscendC::LoadData2DParamsV2—LoadData接口使用描述Ascend 950PR/Ascend 950DT产品中A矩阵L1到L0A和B矩阵L1到L0B的数据搬运参数struct LoadData2DParamsV2 { uint32_t mStartPosition; // M方向起始位置单位512B uint32_t kStartPosition; // K方向起始位置单位512B uint16_t mStep; // M方向搬运分形数 uint16_t kStep; // K方向搬运分形数 int32_t srcStride; // 源端相邻K方向分形间隔单位512B uint16_t dstStride; // 目的端相邻K方向分形间隔单位512B bool ifTranspose; // 是否转置每个分形默认false uint8_t sid; // 预留配置为0 };Ascend 950PR/Ascend 950DT产品中L0A上的排布格式为Nz搬运A矩阵时使用{0, 0, baseM / CUBE_BLOCK, baseK / CUBE_BLOCK, baseM / CUBE_BLOCK, baseM / CUBE_BLOCK, false, 0}一次完成A矩阵Nz到Nz搬运搬运B矩阵时使用{0, 0, baseK / CUBE_BLOCK, baseN / CUBE_BLOCK, baseK / CUBE_BLOCK, baseN / CUBE_BLOCK, true, 0}一次完成B矩阵Nz到Zn搬运。AscendC::MmadParams—Mmad接口使用描述矩阵乘参数struct MmadParams { uint16_t m; // 左矩阵HeightM维[0, 4095] uint16_t n; // 右矩阵WidthN维[0, 4095] uint16_t k; // 左矩阵Width/右矩阵HeightK维[0, 4095] uint16_t unitFlag; // Mmad与Fixpipe细粒度并行控制默认0 bool cmatrixSource; // C矩阵初始值来源falseCO1trueC2默认false bool cmatrixInitVal; // C矩阵初始值是否为0默认true };例如{baseM, baseN, baseK, 0, false, true}计算baseM×baseN输出块并在K方向累加baseK长度。AscendC::FixpipeParamsV220—Fixpipe接口使用描述L0C到GM的数据搬运和精度转换参数struct FixpipeParamsV220 { int32_t nSize; // 源Nz矩阵N方向大小[1, 4095] uint16_t mSize; // 源Nz矩阵M方向大小Nz2ND时[1, 8192] uint16_t srcStride; // 源Nz相邻Z排布起始偏移单位C0_SIZE[0, 65535] int32_t dstStride; // Nz2ND时目的ND矩阵每行元素数单位element bool reluEn; // 是否使能ReLU QuantMode_t quantPre; // 量化模式F322F16表示float→half uint64_t deqScalar; // scalar量化参数单个scale值 int32_t ndNum; // 源Nz矩阵数目[1, 65535] int32_t srcNdStride; // 不同Nz矩阵起始地址间隔单位16×C0_SIZE[1, 512] int32_t dstNdStride; // 目的相邻ND矩阵偏移单位element[1, 65535] int32_t unitFlag; // Mmad与Fixpipe并行控制 };例如{baseN, baseM, baseM, N, false, F322F16, 0, 1, 0, 0, 0}将L0C中的baseM×baseN float32结果转为half并写回GM。编译运行在本样例根目录下执行如下步骤编译并执行样例。配置环境变量请根据当前环境上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 ..;make -j; # 编译工程默认npu模式 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 架构Ascend 950PR/Ascend 950DTAtlas A3 训练系列产品/Atlas A3 推理系列产品Atlas A2 训练系列产品/Atlas A2 推理系列产品执行结果执行结果如下说明精度对比成功。test pass!创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考