【合入来源】

• 问题单

背景

aten.cumprod 在 Inductor lowering 后会走 ops.scan（scan IR）。在 torch_npu 的 NPU Triton codegen 中，scan kernel 可能会被标记为 inside_reduction=True，但它不一定存在 torch._inductor.ir.Reduction 节点。

旧逻辑在以下两个地方会错误地触发 ReductionAnalysis，从而在 ReductionAnalysis.__init__() 内部抛出：RuntimeError: failed to get one reduction node。

• scheduling 阶段：对 inside_reduction 的 kernel 直接/间接做 reduction 分析
• triton codegen 阶段：dense_size_list()/dense_size_str() 在 inside_reduction 时无条件构建 ReductionAnalysis

拿aten.cumprod用例举例：

import os
import torch


os.environ.setdefault("TORCHDYNAMO_DEBUG", "1")
os.environ.setdefault("TORCHDYNAMO_LOG_LEVEL", "debug")
os.environ.setdefault("TORCHINDUCTOR_DEBUG", "1")
os.environ.setdefault("TORCH_COMPILE_DEBUG", "1")
os.environ.setdefault("TORCHINDUCTOR_FORCE_DISABLE_CACHES", "1")


def aten_add_cumprod_chain(x):
	a = x + 1
	b = torch.ops.aten.cumprod(a, dim=-1)
	c = b + 1
	return c

def main():
	device = torch.device("npu")
	a = torch.ones(3, 4, 5, device=device)
	for name, fn in [
		("dim-1", aten_add_cumprod_chain),
	]:
		compiled = torch.compile(fn, backend="inductor")
		b = compiled(a)
		c = fn(a)
		max_abs_diff = (b - c).abs().max().item()
		print(f"=== {name} ===")
		print(f"max_abs_diff: {max_abs_diff}")
		print(f"allclose: {torch.allclose(b, c, rtol=1e-4, atol=1e-4)}")
		print(f"inductor output: {b}")
		print(f"eager output: {c}")


if __name__ == "__main__":
	main()

错误日志片段如下：

