Pull Request已成功合入, 合并人@ascend-robot
(感谢 luqichao 的贡献)变更摘要
本 PR 主要优化 inductor 中默认 pointwise/elementwise 动态分组的策略与运行时行为。在 torch_npu/_inductor/codegen/split_tiling.py 中新增默认动态轴分组特征的构建逻辑(_default_dynamic_axis_feature),为分组补充 bucket_factor,并通过 _add_wide_group_midpoints 为跨度较大的分组边界增加中间档位,避免 representative autotune 覆盖过大的动态范围;同时重构 split 轴处理(_downgrade_split_axes、_downgrade_suffix_static_split_axes),减少动态轴及静态后缀轴不必要的 split。在 torch_npu/_inductor/runtime/symbolic_grouping.py 中为 GroupFeatureSpec 增加 bucket_factor 字段并在 GroupedKernelMeta 元数据中序列化,分组 ID 计算(make_group_id)与 representative 构建(build_group_representatives)均按 bucket_factor 归一化取值。在 torch_npu/_inductor/runtime/triton_heuristics.py 中,NPUSymbolicGroupedAutotuner 的运行时分组同样应用 bucket_factor,并为 pointwise/elementwise_numel 特征在运行时规则中增加 representative_numel 与 max_grid,使 resolve_grouped_runtime_block 在较小 shape 下可用更多 program 并保持 BLOCK 对齐,提升不同 shape 下的执行稳定性。
主要改动
- 默认动态分组特征构建:在
SplitTiling._build_group_features中接入新增的_default_dynamic_axis_feature,为默认 pointwise/elementwise 分组生成带bucket_factor(前缀静态轴乘积)的GroupFeatureSpec,并新增_add_wide_group_midpoints在相邻边界比值过大时插入中间档位,收敛 representative 覆盖的动态范围。 - split 轴处理优化:将
_downgrade_secondary_runtime_split_axes泛化为_downgrade_split_axes,并新增_downgrade_suffix_static_split_axes,当动态轴及之前的 split 总量不小于num_vector_core时降级后缀静态 split 轴,减少不必要的 split;_build_group_features相应新增dynamic_split_axes/static_split_axes参数传递。 bucket_factor元数据扩展:在symbolic_grouping.py中为GroupFeatureSpec增加bucket_factor字段(默认 1),随GroupedKernelMeta序列化/反序列化,并新增_require_positive_int校验;make_group_id与build_group_representatives均按value // bucket_factor归一化后参与 bucketize。- 运行时分组与 block 配置适配:
NPUSymbolicGroupedAutotuner的 group_id 计算读取bucket_factor并校验其合法性;对pointwise/elementwise_numel主特征,build_grouped_launch_policy生成的运行时规则新增representative_numel与max_grid(基于npu_num_vector_core与prior_programs),resolve_grouped_runtime_block在axis_numel < representative_numel时改用max_grid推导 program 数,保证 BLOCK 仍对齐block_sub。


/lgtm
/approve




进入 grouped autotune 的条件保持不变。默认分组特征从所有轴的 outer product
调整为:
A = 动态拆分轴 × 该轴之前的静态拆分轴乘积
B = 排除 split 轴和动态轴后的 tiling 轴乘积
运行时直接根据 A 选择分组,分组边界根据数据规模和核间并行度计算:
lower = max(
4 KiB / dtype_bytes / B,
next_power_of_2(2 * vector_core),
)
upper = max(
前置静态拆分轴乘积 * 128,
next_power_of_2(8 * vector_core),
ceil(4096 * vector_core / B),
)
各阈值的作用如下:
4 KiB / dtype_bytes / B:保证A × B × dtype_bytes至少覆盖约4 KiB 数据量,用于区分小数据场景。
2 * vector_core:保证 lower 至少能够为每个 vector core 提供两个A 维度上的处理单元,并使用 2 的幂对齐边界。
前置静态拆分轴乘积 * 128:保证大组边界对应的动态轴取值至少为 128,避免前置静态轴较大时 open group 的 representative 过小。
8 * vector_core:保证大组在 A 维度上具有足够的核间并行规模。4096 * vector_core / B:保证A × B达到平均每个 vector core约 4096 个元素的工作量。
upper 取这些条件的最大值,使大组同时满足动态轴规模、核间并行度和单核工作量
要求。
当相邻边界跨度超过 8 倍时增加中间档位,避免同一个 representative 覆盖过大
的动态范围。bucket 上界无法被前置静态轴乘积整除时,动态轴 representative
向上取整。
broadcast、transpose 等已有专用分组策略保持不变。
对于只有一个动态 split 轴的 pointwise kernel,若动态轴及其之前的 split 轴
已经能够提供不少于 num_vector_core 的并行度,则将动态轴之后的静态 split
轴降级为普通 kernel 轴,不再生成对应的 runtime BLOCK 和 grid 维度。
并行度不足或无法获得可靠 size hint 时,继续保留这些静态 split 轴。
每个 group 使用 representative shape 执行 autotune,确定 BLOCK_SUB、
静态 split 轴的 BLOCK 及其他编译配置。
运行时根据实际 shape 选择 group,并在对应 BLOCK_SUB 的约束下重新计算
动态轴的 runtime BLOCK。当实际 shape 小于 representative 时,允许使用更多
program,减少因沿用 representative grid 导致的核间并行度不足。