PTO 自动模式下的 TopK 算子实现：基于 TSort32 与 TMrgsort 的多核排序流水实战（Ascend A2/A3）

PTO 自动模式下的 TopK 算子实现：基于 TSort32 与 TMrgsort 的多核排序流水实战（Ascend 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/automode/a2a3/topk 示例为线索，完整讲解如何用 PTO（Parallel Tile Operation）虚拟指令集在 Ascend A2/A3 上实现一个高性能 TopK 算子。你将掌握从算子规格与 Tiling 划分、TSort32块排序与TMrgsort归并排序的算法组装、双缓冲流水调度（MTE2→V→MTE1→V→MTE3），到数据生成、构建运行与 golden 校验的完整闭环，并理解 PTO auto 模式（--cce-pto-auto-enable）下编译器自动管理 Tile 内存与同步所带来的编程简化。

一、示例定位：PTO auto 模式的 TopK 最小可运行实现

本示例展示如何使用 PTO 在A2/A3 芯片上实现一个固定维度的 TopK 算子，覆盖工程搭建、kernel 编写、构建与执行的完整流程。它属于 kernels/automode 目录，即采用 PTO 的auto 编译模式：与 manual 模式相比，开发者不需要手动调用TASSIGN为 Tile 分配硬件 buffer 内存，也不需要手动插入 pipe 间的同步事件，编译器会自动完成 Tile 的 liveness 分析、内存分配与同步指令插入（详见 docs/auto_mode/Auto_Mode_Overview_zh.md）。这大幅降低了开发门槛，同时通过合理的 kernel 结构设计（见第四节流水调度）仍能获得接近手写优化的性能。

示例规格固定为[rows, cols] = [4800, 1024]，从每行 1024 个元素中选出最大的topk = 1000个元素，输出排序后的data（值）与index（原始列下标）两个结果张量。

1.1 算子规格

项目	值
OpType	topk
输入	[rows, cols] = [4800, 1024]，以及每列对应的初始索引
输出	data（topk 个值）、index（topk 个原始下标）
Kernel 名称	topk_kernel（源文件为 topk_kernel.cpp）
支持处理器	A2/A3

1.2 目录结构

kernels/automode/a2a3/topk/ ├── scripts/ │ └── gen_data.py # 生成输入与 golden 输出 ├── CMakeLists.txt # 构建配置（含 auto 模式编译选项） ├── topk_kernel.cpp # Kernel 实现 ├── main.cpp # Host 侧入口（ACL 数据搬运与结果校验） ├── multiBuffer.hpp # 双缓冲流水辅助模板 └── run.sh # 一键构建运行脚本

注意：原 README 中构建命令书写的路径为kernels/manual/a2a3/topk，而本示例实际位于 auto 模式目录kernels/automode/a2a3/topk，下文均以仓库实际路径为准。

二、Tiling 策略：48 核上的数据切分

示例面向的验证平台拥有48 个 VEC 核，工作负载按行切分到各个核上。从 topk_kernel.cpp 的启动参数可以看到：

constexpr int blockDim = 48; constexpr int gShape3 = 4800; // 逻辑有效行数 validRow constexpr int gShape4 = 1024; // 逻辑有效列数 validCol constexpr int gWholeShape3 = 4800; // 物理行数（含 padding） constexpr int gWholeShape4 = 1280; // 物理列数（含 padding） constexpr int topk = 1000; Topk<T, 1, 1, 1, gShape3, gShape4, 1, 1, 1, gWholeShape3, gWholeShape4, topk, blockDim> <<<blockDim, nullptr, stream>>>(out, index, src, inIdx);

每个核通过get_block_idx()定位自己的数据分片：

constexpr int validRow = gShape0 * gShape1 * gShape2 * gShape3 / blockDim; // 4800 / 48 = 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;

于是每核处理的分片形状为：

• rows = 100，cols = 1024（有效列）；
• 物理列宽cols = 1280（gWholeShape4），即全局内存按 1280 列 stride 存放，有效数据只占前 1024 列，多出的 256 列用于满足内存对齐与块边界要求；
• 每核输出100 × topk个值以及100 × topk个索引。

在Check<>()编译期检查中（topk_kernel.cpp），还通过static_assert保证：

