人工智能指令集算子库CANNAscend【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址https://gitcode.com/cann/pto-isa点击查看免费下载导读TMATMUL_MX 是 Ascend CANN PTOParallel Tile Operation虚拟指令集中的矩阵乘法GEMM指令其核心特征是额外携带 A/B 两侧的缩放 TileScale Tile用于在支持的昇腾目标上实现 FP8/FP4 混合精度与量化矩阵乘法。本文将基于 docs/isa/TMATMUL_MX_zh.md 展开结合仓库中 A5/A6 的底层实现与 ST 测试用例系统讲解该指令的数学语义、三级汇编语法、C 内建接口、类型/形状合法性约束以及从数据加载、缩放 Tile 地址绑定到发射指令的完整实战流程。读完本文你将能够在 Ascend 950PR/950DT 与 Ascend 960 平台上正确编写并校验基于 TMATMUL_MX 的量化 GEMM Kernel。一、指令概述带额外缩放 Tile 的矩阵乘法TMATMUL_MX 是一条带额外缩放 Tile 的矩阵乘法GEMM指令设计目的就是在支持的昇腾目标上执行混合精度/量化矩阵乘法。与常规的 TMATMUL 相比它额外接收两个缩放 Tile 参数aScaleMatrix作用于左矩阵 A与bScaleMatrix作用于右矩阵 B从而在单条指令内完成低比特数据乘累加 缩放校正的量化语义。该指令目前仅在Ascend 950PR/Ascend 950DTA5与Ascend 960A6上实现对应的实现文件为 include/pto/npu/a5/TMatmul.hpp 与 include/pto/npu/a6/TMatmul.hpp。而其他昇腾目标如 Kirin9030则明确不支持其实现中通过static_assert(sizeof(TileRes::DType) 0, no support instruction.)在编译期直接报错参见 include/pto/npu/kirin9030/TMatmul.hpp#L254-L277因此在使用前务必确认目标平台。二、数学语义缩放 Tile 的作用边界设M aMatrix.GetValidRow()K aMatrix.GetValidCol()N bMatrix.GetValidCol()概念上结果对应于有效矩阵乘法域0 i M0 j N上的标准矩阵乘法$$ \mathrm{C}{i,j} \sum{k0}^{K-1} \mathrm{A}{i,k} \cdot \mathrm{B}{k,j} $$缩放 TileaScaleMatrix/bScaleMatrix用于配置实现定义的混合精度行为。需要特别注意的是aScaleMatrix/bScaleMatrix的确切作用以及任何反量化/量化语义由目标硬件定义PTO 指令集层面只保证存在缩放 Tile 参与运算这一事实而不规定具体的缩放算法。从仓库的实现看缩放 Tile 的类型固定为float8_e8m0_t8 位指数、0 位尾数、偏置 127 的 MX 格式缩放因子定义于 include/pto/cpu/MXTypes.hpp#L326-L334由硬件按 MX 规范解释。换句话说PTO 层保证的是指令形态与 Tile 契约实际缩放语义以目标硬件的指令手册为准。这也解释了为什么缩放 Tile 的尺寸与数据 Tile 不同见下文示例中的ScaleA形状为16x2、ScaleB形状为2x32而数据 Tile 为16x64与64x32——缩放 Tile 按每 32 个 K 元素共享一个缩放因子的 MX 分块粒度组织见测试用例中kMX CeilDiv(kAlign, 32)的推导tests/npu/a5/src/st/testcase/tmatmul_mx/tmatmul_mx_kernel.cpp#L73-L75。三、汇编语法从概念形式到三级 AS 表示3.1 同步形式概念性指令提供三种形式基础乘法、累加.acc与偏置.bias%c tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile... %c_out tmatmul.mx.acc %c_in, %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile... %c tmatmul.mx.bias %a, %a_scale, %b, %b_scale, %bias : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile...基础形式c A * B缩放 Tile 参与量化语义.acc累加形式c_out c_in A * B用于 K 维切分Split-K或连续累加.bias偏置形式c A * B bias其中bias为单行偏置 Tile。3.2 AS Level 1SSA 形式SSA 形式使用!pto.tile...类型操作数以 SSA 值%a、%b等呈现%c pto.tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile... %c_out pto.tmatmul.mx.acc %c_in, %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile... %c pto.tmatmul.mx.bias %a, %a_scale, %b, %b_scale, %bias : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) - !pto.tile...3.3 AS Level 2DPS 形式DPSDestination-Passing Style形式将结果写入显式的outs()缓冲区输入为!pto.tile_buf...pto.tmatmul.mx ins(%a, %a_scale, %b, %b_scale : !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf...) outs(%c : !pto.tile_buf...) pto.tmatmul.mx.acc ins(%c_in, %a, %a_scale, %b, %b_scale : !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf...) outs(%c_out : !pto.tile_buf...) pto.tmatmul.mx.bias ins(%a, %a_scale, %b, %b_scale, %bias : !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf...) outs(%c : !pto.tile_buf...)两种 AS 层级的关系Level 1SSA便于编译器分析与调度Level 2DPS更贴近最终硬件指令显式指定 ins/outs 缓冲区。四、C 内建接口六个模板重载4.1 接口签名TMATMUL_MX 的 C 内建接口声明于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hpp。共六个重载对应三种功能形态 × 是否显式指定累加相位template typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, WaitEvents ... events); template AccPhase Phase, typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, WaitEvents ... events); template typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cOutMatrix, TileRes cInMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, WaitEvents ... events); template AccPhase Phase, typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cOutMatrix, TileRes cInMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, WaitEvents ... events); template typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename TileBias, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, TileBias biasData, WaitEvents ... events); template AccPhase Phase, typename TileRes, typename TileLeft, typename TileLeftScale, typename TileRight, typename TileRightScale, typename TileBias, typename... WaitEvents PTO_INST RecordEvent TMATMUL_MX(TileRes cMatrix, TileLeft aMatrix, TileLeftScale aScaleMatrix, TileRight bMatrix, TileRightScale bScaleMatrix, TileBias biasData, WaitEvents ... events);六个重载的参数语义可归纳为重载形态关键参数语义基础无 PhasecMatrix, aMatrix, aScaleMatrix, bMatrix, bScaleMatrixc A·B量化缩放基础带AccPhase Phase同上显式指定累加相位如清零/累加累加无 PhasecOutMatrix, cInMatrix, aMatrix, aScaleMatrix, bMatrix, bScaleMatrixc_out c_in A·B累加带AccPhase Phase同上Split-K 等场景显式控制相位偏置无 PhasecMatrix, aMatrix, aScaleMatrix, bMatrix, bScaleMatrix, biasDatac A·B bias偏置带AccPhase Phase同上显式指定累加相位所有重载均以变参WaitEvents ... events结尾用于声明依赖事件如等待 MTE2 数据搬运完成的同步事件。返回值类型为PTO_INST RecordEvent用于记录该指令产生的同步事件供下游指令等待。4.2 从源码看实现调用链A5/A6 共享的架构无关实现位于 include/pto/npu/a5/TMatmulImpls.hpp该文件被 A5/A6 各自的TMatmul.hpp包含。以偏置形式为例其内部调用链为// TMATMUL_MX_IMPLbias 形式 CheckMadMxValidTileRes, TileLeft, TileLeftScale, TileRight, TileRightScale(); static_assert(std::is_same_vtypename TileBias::DType, float, TMatmulMX:No supported bias data type.); static_assert((TileBias::Loc TileType::Bias) (TileBias::Rows 1), TMatmulMX:TileBias must be single row.); uint16_t m aMatrix.GetValidRow(); uint16_t k aMatrix.GetValidCol(); uint16_t n bMatrix.GetValidCol(); CheckDynamicMmad(m, k, n); TMatmulMxBiasPhase, TileRes, TileLeft, TileRight, true, false, true( cMatrix.data(), aMatrix.data(), bMatrix.data(), biasData.data(), m, k, n);从源码结构看include/pto/npu/a5/TMatmulImpls.hpp#L120-L138调用链可分为四步静态合法性检查CheckMadMxValid...()编译期校验数据类型三元组、分形布局与 K 对齐偏置静态断言TileBias::DType必须为floatTileBias::Loc TileType::Bias且TileBias::Rows 1动态形状检查CheckDynamicMmad(m, k, n)从 Tile 的有效形状ValidRow/ValidCol提取 m/k/n 并做运行期范围校验底层发射调用TMatmulMxBias内部最终落到mad类硬件指令见 include/pto/npu/a5/TMatmul.hpp#L40-L49 中mad(c, a, b, m, k, n, ...)的底层封装。值得注意的是缩放 TileaScaleMatrix/bScaleMatrix在TMatmulMx底层调用中并未作为显式数据参数传入——缩放因子是通过与数据 Tile 绑定的方式GetScaleAddr获得的地址由硬件识别这正是实现定义缩放语义的体现。五、类型、形状与合法性约束5.1 支持的 (C, A, B) 数据类型三元组C始终为float两侧缩放 Tile 固定为float8_e8m0_t。支持的 A/B 数据组合A5/A6FP84 种组合A 类型B 类型float8_e4m3_tfloat8_e4m3_tfloat8_e4m3_tfloat8_e5m2_tfloat8_e5m2_tfloat8_e4m3_tfloat8_e5m2_tfloat8_e5m2_tFP44 种组合A 类型B 类型float4_e1m2x2_tfloat4_e1m2x2_tfloat4_e1m2x2_tfloat4_e2m1x2_tfloat4_e2m1x2_tfloat4_e2m1x2_tfloat4_e2m1x2_tfloat4_e1m2x2_t这些组合在CheckMadMxValid中通过isSupportedFp4ComboAType, BType/isSupportedFp8ComboAType, BType与static_assert((isFp4 || isFp8) std::is_same_vCType, float, TMatmulMX:No supported data type combination.)强制保证include/pto/npu/a5/TMatmul.hpp#L51-L75。5.2 形状与分形布局约束A5 实现中CheckMadMxValid还强制执行以下静态约束include/pto/npu/a5/TMatmul.hpp#L52-L75K 对齐TileLeft::Cols必须是 64 的倍数BASEK 64对 FP4 数据类型TileLeft::Cols还必须是偶数分形布局左矩阵TileLeft::Loc TileType::Left、非行主序!isRowMajor、SFractal SLayout::RowMajor右矩阵TileRight::Loc TileType::Right、行主序isRowMajor、SFractal SLayout::ColMajor累加器TileRes::Loc TileType::Acc、非行主序、SFractal SLayout::RowMajorL0C 容量累加器字节数Rows * Cols * sizeof(CType)不得超过PTO_L0C_SIZE_BYTESA5 的 L0C 缓冲容量。5.3 偏置形式约束偏置形式额外要求A5/A6 通过static_assert检查见 include/pto/npu/a5/TMatmulImpls.hpp#L128-L129TileBias::DType必须为floatTileBias::Loc TileType::BiasTileBias::Rows 1单行偏置。六、代码示例自动模式与手动模式6.1 自动Auto模式自动模式下Tile 的物理地址与调度由编译器/运行时统一管理用户只需声明 Tile 类型与形状#include pto/pto-inst.hpp using namespace pto; void example_auto() { using A TileLeftfloat8_e5m2_t, 16, 64; using B TileRightfloat8_e5m2_t, 64, 32; using ScaleA TileLeftScalefloat8_e8m0_t, 16, 2; using ScaleB TileRightScalefloat8_e8m0_t, 2, 32; using Bias TileTileType::Bias, float, 1, 32; using C TileAccfloat, 16, 32; A a; B b; ScaleA scaleA; ScaleB scaleB; Bias bias; C c; TMATMUL_MX(c, a, scaleA, b, scaleB, bias); }形状语义说明A TileLeftfp8, 16, 6416 行 × 64 列M16, K64K 满足 64 对齐B TileRightfp8, 64, 3264 行 × 32 列K64, N32ScaleA TileLeftScalefp8_e8m0_t, 16, 2A 侧缩放 Tile16 行 × 2 列K/32 2ScaleB TileRightScalefp8_e8m0_t, 2, 32B 侧缩放 Tile2 行 × 32 列Bias TileTileType::Bias, float, 1, 32单行float偏置满足Rows 1C TileAccfloat, 16, 32float累加结果。6.2 手动Manual模式手动模式下所有 Tile 必须先用TASSIGN显式绑定物理地址再发射指令。缩放 Tile 的地址通过GetScaleAddr()从对应数据 Tile 推导GetScaleAddr定义于 include/pto/npu/a5/utils.hpp#L82#include pto/pto-inst.hpp using namespace pto; void example_manual() { using A TileLeftfloat8_e5m2_t, 16, 64; using B TileRightfloat8_e5m2_t, 64, 32; using ScaleA TileLeftScalefloat8_e8m0_t, 16, 2; using ScaleB TileRightScalefloat8_e8m0_t, 2, 32; using Bias TileTileType::Bias, float, 1, 32; using C TileAccfloat, 16, 32; A a; B b; ScaleA scaleA; ScaleB scaleB; Bias bias; C c; TASSIGN(a, 0x1000); TASSIGN(b, 0x2000); TASSIGN(scaleA, GetScaleAddr(a.data())); TASSIGN(scaleB, GetScaleAddr(b.data())); TASSIGN(bias, 0x3000); TASSIGN(c, 0x4000); TMATMUL_MX(c, a, scaleA, b, scaleB, bias); }手动模式的关键点数据 Tile 的物理地址由用户规划示例中的0x1000/0x2000/0x3000/0x4000需根据片上内存布局自行计算而缩放 Tile 地址应通过GetScaleAddr(a.data())/GetScaleAddr(b.data())由硬件工具函数推导不要手写以保证缩放因子与数据块的正确绑定关系。七、汇编示例7.1 自动模式# 自动模式由编译器/运行时负责资源放置与调度。 %c pto.tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...)7.2 手动模式# 手动模式先显式绑定资源再发射指令。 # 可选当该指令包含 tile 操作数时 # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %c pto.tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...)7.3 PTO 汇编形式DPS%c pto.tmatmul.mx %a, %a_scale, %b, %b_scale : (!pto.tile..., !pto.tile..., !pto.tile..., !pto.tile...) # AS Level 2 (DPS) pto.tmatmul.mx ins(%a, %a_scale, %b, %b_scale : !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf..., !pto.tile_buf...)八、实战流程从 Global Memory 到量化 GEMM 结果的完整 Kernel仓库中的 ST 测试 tests/npu/a5/src/st/testcase/tmatmul_mx/tmatmul_mx_kernel.cpp 给出了一个可直接参考的完整量化 GEMM Kernel 流程其中蕴含了 TMATMUL_MX 实战中的关键经验1. 形状与对齐规则对应测试中的RunTMATMULMXtmatmul_mx_kernel.cpp#L66-L75constexpr int blockAlign isFp4 ? 64 : 32; // need to be 32B aligned constexpr int M CeilAlignint(validM, 16); constexpr int kAlign CeilAlignint(validK, 64); constexpr int N CeilAlignint(validN, blockAlign); constexpr uint8_t kMX CeilDiv(kAlign, 32);M 对齐到 16A5 矩阵分形行粒度K 对齐到 64与BASEK 64的静态约束一致N 对齐到 32FP8/ 64FP4保证 32B 对齐缩放 Tile 的 K 维为kMX K / 32即每 32 个 K 元素共享一个float8_e8m0_t缩放因子。2. 缩放 Tile 的布局A 侧缩放 Tile 使用Layout::MX_A_ZZ16 行对齐B 侧使用Layout::MX_B_NNN 对齐与TLOAD支持的 MX 布局一致参见 include/pto/npu/a5/TLoad.hpp#L138 中 MX_A_ZZ / MX_B_NN 的加载支持。3. 完整数据通路Global MemoryGlobalTensor→TLOAD到 L1/UB数据 Tile 缩放 Tile 偏置 Tile→TMOV搬入计算侧 → 手动模式绑定缩放地址或自动模式调用AssignMxScaleAddr内部为__cce_pto_get_mx_scale_tile(scaleTile.data(), dataTile.data())→TMATMUL_MX发射 →TSTORE写回 Global Memory。4. 同步与事件数据依赖通过set_flag/wait_flag如PIPE_MTE2 → PIPE_MTE1 → PIPE_M → PIPE_FIX在流水线阶段间同步确保缩放/数据 Tile 就绪后再发射矩阵乘。5. Split-K 变体测试还提供了RunTMATMULMX_SPLIT_Ktmatmul_mx_kernel.cpp#L216 起将 K 维切分为BASEK 64的块多次调用TMATMUL_MX的.acc形式累加这正对应接口中cOutMatrix/cInMatrix双 C 参数重载的使用场景——K 维较大时可通过 Split-K 提升并行度。九、平台支持范围小结平台TMATMUL_MX 支持情况依据Ascend 950PR / Ascend 950DTA5✅ 支持include/pto/npu/a5/TMatmul.hpp、TMatmulImpls.hppAscend 960A6✅ 支持include/pto/npu/a6/TMatmul.hpp复用 A5 的 TMatmulImpls 包装层Kirin9030❌ 不支持include/pto/npu/kirin9030/TMatmul.hpp#L254-L277static_assert no support instruction.从源码结构看A6 通过复用 A5 的架构无关实现层TMatmulImpls.hpp 头部注释明确shared by a5 and a6获得同样的 MX 矩阵乘能力。因此若目标平台不在支持列表内应改用 TMATMUL标准混合精度无缩放 Tile或等待对应平台补齐该指令否则会直接触发编译期报错。十、总结TMATMUL_MX 是 PTO 指令集中面向混合精度/量化 GEMM的核心指令它通过aScaleMatrix/bScaleMatrix两个float8_e8m0_t缩放 Tile 将量化语义封装进单条矩阵乘指令支持 FP8 与 FP4 共 8 种 (A, B) 数据类型组合并提供基础/累加/偏置三种形式、六个 C 模板重载与三级汇编表示。编写 TMATMUL_MX Kernel 时需要重点把握K 维 64 对齐、缩放 Tile 按 K/32 粒度组织、偏置必须为单行float、分形布局遵循 Left(RowMajor)/Right(ColMajor)/Acc(RowMajor) 约定以及在手动模式下通过GetScaleAddr正确绑定缩放地址。结合仓库的 ST 测试用例即可在 Ascend 950PR/950DT 与 Ascend 960 上构建出正确、可校验的量化 GEMM 实现。进一步阅读指令总览PTO Virtual ISA Manual 与 PTOISA同族指令TMATMUL、TMATMUL_BIAS、TMATMUL_ACC缩放数据通路TLOAD、TMOV、TSTORE相关测试tests/npu/a5/src/st/testcase/tmatmul_mx/含 kernel 与 main.cpp 驱动赞分享人工智能指令集算子库CANNAscend【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址https://gitcode.com/cann/pto-isa点击查看免费下载相关推荐PTO 指令详解TGEMV_MX——Ascend 混合精度/量化矩阵向量乘GEMV with MX Scaling的 Tile 级实现PTO 指令详解TGEMV_MX——Ascend 混合精度/量化矩阵向量乘GEMV with MX Scaling的 Tile 级实现 TGEMV_MX人工智能指令集算子库CANNAscendCANN PTO-ISA TMATMUL_MX 指令详解基于缩放 Tile 的混合精度/量化 GEMM 编程指南CANN PTO ISA TMATMUL_MX 指令详解基于缩放 Tile 的混合精度/量化 GEMM 编程指南 导读 TMATMUL_MX 是 CANN P人工智能指令集算子库CANNAscendgmaps开发者指南如何扩展和自定义地图插件功能gmaps开发者指南如何扩展和自定义地图插件功能 gmaps是Jupyter Notebook中功能强大的Google Maps交互式地图插件它为数据科学家人工智能指令集算子库CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考