PTO-ISA TCMP 指令详解：Tile 比较与打包谓词掩码的跨平台实现指南

PTO-ISA TCMP 指令详解Tile 比较与打包谓词掩码的跨平台实现指南【免费下载链接】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-isaTCMPTile Compare是 Ascend CANN 并行 Tile 操作虚拟指令集PTO-ISAParallel Tile Operation中的核心矢量比较指令它逐元素比较两个 Tile 的数据将比较结果以**打包的谓词掩码packed predicate mask**形式写入第三个 Tile是 Tile 级条件分支、数据过滤、掩码生成等算法的基础构件。本篇文章以 TCMP 指令文档 为骨架结合仓库内include/pto下的 A2A3、A5、CPU 三套实现与tests/cpu/st/testcase/tcmp中的测试用例完整讲解 TCMP 的数学语义、汇编语法、C 内建接口、类型约束、掩码编码规则与实战用法读者读完后可以在自己的 PTO kernel 中正确使用TCMP生成掩码并理解其在各平台上的行为差异。指令概述与数学语义TCMP 的语义非常直接比较两个 Tile并写入一个打包的谓词掩码。概念上对于有效区域中的每个元素(i, j)指令定义一个谓词$$ p_{i,j} \left(\mathrm{src0}{i,j}\ \mathrm{cmpMode}\ \mathrm{src1}{i,j}\right) $$其中cmpMode是本次比较采用的模式p_{i,j}为布尔值。所有谓词位按照实现定义的打包布局packed layout存入目标 Tiledst。也就是说dst中存放的不是逐元素的 0/1 数值而是将 8 个比较结果压缩进 1 个字节的高密度位掩码——这正是它在内存带宽上优于逐元素写出的原因。CmpMode六种比较模式CmpMode枚举定义在 include/pto/common/type.hpp#L189-L196取值如下enum class CmpMode : uint8_t { EQ 0, NE 1, LT 2, LE 3, GT 4, GE 5, };六个模式分别对应相等、不等、小于、小于等于、大于、大于等于覆盖了绝大多数数据筛选场景。在后续的汇编形式中比较模式通过属性cmpMode #pto.cmpEQPTO 汇编层或cmpMode #ptocmp xxAS Level 1/2指定。汇编语法TCMP 提供三种粒度的汇编形式从 PTO 指令集汇编同步形式到 MLIR 风格的 SSA 再到 DPS 数据流形式。同步形式PTO 汇编%dst tcmp %src0, %src1 {cmpMode #pto.cmpEQ} : !pto.tile... - !pto.tile...AS Level 1SSA 形式——操作数与结果都是!pto.tile...值%dst pto.tcmp %src0, %src1{cmpMode #ptocmp xx}: (!pto.tile..., !pto.tile...) - !pto.tile...AS Level 2DPS 数据流形式——操作数改用带存储属性的!pto.tile_buf...明确ins输入与outs输出区分pto.tcmp ins(%src0, %src1{cmpMode #ptocmp xx}: !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)C 内建接口在 C kernel 中TCMP 以模板内建函数的形式暴露公共头文件为pto/pto-inst.hpp声明位于 include/pto/common/pto_instr.hpp#L333-L339template typename TileDataDst, typename TileDataSrc0, typename TileDataSrc1, typename... WaitEvents PTO_INST RecordEvent TCMP(TileDataDst dst, TileDataSrc0 src0, TileDataSrc1 src1, CmpMode cmpMode, WaitEvents ... events);接口要点dst、src0、src1三个 Tile 分别对应目标掩码、左操作数、右操作数cmpMode传入CmpMode枚举值如CmpMode::GT变参WaitEvents支持事件同步接口内部先调用detail::PtoWaitEvents(events...)等待前序事件再通过MAP_INSTR_IMPL(TCMP, ...)宏分发到具体平台的实现返回RecordEvent可作为后续指令的等待事件参与流水编排。约束条件A2A3 与 A5 的实现差异TCMP 在不同硬件代际上的支持面差异明显编写可移植 kernel 前务必对照约束。A2A3Atlas A2/A3 训练/推理系列检查项实现位于 include/pto/npu/a2a3/TCmp.hpp编译期与运行期检查包括输入类型仅支持int32_t、half、float三种输出类型必须为uint8_tTile 位置src0/src1/dst的TileType::Loc必须为TileType::Vec矢量单元静态有效边界TileDataSrc::ValidRow TileDataSrc::Rows且TileDataSrc::ValidCol TileDataSrc::Cols编译期用static_assert保证运行期形状src0与src1的有效行数、有效列数必须分别相等且src0.GetValidRow() dst.GetValidRow()目标列语义dst的有效列数描述的是打包容量不要求等于源有效列数int32_t 特例仅支持EQ与NENE对相等比较结果取反其余模式一律走EQ路径。从源码 include/pto/npu/a2a3/TCmp.hpp#L27-L65 可以看到A2A3 的底层是vcmpv_eq/lt/gt/ge/le系列矢量比较指令且定义了单次发射最大 repeat 数TCMP_REPEAT_MAX 240超过部分通过外层numLoop循环切分。NE的实现方式是在比较之后额外插入pipe_barrier(PIPE_V)并执行一次vnot对掩码取反include/pto/npu/a2a3/TCmp.hpp#L102-L111。A5Ascend 950PR / 950DT检查项实现位于 include/pto/npu/a5/TCmp.hpp支持面显著扩大输入类型支持uint32_t、int32_t、int64_t、uint64_t、uint16_t、int16_t、uint8_t、int8_t、float、half、bfloat16_t共 11 种输出为打包谓词字节可使用 RowMajor 的uint8_t掩码 Tile迭代域以src0.GetValidRow()/src0.GetValidCol()为比较次数src1必须提供对应的有效元素dst有效列只描述打包容量不决定比较次数编译期检查要求三个 Tile 均为TileType::Vec且为 RowMajor 布局见 include/pto/npu/a5/TCmp.hpp#L301-L325 的TcmpCheck。A5 实现按元素宽度分派到三条路径include/pto/npu/a5/TCmp.hpp#L327-L344sizeof(T) 8int64/uint64走Int64Compare将 64 位数据拆成高/低 32 位寄存器对先比较高位、再按模式用psel/por组合低位结果最终pdintlv_b8psts(..., PK)完成位打包sizeof(T) 4走TCmp_32B用双发射 pdintlv_b8交叉解交织完成 32 位元素的位压缩sizeof(T) 2/1走TCmp_8B_16B单发比较后按PK/NORM分布写入。掩码编码规则掩码的具体位布局由目标平台定义文档与实现给出了 A5 平台64 位输入的明确规则对 64 位输入第j列的比较结果存放在该行第j / 8个字节的第j % 8位最低有效位在前little-endian bit order对于uint8_t掩码有效形状可设为[R, ceil(C / 8)]其中[R, C]是源 Tile 的有效形状物理Cols需要按 32 字节对齐行地址按目标 Tile 的物理行步长RowStride计算最后一个有效位之后的填充位取值未指定不应依赖其值。CPU 模拟实现 include/pto/cpu/TCmp.h#L48-L66 印证了这一规则它按字宽uint8_t为 8 位、uint32_t为 32 位将源列切分为validWords ceil(srcValidCol / kBitsPerWord)个打包字逐位packedWord | (cmp bit)组装再写入目标 Tile。CPU 侧还额外支持uint32_t输出类型此时 dst 数据类型断言放宽为uint8_t/uint32_t见 include/pto/cpu/TCmp.h#L100-L102并在写掩码前调用ZeroTileData将目标清零避免脏数据。实战示例Auto 模式Auto 模式下 Tile 的资源放置与调度由编译器/运行时管理只需声明 Tile 类型并调用接口#include pto/pto-inst.hpp using namespace pto; void example_auto() { using SrcT TileTileType::Vec, float, 16, 16; using MaskT TileTileType::Vec, uint8_t, 16, 32, BLayout::RowMajor, -1, -1; SrcT src0, src1; MaskT mask(16, 2); TCMP(mask, src0, src1, CmpMode::GT); }注意MaskT的声明物理Cols 32满足 32 字节对齐运行期构造时传入有效形状(16, 2)即 16 行 × 2 个打包字节ceil(16 / 8) 2正好容纳 16×16 源区域 16 行每行 16 个比较位的掩码。Manual 模式Manual 模式下需要先用TASSIGN显式绑定各 Tile 的物理地址再发射 TCMP#include pto/pto-inst.hpp using namespace pto; void example_manual() { using SrcT TileTileType::Vec, float, 16, 16; using MaskT TileTileType::Vec, uint8_t, 16, 32, BLayout::RowMajor, -1, -1; SrcT src0, src1; MaskT mask(16, 2); TASSIGN(src0, 0x1000); TASSIGN(src1, 0x2000); TASSIGN(mask, 0x3000); TCMP(mask, src0, src1, CmpMode::GT); }对应的 PTO 汇编形式同样体现了自动/手动两种风格# Auto mode: compiler/runtime-managed placement and scheduling. %dst pto.tcmp %src0, %src1{cmpMode #ptocmp xx}: (!pto.tile..., !pto.tile...) - !pto.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.tcmp %src0, %src1{cmpMode #ptocmp xx}: (!pto.tile..., !pto.tile...) - !pto.tile...端到端验证CPU 仿真测试仓库在 tests/cpu/st/testcase/tcmp/tcmp_kernel.cpp 提供了完整的 CPU 仿真验证用例展示了 TCMP 与TLOAD/TSTORE的组合使用范式template typename T, typename TDst, int kGRows_, int kGCols_, int kTRows_, int kTCols_ AICORE void runTCmp(__gm__ TDst __out__* out, __gm__ T __in__* src0, __gm__ T __in__* src1, pto::CmpMode mode) { constexpr int kBitsPerDst sizeof(TDst) * 8; constexpr int kPackedCols (kTCols_ kBitsPerDst - 1) / kBitsPerDst; using SrcTile TileTileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1; using DstTile TileTileType::Vec, TDst, kTRows_, kTCols_, BLayout::RowMajor, -1, -1; SrcTile src0Tile(kTRows_, kTCols_); SrcTile src1Tile(kTRows_, kTCols_); DstTile dstTile(kTRows_, kPackedCols); TASSIGN(src0Tile, 0); TASSIGN(src1Tile, kTRows_ * kTCols_ * sizeof(T)); TASSIGN(dstTile, 2 * kTRows_ * kTCols_ * sizeof(T)); TLOAD(src0Tile, src0Global); TLOAD(src1Tile, src1Global); TCMP(dstTile, src0Tile, src1Tile, mode); TSTORE(dstGlobal, dstTile); }该用例的显式模板实例化覆盖了 CPU 侧支持的类型组合tests/cpu/st/testcase/tcmp/tcmp_kernel.cpp#L61-L90float、int32_t、int64_t、uint64_t、aclFloat16搭配uint8_t输出以及uint32_t、int32_t、uint16_t、int16_t、uint8_t、int8_t、float、aclFloat16含可选bfloat16_t搭配uint32_t输出与上述约束章节完全对应。测试目录同时提供了 gen_data.py 生成对比数据并通过tests/cpu/st/testcase/CMakeLists.txt挂入 ST 测试框架配合tests/run_cpu_tests.sh即可在 CPU 仿真环境下验证掩码输出的正确性。小结功能定位TCMP 是 Tile 级逐元素比较指令输出按位打包的谓词掩码而非逐元素的 0/1 数据适合作为后续TSEL、TGATHER等掩码驱动指令的数据源模式选择CmpMode提供EQ/NE/LT/LE/GT/GE六种模式A2A3 上int32_t输入仅支持EQ/NE其余类型需关注平台差异掩码布局A5 上 64 位输入按行内“第j/8字节第j%8位、LSB 在前”编码uint8_t掩码有效形状取[R, ceil(C/8)]、物理列 32 字节对齐填充位不可依赖跨平台实现A2A3 走vcmpv_*原生比较 repeat 切分240 上限A5 按 8/4/2/1 字节宽度分派三条打包路径并支持 64 位拆比较CPU 仿真侧则提供位打包参考实现与完整 ST 用例三套实现共同保障了 PTO-ISA 的跨平台一致性。【免费下载链接】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），仅供参考