• totalRow % blockDim == 0：总行数可被核数整除；
• validRow % (SINGLE_LOOP_ROW * 2) == 0：保证双缓冲流水迭代数完整；
• 所有 Tile 的 UB 内存占用总和不超过 192 KiB（192 * 1024字节），避免 kernel 直接内存耗尽。

三、算法原理：TSort32 块排序 + TMrgsort 归并排序

3.1 整体数据流

TopK 的实现思路非常典型：先在 UB 上对每个 32 元素块做基础排序，再通过多轮归并把有序块合并为整行有序，最后截取前 topk 个值/索引写回 GM。README 中的描述可以映射到 topk_kernel.cpp 的实际调用链：

TLOAD(srcTile, srcGlobal) // GM -> UB 加载数据行 → TSORT32(sort32DstTile, srcTile, inIdx) // 每 32 个元素 + 索引排序（SFU VBS32） → TMRGSORT(mrgDstTile, sort32DstTile) // 多轮 4 路归并（VEC vmrgsort4） → TGATHER 提取 value 与 index // 从 value-index 交错对中分别抽取 → TSTORE 写回 GM（data、index 分别存储）

关键点：PTO 的排序指令处理的是“值-索引对”。TSort32和TMrgsort底层硬件指令（vbitsort/vmrgsort4）都按 8 字节的[value, index]元组工作，因此示例先把原始数据与初始索引交错摆放（见 gen_data.py），排序完成后再用TGATHER把两路结果拆开。

3.2 类型定义：5 维 Shape/Stride 描述 GM 布局

README.md 中给出的类型定义，是理解数据布局的关键。GM 侧的张量用 5 维Shape+Stride描述，前 4 维用于承载多核/多行逻辑，最后一维是列：

// data（源数据）：单次处理 1 行 × validCol 列，列 stride 为物理 Cols using DynShapeDim5 = Shape<1, 1, 1, singleLoopRow, validCol>; using DynStridDim5 = Stride<singleLoopRow * Cols, singleLoopRow * Cols, singleLoopRow * Cols, Cols, 1>; using GlobalData = GlobalTensor<T, DynShapeDim5, DynStridDim5>; // index（初始列索引）：1 × validCol，每行内 stride 为 1 using IndexShapeDim5 = Shape<1, 1, 1, 1, validCol>; using IndexStridDim5 = Stride<validCol, validCol, validCol, validCol, 1>; using IndexGlobalData = GlobalTensor<indexT, IndexShapeDim5, IndexStridDim5>; // 排序后的 data / index 输出：每行只保留 topk 个 using DstShapeDim5 = Shape<1, 1, 1, singleLoopRow, topk>; using DstStridDim5 = Stride<singleLoopRow * topk, singleLoopRow * topk, singleLoopRow * topk, topk, 1>; using DstDataGlobalData = GlobalTensor<T, DstShapeDim5, DstStridDim5>; using DstIdxGlobalData = GlobalTensor<indexT, DstShapeDim5, DstStridDim5>;

其中T为float或half，indexT为uint32_t。注意源数据的列 stride 是物理Cols = gWholeShape4 = 1280，与有效列validCol = 1024不同，这正是 GM 上“有效数据 + padding 列”布局的体现。

UB 侧则用Tile<TileType::Vec, T, Rows, Cols, BLayout::RowMajor, ...>描述，例如：

using SrcTileData = Tile<TileType::Vec, T, SINGLE_LOOP_ROW, validCol, BLayout::RowMajor, SINGLE_LOOP_ROW, validCol>; using DstTileData = Tile<TileType::Vec, T, SINGLE_LOOP_ROW, dstCols, BLayout::RowMajor, SINGLE_LOOP_ROW, dstCols>; using IndexTileData = Tile<TileType::Vec, indexT, 1, validCol, BLayout::RowMajor, 1, validCol>;

其中dstCols = validCol * 2 * TYPE_COEF，TYPE_COEF = sizeof(float)/sizeof(T)：对 float 是 ×2（每个元素变 2 个 float 槽位），对 half 是 ×4（每个元素变 4 个 half 槽位），对应排序指令输出的 value-index 对展开宽度。

3.3 阶段一：TSort32 基础排序（每 32 元素一块）

TSort32的硬件底层是 SFU 指令VBS32（vbitsort），一次调用可对若干个相互独立的 32 元素块排序，排序规则为按值降序，值相同按下标升序。输入为值数组 + 索引数组，输出为排好序的 value-index 交错对，输出宽度按类型展开：float 每元素占 2 个 float 槽位，half 每元素占 4 个 half 槽位（详见 docs/isa/TSORT32.md）。

