Pull Request已成功合入, 合并人@CANN-robot
(感谢 黄晓彬 的贡献)变更摘要
本 PR 将累积的 ping-pong 归并优化适配到 master 分支,核心是新增统一的 UB 内 ping-pong 归并工具 PingPongMergeSortCommon::SortToProposal,用它替换 sort 算子各路径原先的 Concat + Sort 流程,并对 merge more-core 调度做重构:支持 direct(单轮直接并行)、sync-merge(keyParams1 逻辑块)与 multi-round(多轮)三种模式,同时新增 SORT_SCHID_12 直接调度内核实例。改动覆盖 tiling 侧(sort_tiling_common.cpp/.h、sort_tiling_arch35.cpp)与内核侧(merge_more_core_base.h、merge_intra_core_base.h、sort_merge_sort.h、non_last_small_axis_base.h、merge_sort_big_size.h),并配套更新了单测 test_sort_tiling.cpp。
主要改动
- 新增 ping-pong 归并工具
ping_pong_merge_sort.h:新增SortToProposal与MergeStage,先用Sort32生成DEALING_SORT_NUM_ONCE长度的有序 run,再在 ping/pong 两个 proposal buffer 间迭代做 4 路归并,返回最终结果所在 buffer,替代原先依赖Concat临时缓冲的排序流程。 - 重构 merge more-core 调度:tiling 侧以
IsMergeMoreCoreProfitable、SelectMergeMoreCoreDataSize、SelectMergeSyncMergeBlockSize取代IsMergeMoreCoreSupported,ComputeMergeMoreCorePlan按 direct → sync-block → multi-round 顺序选择,新增常量MERGE_MORE_CORE_DATA_SIZE_BASE/LARGE与MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND;内核侧merge_more_core_base.h重写Process()支持多轮循环,新增ProcessDirect()、SortLogicalBlock/MergeLogicalBlock、MergeAcrossCores,keyParams1承载 sync-merge 逻辑块大小(0 表示普通 more-core),并新增SORT_SCHID_12直接调度模板参数。 - 重写
sort_merge_sort.h主归并内核:移除Concat相关缓冲,改用SortToProposal生成 proposal 后Extract输出 value/index;新增ProcessDoubleBuffered双缓冲流水(MTE2/MTE3 与 Vector 重叠),并为 bf16 增加SortRowsBf16在 cast buffer 内就地输出结果。 - 增加单块快速路径与紧凑输出:
merge_intra_core_base.h在blocksPerRow_ == 1时跳过 merge 阶段(省去concatTmpBuf_及 Phase2 UB);tiling 侧新增ComputeMergeIntraCoreSingleBlockSortSize、ComputeMergeSortBatchCapacity,并为FillMergeSortInfo/ComputeMergeSortTiling增加compactOutput选项;non_last_small_axis_base.h的 merge 路径改用SortToProposal+Extract,并将sortInput_复用为 value 输出缓冲。 - 调整调度选择顺序与数据类型的阈值:
sort_tiling_arch35.cpp中TryMerge引入 fp16/bf16 的 merge axis 阈值(2048/1536),新增PreferSortMergeMoreCore决策函数,整体选择顺序改为TryMerge→TryMergeIntraCore→TryRadixOneCore,并更新了 tiling 单测用例与期望 tiling key。


🤖 AI 代码检视正在进行中,请稍候…


Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| math/sort | ✅ yue-ma, 王瑞 (2/2) | ✅ 王瑞, yue-ma (2/1) |
| repo-cann/ops-math | ✅ 王瑞, yue-ma (2/2) | ✅ 王瑞, yue-ma (2/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
ConanHuang, thanks for your pull request. All authors of the commits have signed the CLA. 👍


描述
优化 Sort 在多行归并、非末轴小轴及 UB 可容纳整行的大轴场景中的执行开销,保持稳定排序、升降序及 values/indices 语义。
关联的Issue
https://gitcode.com/cann/ops-math/issues/3230
测试
文档更新
不涉及
类型标签