问题描述

torch.compile 使用 npu_backend=triton_experimental 时，输出为单个标量的全量归约融合 kernel（典型：CausalLM 的 log_softmax + nll_loss(ignore_index) 路径，如 triton_unk_fused_clone_nll_loss_forward_slice_view_*）编译失败：

ValueError('Cannot broadcast, rank mismatch: [1, 1, 1], [1, 1]')
tl.store(in_out_ptr0 + (tl.full([1, 1], 0, tl.int32).broadcast_to(XBLOCK, 1)), tmp15, None)

HuggingFace runner 的 ElectraForCausalLM / RobertaForCausalLM / XGLMForCausalLM 三模型均因同一 kernel 模板触发。

复现环境

Ascend NPU（aarch64）+ PyTorch 2.13.0 + torch_npu 2.13.0（triton_experimental 后端）

根因分析

1. 上游 store codegen（integer-index store 分支）会手动追加 .broadcast_to(<value.shape>)，其中 value.shape 取自发射时刻 CSE 记录的形状——此刻归约值与索引秩一致，代码正确；
2. triton_experimental 随后做 r 树提升（promoted r-tree）改写，把归约值表达式整形为更高秩（如 [None, None, :] 秩 3），但 store 行中手动追加的 broadcast_to 参数与基座 tl.full 仍停留在提升前的秩；
3. Triton 要求 tl.store 两侧秩相等，直接编译失败。

修复

PR #44955：在 _rewrite_reduction_store_shape 内补全陈旧秩分支——当 broadcast_to 参数个数小于 real_ndim 时重建参数（r 槽强制置 1、旧非 1 参数按序映射、tl.full shape 补齐到 real_ndim）。写入地址、数值与冗余写次数均不变，语义严格等价。