在示例中，每行 1024 个元素恰好是 32 的整数倍，因此可以干净地切分为 32 个 32 元素块。kernel 中通过SortEachGroup调用 4 参形式的TSort32：

IdxRowTile tmpTile(1, validCol); TRESHAPE(dstRowTile, dst); TRESHAPE(srcRowTile, src); detail::PtoSubTileView(tmpTile, inIdx, 0, kTCols_); TSORT32(dstRowTile, srcRowTile, inIdx, tmpTile);

4 参形式（带tmp临时 Tile）用于支持非 32 对齐的尾部块补齐；本示例validCol=1024已对齐，但调用形式与 ISA 文档中的约束一致：src/dst必须同为half或float，idx为uint32_t，全部使用TileType::Vec且行主序。

3.4 阶段二：TMrgsort 多轮归并

TMrgsort的硬件底层是 VEC 指令vmrgsort4，可将最多 4 条已排序链表归并为一条降序输出（详见 docs/isa/TMRGSORT.md）。该指令有两种形态：

• 单列表形态TMRGSORT(dst, src, blockLen)：把src视为 4 个等长且各自有序的块，做一次 4 路归并，blockLen必须是 64 的倍数；
• 多列表形态TMRGSORT<..., exhausted>(dst, executedNumList, tmp, src0, src1, ...)：归并 2~4 条独立的预排序链表，结果先写入tmp再 MOV 到dst，exhausted=true时可通过MrgSortExecutedNumList得知每条链表实际消费的元素数。

示例的归并逻辑分为两个函数：

1）MrgsortSingleRow——逐轮 4 路归并（topk_kernel.cpp）

uint32_t blockLen = 64 * TYPE_COEF; SrcTileData tmpTile(1, valid_col); // 每轮把 4 个 blockLen 长的有序块合并 for (; blockLen * 4 <= valid_col; blockLen *= 4) { uint16_t cols = valid_col / (blockLen * 4) * (blockLen * 4); SrcTileData srcSortedTile(1, cols); SrcTileData tmpSortedTile(1, cols); TRESHAPE(srcSortedTile, srcTile); TRESHAPE(tmpSortedTile, tmpTile); TMRGSORT<SrcTileData, SrcTileData>(tmpSortedTile, srcSortedTile, blockLen); TMOV(srcSortedTile, tmpSortedTile); // 归并结果搬回 src，作为下一轮输入 }

每轮把 4 个长度为blockLen的有序段归并成长度为4*blockLen的有序段，blockLen按 4 倍增长（64 → 256 → 1024 …），直到覆盖整行。注意blockLen同样乘了TYPE_COEF，这是因为归并指令以 value-index 对为单位计数，float 下每对占 2 个元素、half 下占 4 个元素。

2）SortTailBlock——尾部块裁剪归并（topk_kernel.cpp）

当blockLen * 4 > valid_col无法整轮归并时，剩余的有序块使用 2 列表形态归并，并在每一步把参与归并的长度裁剪到不超过topk——因为最终只需要前 k 个最大值，超出部分无需参与后续归并，这是 TopK 相对全排序的关键剪枝：

if (tmpMrgSortedLen > topk) { tmpMrgSortedLen = topk; } if (tmpMrgArray > topk) { tmpMrgArray = topk; } ... TMRGSORT<DstTileData, TmpTileData, SrcTileData, SrcTileData, 0>( curDstTile, executedNumList, tmp1Tile, src0Tile, src1Tile);

FillMrgArray负责把行内剩余元素按 4 的幂次（i从blockLen起每次除以 4，且不小于 64）分解成可归并的段长序列。归并段的选择逻辑与 docs/isa/TMRGSORT.md 中 “Typical Usage: TopK” 一节描述的两阶段流程完全对应：Phase 1 用单列表形态逐轮放大blockLen，Phase 2 用多列表形态合并尾部并只保留 top K。

3）结果提取ExtractDataOrIndex（topk_kernel.cpp L145-L173）

归并完成后，mrgDstTile里是前topk*2*TYPE_COEF个元素长度的 value-index 交错对。通过TGATHER配合不同MaskPattern把值与索引分开：

