PTO ISA 行指数差指令 TROWEXPANDEXPDIF 全面解析语义、双模式、临时缓冲与跨平台实现【免费下载链接】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本指南基于 CANN pto-isa 仓库的 TROWEXPANDEXPDIF 官方文档围绕行指数差这一核心指令展开它计算exp(src0 - src1)其中src1为每个数据行提供一个标量或 32 字节块。读完本文你将掌握该指令的数学语义、双模式数据布局、汇编三级语法SSA/DPS、C 内建接口的两种重载以及 A2/A3 与 A5 平台在临时缓冲和底层指令调度上的实现差异并能结合仓库测试用例独立写出可运行的调用代码。指令概览与数学语义TROWEXPANDEXPDIFRow-wise Exp-Diff是 PTOParallel Tile Operation平铺并行操作虚拟指令集中用于行方向广播指数差运算的指令。它属于TROWEXPAND*行扩展二元运算指令族与 TROWEXPANDSUB、TROWEXPANDADD、TROWEXPANDMUL 等指令共享相同的扩展操作数 逐元素运算框架其特殊之处在于最终环节叠加了一次指数exp运算。设R dst.GetValidRow()、C dst.GetValidCol()s_i为从扩展操作数src1中取出的每行标量每行一个值则对于0 i R、0 j C$$ \mathrm{dst}{i,j} \exp(\mathrm{src0}{i,j} - s_i) $$即先把第i行的标量s_i沿列方向广播到整行用src0逐元素减去该广播值再对每个差逐元素求指数。该语义天然适用于归一化如 softmax 中的exp(x - max)、概率计算、数值稳定化等以行为粒度的指数差场景。双模式数据布局指令的行为由扩展操作数的布局决定当src0与dst形状一致时扩展操作数是src1当src1与dst形状一致时扩展操作数是src0模式 1每行一个标量ColMajor扩展操作数采用ColMajor布局、单列有效形状覆盖R个值每个标量广播到整行。模式 2每行 32 字节块RowMajor扩展操作数采用RowMajor布局每行32 / sizeof(T)个元素即一个 32 字节块该块在行内按j mod (32 / sizeof(T))周期性重复。从源码可见这一判定逻辑。在 include/pto/npu/a2a3/TRowExpandBinOp.hpp 中TRowExpandBin通过编译期if constexpr (TileDataSrc1::isRowMajor)分流RowMajor 的src1走TRowExpandBinaryInstr32B模式 2ColMajor 的src1走TRowExpandBinaryInstr模式 1需要临时广播缓冲。A5 后端 include/pto/npu/a5/TRowExpandExpdif.hpp 同样用PTO_ASSERT校验扩展操作数形状RowMajor 时validCol 32 / sizeof(T)ColMajor 时validCol 1且有效行数必须等于validRow。汇编语法从 SSA 到 DPSTROWEXPANDEXPDIF 与 PTO 其他指令一样提供三级汇编表示同步形式可直接书写%dst trowexpandexpdif %src0, %src1 : !pto.tile..., !pto.tile... - !pto.tile...AS Level 1SSA 形式操作数类型为抽象 Tile 类型%dst pto.trowexpandexpdif %src0, %src1 : !pto.tile..., !pto.tile... - !pto.tile...AS Level 2DPS 形式操作数类型为具体缓冲 Tilepto.trowexpandexpdif ins(%src0, %src1 : !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)自动模式与手动模式在自动模式下Tile 的资源放置与调度由编译器/运行时统一管理直接发射指令即可# Auto mode: compiler/runtime-managed placement and scheduling. %dst pto.trowexpandexpdif %src0, %src1 : !pto.tile..., !pto.tile... - !pto.tile...在手动模式下必须先通过pto.tassign显式绑定 tile 的 UB 地址再发射指令可选的tassign用于为 tile 操作数指定地址# Manual mode: resources must be bound explicitly before issuing the instruction. # Optional for tile operands: # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %dst pto.trowexpandexpdif %src0, %src1 : !pto.tile..., !pto.tile... - !pto.tile...PTO 汇编形式的完整写法同时给出两级表示%dst trowexpandexpdif %src0, %src1 : !pto.tile..., !pto.tile... - !pto.tile... # AS Level 2 (DPS) pto.trowexpandexpdif ins(%src0, %src1 : !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)C 内建接口TROWEXPANDEXPDIF 的 C 内建接口声明于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hpp。接口以RecordEvent为返回类型并提供两个重载template typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1, typename... WaitEvents PTO_INST RecordEvent TROWEXPANDEXPDIF(TileDataDst dst, TileDataSrc0 src0, TileDataSrc1 src1, WaitEvents ... events); template typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1, typename TileDataTmp, typename... WaitEvents PTO_INST RecordEvent TROWEXPANDEXPDIF(TileDataDst dst, TileDataSrc0 src0, TileDataSrc1 src1, TileDataTmp tmp, WaitEvents ... events);两个重载的内部实现均先调用detail::PtoWaitEvents(events...)处理事件同步再经MAP_INSTR_IMPL宏见 include/pto/common/pto_instr.hpp展开为对应平台的TROWEXPANDEXPDIF_IMPL。变参WaitEvents允许传入事件对象实现跨 PIPE 或跨核的流水线同步这也是 PTO 支持异步并发的关键设计。两个重载的区别在于3 参数重载同时支持模式 1 与模式 2模式 1 时内部使用 8KB 固定缓冲TMP_UB_OFFSET4 参数重载显式传入TileDataTmp tmp但仅支持模式 1ColMajor 扩展操作数、每行标量。约束与数据类型要求根据 docs/isa/TROWEXPANDEXPDIF.md 与 A5 后端源码中的static_assertinclude/pto/npu/a5/TRowExpandExpdif.hpp调用前必须满足三个操作数数据类型一致TileDataDst::DType TileDataSrc0::DType TileDataSrc1::DType数据类型只能是half或float编译期 Tile 形状/布局约束TileDataDst::isRowMajor目标必须行主序模式 1扩展操作数的有效形状必须覆盖R个值每行一个标量模式 2扩展操作数每行提供 32 字节数据确切的分形/布局约束与目标平台相关可参考后端头文件 include/pto/npu/a2a3/TRowExpandExpdif.hpp 与 include/pto/npu/a5/TRowExpandExpdif.hpp。此外A5 后端还要求src0或src1中至少有一个的有效形状与dst相等PTO_ASSERT校验并据此判定扩展操作数是哪一侧。临时 Tile 机制与缓冲大小计算A2/A3TROWEXPANDEXPDIF在内部由TROWEXPANDSUB后接TEXP复合实现因此 4 参数重载中的tmpTile 实际是减法步骤的广播缓冲区。这一结论可直接从 include/pto/npu/a2a3/TRowExpandExpdif.hpp 的实现代码得到印证template typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1 PTO_INTERNAL void TROWEXPANDEXPDIF_IMPL(TileDataDst dst, TileDataSrc0 src0, TileDataSrc1 src1) { TROWEXPANDSUB_IMPL(dst, src0, src1); pipe_barrier(PIPE_V); TEXP_IMPL(dst, dst); }即先做行扩展减法随后以pipe_barrier(PIPE_V)同步向量流水线再对结果就地求指数TEXP_IMPL(dst, dst)。在Atlas A2/A3 系列上ColMajor 扩展操作数模式 1的每行标量通过vbrcb指令广播到tmp缓冲区为每行创建一个 32 字节块随后作为减法步骤的扩展操作数使用。vbrcb的 repeat stride 为 8 个块256 字节每个 repeat 处理 8 行。据此模式 1 的最小tmp大小计算如下公共参数R dst.GetValidRow()T TileDataDst::DType当R 256$$ \text{tmpSize} \left\lceil\frac{R}{8}\right\rceil \times 256 \text{ bytes} $$当R 256操作按循环执行每轮循环最多 30 个 repeat240 行tmp缓冲在各轮之间复用因此单轮需求为$$ \text{tmpSize} 30 \times 256 7680 \text{ bytes} $$对于任意模式 1 调用一个紧凑且与形状无关的上界是8KB8192 字节3 参数重载无tmp在模式 1 下使用内部 8KB 缓冲TMP_UB_OFFSET模式 2 不需要广播缓冲。上述行数与 repeat 的对应关系在 include/pto/npu/a2a3/TRowExpandBinOp.hpp 的TRowExpandBinaryNormMode中可逐一验证repeatTimes CeilDivision(validRow, 8)常量BRCB_REPEAT_MAX 30、MAX_ROW 240vbrcb(tmpPtr_, src1Ptr, 1, 8, BRCB_REPEAT_MAX)一次广播 30 个 repeat、覆盖 240 行然后处理余数validRow % MAX_ROW。CountMode / NormMode 两种执行路径A2/A3 的广播减法内部还根据形状动态选择两种执行策略include/pto/npu/a2a3/TRowExpandBinOp.hppCountMode计数模式基于set_mask_count()SetVectorCount(validCol)逐行for循环调用RowExpandBinInstr。当 repeat stride 超过 255 或validCol / elementsPerRepeat validRow时启用避免 uint8 类型的 repeatStride 越界NormMode常规模式基于掩码 repeat 展开DstRepeatStride/SrcRepeatStride以uint8_t存储repeatStride / blockSizeElem单次可覆盖整批行行数超过 240 时按 30-repeat 循环。两种模式最终都会在行尾通过SetFullVecMaskByDTypeT()恢复完整向量掩码保证后续指令状态一致。A5 平台的实现差异Ascend 950PR/Ascend 950DTA5的硬件通过vlds指令的广播模式原生支持行广播因此不需要任何临时广播缓冲区。4 参数重载中的tmp被接受但忽略标记为[[maybe_unused]]见 include/pto/npu/a5/TRowExpandExpdif.hpp。A5 在运算指令层面的优化同样显著对float类型直接使用单条向量指令pto_vexpdif(reg_dst, reg_src0, reg_src1, preg, PART_ODD)完成差 指数include/pto/npu/a5/TRowExpandExpdif.hpp无需先减后 exp 两步对half类型则退化为vsubvexp。同时A5 实现还区分src0eqdstsrc0与dst形状一致与相反情形前者按exp(src0 - src1)方向运算后者自动交换操作数顺序按exp(src1 - src0)方向运算保证无论哪一侧是扩展操作数语义都正确。CPU 仿真支持PTO 提供 CPU 仿真后端便于在无昇腾硬件环境下的功能验证。TROWEXPANDEXPDIF 的 CPU 实现位于 include/pto/cpu/TRowExpandOp.hpp通过统一的TRowExpandOp模板与ElementOp::OP_EXPDIF枚举驱动与 TROWEXPANDSUB/ADD/MUL/MAX/MIN 共用同一套行扩展框架4 参数重载同样存在include/pto/cpu/TRowExpandOp.hpp。测试与正确性验证仓库为该指令提供了完备的 NPU 与 CPU 测试覆盖A2/A3 单算子测试tests/npu/a2a3/src/st/testcase/trowexpandexpdif/ 目录包含 kernel、main 与数据生成脚本A5 单算子测试tests/npu/a5/src/st/testcase/trowexpandexpdif/CPU 仿真测试tests/cpu/st/testcase/trowexpandop/。Golden 数据生成逻辑gen_data.py 展示了标准验证流程随机生成input1形状validRow × validCol与input2形状validRow × src1vc模式 1 时src1vc 1、模式 2 时src1vc 32 / itemsize随后按语义计算 goldensrc0eqdst为真golden[i*vcj] input1[i*vcj] - input2[i*src1vc j % src1vc]再整体np.exp(golden)src0eqdst为假交换顺序为golden[i*vcj] input2[...] - input1[...]再取指数。18 个用例覆盖了float/float16、16×16、32×32、1×16384极端瘦长、2048×8极端宽扁、两种布局RowMajor/ColMajor以及src0eqdst两种方向gen_data.py。Kernel 中的典型调用方式A2/A3 kernel 测试trowexpandexpdif_kernel.cpp展示了完整的实战写法以TileTileType::Vec, T, Row, Col, BLayout::RowMajor定义 dst/src0以TileTileType::Vec, T, src1Row, 1, BLayout::ColMajor定义模式 1 的 src1通过TASSIGN绑定 UB 地址、TLOAD装载数据、set_flag/wait_flag完成 MTE2→V 的依赖同步后调用指令最后TSTORE写回并同步 V→MTE3。4 参数重载runTRowExpandExpdif3中tmp只需取初始地址分配空间 ≥src1Row*32B或 ≥8KB 即可行数 256 时甚至可与src1共用空间也印证了前文tmp 仅作广播缓冲、地址有效即可的结论。典型应用场景TROWEXPANDEXPDIF 将行级标量/块广播、减法、指数三步融合为一条指令在以下场景可显著减少指令发射与流水线开销行归一化与数值稳定化softmax 中exp(x - row_max)类算子src1携带每行最大值概率与注意力计算注意力分数exp(q·kᵀ - mask)的行掩码减法指数族统计运算按行计算exp(x - s)的似然、权重更新等。配合TROWEXPANDSUB、TROWEXPANDMUL等同族指令参考 docs/isa/TROWEXPANDSUB.md 与 docs/isa/TROWEXPANDADD.md可在 PTO 编程模型下统一构建行方向广播运算流水线。【免费下载链接】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），仅供参考