变更摘要

该 PR 针对 inductor 的完全静态 Welford 优化在存在原地内存更新（inplace update）场景下可能产生的正确性问题进行了修复。核心改动是在 torch_npu/_inductor/codegen/triton.py 中新增缓冲区变异风险分析逻辑，并在索引最终确定阶段重新评估 full_static_welford_reduction，从而在检测到跨归约且被原地修改的缓冲区时禁用完全静态 Welford 优化；同时在 test/_inductor/test_var_mean.py 中新增了对应的回归测试用例。

主要改动

• 新增 _full_static_welford_mutation_hazards 分析方法: 在 NPUIndexTritonKernel 中遍历调度节点，分别收集归约前/后使用的缓冲区、Welford 源/后置缓冲区、别名组与被修改缓冲区，并通过别名闭包展开计算跨归约且被修改的缓冲区集合，作为完全静态 Welford 的变异风险来源。
• 新增 finalize_indexing 覆盖逻辑: 在索引最终确定阶段，当缓冲区复用与原地更新已经确定但加载代码尚未生成时，重新调用 _full_static_welford_mutation_hazards，若存在变异风险则关闭 full_static_welford_reduction。
• 调整完全静态 Welford 的判定条件: 原逻辑在 self.mutations 非空时直接禁用完全静态 Welford，现改为仅依据归约轴长度判断其静态性，再结合后续变异风险分析决定是否真正启用。
• 新增回归测试 test_welford_full_static_mutation_hazard: 构造一个归约输入被原地复用/更新的 layer norm 场景，验证结果正确性，并断言生成代码中不存在 'vectorized_welford_axis':，同时包含 mutated_arg_names': ['in_out_ptr。

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

代码审查

审查总结

我已逐一审查了本次变更的两个文件：

• test/_inductor/test_var_mean.py：新增的 test_welford_full_static_mutation_hazard 测试逻辑合理（构造了 pre_norm 同时作为 reduction 输入、后置计算输入并发生 in-place mutation 的场景），断言覆盖了优化被禁用（vectorized_welford_axis 不存在）与数值一致性检查。未发现需要报告的问题。
• torch_npu/_inductor/codegen/triton.py：新增的 _full_static_welford_mutation_hazards 分析逻辑本身（inside_reduction 初始值、别名传递闭包、cross_reduction_buffers & mutated_buffers 判定）经核对对目标场景分类正确；主要风险在于 finalize_indexing 中重算 full_static_welford_reduction 的执行时机相对 create_inductor_meta（line 2356 读取该标志）不确定。

问题统计：

• P0：0
• P1：0
• P2：1
• P3：0

整体风险判断：本次修复方向正确、分析逻辑基本完备，但存在一个中等风险点——hazard 重算依赖基类 finalize_indexing 的调用时机，而该时机在当前文件中不可见且 codegen_body/codegen_kernel 均被完全重写、未调用 super()，可能导致 meta 与循环消除决策不一致或优化未被真正禁用。建议显式保证重算发生在 create_inductor_meta 之前并补充对应验证。

⛔ 需要修改

The MR can not be merged, because of CodeReview discussion not resolved

If you want to solve this problem, you can click here to do it in the FAQs.

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.