已关闭
[torch.compile][triton_experimental] 归约 kernel 元素总数超 2^31 时 int32 索引溢出：结果错误/设备故障，全量 int64 提升则编译挂死 #4373
huyuchao创建于 15 天前关闭于 1 天前
15 天前 添加了label：triage-review
TorchNPU-Bot
15 天前 评论：
15 天前 评论：
issue待分派，添加triage-review标签
15 天前 添加了label：bot-triaged；删除了label：triage-review
TorchNPU-Bot
15 天前 评论：
15 天前 评论：
检测到当前 issue 已关联 PR，自动添加标签：bot-triaged
huyuchao
12 天前 评论：
12 天前 评论：
修复进展（关联 PR：#43736）：本 issue 标题所述两类失效均已在 triton_experimental 后端解决——
1. int32 索引溢出（结果错误 / 507035 设备故障）
int64 模式下地址表达式全部轴因子无条件升宽，正确性由 triton 类型提升规则构造性保证（每个地址算术含 i64 操作数，任意形状/项形态/求值顺序不回卷），不依赖求界分析、形状快照或运行时 guard：
tl.load(in_ptr0 + (256*x0.to(tl.int64) + r0_1.to(tl.int64)), r0_mask & x0mask, ...)
int32 模式亦有 #186057 类防御武装（大常量 / 含 Mod / 整表达式 bound 越界即升宽，平局全部 fail-safe）。
2. 全量 int64 提升编译挂死（507034）
lane tile（tl.arange / tl.full）保持 int32（逐行断言钉死），i64 只出现在地址使用点轴因子与 i64 运行时形参——UB 不膨胀。另修复静态 numel ≥2^31 落 triton 字面量的三条编译失败路径（[2^31,2^32) uint32 符号性冲突 / constexpr 轴同毒 / equal_to 特化被 BiShengIR 拒绝的 uint32→i64 vcast）。
验证：2^31 边界端到端 9/9（含 8.6GB 真实跨界 kernel ×3、总块数 >2^31 的 expand 用例）、enable 套件 13/13（三入口 + 混跑隔离）、模型级 A/B 生成代码相同（升宽武装率 0）无回归。机制与证据链详见 PR 描述。
1 天前 关闭了 issue
1 天前 issue状态由 TODO 改变为 DONE
1 天前 添加了label：resolved
环境
torch.compile(..., options={"npu_backend": "triton_experimental"})）问题现象
参与
torch.compile归约的 tensor 元素总数超过 int32 上限（2^31 = 2147483648）时：复现
import torch import torch_npu # noqa x = torch.randn(8388609, 256, device="npu") # numel = 2^31 + 257 cfn = torch.compile(lambda t: t.sum(dim=1), options={"npu_backend": "triton_experimental"}) y = cfn(x) # 修复前：结果错误 / 设备故障根因
上游
select_index_dtype()在numel * reduction_numel > int32_max时将index_dtype置为tl.int64。旧实现将该信号传导到整棵索引 codegen，导致大向量 arange tile 全量 upcast；而实际溢出只发生在个别乘加项（如268435456*x1，系数 × (轴长-1) 超 2^31）上——全量 upcast 属过度修复（UB 翻倍挂死），全不提升则地址回绕。修复
PR：https://gitcode.com/Ascend/pytorch/pull/43736
「贪心提升 + 局部 guard」：静态/字面系数跨界仅提升溢出 addend（生成
268435456*x1.to(tl.int64)，大 tile 保持 int32，UB 不变）；动态轴长以 trace-time hint 快照判定，int32 保留项安装符号 guard 闭环（运行期轴长超快照 → guard 失败重编译）。