• 提取值：float 用MaskPattern::P0101、half 用MaskPattern::P0001（跳过索引槽位）；
• 提取索引：把uint32_t索引按 2 元素宽的CopySrcTileData重解释后，用MaskPattern::P1010抽取出偶槽位的索引。

if constexpr (isIndex == false) { ... if constexpr (std::is_same_v<T, half>) { TGATHER<DstRowTile, SrcRowTile, MaskPattern::P0001>(rowDTile, rowTile); } else { TGATHER<DstRowTile, SrcRowTile, MaskPattern::P0101>(rowDTile, rowTile); } } else { ... TGATHER<IndexRowTileData, CopySrcTileData, MaskPattern::P1010>(rowITile, copyTile); }

3.5 单行处理流水ProcessSingleRow

把上述步骤串起来即 ProcessSingleRow：

TLOAD(srcTile, srcGlobal); SortEachGroup<...>(sort32DstTile, srcTile, indexTile); // TSORT32 块排序 MrgsortSingleTile<...>(mrgDstTile, sort32DstTile); // 多轮 TMRGSORT + 尾部裁剪 ExtractDataOrIndex<..., 0>(dTile, mrgDstTile); // 提取 topk 值 ExtractDataOrIndex<..., 1>(iTile, mrgDstTile); // 提取 topk 索引 TSTORE(dstDataGlobal, dTile); TSTORE(dstIdxGlobal, iTile);

单行数据流为TLOAD → TSORT32 → TMRGSORT → TGATHER → TSTORE，即 README 中描述的TLOAD->TSORT32->TMRGSORT(含 MRGSORT 和 MOV 操作)->TSTORE。

四、流水线调度：UB 双缓冲与跨 pipe 依赖

4.1 为什么需要双缓冲

单行处理链路涉及多个硬件 pipe：TLOAD/TSTORE走MTE2/MTE3（数据搬运），TSORT32走SFU，TMRGSORT、TMOV、TGATHER走VEC（V）。如果严格串行执行，数据搬运期间计算单元会空转。示例通过在UB 上使用双缓冲（BUFFER_NUM = 2）把数据搬运与计算重叠起来，提高硬件利用率。

4.2 MultiBuffered 模板与流水阶段

multiBuffer.hpp 提供了pto_auto::MultiBuffered<NumBuffs>模板，把循环体自动拆分为Prologue（序言）/ Main（主体）/ Epilogue（尾声）三个阶段执行，并通过#pragma pto v_loop_barrier标注 V 指令屏障。kernel 中的用法是：

constexpr uint32_t loopNum = validRow / SINGLE_LOOP_ROW; // 100 行 / 1 = 100 次迭代 MultiBuffered<BUFFER_NUM> double_buffer; double_buffer.loop<Range<loopNum>>(& { int iter = context.iter; DstTileData sort32DstTile; SrcTileData srcTile; DstTileData mrgDstTile; DstDataTileData dTile; DstIndexTileData iTile; ProcessIteration<...>(out, src, index, iter, sort32DstTile, srcTile, indexTile, mrgDstTile, dTile, iTile); });

ProcessIteration在一次迭代中同时处理两行（i*SINGLE_LOOP_ROW与i*SINGLE_LOOP_ROW+1两段 GM 地址），两套 Tile 分别占用双缓冲的两个 buffer，从而让“第 2 行的 TLOAD”可以和“第 1 行的 TSTORE”并行。这与 README 中“每次循环执行两组操作”的描述一致：

每次循环执行两组操作，TLOAD->TSORT32->TMRGSORT(含 MRGSORT 和 MOV)->TSTORE。单组操作的依赖顺序是MTE2->V->MTE1->V->MTE3。

注意这里MTE2→V→MTE1→V→MTE3中的MTE1对应归并内部的 UB-to-UB 搬运（TMOV）。

4.3 跨迭代的反向依赖

双缓冲的收益来自“第二组的 TLOAD 不必等待第一组全部完成”，但必须保证下一轮迭代的 TLOAD 不会覆盖上一轮还在被 VEC 指令读取的 buffer。因此示例额外插入了一条从V 到 MTE2 的反向依赖（V->MTE2）：即下一轮 TLOAD 必须等对应 buffer 上的 VEC 操作（TSORT32/MRGSORT）全部结束才开始。README 对此的表述是：

