PTO Auto Mode 示例解析:TADD 与 TMATMUL 的自动内存分配、自动同步写法及编译启用
PTO Auto Mode 示例解析TADD 与 TMATMUL 的自动内存分配、自动同步写法及编译启用【免费下载链接】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-isa 仓库中 Auto Mode Examples 为主体通过 TADD向量逐元素加法与 TMATMUL矩阵乘法两组auto mode vs manual mode完整代码对照讲清 PTO 自动编译模式到底替开发者省掉了哪些工作——TASSIGN显式显存UB地址分配和Event显式流水线同步——并给出在 CANN 工具链中启用 auto mode 编译的完整命令与 CMake 配置。读完本文你可以直接照抄代码骨架编写 auto mode kernel并理解 manual mode 中每个Event所对应的硬件流水线边界MTE → Cube → Store。文档范围与核心结论docs/auto_mode/Examples.md 的范围声明很明确展示以auto mode编写与以manual mode编写的 kernel 对照示例。结合 Auto_Mode_Overview.md可以提炼出 auto mode 的两个核心抽象本文两组示例正是围绕它们展开Tile 内存自动分配manual mode 下实例化Tile变量后必须用TASSIGN手动分配一段专属 UBUnified Buffer缓冲区地址auto mode 下只需实例化Tile变量编译器自动完成缓冲区地址分配自动同步auto-syncmanual mode 需要用 PTO 的 event model 在精确的代码位置插入同步兼顾功能正确性与性能auto mode 下编译器自动决定同步插入位置。一个关键事实TASSIGN和显式Event同步在 auto mode 中并不是被禁止的非法调用而是什么都不做no-op——这保证了同一份 manual mode 代码可以不加修改地在 auto mode 下编译。这也是为什么下文 manual mode 示例中TASSIGN/Event标注着 only in manual mode而 auto mode 版本中它们被整体删除。示例一TADD 逐元素加法Auto mode 版本文档给出的 auto mode 版本如下来自 Examples.md#include pto/pto-inst.hpp #include pto/common/constants.hpp using namespace pto; AICORE void runTAdd(__gm__ float __out__ *out, __gm__ float __in__ *src0, __gm__ float __in__ *src1) { using DynShapeDim5 Shape1, 1, 1, 64, 64; using DynStridDim5 Stride1, 1, 1, 64, 1; using GlobalData GlobalTensorfloat, DynShapeDim5, DynStridDim5; using TileData TileTileType::Vec, float, 64, 64, BLayout::RowMajor, 64, 64; TileData src0Tile(64, 64); TileData src1Tile(64, 64); TileData dstTile(64, 64); GlobalData src0Global(src0); GlobalData src1Global(src1); GlobalData dstGlobal(out); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); TADD(dstTile, src0Tile, src1Tile); TSTORE(dstGlobal, dstTile); }要点拆解Shape1, 1, 1, 64, 64与Stride1, 1, 1, 64, 1是 PTO 的 5 维张量描述前 3 维退化为 1/1实际有效的是末两维 64 行 × 64 列的行主序布局详见 GlobalTensorTileTileType::Vec, float, 64, 64, BLayout::RowMajor, 64, 64声明了一个向量核AIV上的 64×64 float TileTileType::Vec表示该 Tile 驻留在向量核的 UB 中Tile 抽象见 Tile指令流是最短闭环TLOAD两次从 GM 搬运到 UB →TADD做逐元素加法见 TADD.md→TSTORE写回 GM见 TSTORE.md。没有任何地址分配和同步代码。Manual mode 版本同一功能在 manual mode 下需要补充TASSIGN与Event链#include pto/pto-inst.hpp #include pto/common/constants.hpp using namespace pto; AICORE void runTAdd(__gm__ float __out__ *out, __gm__ float __in__ *src0, __gm__ float __in__ *src1) { using DynShapeDim5 Shape1, 1, 1, 64, 64; using DynStridDim5 Stride1, 1, 1, 64, 1; using GlobalData GlobalTensorfloat, DynShapeDim5, DynStridDim5; using TileData TileTileType::Vec, float, 64, 64, BLayout::RowMajor, 64, 64; TileData src0Tile(64, 64); TileData src1Tile(64, 64); TileData dstTile(64, 64); /* TAssign only in manual mode */ TASSIGN(src0Tile, 0x0); TASSIGN(src1Tile, 0x10000); TASSIGN(dstTile, 0x20000); GlobalData src0Global(src0); GlobalData src1Global(src1); GlobalData dstGlobal(out); /* event model only in manual mode */ EventOp::TLOAD, Op::TADD event0; EventOp::TADD, Op::TSTORE_VEC event1; TLOAD(src0Tile, src0Global); event0 TLOAD(src1Tile, src1Global); event1 TADD(dstTile, src0Tile, src1Tile, event0); TSTORE(dstGlobal, dstTile, event1); }manual mode 多出来的部分正是 auto mode 自动化的两件事机制Manual mode 代码作用Auto mode 处理Tile 内存分配TASSIGN(src0Tile, 0x0)、TASSIGN(src1Tile, 0x10000)、TASSIGN(dstTile, 0x20000)把三个 Tile 分别绑定到 UB 中的不同基地址见 TASSIGN.md编译器基于 Tile 生命周期分析自动分配源码删除流水线同步EventOp::TLOAD, Op::TADD event0约束第二个TLOAD完成后才能开始TADDEventOp::TADD, Op::TSTORE_VEC event1约束TADD完成后才能TSTORE向量核写回路径为Op::TSTORE_VEC表达 GM→UBMTE2 路径与计算VF 路径之间的异步依赖编译器自动插入同步Event传参在 auto mode 下为 no-op注意event0 TLOAD(...)这种写法manual mode 的指令调用可以返回/接收Event句柄把前驱操作完成这一事实显式传递到后一条指令。auto mode 版本中同样的调用全部退化为一参或两参形式。仓库中真实落地的 auto mode TADD kernelExamples.md 的 TADD 是单核最小示例仓库演示工程 demos/auto_mode/baseline/add/csrc/kernel/add_custom.cpp 给出了面向 A2A3 的完整工程化写法结构上可以直接对照template typename T, unsigned tileRows, unsigned tileCols AICORE void runTAdd(__gm__ T* z, __gm__ T* x, __gm__ T* y, uint32_t totalLength) { set_mask_norm(); set_vector_mask(-1, -1); static_assert(BLOCK_ROWS * BLOCK_COLS BLOCK_DIM, Wrong block tilling!); // inter-vector cores block tiling constexpr unsigned bTileRows tileRows / BLOCK_ROWS; constexpr unsigned bTileCols tileCols / BLOCK_COLS; static_assert(bTileRows * bTileCols * sizeof(T) MAX_TILE_SIZE, UB buffer overflow.); // define GlobalData on global memory with shape and stride using ShapeDim5 pto::Shape1, 1, 1, bTileRows, bTileCols; using StridDim5 pto::Stride1, 1, 1, tileCols, 1; using GlobalData pto::GlobalTensorT, ShapeDim5, StridDim5; const unsigned offset block_idx * bTileRows * bTileCols; GlobalData xGlobal(x offset); GlobalData yGlobal(y offset); GlobalData zGlobal(z offset); using TileData TileTileType::Vec, T, bTileRows, bTileCols, BLayout::RowMajor, -1, -1; TileData xTile(bTileRows, bTileCols), yTile(bTileRows, bTileCols), zTile(bTileRows, bTileCols); TLOAD(xTile, xGlobal); TLOAD(yTile, yGlobal); TADD(zTile, xTile, yTile); TSTORE(zGlobal, zTile); } // kernel entry extern C __global__ AICORE void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength) { // Define the tile size constexpr unsigned tileRows 20; constexpr unsigned tileCols 2048; // main kernel, totalLength is dynamic input runTAddhalf, tileRows, tileCols((__gm__ half*)z, (__gm__ half*)x, (__gm__ half*)y, totalLength); }相对文档示例的增量在于按BLOCK_DIM 20个 AIV 做核间 tilingblock_idx切分、static_assert防止 UB 溢出、half数据类型与多核入口add_custom。指令闭环TLOAD → TLOAD → TADD → TSTORE与文档示例完全一致——这印证了 auto mode 下 kernel 作者只需关心Tile 级数据流地址与同步全部交给编译器。该示例的工程说明见 demos/auto_mode/baseline/add/README.md另有 PyTorch JIT 变体 demos/auto_mode/torch_jit/add/jit_util_add.py 展示通过 JIT 编译同款 kernel 的路径。示例二TMATMUL 矩阵乘法含 FbTile 缩放这是文档中更复杂的示例数据流跨越 MTE2GM→UB 搬运、MTE1UB→Cube/L0 搬运即TMOV、CubeTMATMUL矩阵乘法、以及带逐列缩放因子的 FbTile 写回路径是理解 auto-sync 价值的最佳样本。Auto mode 版本#include pto/pto-inst.hpp #include pto/common/constants.hpp using namespace pto; template typename cType, typename aType, typename bType, typename fbType, typename l0cType, int M, int K, int N, int ValidM, int ValidK, int ValidN __global__ AICORE void runTMatMul(__gm__ cType *out, __gm__ aType *src0, __gm__ bType *src1, __gm__ fbType *src2) { using GlobalDataSrc0 GlobalTensoraType, pto::Shape1, 1, 1, ValidM, ValidK, pto::StrideValidM * ValidK, ValidM * ValidK, ValidM * ValidK, ValidK, 1; using GlobalDataSrc1 GlobalTensorbType, pto::Shape1, 1, 1, ValidK, ValidN, pto::StrideValidK * ValidN, ValidK * ValidN, ValidK * ValidN, ValidN, 1; using GlobalDataSrc2 GlobalTensorfbType, pto::Shape1, 1, 1, 1, ValidN, pto::StrideValidN, ValidN, ValidN, ValidN, 1; using GlobalDataOut GlobalTensorcType, pto::Shape1, 1, 1, ValidM, ValidN, pto::StrideValidM * ValidN, ValidM * ValidN, ValidM * ValidN, ValidN, 1; GlobalDataSrc0 src0Global(src0); GlobalDataSrc1 src1Global(src1); GlobalDataSrc2 src2Global(src2); GlobalDataOut dstGlobal(out); using TileMatAData TileTileType::Mat, aType, M, K, BLayout::ColMajor, ValidM, ValidK, SLayout::RowMajor, 512; using TileMatBData TileTileType::Mat, bType, K, N, BLayout::ColMajor, ValidK, ValidN, SLayout::RowMajor, 512; using TileMatFbData TileTileType::Mat, fbType, 1, N, BLayout::RowMajor, 1, ValidN, SLayout::NoneBox; using LeftTile TileLeftaType, M, K, ValidM, ValidK; using RightTile TileRightbType, K, N, ValidK, ValidN; using AccTile TileAccl0cType, M, N, ValidM, ValidN; using FbTile TileTileType::Scaling, fbType, 1, N, BLayout::RowMajor, 1, ValidN, SLayout::NoneBox; TileMatAData aMatTile; TileMatBData bMatTile; TileMatFbData fbMatTile; LeftTile aTile; RightTile bTile; AccTile cTile; FbTile fbTile; TLOAD(aMatTile, src0Global); TLOAD(bMatTile, src1Global); TLOAD(fbMatTile, src2Global); /**************************TMOV TMATMUL**************************/ TMOV(aTile, aMatTile); TMOV(bTile, bMatTile); TMATMUL(cTile, aTile, bTile); TMOV(fbTile, fbMatTile); /********************************TSTORE****************************/ TSTORE_FPAccTile, GlobalDataOut, FbTile(dstGlobal, cTile, fbTile); }Manual mode 版本#include pto/pto-inst.hpp #include pto/common/constants.hpp using namespace pto; template typename cType, typename aType, typename bType, typename fbType, typename l0cType, int M, int K, int N, int ValidM, int ValidK, int ValidN __global__ AICORE void runTMatMul(__gm__ cType *out, __gm__ aType *src0, __gm__ bType *src1, __gm__ fbType *src2) { using GlobalDataSrc0 GlobalTensoraType, pto::Shape1, 1, 1, ValidM, ValidK, pto::StrideValidM * ValidK, ValidM * ValidK, ValidM * ValidK, ValidK, 1; using GlobalDataSrc1 GlobalTensorbType, pto::Shape1, 1, 1, ValidK, ValidN, pto::StrideValidK * ValidN, ValidK * ValidN, ValidK * ValidN, ValidN, 1; using GlobalDataSrc2 GlobalTensorfbType, pto::Shape1, 1, 1, 1, ValidN, pto::StrideValidN, ValidN, ValidN, ValidN, 1; using GlobalDataOut GlobalTensorcType, pto::Shape1, 1, 1, ValidM, ValidN, pto::StrideValidM * ValidN, ValidM * ValidN, ValidM * ValidN, ValidN, 1; GlobalDataSrc0 src0Global(src0); GlobalDataSrc1 src1Global(src1); GlobalDataSrc2 src2Global(src2); GlobalDataOut dstGlobal(out); using TileMatAData TileTileType::Mat, aType, M, K, BLayout::ColMajor, ValidM, ValidK, SLayout::RowMajor, 512; using TileMatBData TileTileType::Mat, bType, K, N, BLayout::ColMajor, ValidK, ValidN, SLayout::RowMajor, 512; using TileMatFbData TileTileType::Mat, fbType, 1, N, BLayout::RowMajor, 1, ValidN, SLayout::NoneBox; using LeftTile TileLeftaType, M, K, ValidM, ValidK; using RightTile TileRightbType, K, N, ValidK, ValidN; using AccTile TileAccl0cType, M, N, ValidM, ValidN; using FbTile TileTileType::Scaling, fbType, 1, N, BLayout::RowMajor, 1, ValidN, SLayout::NoneBox; TileMatAData aMatTile; TileMatBData bMatTile; TileMatFbData fbMatTile; /* TAssign only in manual mode */ TASSIGN(aMatTile, 0x0); TASSIGN(bMatTile, 0x10000); TASSIGN(fbMatTile, 0x20000); LeftTile aTile; RightTile bTile; AccTile cTile; FbTile fbTile; /* TAssign only in manual mode */ TASSIGN(aTile, 0x0); TASSIGN(bTile, 0x0); TASSIGN(cTile, 0x0); TASSIGN(fbTile, 0x0); /* event model only in manual mode */ EventOp::TLOAD, Op::TMOV_M2L evtLoad_Mov; EventOp::TMOV_M2B, Op::TMATMUL evtMov_Matmul; EventOp::TMATMUL, Op::TMOV_M2S evtMatmul_MovM2s; TLOAD(aMatTile, src0Global); TLOAD(bMatTile, src1Global); evtLoad_Mov TLOAD(fbMatTile, src2Global); /**************************TMOV TMATMUL**************************/ TMOV(aTile, aMatTile, evtLoad_Mov); evtMov_Matmul TMOV(bTile, bMatTile); evtMatmul_MovM2s TMATMUL(cTile, aTile, bTile, evtMov_Matmul); TMOV(fbTile, fbMatTile, evtMatmul_MovM2s); /********************************TSTORE****************************/ TSTORE_FPAccTile, GlobalDataOut, FbTile(dstGlobal, cTile, fbTile); }Manual mode 中的同步依赖图auto mode 自动化的对象对照两个版本manual mode 版本显式维护了三条事件边恰好覆盖了 Cube 计算单元的全部上游/下游EventOp::TLOAD, Op::TMOV_M2L evtLoad_MovTLOADMTE2 完成 GM→UB 搬运aMatTile→TMOV_M2LUB→Cube 左操作数搬运。注意evtLoad_Mov TLOAD(fbMatTile, src2Global)只绑定了第三个 TLOAD即作者刻意让 a 路先行、f 路最后加载为 Cube 争取更早的启动窗口EventOp::TMOV_M2B, Op::TMATMUL evtMov_MatmulTMOV_M2BUB→Cube 右操作数bTile搬运→TMATMUL真正执行即 Cube 必须等到左右操作数都到位EventOp::TMATMUL, Op::TMOV_M2S evtMatmul_MovM2sTMATMUL完成 →TMOV_M2SfbMatTile向fbTile的缩放因子搬运复用/写回路径把写回依赖挂到计算完成之后。这条搬运→计算→写回的依赖链与仓库文档中的 GEMM 流水线示意图见文首配图一致MTE2 上的aMatTile/bMatTile经 TMOV 进入 MTE1 的aTile/bTile再交给 mmadCube 矩阵乘累加计算反向边如 mmad→MTE1用于在计算完成后再允许后续搬运。图中每一条同步箭头在 manual mode 下都由开发者手写Event表达在 auto mode 下这些边全部由编译器依据 Tile 生命周期与指令依赖自动推导插入源码中只保留TLOAD → TMOV → TMATMUL → TSTORE_FP的纯数据流。此外可以观察到 Tile 层次的两层分配差异manual mode 中TASSIGN分成了两组——UB 层aMatTile/bMatTile/fbMatTile分配到0x0/0x10000/0x20000等 UB 地址与 Cube/L0 层aTile/bTile/cTile/fbTile统一在0x0起的 L0 区。auto mode 中这两层缓冲区的分配全部隐去。TileMatAData等模板参数M/K/N与ValidM/ValidK/ValidN分离、SLayout::RowMajor与512的 L0 布局参数表达的是 Tile 的物理布局与有效尺寸与地址分配无关因此两个版本中完全一致——这也提醒读者auto mode 自动化的是地址 同步不改变任何 Tile 布局语义。指令语义可分别参见 TMATMUL.md、TMOV.md、TSTORE_FP.md。仓库中的规模化 auto mode GEMM/FA kernel文档示例是单 tile 教学骨架仓库kernels/automode/下有一批真实规模较大的 auto mode kernel可作进阶阅读kernels/automode/a2a3/gemm/A2A3 平台的 GEMM 性能 kernel与本文 TMATMUL 示例同构kernels/automode/a2a3/flash_atten/含multiBuffer.hpp与 softmax/matmul 宏封装kernels/automode/a5/flash_atten/A5 平台的 FlashAttention 性能 kernelGU/DN 双路径。从这些目录结构看规模化 kernel 普遍采用pto_macro_*.hpp宏封装 Tile 级数据流 multiBuffer.hpp处理多缓冲的组织方式核心指令序列仍是示例中的 TLOAD/TMOV/TMATMUL/TSTORE 闭环。启用 Auto Mode 编译示例代码的价值取决于能否正确编译。结合 Auto_Mode_Overview.md 与仓库演示工程启用方式如下适用前提已安装含 Bisheng CCE 工具链的 CANN 环境--cce-aicore-arch需按目标 SoC 调整本仓库示例使用dav-c220-vec、dav-c310-vec等。单文件命令行编译device 编译source /usr/local/Ascend/ascend-toolkit/latest/bin/setenv.bash bisheng -c -x cce -O2 --cce-aicore-only \ --cce-aicore-archdav-c310-vec \ -stdc17 \ --cce-pto-enable \ --cce-pto-auto-enable \ kernel.cpp -o kernel.o两个关键开关--cce-pto-enable打开 PTO 支持--cce-pto-auto-enable打开 auto mode pass-O2是 auto mode 的硬性要求见下方 CMake 注释。CMake 工程编译demos/auto_mode/baseline/add/CMakeLists.txt 展示了通过ascendc_library集成的写法ascendc_library(no_workspace_kernel STATIC csrc/kernel/add_custom.cpp ) ascendc_include_directories(no_workspace_kernel PRIVATE ${PTO_LIB_PATH}/include ${PTO_LIB_PATH}/include/pto/common ) # --cce-pto-enable --cce-pto-auto-enable enables auto mode of compiler # auto mode only works with -O2 ascendc_compile_options(no_workspace_kernel PRIVATE --cce-pto-enable --cce-pto-auto-enable -O2)配套的构建与运行脚本见 demos/auto_mode/baseline/add/run.shpython3 setup.py bdist_wheel打 wheel → 安装 →test/test.py验证其中要求先export PTO_LIB_PATH指向 pto-isa 仓库根目录并在 CMakeLists.txt 中通过SOC_VERSION如Ascend910B1/ascend910b4指定目标芯片。Auto Mode 编写约束写示例代码时的常见坑上述示例代码之所以看起来这么干净是因为遵守了 Kernel_Developer_Rules_And_Limitations.md 中列出的约束违反它们可能导致编译失败、功能错误或性能退化。与本文两组示例直接相关的要点不要依赖TASSIGN生效auto mode 下它是 no-op。需要表达两个 Tile 共享基址别名时用TRESHAPE见 TRESHAPE.md需要表达子视图时用 sub-tile aliasing 接口注意 auto mode 下 Tile 地址一经分配终生不变——可以把 Tile 心智模型理解为C 引用。不要对纯输出 Tile 调TLOADauto mode 基于生命周期分析复用 UB若一个只写不读的 Tile 被TLOAD可能与同地址的其他TLOAD发生数据竞争。不要直接调用 CCE intrinsics 或Tile::data()自动分配/同步只建立在 PTO 指令分析之上且 auto mode 下.data()返回的是向量类型而非指针见 include/pto/common/memory.hpp 与 include/pto/common/pto_tile.hpp 中 auto mode only supports static shapes 的注释在 kernel 中当指针用会直接编译失败。库开发者侧对应规则见 Library_Developer_Rules_And_Limitations.md。暂不要使用 double/multi bufferingauto mode 对双缓冲支持尚不完整规则 1.4 与 demos README 注记 均明确提示复杂控制流会使 auto-sync 退化为保守插入带来性能损失。控制流保持静态可判定首/末迭代守卫应写成可静态求值的形式以便编译器 peel 循环复杂的 if 条件先求值到变量再使用tile function 内部TF 之下的同步仍需库开发者手写set_flag/wait_flag/pipe_barrier因为 auto-sync 不会穿透 tile function 内部。参考路径路径内容docs/auto_mode/Examples.md本文主体文档TADD / TMATMUL 的 auto vs manual 对照示例docs/auto_mode/Auto_Mode_Overview.mdAuto mode 概念、抽象层次与 bisheng 编译命令docs/auto_mode/Kernel_Developer_Rimitations.mdkernel 开发者规则与限制控制流、内存分配、通用规则docs/auto_mode/Library_Developer_Rules_And_Limitations.md库开发者侧约束tile function、*_IMPL规范docs/coding/Event.mdmanual mode event model 语义demos/auto_mode/baseline/add/auto mode add kernel 完整工程CMake、PyTorch 集成、测试kernels/automode/规模化 auto mode kernelgemm / flash_atten / topkA2A3 与 A5include/pto/common/pto_tile.hppTile抽象实现含 auto mode 静态 shape 限制注释【免费下载链接】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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考