Pull Request已成功合入, 合并人@CANN-robot
(感谢 黄晓彬 的贡献)变更摘要
本 PR 将累积直方图(cumulative histogram)等 kth-value / median 性能优化从开发分支移植到 master,核心思路是:radix select 改用累积字节直方图 + 二分定位桶、跳过相等的首字节区间、以向量化压缩替代串行 SIMT 扫描;排序阶段引入 UB 内的 ping-pong 归并(SortToProposal/MergeStage)取代原有的 Concat+Sort 流程;tiling 侧新增驻留(resident)行选择、int16 merge 等经验路由策略及 merge-more-core 多轮/sync-merge 调度;同时为单核路径拆分出紧凑 tiling 数据结构(KthValueMergeOneCoreTilingData、KthValueRadixOneCoreTilingData),并在 kth_value、median、nan_median 内核中按 schId 注册对应结构。
主要改动
- radix select 累积直方图与快速选择:
kth_value_radix_select.h将桶直方图从FREQUENCY改为ACCUMULATE(含 SIMT 路径的累积化),新增SelectCumulativeBucket二分查找首超桶、FindSelectedKeyRangeVec/FindSelectedKeyRangeB64基于最小/最大键跳过相等的低字节轮次,以及CompactMatchingPositions向量化候选压缩;整数键改为原地可逆变换(复用输入 UB,不再额外申请keyBuf_),并通过RestoreInputValue还原符号位。 - 新增 UB ping-pong 归并排序:
ping_pong_merge_sort.h提供SortToProposal(Sort32 产块 + 4 路MergeStage交替 ping/pong 归并)与MergeStage,merge_intra_core_base.h、merge_more_core_base.h、kth_value_merge_sort_one_core.h、non_last_small_axis_base.h均改用它并移除concatTmpBuf_/Concat 相关成员与DEALING_CONCAT_NUM_ONCE依赖。 - tiling 路由策略与 merge-more-core 调度扩展:
kth_value_tiling_arch35.cpp新增KthRoutePolicy(驻留轴范围、int16 merge 轴区间等)、IsResidentIntegerSelect/IsResidentFloatSelect并在SelectKthValueRoute中赋予驻留 radix select 优先级;新增useInt16Merge/FitsInt16MergeUb、PreferKthMergeMoreCore;kth_value_tiling_common.cpp将 merge-more-core 重构为SelectMergeMoreCoreDataSize/SelectMergeSyncMergeBlockSize/IsMergeMoreCoreProfitable,支持 2048/4096 分块、sync-merge 逻辑块与多轮(新增常量MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND)调度。 - 按 schId 拆分的紧凑 tiling 数据与内核注册:
kth_value_tiling_data.h新增KthValueMergeOneCoreTilingData、KthValueRadixOneCoreTilingData,host 侧StoreKthValueTiling按schId序列化紧凑布局;kth_value.cpp、median.cpp、nan_median.cpp通过REGISTER_TILING_FOR_TILINGKEY分别注册对应结构,Dispatch改为按 tiling 数据类型模板化分发。 - merge-more-core 基类多轮/同步归并流程:
merge_more_core_base.h新增Process()(含行活跃度判断与SyncAll防死锁)与单轮快速路径ProcessDirect(),支持SortLogicalBlock/MergeLogicalBlock同步归并模式及MergeAcrossCores多轮归并;kth_value_merge_sort_more_core.h按轮次通过PrepareRowWorkspace重绑行工作区并复用物理行组的 median 计数组,ResolveRowK改用staticKthIndex_。


🤖 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 |
|---|---|---|
| index | ✅ 王瑞, 唐玮玮 (2/2) | ✅ 王瑞 (1/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. 👍


/lgtm
/approve


The following users do not have permission to comment /lgtm or /approve on any module in this PR:
/approve: 唐玮玮


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


描述
优化 KthValue、Median、NanMedian 的选择与共享归并路径,减少完整排序、重复扫描和逐行控制开销,保持各算子的目标秩、NaN 和索引语义。
关联的Issue
https://gitcode.com/cann/ops-nn/issues/5977
测试
文档更新
不涉及
类型标签
AI/Agent生成声明