增加了循环之间的从V->MTE2的反向依赖，以保证下一个循环的 TLOAD 是在对应的 VEC 操作执行完后再开始的。

这一依赖由编译器依据 Tile 的 liveness 自动生成（auto 模式特性），开发者只需按双缓冲的结构组织代码，无需手写WaitEvents。

4.4 循环展开与内存上限

SINGLE_LOOP_ROW = 1（源码注释标明 “Other values not currently working for auto mode”，即当前 auto 模式下该参数取其他值不可用），每次迭代处理 1 行、两套操作，loopNum = validRow / 1 = 100次迭代。MultiBuffered会把主体循环按 buffer 数展开，使两条链路的指令交错发射。同时Check<>()中的static_assert(sort32DstSize * 3 + validCol * sizeof(uint32_t) * 5 + srcSize < 192 * 1024)保证了所有 Tile（sort32 输出、索引、源数据等）在 UB 上排布后总量不超过 192 KiB。

五、实测性能参考（A3，48 VEC 核）

README 给出了在 Ascend A3（48 个 VEC 核）上测得的参考数据，覆盖不同validRow与数据类型。表中aiv_vec_ratio（VEC 占用率）、aiv_scalar_ratio（标量占用率）、aiv_mte2_ratio（MTE2 搬运占用率）、aiv_mte3_ratio（MTE3 搬运占用率）均为各 pipe 的占用比例，task_duration(us)为任务耗时：

参数	aiv_vec_ratio	aiv_scalar_ratio	aiv_mte2_ratio	aiv_mte3_ratio	task_duration(us)
type=floatvalidRow=rows=4800validCol=1024cols=1280topk=1000	94%	3.2%	11.7%	10.4%	324.106
type=floatvalidRow=rows=3456validCol=1024cols=1280topk=1000	91.5%	4.6%	12.3%	10.5%	238.819
type=floatvalidRow=rows=2304validCol=1024cols=1280topk=1000	88.7%	6%	12.4%	10.1%	161.375
type=halfvalidRow=rows=4800validCol=1024cols=1280topk=1008	93.7%	2.4%	11.5%	9.6%	326.886

从数据可以看出：

• VEC 占用率高达 88%~94%，说明归并排序（TMrgsort）是计算主体，且双缓冲已使计算与搬运高度重叠；
• MTE2/MTE3 占用率约 10%~12%，远低于 VEC，搬运不是瓶颈；
• 标量占用率随行数减少而上升（3.2%→6%），说明小规模任务下标量开销占比相对增大；
• 任务耗时与validRow近似线性（4800→2304 行，耗时 324→161us），符合逐行处理的计算模型。

需要说明：以上为示例环境下的参考数据，实际性能会随 SOC 型号、CANN 版本、物理列 padding 宽度与 topk 取值变化，不应作为跨平台绝对性能结论。

六、数据生成与结果校验

6.1 gen_data.py：构造输入与 golden

scripts/gen_data.py 使用固定随机种子np.random.seed(19)保证结果可复现，生成流程：

1. 构造有效数据：对每个有效行i，从np.random.uniform(i, i + valid_col, size=valid_col)采样 1024 个 float32 值，使各行数据范围不同、便于肉眼核对；
2. 交错摆放 value-index 对：x1_gm[i, j*2] = value、x1_gm[i, j*2+1] = counter（counter 从 0 递增），形成排序指令需要的[value, index]8 字节元组布局，落盘为input/x1_gm.bin；
3. 生成初始索引：idx = np.arange(valid_col).astype(np.uint32)，落盘为input/x1_idx.bin；
4. 生成 golden 结果：对每行用np.lexsort((idx, -row))按“值降序、下标升序”排序，取前topk个，值写入output/golden_d.bin、索引写入output/golden_i.bin。

TopkParams类的默认参数与 kernel 完全一致：src_type=np.float32、g_shape3=4800、g_shape4=1024、g_whole_shape4=1280、topk=1000。

6.2 main.cpp：Host 侧数据搬运与校验

main.cpp 是标准的 ACL host 流程：

• aclInit → aclrtSetDevice(0) → aclrtCreateStream初始化运行环境；
• 用aclrtMallocHost/aclrtMalloc分配 host 与 device 内存，把input/x1_gm.bin、input/x1_idx.bin拷入 device；
• 调用launchTopk<T>下发 kernel（T=float，与 gen_data 的np.float32匹配）；
• 同步后把output_z.bin（值）、index_z.bin（索引）拷回 host 落盘；
• 分别用ValidateDataResults/ValidateIndexResults与 golden 做ResultCmp(..., 0.001f)容差比较，全部通过时打印：

