PTO 基础 Topk 算子实战解析：基于 TSORT32 与 TMRGSORT 的 A2/A3 排序实现

PTO 基础 Topk 算子实战解析基于 TSORT32 与 TMRGSORT 的 A2/A3 排序实现【免费下载链接】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 仓库中 kernels/manual/a2a3/topk 示例为对象深入剖析如何用 Parallel Tile OperationPTO虚拟指令集实现一个固定维度的 Topk 算子从值-索引对的 Tile 布局、TSORT32 与 TMRGSORT 的分级排序算法、UB 双缓冲流水线调度到数据生成、golden 校验与 NPU/仿真两种运行模式。读完本文你将掌握 PTO 手写 Topk kernel 的完整工程骨架与关键实现技巧并能在 Ascend A2/A3 平台上独立构建、运行和验证该示例。示例概览与工程结构该示例位于kernels/manual/a2a3/topk/用 PTO 实现了固定形状[rows, cols] [4800, 1024]的 Topk 算子从每行 1024 个元素中选出前topk 1000个最大值及其索引。完整的工程结构如下kernels/manual/a2a3/topk/ ├── scripts/ │ └── gen_data.py # 生成输入与 golden 输出 ├── CMakeLists.txt # 构建配置 ├── main.cpp # Host 侧入口ACL 初始化、数据搬运与结果校验 ├── README.md # 英文说明 ├── README_zh.md # 中文说明 ├── run.sh # 一键构建运行脚本 └── topk_kernel.cpp # Kernel 实现AICORE 侧各部分职责明确topk_kernel.cpp是算子核心__global__kernel 与模板化的辅助函数main.cpp通过 ACL 完成 Host/Device 内存管理、kernel 启动与 golden 比对gen_data.py负责生成输入与参考输出run.sh与CMakeLists.txt打通了从源码到可执行程序的完整构建链路。算子规格与 Tiling 设计算子规格如下项目值OpTypetopk输入[rows, cols] [4800, 1024]输出data、indexKernel 名称topk_kernel验证平台有48 个核因此每核承担的形状为rows 100、cols 1024从源码看这一 Tiling 关系被直接编译进 kernel 模板参数。在 topk_kernel.cpp 的launchTopk中constexpr int blockDim 48; // 48 个核 constexpr int gShape3 4800; // 有效行数validRow constexpr int gShape4 1024; // 有效列数validCol constexpr int gWholeShape3 4800; constexpr int gWholeShape4 1280; // 物理列数含 padding constexpr int topk 1000;注意这里存在“有效形状gShape”与“物理形状gWholeShape”的区分物理列宽cols 1280但实际参与排序的有效列只有validCol 1024。这种设计使得 GM 上的行数据可以按 1280 列对齐访问满足TLOAD的搬运对齐要求而计算只处理前 1024 列是典型的手写 kernel 中按物理形状搬运、按有效形状计算的布局技巧。在runTOPK中每核通过get_block_idx()计算自己的数据切片topk_kernel.cppconstexpr int validRow gShape0 * gShape1 * gShape2 * gShape3 / blockDim; // 100 __gm__ T* src origSrc get_block_idx() * validRow * gWholeShape4; __gm__ T* out origOut get_block_idx() * validRow * topk; __gm__ uint32_t* index origIndex get_block_idx() * validRow * topk;每核读取自己的 100 行数据经排序后各自输出 100 × topk 的结果互不重叠。核心算法TSORT32 TMRGSORT 分级排序Topk 的整体思路原文概述是从 GM 加载数据和索引到 UB用 TSORT32 对每 32 个数据进行排序用 TMRGSORT 对每个 Tile 内部做归并排序最后分别取出 topk 个数据和索引存回 GM。下面结合指令语义与源码逐层展开。值-索引对的数据布局Topk 需要同时输出排序后的值和对应的原始下标因此 PTO 采用值-索引对value-index pair作为排序单元float每个 pair 占 8 字节 4 字节 value 4 字节uint32_tindex即2 个 float 槽位half2 字节 value 2 字节 padding 4 字节 index即4 个 half 槽位。这正是 TSORT32_zh.md 与 TMRGSORT_zh.md 中反复强调的 8 字节 tuple 结构。源码中以TYPE_COEF sizeof(float) / sizeof(T)统一表达这一扩展因子constexpr int TYPE_COEF sizeof(float) / sizeof(T); // float1, half2 constexpr int dstCols validCol * 2 * TYPE_COEF; // 排序后 pair 占用的槽位数排序方向为值降序、同值时索引小者优先天然满足 Topk 取前 K 个最大值的需求。GM 侧类型定义原文给出的类型定义展示了数据在 GM 上的 5 维张量描述topk_kernel.cpp// data using DynShapeDim5 Shape1, 1, 1, singleLoopRow, validCol; using DynStridDim5 StridesingleLoopRow * Cols, singleLoopRow * Cols, singleLoopRow * Cols, Cols, 1; using GlobalData GlobalTensorT, DynShapeDim5, DynStridDim5; // index using IndexShapeDim5 Shape1, 1, 1, 1, validCol; using IndexStridDim5 StridevalidCol, validCol, validCol, validCol, 1; using IndexGlobalData GlobalTensorindexT, IndexShapeDim5, IndexStridDim5; // sorted data and index using DstShapeDim5 Shape1, 1, 1, singleLoopRow, topk; using DstStridDim5 StridesingleLoopRow * topk, singleLoopRow * topk, singleLoopRow * topk, topk, 1; using DstDataGlobalData GlobalTensorT, DstShapeDim5, DstStridDim5; using DstIdxGlobalData GlobalTensorindexT, DstShapeDim5, DstStridDim5;其中Shape/Stride是 PTO 描述 GlobalTensor 布局的元编程工具Shape声明每维长度Stride声明每维的字节步长。GlobalData每维步长为singleLoopRow * Cols即按物理列宽Cols 1280取行ND2ND 形式输出张量DstDataGlobalData/DstIdxGlobalData则按topk列紧凑排列。该示例中singleLoopRow 2即每次循环处理两行详见流水线一节。第一级TSORT32 对每 32 元素排序SortEachGroup负责第一级排序topk_kernel.cppRowTile dstRowTile(1, validCol * 2 * TYPE_COEF); // 排序后的值-索引对 RowTile srcRowTile(1, validCol); RowTile tmpTile(1, validCol); TASSIGN(dstRowTile, (uint64_t)dst.data() i * DstTileData::Cols * sizeof(T)); TASSIGN(srcRowTile, (uint64_t)src.data() i * SrcTileData::Cols * sizeof(T)); TASSIGN(tmpTile, (uint64_t)inIdx.data() kTCols_ * sizeof(indexT)); TSORT32(dstRowTile, srcRowTile, inIdx, tmpTile); pipe_barrier(PIPE_V);根据 TSORT32_zh.mdTSORT32 的底层 SFU 指令为VBS32vbitsort对 src 的每个 32 元素块连同 idx 中对应的索引一起排序并将排序后的值-索引对写入 dst。本示例调用的是 4 参数形式带tmp它支持非 32 对齐的尾部validCol % 32 ! 0通过 tmp 填充补齐。排序后的输出 dst 元素数为输入的2 * TYPE_COEF倍正是值-索引对展开的结果。第二级TMRGSORT 归并排序TSORT32 只保证每 32 个元素块内有序1024 列需要进一步归并。MrgsortSingleRow实现了分层的 4 路归并topk_kernel.cppuint32_t blockLen 64 * TYPE_COEF; // Merge sort data for every 4 blockLen lengths. for (; blockLen * 4 valid_col; blockLen * 4) { uint16_t cols valid_col / (blockLen * 4) * (blockLen * 4); SrcTileData srcSortedTile(1, cols); SrcTileData tmpSortedTile(1, cols); TASSIGN(srcSortedTile, (uint64_t)srcTile.data()); TASSIGN(tmpSortedTile, tmpAddr); TMRGSORTSrcTileData, SrcTileData(tmpSortedTile, srcSortedTile, blockLen); pipe_barrier(PIPE_V); TMOV(srcSortedTile, tmpSortedTile); pipe_barrier(PIPE_V); }根据 TMRGSORT_zh.md这对应变体 A单列表排序把 src Tile 视为 4 个连续等长的已排序块在单个 Tile 内完成 4 路归并底层为硬件指令vmrgsort4。该变体的关键约束是blockLen必须是64的倍数src.GetValidCol()必须是blockLen * 4的整数倍repeatTimes src.GetValidCol() / (blockLen * 4)必须在[1, 255]范围内。因此循环以 64 为起点、每次乘 4 扩张块长直到无法再整除为止。TMRGSORT输出写入 tmptmpAddr随后用TMOV搬回原位置形成归并-搬回交替。当剩余块无法继续 4 路归并时SortTailBlock处理尾部topk_kernel.cpp它调用的是变体 B多列表归并TMRGSORTDstTileData, TmpTileData, SrcTileData, SrcTileData, 0( curDstTile, executedNumList, tmp1Tile, src0Tile, src1Tile);变体 B 可将 2~4 个独立已预排序列表归并为一个降序输出并借助MrgSortExecutedNumList对应硬件寄存器VMS4_SR记录每个列表实际处理的元素数。FillMrgArray预先按块长等比递减计算出每次归并的块长序列topk_kernel.cpptemplate int Cols PTO_INTERNAL int32_t FillMrgArray(int32_t* mrgArray, int blockLen) { int32_t arrayCount 0; int32_t tmpInner Cols; for (int32_t i blockLen; i 64; i / 4) { int32_t count; for (count 0; count tmpInner / i; count) { mrgArray[arrayCount] i; } tmpInner - count * i; } return arrayCount; }注意SortTailBlock中每次归并都会把已排序长度和下一块长度截断到topktmpMrgSortedLen topk时取topk。这是 Topk 的关键优化只关心前 K 个最大值超过 topk 的部分无需继续参与归并从而大幅减少无效计算。MrgsortSingleTile再对每行循环调用MrgsortSingleRow把整个 TilevalidRow行逐行完成归并topk_kernel.cpp。第三级TGATHER 提取 data 与 index排序完成后值-索引对以交错布局存于 UB需要分别抽出前 topk 个纯值和纯索引。ExtractDataOrIndex用TGATHER完成抽取topk_kernel.cppif constexpr (std::is_same_vT, half) { TGATHERRowTile, RowTile, MaskPattern::P0001(rowDTile, rowTile); // half从 4 槽位 pair 中抽 value } else { TGATHERRowTile, RowTile, MaskPattern::P0101(rowDTile, rowTile); // float从 2 槽位 pair 中抽 value } // index 抽取使用 P1010 TGATHERIndexRowTileData, CopySrcTileData, MaskPattern::P1010(rowITile, copyTile);根据 TGATHER_zh.md 的 mask 语义P0101每 2 个元素取第 1 个索引 0、2、4…P1010取第 2 个索引 1、3、5…P0001则是 4 元素步长取第 4 个。因此float 输入中 value 占 pair 的第 0 槽位 →P0101抽取 valueP1010抽取 indexhalf 输入中 value 占 pair 的第 0 槽位、其后有 2 字节 padding再跟 index →P00014 槽位步长抽取 value。这是 PTO 中通过 mask 模式在交错数据流中挑元素的典型用法。抽出的dTiletopk 个值与iTiletopk 个索引随后由TSTORE写回 GM 的两个独立输出张量。UB 双缓冲流水线调度原文指出本示例通过 UB 上的双缓冲来重叠数据搬运与计算。每次循环执行两组操作TLOAD - TSORT32 - TMRGSORT(含 MRGSORT 和 MOV) - TSTORE单组操作的依赖顺序是MTE2 - V - MTE1 - V - MTE3即MTE2 搬运TLOAD→ V 向量计算TSORT32→ MTE1TMRGSORT 的 tmp 归并→ V 计算TMOV/TGATHER→ MTE3 搬出TSTORE。第二组的 TLOAD 不必等第一组全部执行完再开始从而提高了流水并行度。BUFFER_NUM 2定义于文件开头topk_kernel.cppInitBuffers在 UB 基址上为两组 buffer 依次划分了sort32DstTile[2]、mrgDstTile[2]、indexTile、dTile[2]、iTile[2]、srcTile[2]及两个临时区tmpAddr/nextTmpAddrtopk_kernel.cpp。UB 容量约束由编译期断言把关topk_kernel.cppstatic_assert(totalRow % blockDim 0, expect totalRow % blockDim 0); static_assert(sort32DstSize * 3 validCol * sizeof(uint32_t) * 5 srcSize 192 * 1024, memory is exhausted.); static_assert(validRow % (SINGLE_LOOP_ROW * 2) 0, expect validRow % (SINGLE_LOOP_ROW * 2) 0.);其含义是所有 UB buffer 总和必须小于 192KB 的 UB 空间且每核行数须能被SINGLE_LOOP_ROW * 2 4整除以支持两级双缓冲循环。在ProcessSingleRow中可以看到事件同步的完整细节topk_kernel.cppwait_flag(PIPE_V, PIPE_MTE2, loadEvent); // 等上一个循环的 V 完成反向依赖 TLOAD(srcTile[cur], srcGlobal); set_flag(PIPE_MTE2, PIPE_V, (event_t)cur); // MTE2 - V wait_flag(PIPE_MTE2, PIPE_V, (event_t)cur); SortEachGroup...(sort32DstTile[cur], srcTile[cur], indexTile); set_flag(PIPE_V, PIPE_MTE2, loadEvent); // V - MTE2反向依赖串起下一循环 ... set_flag(PIPE_V, PIPE_MTE3, storeEvent); // V - MTE3 ... TSTORE(dstDataGlobal, dTile[cur]); TSTORE(dstIdxGlobal, iTile[cur]);其中增加了循环之间的从 V-MTE2 的反向依赖以保证下一个循环的 TLOAD 是在对应的 VEC 操作执行完后再开始的即第 331-332 行在启动循环前预先set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0/EVENT_ID1)循环末尾wait_flag闭合构成跨迭代的依赖链。主循环每次步进 2for (i 0; i loopNum; i 2)ProcessIteration内依次调用ProcessSingleRow处理第 0/1 两组 buffer实现一组在算、一组在搬的双缓冲流水。数据生成与 golden 校验gen_data.py输入与参考输出scripts/gen_data.py 使用固定随机种子np.random.seed(19)生成可复现的数据构造rows × cols的输入矩阵每行数据取np.random.uniform(i, i valid_col)保证行间值域分离、便于核对生成x1_gm.bin将[value, index]交错写入cols * 2列模拟 kernel 中值-索引对的输入布局gen_data.py生成x1_idx.bin即每列的原始下标idx arange(valid_col)用 NumPy 计算 goldennp.lexsort((idx, -row))实现值降序、同值索引升序的稳定排序取前 topk 写入golden_i.bin索引与golden_d.bin值gen_data.py。默认用例参数与 kernel 侧严格对应gen_data.pyTopkParams(np.float32, np.int32, 1, 1, 1, 4800, 1024, 1, 1, 1, 4800, 1280, 1000)main.cppHost 侧执行与校验main.cpp 是典型的 ACL Host 流程aclInit→aclrtSetDevice(0)→aclrtCreateStream初始化运行环境aclrtMallocHost/aclrtMalloc分配 Host/Device 内存输入、输出、索引各 4 块读入../input/x1_gm.bin与../input/x1_idx.binaclrtMemcpy拷贝到 Device调用launchTopkT启动 kernelaclrtSynchronizeStream等待完成结果拷回 Host写入../output/output_z.bin与../output/index_z.binValidateDataResults/ValidateIndexResults用 tests/common/test_common.h 提供的ResultCmp与 golden 比对容差0.001f分别打印test data success/test index success全部通过则输出All tests passed!main.cpp。构建与运行环境准备需要已安装 Ascend CANN 工具链并配置环境变量示例路径source ${ASCEND_INSTALL_PATH}/bin/setenv.bashCMakeLists.txt 会校验ASCEND_HOME_PATH环境变量缺失时直接报错Cannot find ASCEND_HOME_PATH, please run set_env.sh.。构建依赖笔式编译器set(CMAKE_COMPILER bisheng)kernel 部分以--cce-pto-enable开启 PTO 编译支持并以--cce-aicore-archdav-c220-vec指定 A2/A3 的 VEC 核架构CMakeLists.txt。运行示例在克隆仓库后进入示例目录执行cd ${git_clone_path}/kernels/manual/a2a3/topk bash run.sh -r npu -v Ascend910B1成功时输出test successrun.sh 支持两个参数参数说明-r/--run-mode运行模式npu真实设备或simCPU 仿真需对应仿真库-v/--soc-versionSoC 版本必须以Ascend开头如Ascend910B1脚本行为先用python ./scripts/gen_data.py生成输入与 golden再清理并重建build/目录按RUN_MODE/SOC_VERSION配置 CMake 后make -j16最后直接执行生成的./topk可执行程序。仿真模式下会额外设置LD_LIBRARY_PATH指向tools/simulator/${SOC_VERSION}/libAscend910B4-1版本不支持 sim 模式脚本会提前报错run.sh。从 CMakeLists.txt 可以看到两种模式的区别在于链接库sim链接runtime_camodelnpu链接runtime同时统一链接ascendcl、m、tiling_api、platform、nnopbase等运行时库。实测性能参考以下数据在 Ascend A348 个 VEC 核上测得覆盖多个尺寸以及不同类型可作为调优基线数据来自原文档为参考值而非承诺指标参数aiv_vec_ratioaiv_scalar_ratioaiv_mte2_ratioaiv_mte3_ratiotask_duration(us)typefloatvalidRowrows4800validCol1024cols1280topk100094%3.2%11.7%10.4%324.106typefloatvalidRowrows3456validCol1024cols1280topk100091.5%4.6%12.3%10.5%238.819typefloatvalidRowrows2304validCol1024cols1280topk100088.7%6%12.4%10.1%161.375typehalfvalidRowrows4800validCol1024cols1280topk100893.7%2.4%11.5%9.6%326.886可以观察到两个有价值的规律VEC 占用率极高88%排序类算子计算密度大双缓冲流水让向量流水基本饱和而 MTE2/MTE3 占比10% 上下说明搬运与计算重叠良好half 场景比值与 float 接近因为值-索引对统一为 8 字节结构half 的 TYPE_COEF2 使排序宽度翻倍但整体瓶颈仍集中在 VEC 排序本身。小结与延伸阅读本文从工程结构、Tiling 设计、三级排序算法TSORT32 → TMRGSORT → TGATHER、UB 双缓冲流水到数据生成与构建运行完整还原了 PTO 基础 Topk 算子的实现。该示例是理解 PTO 手写算子的优秀范本它同时覆盖了 GlobalTensor/Tile 类型系统、值-索引对布局、mask 模式数据抽取以及跨流水线事件同步四大 PTO 核心主题。若想进一步深入建议按以下路径阅读仓库资料指令语义TSORT32_zh.mdVBS32 底层实现与 32 元素块排序规则、TMRGSORT_zh.mdvmrgsort4 两种变体与值-索引对格式、TGATHER_zh.mdMaskPattern 抽取语义工程入口topk_kernel.cpp完整 kernel 实现、main.cppHost 侧流程、run.sh构建运行脚本编程模型ProgrammingModel_zh.md 与 Tile_zh.md 可帮助你理解 Shape/Stride 与 Tile 的底层抽象手工算子总览kernels/manual/README_zh.md 展示了更多 A2/A3 手写算子示例如 allgather_gemm 等可作为后续实战参考。【免费下载链接】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），仅供参考