CANN PTO-ISA 指令详解TREM 逐元素余数运算指令的数学语义、汇编形式与多平台实现【免费下载链接】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 仓库中 TREM 指令文档 展开系统讲解 PTOParallel Tile Operation并行 Tile 操作虚拟指令集体系下 TREM 指令的数学语义、同步/SSA/DPS 三级汇编语法、C 内建接口、平台约束与实现原理。TREM 是对两个 Tile 执行逐元素余数remainder运算的向量指令其结果的符号与除数一致是与 TFMOD 在符号规则上存在关键差异的算术指令。读者阅读本文后将掌握 TREM 在不同 Ascend 平台Atlas A2/A3 系列、Ascend 950PR/Ascend 950DT上的数据类型支持差异、tmp 临时 Tile 的约束与用法、汇编与 C 两种调用方式并能结合仓库中的 CPU 仿真与 NPU 实现理解其底层算法细节。指令概述与定位TREMTile Remainder是 PTO 指令集中用于计算两个 Tile 逐元素余数的指令。在include/pto/common/pto_instr.hpp中TREM 与 TFMOD 一起被定义在算术指令族中二者共享相似的模板签名但在数学语义上存在根本区别。TREMremainder(src0, src1)结果符号与**除数src1**相同数学上等价于src0 - floor(src0/src1) * src1TFMODfmod(src0, src1)结果符号与**被除数src0**相同。这一差异意味着两者适用于不同的数值语义场景TREM 对应 Python 的math.remainder/ C 的remainder()风格语义截断向负无穷方向取整后的余数而 TFMOD 对应 C 的fmod()语义。在 CPU 仿真实现 中可以看到TREM 的 CPU 版本正是先调用std::fmod计算基础余数再通过符号判断((src0 0) ! (src1 0)) dst ! 0时执行dst src1来把结果符号修正为与除数一致这一逻辑与文档中的数学定义完全吻合。数学语义TREM 对有效区域由dst.GetValidRow()与dst.GetValidCol()决定的迭代域内的每个元素(i, j)执行$$\mathrm{dst}{i,j} \mathrm{remainder}(\mathrm{src0}{i,j}, \mathrm{src1}{i,j}) \mathrm{src0}{i,j} - \mathrm{floor}\left(\frac{\mathrm{src0}{i,j}}{\mathrm{src1}{i,j}}\right) \times \mathrm{src1}_{i,j}$$其核心特征是结果符号会被修正为与除数src1的符号相同。例如remainder(-7, 3) -7 - floor(-7/3) * 3 -7 - (-3) * 3 2结果为正与除数 3 同号而fmod(-7, 3) -1符号与被除数 -7 相同。在 Atlas A2/A3 训练系列产品/Atlas A2/A3 推理系列产品的 NPU 实现 的注释中实现逻辑被明确描述为对于 fp32计算出余数后检查result * divider 0若符号相反则result divider以确保结果与除数同号。这正是上述数学定义在硬件层面的直接落地。汇编语法TREM 提供三种层级的汇编表示覆盖从抽象 SSA 形式到具体 DPSData Parallel Semantics形式的描述。同步形式PTO 汇编%dst trem %src0, %src1 : !pto.tile...AS Level 1SSA%dst pto.trem %src0, %src1 : (!pto.tile..., !pto.tile...) - !pto.tile...AS Level 2DPSpto.trem ins(%src0, %src1 : !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)C 内建接口TREM 的 C 内建接口声明于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hpp内部声明位于pto/common/pto_instr.hpp。template auto PrecisionType RemAlgorithm::DEFAULT, typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1, typename TileDataTmp, typename... WaitEvents PTO_INST RecordEvent TREM(TileDataDst dst, TileDataSrc0 src0, TileDataSrc1 src1, TileDataTmp tmp, WaitEvents ... events);接口说明模板参数PrecisionType默认取RemAlgorithm::DEFAULT用于控制高精度算法选择dst为输出 Tilesrc0、src1为两个输入 Tiletmp为临时 Tile用于存放中间结果可变参数events用于事件同步接口内部先调用detail::PtoWaitEvents(events...)等待前置事件再执行TREM_IMPLPrecisionType最后返回RecordEvent以便后续指令依赖与 TFMOD 不同TREM 多出一个tmp参数——这是因为余数运算需要借助除法、乘法、减法、比较与选择等多个中间步骤需要一块临时缓冲区存放候选结果与比较掩码。平台约束与类型支持TREM 在两类平台上有不同的实现检查约束使用时需特别注意。Atlas A2/A3 训练系列产品 / Atlas A2/A3 推理系列产品TileData::DType必须是float、float32_t、int32_t之一Tile 布局必须是行主序TileData::isRowMajorTile 位置必须是向量TileData::Loc TileType::Vec运行时约束src0、src1和dst应具有相同的validRow/validColtmpTile 必须至少有 2 行和validCols列第 0 行用于中间结果第 1 行用于比较掩码。上述约束在 a2a3 实现 的TRemCheck中以static_assert与PTO_ASSERT的形式落实编译期检查数据类型与行主序布局运行期检查三个 Tile 的validRow/validCol一致性以及tmp的容量。Ascend 950PR / Ascend 950DTTileData::DType支持范围更广half、float、int16_t、uint16_t、int32_t、uint32_t、int64_t、uint64_t同样要求行主序布局与向量位置运行时约束与 A2/A3 相同三个 Tile 的validRow/validCol一致特别说明tmp参数在 Ascend 950PR/Ascend 950DT 上被接受但不会使用接口签名保持一致是为了统一调用方式。通用行为约定有效区域操作使用dst.GetValidRow()/dst.GetValidCol()作为迭代域因此输出 Tile 的有效区域决定了实际计算范围除零行为由目标定义CPU 仿真在调试构建assert中会触发断言。在 CPU 的 ElementOpCal 实现 中可以看到assert(src1 ! static_castDType(0) Divider cannot be equal to zero)的显式检查高精度算法仅在 Ascend 950PR/Ascend 950DT 上对float类型有效PrecisionType选项在 Atlas A2/A3 系列产品上会被忽略。平台实现与算法原理CPU 仿真实现CPU 仿真是理解指令语义的最佳参考。ElementTileOp.h 中的TREM_IMPL忽略tmp参数直接调用BinaryElementTileOp_ImplElementOp::OP_REM即逐元素二元运算框架。而 ElementOp.h 中的ElementOpCalDType, ElementOp::OP_REM::apply展示了核心算法对整型执行src0 % src1对浮点执行std::fmod(src0, src1)若(src0 0) ! (src1 0)且结果不为 0则dst src1完成符号修正。这也解释了 TREM 与 TFMODOP_FMOD直接返回std::fmod结果在 CPU 语义上的区别来源。Atlas A2/A3 NPU 实现include/pto/npu/a2a3/TRem.hpp 展示了 fp32 路径的完整指令序列其本质上是用一组向量指令组合出余数语义vdiv(tmp, src0, src1, ...)计算商随后vconv_f322f32做向负无穷方向的截断floor 语义vmul(dst, tmp, src1, ...)计算floor(src0/src1) * src1vsub(dst, src0, dst, ...)得到候选余数通过vcmpvs_lt在 norm mask 模式下比较候选余数 * src1 0生成符号修正掩码再用vselVSEL_TENSOR_TENSOR_MODE读打包位掩码在“余数 src1”与“原候选”之间选择完成与除数同号的修正。int32 路径RemInt32Instr在此基础上额外包含vconv_s322f32/vconv_f322s32系列整数-浮点转换指令并处理了 f32 表示精度问题。值得注意的细节是比较指令必须显式传入repeatTimes并切换到 norm mask 模式否则仅发射单次 repeat 会遗留过期的 UB 掩码数据掩码缓冲区之后还跟随一个 32 字节对齐的地址缓冲区供set_cmpmask的两级地址间接寻址使用。这些实现细节体现了tmp缓冲区在 A2/A3 上需要“2 行”的原因——第 0 行存中间结果第 1 行存比较掩码及地址缓冲。Ascend 950 系列寄存器级实现对于 Kirin9030、KirinDev0000Ascend 950PR/950DT 相关架构include/pto/common/arch/register/trem_common.hpp 提供了架构无关的浮点/半精度余数内核RemOpBaseRemFloat以寄存器张量执行vdiv → vtrc(ROUND_F) → vmul → vsub得到候选余数再用vcmps_lt生成符号差异掩码通过vadd(diffSign, dst, src1, diffSignMask, MODE_MERGING)完成符号修正同时处理了 |src0| inf 时的 NaN 结果vcmps_eq检测后写入 NaN 位模式保证边界语义RemHalf则先把半精度拆分为偶数/奇数通道提升为 fp32 运算vcvt PART_EVEN/PART_ODD计算完成后再合并回落规避半精度中间运算的精度损失整型余数内核是架构相关的Kirin9030 使用硬件整型除法而 KirinDev0000 缺少该能力回退到软件除法因此各架构外壳shell在基类之上自行提供RemInt与分发逻辑。这也印证了文档中“tmp参数在 Ascend 950PR/Ascend 950DT 上被接受但不使用”的说明——该平台走的是寄存器级RegTensor实现无需 UB 临时缓冲区。使用示例C 内建接口示例以下示例直接取自 TREM 指令文档展示了最精简的调用方式#include pto/pto-inst.hpp using namespace pto; void example() { using TileT TileTileType::Vec, float, 16, 16; using TmpT TileTileType::Vec, float, 2, 16; TileT dst, src0, src1; TmpT tmp; TREM(dst, src0, src1, tmp); }要点输出/输入 Tile 使用行主序向量 TileTileType::Vectmp在 Atlas A2/A3 上按文档约束声明为至少 2 行 × 有效列数的 Tile此处TmpT为 2×16与TileT的 16 列对应。完整加载-计算-存储流程参考 CPU 测试用例仓库的 CPU 单测 tests/cpu/st/testcase/trem/trem_kernel.cpp 给出了一个完整的TLOAD → TREM → TSTORE流程可作为实际算子编写的参考骨架#include pto/pto-inst.hpp #include pto/common/constants.hpp using namespace pto; template typename T, int kDRows_, int kDCols_, int kTRows_, int kTCols_ AICORE void runTRem(__gm__ T __out__* out, __gm__ T __in__* src0, __gm__ T __in__* src1) { using DynShapeDim5 Shape1, 1, 1, kTRows_, kTCols_; using DynStridDim5 Stride1, 1, 1, kTCols_, 1; using GlobalData GlobalTensorT, DynShapeDim5, DynStridDim5; using TileDataDst TileTileType::Vec, T, kDRows_, kDCols_, BLayout::RowMajor, -1, -1; using TileDataSrc TileTileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1; TileDataSrc src0Tile(kTRows_, kTCols_); TileDataSrc src1Tile(kTRows_, kTCols_); TileDataDst dstTile(kTRows_, kTCols_); TileDataDst tmpTile(1, kTCols_); // 临时 TileCPU 仿真下忽略其内容 GlobalData src0Global(src0); GlobalData src1Global(src1); GlobalData dstGlobal(out); TASSIGN(src0Tile, 0); TASSIGN(src1Tile, kTRows_ * kTCols_ * sizeof(T)); TASSIGN(dstTile, 2 * kTRows_ * kTCols_ * sizeof(T)); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); TREM(dstTile, src0Tile, src1Tile, tmpTile); TSTORE(dstGlobal, dstTile); }该用例在 CPU 仿真下对float、int64_t、uint64_t、aclFloat16以及启用CPU_SIM_BFLOAT_ENABLED时的bfloat16_t等类型做了实例化覆盖了多种形状如 64×64、16×256、64×512 等可作为理解 TREM 数据通路和进行功能验证的参考。需要说明的是CPU 仿真支持的类型范围比 NPU 平台的实现检查更广实际部署时应以目标平台的约束表为准。汇编示例ASM自动模式自动模式下资源的放置与调度由编译器/运行时负责开发者只需描述数据流# 自动模式由编译器/运行时负责资源放置与调度。 %dst pto.trem %src0, %src1 : (!pto.tile..., !pto.tile...) - !pto.tile...手动模式手动模式下需要先显式绑定 Tile 资源通过pto.tassign将虚拟 Tile 绑定到物理地址再发射指令# 手动模式先显式绑定资源再发射指令。 # 可选当该指令包含 tile 操作数时 # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %dst pto.trem %src0, %src1 : (!pto.tile..., !pto.tile...) - !pto.tile...PTO 汇编形式%dst trem %src0, %src1 : !pto.tile... # AS Level 2 (DPS) pto.trem ins(%src0, %src1 : !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)与 TFMOD 的对比与选型建议维度TREMTFMOD数学定义src0 - floor(src0/src1) * src1C 语言fmod语义结果符号与除数 src1相同与被除数 src0相同典型语义math.remainder向负无穷取整fmod向零取整临时 Tile 参数需要tmpA2/A3 上必须提供不需要高精度选项RemAlgorithm::DEFAULT950 系列 float 有效FmodAlgorithm::DEFAULT选型建议当算子语义要求“余数与被除数同号”例如某些降采样、对齐计算中的边界修正时使用 TFMOD当语义要求“余数与除数同号”、或需要与math.remainder/fmodIEEE 754remainder的 floor 变体对齐时使用 TREM。在实现归一化、循环寻址、量化误差修正等算法时TREM 的“与除数同号”特性可避免结果跨过 0 点出现跳变使余数始终落在[0, |src1|)区间方向上的稳定一侧。进一步阅读指令总览与体系可参考 docs/isa/PTO-Virtual-ISA-Manual.md 与 docs/isa/README.md相关算术指令TFMOD余数、被除数符号见 TFMOD.md、TDIV除法见 TDIV.md、TMOD 相关整数取模语义可对照 TREM_zh.md 所在目录下的指令族文档标量形式 TREMSsrc1 为标量的余数运算对应文档 TREMS.md源码实现CPU 仿真见 include/pto/cpu/ElementTileOp.h 与 include/pto/cpu/ElementOp.hA2/A3 NPU 实现见 include/pto/npu/a2a3/TRem.hpp950 系列寄存器实现见 include/pto/common/arch/register/trem_common.hpp测试用例CPU 单测见 tests/cpu/st/testcase/trem/A2/A3 NPU 测试见 tests/npu/a2a3/src/st/testcase/trem/其中main.cpp负责组装数据、gen_data.py生成输入输出基准数据可用于验证本文所述的符号语义。参考docs/isa/TREM_zh.md本文主依据文档include/pto/common/pto_instr.hppTREM/TFMOD 内建接口声明include/pto/npu/a2a3/TRem.hppAtlas A2/A3 NPU 实现与约束检查include/pto/common/arch/register/trem_common.hppAscend 950 系列寄存器级余数内核include/pto/cpu/ElementOp.hCPU 仿真逐元素余数语义tests/cpu/st/testcase/trem/trem_kernel.cppCPU 仿真完整调用流程示例【免费下载链接】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），仅供参考