已合并
perf(inductor): refine pointwise default grouping #44926
perf(inductor): refine pointwise default grouping #44926
已合并
luqichao创建于 5 天前
luqichao
luqichao成员
5 天前
  1. 优化默认 pointwise/elementwise 动态分组策略

进入 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 等已有专用分组策略保持不变。

  1. 减少不必要的静态轴 split
    对于只有一个动态 split 轴的 pointwise kernel,若动态轴及其之前的 split 轴
    已经能够提供不少于 num_vector_core 的并行度,则将动态轴之后的静态 split
    轴降级为普通 kernel 轴,不再生成对应的 runtime BLOCK 和 grid 维度。
    并行度不足或无法获得可靠 size hint 时,继续保留这些静态 split 轴。
  2. 适配分组配置与运行时 tiling
    每个 group 使用 representative shape 执行 autotune,确定 BLOCK_SUB、
    静态 split 轴的 BLOCK 及其他编译配置。
    运行时根据实际 shape 选择 group,并在对应 BLOCK_SUB 的约束下重新计算
    动态轴的 runtime BLOCK。当实际 shape 小于 representative 时,允许使用更多
    program,减少因沿用 representative grid 导致的核间并行度不足。
likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 luqichao 的贡献)
luqichaoluqichao成员
5 天前 创建了 pull request,commit 983928cf
atomgit-bot
atomgit-bot
5 天前 评论:

变更摘要

本 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_numelmax_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_idbuild_group_representatives 均按 value // bucket_factor 归一化后参与 bucketize。
  • 运行时分组与 block 配置适配NPUSymbolicGroupedAutotuner 的 group_id 计算读取 bucket_factor 并校验其合法性;对 pointwise/elementwise_numel 主特征,build_grouped_launch_policy 生成的运行时规则新增 representative_numelmax_grid(基于 npu_num_vector_coreprior_programs),resolve_grouped_runtime_blockaxis_numel < representative_numel 时改用 max_grid 推导 program 数,保证 BLOCK 仍对齐 block_sub
likedislike
atomgit-bot
atomgit-bot
5 天前 评论:

代码审查

✅ 未发现问题

likedislike
ascend-robotascend-robot成员
5 天前 添加了label:ascend-cla/yes
此处折叠了109条消息 查看更多
dezheng889成员
3 天前 评论:

/lgtm
/approve

likedislike
ascend-robotascend-robot成员
3 天前 添加了label:approvedlgtm
ascend-robotascend-robot成员
3 天前 关闭了关联的issue
ascend-robotascend-robot成员
3 天前 合入了pull request
ascend-robot
ascend-robot成员
3 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#14425 [ commitID:4f9f5606 ] 已完成
likedislike