| Files | Last commit | Last update |
|---|---|---|
TopkV2算子bitonic算法性能优化 Co-authored-by: caoyan_huawei<caoyan11@huawei.com> Co-authored-by: zhanghedong<zhanghedong@huawei.com> # message auto-generated for no-merge-commit merge: !6320 merge feat/topk-bitonic-perf-optimize into master TopkV2算子bitonic算法性能优化 Created-by: zhanghedong Commit-by: caoyan_huawei;zhanghedong Merged-by: cann-robot Description: 本 PR 围绕 RegBase（arch35）下 TopK / Sort 的小轴与宽整型场景进行选择策略与内核路径优化。在 Tiling 侧新增 SMALL_AXIS_DIRECT_MODE 分支并重写 TryTopKSmallAxisSort 的路由判定，放宽小轴阈值并按数据类型、innerSize、lastAxis、kValue 等条件细分路由；在内核侧新增基于寄存器直写的小轴 SIMT GM 选择内核 top_k_small_k_simt_gm_select.h 与 UB radix 选择内核 top_k_small_k_ub_radix_select.h，共享常量与 Host/Kernel 代价契约统一收敛至 top_k_constant_var_simd.h，并为 64 位整型的 Bitonic 排序接入小 k 阈值选择快路径；同时调整多处核数分配以避免空闲核与重复同步，并更新公版用例与 topk 的 AiCore 选择逻辑。 ## 主要改动 - **新增小轴 SIMT GM 直读直写内核路径**: 新增 SMALL_AXIS_DIRECT_MODE = 11 常量与 top_k_small_k_simt_gm_select.h（含 RunTopKSmallKSimtGmSelect、TopKSimtGmScalarRows、TopKThresholdSimtGm、TopKSimtGmWarpRows 等，并按职责拆分 Load/Collect/Search/Store 辅助函数），top_k_v2_apt.cpp 按 modeType == TOPK_DIRECT_MODE 直接调用该内核。 - **新增 UB radix 选择内核并接入单块模板**: top_k_small_k_ub_radix_select.h（原 input-select）提供两遍直方图阈值选择、候选收集及 B8/B16/B32/B64 的 SIMD/SIMT 变体；radix_sort_top_k_single_block.h 依 keyParams5 标志切换行距与双缓冲布局，通过 RunTopKSmallKUbRadixSelect 接入并支持 Prefetch 预取流水。 - **宽整型阈值选择快路径**: radix_sort_top_k_inter_core_template_optimization.h 在 IS_BITONIC_SORT 且 T 为 64 位整型、topkValueInput_ <= 32 时调用所属模板内的 TopKUbThresholdSelect，后者通过共享缓冲区报告是否走窄区间选择，否则回退原有 radix 选择。 - **重写小轴排序路由判定**: TryTopKSmallAxisSort 将小轴上限从 SMALL_AXIS_THRESHOLD 提升为 8192，新增标量选择直达、三类 input-select 让路、Direct 轴支持与最小行数判定，命中时写入 SMALL_AXIS_DIRECT_MODE 相关 tiling 字段，并对跨步浮点小 k 场景保留回落逻辑。 - **核数分配与并行度调整**: sort_tiling_arch35.cpp 对 DT_FLOAT 且 lastAxis 在 4~8、unsortedDim >= 2048 的场景收紧 maxCoreNum 以保留每核行数；TileModeSmallSize 将 maxCoreNum 限制为 unsortedDimNum；TopKV2Tiling 针对 int8/uint8 尾轴及 int64/uint64 多核优化模式调整 lastDimNeedCore、unsortedDimParallel、sortLoopTimes。 - **常量归并与检视意见修复**: 响应检视意见移除 top_k_small_axis_direct_config.h、top_k_input_select_config.h、top_k_threshold_select_ub.h，常量与 UB 预算函数并入 top_k_constant_var_simd.h 并由 Host 侧共享；公共 tiling 辅助函数下沉 top_k_v2_tiling_base.cpp；VF 循环边界按 SIMD 约束在 VF 外预计算。 - **AiCore 选择逻辑与公版用例更新**: topk.cpp 删除 DT_INT64 基于 size/轴长的 AiCPU 回退判断，RegBase 支持的 dtype 统一走 AiCore；golden.py 降低 _CUDA_SMALL_TOPK_PER_ROW_MAX 至 128，新增 aclnnTopk 映射及 signed_zero_values、_integer_full_range、_integer_extremes 等测试构造逻辑；本仓库 ST 精简为 4 条代表性冒烟用例，全量用例由独立仓库看护。 ## 关联的Issue https://gitcode.com/cann/ops-math/issues/4139 ## 测试 全量ST，二级冒烟测试通过 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [ ] 新特性 - [x] 性能优化 - [ ] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!6320 | 16 days ago |