openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=43214&platform=gitcode&codeHostingPlatformFlag=gitcode

变更摘要

此 PR 为 inductor 的 welford 归约引入了 SIMD 向量化支持。核心思路是在持久化归约 kernel 中选取一个非归约轴（靠近归约维度的 split axis）进行向量化拆分，将归约循环折叠为 tl.sum(..., axis=1, keep_dims=True) 的单次 SIMD 操作，从而消除显式的 for loop_r 归约循环。同时新增了后归约代码生成缓冲区机制，用于在归约循环结束后生成均值、M2 和权重的后处理计算。默认将 enable_welford 配置开启。

主要改动

• enable_welford 默认值切换: torch_npu/_inductor/config.py 中 enable_welford 的默认值从 "0" 改为 "1"，使 SIMD welford 路径默认生效。

• 向量化 Welford 轴选择逻辑: 在 NPUIndexTritonKernel._select_vectorized_welford_axis() 中新增轴选择算法：在 welford_reduce 类型且归约维度均为静态长度时，选取最靠近归约维度的非归约 split axis，标记 is_vectorized_split = True，并将 kernel 类型设为 NPUKernelType.SIMD。

• 后归约代码生成缓冲区: 新增 post_reduction_loads、post_reduction_compute、post_reduction_stores 三个独立缓冲区，以及 _activate_post_reduction_codegen() 和重写的 disable_reduction() 方法。在归约循环结束后将 welford 的 mean、m2、weight 后处理计算放入这些缓冲区，与主循环代码分离。

• Welford SIMD 归约代码生成重写: _gen_reduction_code() 中 welford_reduce 分支从直接抛异常改为完整的 SIMD 实现：使用 tl.sum 沿归约轴一次性累积 acc_sum 和 acc_sum_sq，在 post_loop_combine 中计算均值、M2 和权重，并通过 reduction_resize() 广播回原始形状。

• 静态尾部掩码优化: IterationRangesEntryNPUIndex._codegen_mask() 中新增 use_static_vector_mask 路径：当维度静态可整除时直接设 _mask = True，仅不可整除时生成 if N % BLOCK_SUB != 0 条件掩码，避免了运行时的逐元素比较开销。

• 测试覆盖: test_var_mean.py 新增 test_welford_simd_codegen（验证 SIMD 代码特征如 axis=1, keep_dims=True、_acc_sum、X1BLOCK_SUB 等）、test_welford_simd_codegen_disabled（验证关闭时不出现向量化标记）和 test_welford_simd_tail_codegen（验证非整除维度的尾部处理分支）。

代码审查

审查总结

本次审查覆盖了 3 个文件的变更，每份文件均已逐一审查。

审查结果

按优先级统计

• P2: 1 个 — _select_vectorized_welford_axis 未排除 tiling 轴可能导致生成的 Triton 代码引用未定义变量
• P3: 1 个 — enable_welford 默认值已启用但注释仍称"默认禁用"

整体风险评估

本次变更为 welford 归约引入了 SIMD 代码生成路径，涉及约 500 行新增代码，逻辑较复杂。主要风险在于 _select_vectorized_welford_axis 方法的轴选择逻辑缺少对 tiling 轴的排除检查（P2），在某些图配置下可能导致 kernel 编译失败。其余代码（包括测试用例）逻辑一致，无明显缺陷。建议在合入前修复 P2 问题并更新 config.py 的过时注释。

⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。