File "/home/zhudada/miniconda3/envs/torch-npu/lib/python3.11/site-packages/torch_npu/_inductor/codegen/scheduling.py", line 387, in codegen_node
    return self.codegen_node_schedule(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhudada/miniconda3/envs/torch-npu/lib/python3.11/site-packages/torch_npu/_inductor/codegen/scheduling.py", line 112, in codegen_node_schedule
    self.decide_codegen_dims_in_kernel(node_schedule, kernel)
  File "/home/zhudada/miniconda3/envs/torch-npu/lib/python3.11/site-packages/torch_npu/_inductor/codegen/scheduling.py", line 570, in decide_codegen_dims_in_kernel
    kernel.reduce_analysis = ReductionAnalysis(kernel)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhudada/miniconda3/envs/torch-npu/lib/python3.11/site-packages/torch_npu/_inductor/codegen/kernel_analysis.py", line 206, in __init__
    raise RuntimeError("failed to get one reduction node")
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: failed to get one reduction node

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

[ERROR] 2026-02-28-09:13:40 (PID:3518743, Device:0, RankID:-1) ERR99999 UNKNOWN applicaiton exception

目标

• 让 scan kernel 在 NPU+Inductor 下可以通过编译并正常运行。
• 对齐上游 Inductor 行为：scan 场景禁用 cooperative reduction。

修改前/修改后路径对比

修改前：

修改后：

修复方案

1）Scan 禁用 cooperative reduction（同社区）

文件：torch_npu/_inductor/codegen/scheduling.py

• 改动点：NPUTritonScheduling.create_kernel_choices()
• 逻辑：若 kernel_features.contains_op("scan")，则强制 override_cooperative_reduction=False
• 目的：scan 需要跨块传递 prefix 状态，不能按 cooperative reduction 的方式做跨 program 的分工/合并；禁用该优化可避免语义错误，并对齐上游 Inductor 行为

2）只在真实 reduction 时构建 ReductionAnalysis

文件：torch_npu/_inductor/codegen/scheduling.py

• 改动点：decide_codegen_dims_in_kernel()
• 逻辑：kernel.find_reduction_node() 会返回一个 torch._inductor.ir.Reduction 实例或 None（表示未找到）。仅当返回值非 None（且类型为 ir.Reduction）时才创建 ReductionAnalysis(kernel)
• 目的：scan kernel 可能 inside_reduction=True 但并不存在 ir.Reduction 节点，避免错误进入 reduction 专用分析导致崩溃

3）dense_size_list() / dense_size_str() 对 scan 做兜底

文件：torch_npu/_inductor/codegen/triton.py

• 改动点：dense_size_list() / dense_size_str()
• 逻辑：
  • inside_reduction=True 时先判断 find_reduction_node() 是否存在且为 ir.Reduction
  • 只有真实 reduction 才构建 ReductionAnalysis
  • 否则按普通 dense 逻辑生成 tile shape（scan fallback）
• 目的：避免 scan 场景错误构建 ReductionAnalysis 触发 failed to get one reduction node

4）scan() codegen 适配：动态推断 scan axis，并统一 scan 相关维度（与社区有区别，因为社区假设 scan 轴永远在 dense tile 的最后一维（layout 固定 [X, R]），NPU 这边 dense tile 的 layout 可能是 [R, X...] 或 [X..., R]，所以必须动态推断 scan axis，并把所有相关逻辑统一到同一个 dim。）

文件：torch_npu/_inductor/codegen/triton.py

动机：NPU 的 dense tile layout 可能为 [R, X...] 或 [X..., R] 等变化形态，scan axis 不能写死，需要与实际 layout 保持一致。

修复策略：在 NPUIndexTritonKernel.scan() 中实现 axis 统一化的 codegen 适配：

• 动态推断 scan 轴 dim

  • 触发 dense_size_list() / dense_size_str() 生成 dense tile layout
  • 基于 golden_var_list 推断当前 dense tile 中 r*（scan 轴）对应的维度位置，得到 dim
  • 若推断失败，再 fallback 到上游启发式：dim = triton_tensor_ndim() - num_reduction_dims

• 统一使用同一个 dim

  • tl.associative_scan(..., dim, ...)
  • accumulator 形状：reduced_size[dim] = 1（保持 keepdim 一致）
  • triton_helpers.select_one(..., dim=dim, keep_dims=True)

该策略保证 scan 的 axis 与 dense layout 一致，覆盖 dim=0 / dim=1 / dim=-1 等常见场景。

测试

测试用例从如下几个角度对特性进行了验证，使用的scan算子是aten.cumsum，reduction算子是aten.sum：
1、张量维度：验证了二维、三维场景下的scan类算子的codegen正确性
2、算子融合正确性：验证了aten.cumsum算子和pointwise算子、reduction算子的融合后的codegen以及计算结果正确性
3、规约轴切分场景：验证了规约轴过长，需要切分规约轴进行计算的场景。
测试结果如下，所有测试用例全部通过：

关键概念说明

Scan IR / ops.scan

Inductor 在 lowering aten.cumprod / aten.cumsum / aten.logcumsumexp 等“前缀累计”算子时，通常会生成 ops.scan（IR 层可表现为 ir.Scan）。

scan 和 reduction 的关键区别：

• reduction：沿某个轴把值“聚合成更少的元素”（例如 sum/max），输出维度消失或 keepdim。
• scan：沿某个轴输出“每个前缀”的结果（例如 cumprod），输出维度保留，语义对顺序更敏感。

实现上，scan 在 Triton 侧通常会生成 tl.associative_scan，要求 combine 函数满足结合律（associative）。

inside_reduction=True

在 NPU 的 Triton codegen/调度框架里，scan kernel 往往也会被标记为 inside_reduction=True（因为它同样存在一个 r* 轴，并且需要类似 reduction 的轴/分块处理）。

但这并不意味着它一定存在一个 torch._inductor.ir.Reduction 节点。

ReductionAnalysis

ReductionAnalysis 是为 ir.Reduction 节点服务的分析器，用于推导 reduction 相关的 dense shape、reduced_dim 等信息。

scan kernel 可能没有 ir.Reduction 节点，如果在 scan 场景错误地实例化 ReductionAnalysis，就会在内部寻找 reduction node 时直接报错（failed to get one reduction node）。

补充：kernel.find_reduction_node() 返回值为 ir.Reduction 实例或 None，用于判断当前 kernel 是否真的包含 reduction 节点。

Cooperative reduction

cooperative reduction 指“多个 program/block 协作完成同一个 reduction 输出”的优化策略（例如并行计算 partial reduce，然后再做跨 program 合并）。

对 scan 来说，不能直接复用 cooperative reduction 的“拆分 r 轴 + 末端合并”模式：scan 需要跨块传递前缀状态（prefix carry），否则会破坏前缀结果的正确性。因此上游 Inductor 以及本次 NPU 适配都选择：scan 场景禁用 cooperative reduction。

Dense layout / dense_size_list() / dense_size_str()

在 Triton codegen 中，scan 会先把输入值 broadcast 成一个“dense tile”，然后对该 tile 的某个维度调用 tl.associative_scan。

dense_size_list() / dense_size_str() 就是生成这个 dense tile 的形状（例如 [R, X] 或 [X, R]）的工具函数。

golden_var_list

golden_var_list 是 NPU index codegen 用于确定 tile 维度顺序（layout）的关键变量列表。dense tile 的轴顺序会随其推导结果而变化，这也是 scan 不能把 axis 写死的根本原因。

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常