CANN pto-isa 中的 TSORT32基于 VBS32 的 32 元素块值-索引排序指令详解【免费下载链接】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-isaPTOParallel Tile Operation虚拟指令集中的TSORT32指令用于对 Tile 中每个 32 元素块连同其索引一起做降序排序并将排序后的值-索引对写入目标 Tile底层由 SFU 硬件指令VBS32vbitsort支撑。本文以仓库中 TSORT32_zh.md 指令文档为主体结合 A2/A3 实现、A5 实现、CPU 仿真实现 与 ST 测试用例完整讲清该指令的数学语义、Tile 尺寸规则、尾部非对齐处理策略以及跨平台实现差异读完后可直接在 PTO kernel 中正确调用TSORT32并完成tmp缓冲规划与流水线同步。指令语义每 32 元素块的值-索引对排序TSORT32对src的每个 32 元素块与idx中对应的索引一起排序并将排序后的值-索引对写入dst。对每一行rsrc按独立的 32 元素块处理设块b覆盖列32b … 32b31n_b min(32, C - 32b)为其有效元素数则构成值-索引对$$ (v_k, i_k) (\mathrm{src}{r,32bk},; \mathrm{idx}{r,32bk}), \quad 0 \le k n_b $$按值降序排序输出重排后的序列$$ [(v_{\pi(0)}, i_{\pi(0)}),; (v_{\pi(1)}, i_{\pi(1)}),; \ldots] $$其中π为该块的排序置换值相同时索引小者优先tie-break。两个容易误解的要点idx是输入Tile索引随值一起被重排不是输出dst存储的是排序后的值-索引对而非仅排序后的值。这一语义在 CPU 仿真实现 中可得到直接印证其ScoreIndexPair比较器先按分数降序、分数相同再按原始索引升序并使用std::stable_sort保证稳定性与硬件的 tie-break 规则一致。硬件基础VBS32vbitsortVBS32 运行在SFU非向量流水线上。一次调用可排序repeat个连续的 32 元素块每个块由 32 个值 32 个索引组成void vbitsort(__ubuf__ T *dst, // 排序后的值-索引对输出 __ubuf__ T *src0, // 每块 32 个值 × repeat __ubuf__ uint32_t *src1, // 每块 32 个索引 × repeat uint8_t repeat); // 32 元素块的数量1..255repeat上限REPEAT_MAX 255打包到指令字config[63:56]块在内存中连续分布、步长为 32 个元素块b读取src0[b*32 : b*3232]和src1[b*32 : b*3232]写入dst[b*32*coef : ...]其中coef 2float或 4half——即值-索引对的扩展因子排序顺序按值降序相同值时索引小者优先。从源码结构看REPEAT_MAX与BLOCK_SIZE等常量定义在 common/constants.hpp被各平台实现统一 includeMAX_UB_TMP 32 * 255 8160在 a2a3/TSort32.hpp 中定义它同时是单次vbitsort调用可处理元素数的上限255 块 × 32 元素。当repeat超过 255 块时A2/A3 实现 会将一行拆分为多次vbitsort调用前若干组各处理 255 块最后一组处理余下块。C 内建接口接口声明于 include/pto/common/pto_instr.hpp公共包含头为pto/pto-inst.hpp// TSORT32 does not take wait events. Wait on events before calling it if needed. template typename DstTileData, typename SrcTileData, typename IdxTileData PTO_INST RecordEvent TSORT32(DstTileData dst, SrcTileData src, IdxTileData idx); // 角色映射 OIO1 个输出 2 个输入 template typename DstTileData, typename SrcTileData, typename IdxTileData, typename TmpTileData PTO_INST RecordEvent TSORT32(DstTileData dst, SrcTileData src, IdxTileData idx, TmpTileData tmp); // 角色映射 OIOI1 个输出 3 个输入要点3 参数版本要求src列数 32 对齐validCol % 32 0直接排序4 参数版本支持非 32 对齐尾部validCol % 32 ! 0通过tmp做尾部填充该指令不接受WaitEvents...参数、也不做内部事件同步——如需同步调用者须在调用前显式等待相关 event。这在头文件注释中被明确写出也体现在约束表中。各平台的底层实现通过MAP_INSTR_IMPL_ROLES宏分发到TSORT32_IMPLpto_instr_impl.hpp 中按目标平台分别 include 了pto/npu/a2a3/TSort32.hpp、pto/npu/a5/TSort32.hpp与pto/cpu/TSort32.hpp。Tile 尺寸与数据类型对于src形状为 $R \times C$有效区域、块大小 32Tiledtype尺寸元素数说明srchalf或float$T$$R \times C$待排序的值idxuint32_t$R \times C$或 $1 \times C$ 广播随值重排的索引dst$T$$R \times (2C)$ float$R \times (4C)$ half排序后的值-索引对tmp仅4参数$T$见下方 tmp 尺寸公式尾部填充 scratchdst扩展因子typeCoef每个输入元素生成一个 8 Byte 的 tuple[value (4Byte), index (4Byte)]——float的 value 占满 4 Bytehalf的 2 Byte value 零扩展至 4 Byte。因此dst恒为 $C \times 8$ 字节dtype每个src列对应的dst列数dtype 单位tuple 布局字节/tuplefloat×22 个 float 槽位[value_f32, index_u32]8half×44 个 half 槽位[value_f16, 0x0000, index_u32]8源码印证a2a3 实现 定义了FLOAT_DST_STRIDE_COEF 2与HALF_DST_STRIDE_COEF 4写dst时按typeCoef推进偏移CPU 仿真 中t totalByte / sizeof(T)totalByte 8推进输出指针half 时还显式写入一个 0 槽位与上表 tuple 布局完全一致。关于idx广播实现中idxStride idx.GetValidRow() 1 ? 0 : RowStride见 a2a3/TSort32.hpp即当idx只有一行时行步长取 0实现 $1 \times C$ 的列广播。约束汇总约束原因dst/srcdtype half或float须一致idxuint32_tVBS32 类型分派所有 Tile 为TileType::Vec、BLayout::RowMajorSFU 寻址validCol % 32 03参数每块恰为 32 个元素validCol任意4参数尾块通过tmp填充至 32填充值为 $-\infty$repeat validCol/323参数或ceil(validCol/32)4参数VBS32 repeat 计数每次调用 ≤ 255更大的validCol拆分为多次vbitsort调用tmp4参数≥tmpSize元素见下方公式保存填充后的行/尾块副本无WaitEvents.../ 无内部 event synchronization如需同步须显式调用这些约束在编译期由CheckStatic以static_assert强制检查见 a2a3/TSort32.hppdst/src 必须同为half或float、idx 必须为uint32_t、三个 Tile 均须TileType::Vec且行主序。违反任一条件将直接触发编译错误。tmp尺寸公式4参数设 $C$ validCol$b$ sizeof(T)字节数$G$ 32块大小。实现根据整行大小是否满足MAX_UB_TMP 8160进行分支且阈值单位因目标平台而异$$ \mathrm{tmpSize} \begin{cases} \mathrm{ceil}_{G}(C) \text{A2/A3} C \le 8160 \text{元素数} ;; \text{A5/950PR/950DT} C \cdot b \le 8160 \text{字节} \ G 32 \text{A2/A3} C 8160 \text{元素数} ;; \text{A5} C \cdot b 8160 \text{字节} \end{cases} $$ceil_G(C) $C$ 向上取整到 32 的倍数Atlas A2/A3 训练/推理系列产品阈值单位为元素数。源码 的判断是srcShapeBytesPerRow / sizeof(T) MAX_UB_TMP即 $C \le 8160$与 dtype 无关float、half 均为 $C \le 8160$A5Ascend 950PR/950DT阈值单位为字节。源码 的判断是validCol * sizeof(T) MAX_UB_TMP即 $C \cdot b \le 8160$float → $C \le 2040$half → $C \le 4080$。该阈值为pto_copy_ubuf_to_ubufMOV_UB_TO_UB的 repeat 上限 255 块 × 32 Byte尾块 $t C \bmod G$ 个元素末尾不完整块扩展至 $G$ 并以 $-\infty$ 填充Path A小行从行首整行复制到tmp然后原地填充最后 32 个元素Path B大行仅复制尾块到tmp完整块直接从src排序VBS32 硬件上限每次调用repeat ≤ 255块≤ 8160 元素超过 255 块的行拆分为多次vbitsort调用UB 布局tmp应紧接dst之后放置32 Byte 对齐大小为ceil(C·b, 32)字节因 $b \in {2,4}$ 整除 32等价于ceil(ceil(C, 32)·b, 32)——不应使用固定的 8KB 偏移因为 Path AA2/A3在接近阈值时对 float 需要最多约 32KB$C \le 8160$ 元素 32KB。4 参数尾部处理路径当validCol % 32 ! 0时末尾不完整块$t C \bmod 32$ 个元素须填充为完整的 32 元素块后才能送入vbitsortA2/A3$C \le 8160$元素数 / A5$C \cdot b \le 8160$字节小行整行复制到tmp然后原地覆盖最后 32 个元素为 $-\infty$ 填充从tmp排序整行。A2/A3 路径用vector_dup完成填充见 源码A5 路径则通过vlds/vdup/vsts向量指令配合掩码完成见 源码A2/A3$C 8160$元素数 / A5$C \cdot b 8160$字节大行仅复制尾块到tmp并填充完整块直接从src排序仅尾块从tmp排序。该逻辑封装在LargeTmpBufferImpl中a2a3 版本先对前面各满组做vbitsort(..., REPEAT_MAX)最后一组用srcTailRepeatNum - 1处理剩余满块再将尾块拷入tmp、用set_vector_maskvector_dup填充 $-\infty$、最后单独vbitsort(..., 1)。填充值$-\infty$ -(0.0/0.0)即负 NaN等价于std::numeric_limitsT::lowest()落在降序排序的底部不影响有效元素的相对顺序。若validCol 32 × 255行按REPEAT_MAX大小的组拆分每组通过独立的vbitsort调用排序。平台实现差异从源码看三条执行路径PTO 为同一TSORT32内建指令提供了三套实现理解差异有助于排查跨平台行为A2/A3Atlas A2/A3 系列——include/pto/npu/a2a3/TSort32.hpp每个vbitsort调用后都插入pipe_barrier(PIPE_V)保证 SFU 结果在下一调用前就绪行步长直接取Tile::RowStridedstStride/srcStride为编译期常量小行阈值按元素数判断srcShapeBytesPerRow / sizeof(T) MAX_UB_TMP。A5——include/pto/npu/a5/TSort32.hpp不插入pipe_barrier(PIPE_V)流水线语义由编译器/硬件保证行步长按 32 Byte 对齐计算PTO_CEIL(Cols * byteSize, 32) / byteSize尾部填充改走__VEC_SCOPE__内的vlds/vdup/vsts向量指令且阈值按字节数判断dst与src的行步长按各自的列宽独立向上取整到 32 字节这解释了为何dst尺寸公式中 half 的扩展因子是 4一个 tuple 占 8 字节 4 个 half。CPU 仿真——include/pto/cpu/TSort32.hpp用std::stable_sort 自定义比较器分数降序、同分按索引升序逐块模拟硬件行为保证 CPU 结果可作 golden reference静态断言比 NPU 路径更宽松额外允许int32_t/int16_t/bfloat16_t且 4 参数版本直接复用 3 参数逻辑validCol任意值天然被std::min(sortNum, validCol - j)截断无需显式填充。测试用例如何在仓库中验证仓库为多个平台提供了完整的 STSystem Test用例可直接编译运行验证语义。以 A2/A3 为例tests/npu/a2a3/src/st/testcase/tsort32/main.cpp 定义了 4 个用例恰好覆盖了文档的全部关键路径TEST_F(TSort32Test, case1) { tsort32_test1, float(2, 32, 32); } // 32 对齐单块 TEST_F(TSort32Test, case2) { tsort32_test2, uint16_t(4, 64, 64); } // halfuint16_t 即 f16 存储双块 TEST_F(TSort32Test, case3) { tsort32_test3, float(1, 32 * 256, 32 * 256); } // 256 块 REPEAT_MAX(255)验证分组拆分 TEST_F(TSort32Test, case4) { tsort32_test4, float(2, 13, 16); } // 非 32 对齐尾部13 列4 参数 tmp对应 kernel tsort32_kernel.cpp 展示了完整的 UB 规划与 Manual 模式同步写法// Tile 规划ALIGN_C 为 32 对齐后的列宽TYPE_COEF sizeof(float)/sizeof(T) using SrcTileData TileTileType::Vec, T, ROWS_, ALIGN_C, BLayout::RowMajor, -1, -1; using IdxTileData TileTileType::Vec, uint32_t, ROWS_, ALIGN_C, BLayout::RowMajor, -1, -1; using DstTileData TileTileType::Vec, T, ROWS_, TYPE_COEF * 2 * ALIGN_C, BLayout::RowMajor, -1, -1; using TmpTileData TileTileType::Vec, T, 1, ALIGN_C, BLayout::RowMajor; TASSIGN(srcTile, 0x0); // src 0x0 TASSIGN(idxTile, 0x8000); // idx 0x800032 字节对齐偏移 TASSIGN(dstTile, 0x16000); // dst 0x16000 TASSIGN(tmpTile, 0x20000); // tmp 0x20000 TLOAD(srcTile, srcGlobal); TLOAD(idxTile, idxGlobal); TLOAD(tmpTile, tmpGlobal); #ifndef __PTO_AUTO__ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); // Manual 模式显式等待 MTE2 加载完成 wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); #endif TSORT32(dstTile, srcTile, idxTile, tmpTile); // 无 wait event 参数 #ifndef __PTO_AUTO__ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // Manual 模式SFU 完成后通知 MTE3 写回 wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); #endif TSTORE(dstGlobal, dstTile);注意dst的有效列数为VALID_C * TYPE_COEF * 2、分配总列数为TYPE_COEF * 2 * ALIGN_C与文档的扩展因子规则一致tmp分配为ceil32(C)元素测试中tmpByteSize colsAlign * typeSize也印证了 tmp 尺寸公式。A5、kirin9030、kirinDev0000、kirinX90 及 CPU 仿真均有同构用例A5、CPU、kirin9030每个目录都附带gen_data.py生成数据与 golden 校验。汇编语法与示例AS Level 1SSA%dst pto.tsort32 %src, %idx : (!pto.tile..., !pto.tile...) - !pto.tile...AS Level 2DPSpto.tsort32 ins(%src, %idx : !pto.tile_buf..., !pto.tile_buf...) outs(%dst : !pto.tile_buf...)C 调用示例32 对齐与非对齐尾部#include pto/pto-inst.hpp using namespace pto; // 32 对齐每行单个块 using SrcT TileTileType::Vec, float, 1, 32; using IdxT TileTileType::Vec, uint32_t, 1, 32; using DstT TileTileType::Vec, float, 1, 64; // 2× src 列数float SrcT src; IdxT idx; DstT dst; TSORT32(dst, src, idx); // 非 32 对齐尾部4 参数 tmp using SrcT2 TileTileType::Vec, half, 1, 100; using IdxT2 TileTileType::Vec, uint32_t, 1, 100; using DstT2 TileTileType::Vec, half, 1, 400; // 4× src 列数half using TmpT TileTileType::Vec, half, 1, 128; // ≥ ceil32(100)128 TmpT tmp; TSORT32(dst2, src2, idx2, tmp);Manual 模式 ASM 形式Manual 模式下需先用pto.tassign绑定 tile 地址偏移需 32 字节对齐再发射排序指令与上文 kernel 中的TASSIGN序列一一对应# pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) # pto.tassign %arg2, tile(0x3000) %dst pto.tsort32 %src, %idx : (!pto.tile..., !pto.tile...) - !pto.tile...小结与实践要点选接口validCol % 32 0用 3 参数版本否则用 4 参数版本并预留tmp规划 UBdst按 8 字节/元素对扩展float ×2 列、half ×4 列tmp按tmpSize公式放置于dst之后、32 字节对齐勿用固定 8KB 偏移注意阈值单位差异A2/A3 按元素数$C \le 8160$、A5 按字节数float $C \le 2040$、half $C \le 4080$走小行/大行两条路径显式同步TSORT32不接收 wait eventManual 模式下须按set_flag/wait_flag在 MTE2→V、V→MTE3 边界显式同步Auto 模式下编译器自动处理验证手段跨平台行为可用 cpu 仿真用例 生成 golden再用各 NPU ST 用例比对。参考路径指令文档docs/isa/TSORT32_zh.md、docs/isa/TSORT32.md公共声明include/pto/common/pto_instr.hpp平台实现include/pto/npu/a2a3/TSort32.hpp、include/pto/npu/a5/TSort32.hpp、include/pto/cpu/TSort32.hpp测试用例tests/npu/a2a3/src/st/testcase/tsort32/、tests/npu/a5/src/st/testcase/tsort32/、tests/cpu/st/testcase/tsort32/【免费下载链接】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），仅供参考