test data success test index success All tests passed!

（README 中以test success概括成功输出，实际 host 端打印为上述三行。）

七、构建与运行

7.1 环境准备

示例依赖完整的 Ascend CANN 环境，需先配置环境变量（示例路径）：

source ${ASCEND_INSTALL_PATH}/bin/setenv.bash

CMakeLists.txt 会检查ASCEND_HOME_PATH环境变量，未设置时直接报错Cannot find ASCEND_HOME_PATH, please run set_env.sh。同时它把驱动头文件目录固定为/usr/local/Ascend/driver，编译产物默认使用bisheng编译器。

7.2 一键构建运行

cd ${git_clone_path}/kernels/automode/a2a3/topk bash run.sh -r npu -v Ascend910B1

run.sh 的参数解析：

• -r | --run-mode：运行模式，取npu或sim（仿真器）。CMake 中按模式链接不同运行时：sim 链接runtime_camodel，npu 链接runtime；
• -v | --soc-version：SOC 版本，必须以Ascend开头，否则报Unsupported SocVersion；特殊地，Ascend910B4-1不支持 sim 模式，需改用Ascend910B4。

脚本内部流程：

1. python ./scripts/gen_data.py生成输入与 golden；
2. cmake -DRUN_MODE=${RUN_MODE} -DSOC_VERSION=${SOC_VERSION} ..配置工程；
3. make -j16编译 kernel（共享库）与 host 可执行文件；
4. ./topk运行并自动校验，成功时输出test success语义的成功信息。

7.3 auto 模式的关键编译选项

CMakeLists.txt 中定义了一组 CCE 编译选项，其中与 PTO 直接相关的核心开关是：

-xcce --cce-pto-enable # 使能 PTO 编译 --cce-pto-auto-enable # 使能 PTO auto 模式（编译器自动分配 Tile 内存并插入同步） --cce-aicore-arch=dav-c220-vec # A2/A3 的 VEC 核架构

这与 docs/auto_mode/Auto_Mode_Overview_zh.md 中描述的一致：只要加上--cce-pto-enable --cce-pto-auto-enable即可进入 auto 模式，kernel 内无需TASSIGN与手写WaitEvents。其余选项（-mllvm -cce-aicore-stack-size=0x8000、-cce-aicore-record-overflow=true等）用于控制 AICore 栈大小、地址变换与溢出记录等底层行为，DEBUG_MODE开启时还会追加--cce-enable-print支持设备侧打印。

CMake 的pto_example_vec函数把topk_kernel.cpp编译为libtopk_kernel.so（--cce-fatobj-link生成 fatbin 便于 host 加载），main.cpp编译为可执行文件并链接ascendcl、runtime、tiling_api等 Ascend 运行时库。

八、扩展与借鉴

本示例可作为在 PTO auto 模式下实现“基于排序指令的 TopK 类算子”的模板，值得借鉴的设计包括：

1. 排序指令组合：TSort32（SFU 块排序）+TMrgsort（VEC 多路归并）+TGATHER（掩码抽取）的组合，可推广到 median、percentile、argsort 等排序类算子；两份指令的完整语义、约束与汇编形式分别见 docs/isa/TSORT32.md 与 docs/isa/TMRGSORT.md；
2. TopK 剪枝：归并过程中只保留前 topk 个元素（SortTailBlock中的长度裁剪），避免全排序开销；
3. 多核切分：按blockDim整除总行数，每核get_block_idx()定位分片，是昇腾多核编程的标准范式；
4. 双缓冲流水：通过MultiBuffered模板组织循环体，让 MTE2/MTE3 与 VEC/SFU 重叠执行，并借助编译器插入的V->MTE2反向依赖保证正确性；
5. auto 模式开发效率：整个 kernel 没有任何TASSIGN和显式事件同步，内存分配与 pipe 同步全部由编译器完成，代码可读性与可维护性显著优于 manual 模式。

如果你需要把它适配到其他芯片（如 A5）或其他 topk 取值，重点修改launchTopk中的blockDim、gShape*/gWholeShape*、topk常量，并保持static_assert中的 UB 内存约束即可。

【免费下载链接】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