已关闭
[性能优化] arch35 KthValue / Median / NanMedian 选择与归并优化 #5977
黄晓彬创建于  25 天前关闭于  25 天前
黄晓彬成员
25 天前 创建

背景与目标

KthValue、Median、NanMedian 只需要目标秩对应的元素,通用完整排序和反复读取整行会带来额外开销。本次在 arch35 上按 dtype、轴长、行数、目标秩和 UB 容量细化选择算法与归并路由,减少不必要的排序、数据扫描和标量控制,保留各算子的 NaN、索引及中位数语义。

关键方案

  1. 改进 radix select:从高字节到低字节生成累积直方图,逐轮缩小目标 bucket,避免为选择单个秩执行完整排序。显式说明 ACCUMULATE 的 inclusive cumulative 语义,并通过累计计数定位目标 bucket。
  2. 减少重复扫描:可容纳的行/分片在 UB 中保留 key;根据候选 key 的范围跳过相同高位字节,候选集合缩小时压缩匹配位置,只继续处理有效候选。长行保留分块扫描和按行多核协作,最终按照原位置解决重复值索引并回读原始值。
  3. 补充小轴和多行选择路径:新增窄轴 SIMT select、短秩选择及适用场景的驻留直方图路由;保留完整 warp collective 参与,独立封装 NaN、正负零和符号位 key 转换。按 UB 预算与收益条件选择 INT16 cast+merge 等路径,不对所有 dtype 统一扩大阈值。
  4. 同步公共归并优化:将 Math 的 ping-pong 归并、批处理、索引与缓冲布局优化移植到 NN 中的重复实现,使 KthValue、Median、NanMedian 共用的归并路径保持一致。对不同算子模式保留独立语义和路由条件。
  5. 收敛 tiling 与分发成本:为相关路径使用精简 tiling 结构,保持 schId、序列化字段和 kernel 注册一致;Median 静态模式用编译期分支消除不必要的运行时判断。补充 UB 布局、尾块、直方图和同步约束注释,保持跨流水线依赖。

验证与验收口径

  • 本机 Ascend950PR_9579、CANN 9.3.0、TTK v3.0.0:KthValue 338 条、Median 260 条、NanMedian 265 条全量 ST,以及正负零逐位比较 32 条,合计 895 条精度全部通过。
  • 覆盖新增选择/归并路径及边界、NaN、重复值、正负零和不同 dtype;本轮通过不代表穷尽所有输入组合。
  • 对比前一版已验证实现,对超过 5% 且超过 0.5 μs 的设备耗时异常项复测,并对剩余异常做旧包/当前包同机对照,未确认稳定回退。该验收不等同于所有 case 均达到 910B 的 2.45×。
  • 完整提交检查通过;最后一轮常量命名整改的 33 个 kernel 指令段及注册元数据保持一致。

实现 PR:https://gitcode.com/cann/ops-nn/pull/10411
关联的 Math 优化 PR:https://gitcode.com/cann/ops-math/pull/5672

likedislike
黄黄晓彬成员
25 天前 将 ConanHuang 设为负责人
CANN-robotCANN-robot成员
25 天前 关闭了 issue
CANN-robotCANN-robot成员
24 天前 添加了label:resolved