已合并
perf(index): 优化 KthValue、Median 与 NanMedian 选择及归并路径 #10411
黄晓彬创建于 9月14日
perf(index): 优化 KthValue、Median 与 NanMedian 选择及归并路径 #10411
已合并
黄晓彬创建于 9月14日
黄晓彬成员
9月14日

描述

优化 KthValue、Median、NanMedian 的选择与共享归并路径,减少完整排序、重复扫描和逐行控制开销,保持各算子的目标秩、NaN 和索引语义。

  • 改进 radix select:累积字节直方图定位 bucket、UB 驻留 key、相同高位跳过及候选位置压缩,减少无效扫描。
  • 补充窄轴 SIMT select、短秩选择和适用场景的驻留直方图路径,按 dtype、轴长、行数及 UB 容量选择路由。
  • 同步公共 ping-pong 归并和批处理优化,完善 INT16 merge 路径与 host/kernel UB 布局对应关系。
  • 精简相关 tiling 结构与分发,使用编译期 Median 模式分支;补充浮点 key、正负零、尾块及同步约束说明。

关联的Issue

https://gitcode.com/cann/ops-nn/issues/5977

测试

  • Ascend950PR_9579:KthValue 338 条、Median 260 条、NanMedian 265 条全量 ST,加正负零逐位比较 32 条,合计 895 条精度全部通过。
  • 相对前一版已验证实现,性能异常项复测及旧包/当前包同机对照后,未确认稳定超过 5% 且超过 0.5 μs 的设备耗时劣化。
  • 二级冒烟通过。

文档更新

不涉及

类型标签

AI/Agent生成声明

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 黄晓彬 的贡献)
黄黄晓彬成员
9月14日 创建了 pull request,commit f07e1b51
atomgit-bot
atomgit-bot
9月14日 评论:

变更摘要

本 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_。
likedislike
不准确?
atomgit-bot
atomgit-bot
9月14日 评论:

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

likedislike
不准确?
CANN-robotCANN-robot成员
9月14日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
9月14日 评论:

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 /approve or /lgtm
  • Commenting /approve implies 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. 👍

likedislike
此处折叠了203条消息 查看更多
tangweiwei2成员
9月18日 评论:

/lgtm
/approve

likedislike
CANN-robot
CANN-robot成员
9月18日 评论:

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

likedislike
CANN-robotCANN-robot成员
9月18日 添加了label:lgtm
CANN-robotCANN-robot成员
9月18日 合入了pull request
CANN-robot
CANN-robot成员
9月18日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike