已开启

[triton_experimental] fixlinearize rank mismatch + fmod #44448

AACAES创建于 27 天前从AACAES/pytorch： fix/triton-experimental-grid1d-downcast-rank-mismatch 合入到 Ascend/pytorch： master

[triton_experimental] fixlinearize rank mismatch + fmod #44448

已开启

AACAES创建于 27 天前

从AACAES/pytorch fix/triton-experimental-grid1d-downcast-rank-mismatch合入到Ascend/pytorch master

AACAES成员

27 天前

【合入来源】

如有社区issue，请关联issue链接> 请勿携带内部流程信息（需求链接、问题单、内部issue等）

https://gitcode.com/Ascend/pytorch/issues/4314

需求
问题单
issue/工单
重构优化
资料更新

【修改方案】

请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列> 如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）

1. 问题现象

linearize 模式下 promoted r-tree 的 codegen 暴露 rank mismatch 与 load mask 不匹配两类问题，叠加 aten.fmod 错误 fallback，导致多个 reduction / fill / slice / fmod 用例失败：

用例	错误
`test_buffer_use_after_remove_dynamic_shapes_npu`	`ValueError('Cannot broadcast, rank mismatch: [1,1,1,1], [1,1]')`
`test_slice_mutation3_npu`	`ValueError('Cannot broadcast, rank mismatch: [1], [2, 2]')`
`test_fmod_npu`（含 dynamic shapes）	fmod 退化为 extern `aten.fmod` 调用，融合路径精度不达标

2. 问题根因

根因 1 — load mask 不匹配

promoted r-tree 内不同 tl.load 的 index 只涉及部分 r-node，但代码统一套用 combined_mask（所有 r-node mask 的 AND），Triton 在 index 与 mask 秩不一致时拒绝广播。

根因 2 — `triton_tensor_ndim()` 语义混淆

原实现返回"每个子节点 1 维"，实为迭代空间秩，但下游（reduction_resize / dense_size_str / store address）要求的是输出秩。输出秩下每个 r-tree 只贡献 1 个 reduction axis，无论内部被拆成几个子节点。

根因 3 — `post_resize` 与 store address 的秩不一致

store address 的秩可能为 iter_ndim（x-tree 有 register 槽位时）或 real_ndim（full reduction），但 post_resize 统一用 real_ndim 作目标秩，两者不对齐。

根因 4 — `constant()` 形状误配

NPUTritonKernelOverrides.constant() 用 [1]*ndim 作 shape，在非 linearize 的 2D pointwise kernel 下（triton_tensor_ndim() 返回 1，但 [None,:]/[:,None] broadcast 出 2D store address）造成 fill/slice/mutation 场景 rank mismatch。

根因 5 — `index_vars_per_node` 时序问题

triton_tensor_ndim() 在 __init__ 阶段被调用，但此时 index_vars_per_node 尚未赋值。

根因 6 — `aten.fmod` 不在 GENERATE_LIST

aten.fmod 未注册进 triton_experimental 的 GENERATE_LIST，被错误地当作 fallback 处理（extern aten.fmod 调用），无法走 libdevice 融合 lowering。

3. 修改原理

核心思路：输出秩（real_ndim）与迭代空间秩（iter_ndim）分离，store address 秩由上游 codegen 决定、下游 post_resize 去适配它。

#	位置	修改	原理
1	`NPUTritonKernel.__init__()`	在 `super().__init__()` 前初始化 `index_vars_per_node` / `var_ranges_per_node`	消除根因 5 的时序问题
2	`triton_tensor_ndim()`	每个 r-tree 计 1 维（输出秩）	与 upstream `reduction_resize` 语义对齐
3	`_npu_iter_ndim()`（新增）	全局槽位空间取 max（每个子节点计 1 维）	累加器 / broadcast 形状用迭代空间秩
4	`post_resize`	目标秩按 x-tree 是否有 register 槽位选 `iter_ndim` / `real_ndim`	与 store address 秩一致
5	`_npu_apply_promoted_rtree_lines`	传入 `iter_ndim`；`_rewrite_reduction_store_shape` / `_fix_masked_load_ptr` 保留 upstream 07f4d4ccb2 完整实现	与 upstream promoted rtree shape fix 合流（冲突已解，见合并记录）
6	`_npu_mask_for_load_line()`（新增）	解析每条 load 的 index 表达式，只保留实际引用的 r-node mask	精确匹配每条 load 的 mask（根因 1）
7	`constant()`	统一返回 scalar（0D）	依赖 NPU Triton 的 scalar→任意形状隐式广播（根因 4）
8	`_npu_build_tree_node_mapping()`（从 `_apply_linearize` 提取）	纯映射构建，不碰 body	解耦映射与重写
9	partial-flat fold（新增）	对 divisor==1 节点检测其他节点是否构成覆盖它的 divisor-chain	解决 `constant()` 在 body emit 阶段访问不完整 mapping 的问题
10	`_npu_fixup_in_out_downcast()`（新增）	对 `downcast_args` 中 `*i64` 指针的 `tmpN = tl.load(ptr + ...)` 补 `.to(tl.int32)` 后处理	消除 i64 downcast 融合 kernel 读截断/垃圾值（根因：上游对 NPU pointwise kernel 的 load 遗漏自动 .to）

4. 影响范围

仅 promoted r-tree 路径、constant() 路径、i64 downcast 融合 kernel 受影响
非 promoted r-tree 通过 _linearize_applied 守卫或 iter_ndim == real_ndim 自然退化
lowering_override_list.py：aten.fmod 加入 GENERATE_LIST（净变更 +1 行）
改变 kernel 生成代码（明确声明，替代旧描述"不改变 kernel 运行时行为"）：i64 downcast 融合 kernel 的 load 增补 .to(tl.int32)；promoted r-tree kernel 的 store/load shape 与 mask 更精确
不修改 store address 生成逻辑

5. 测试验证

新增两个看护测试文件（9 个用例，均 PASSED，dev 6，2026-09-05/07）：

场景	用例	看护的修复
Promoted r-tree 输出秩	`test_reduction_fused_pointwise_broadcast`	norm→div 广播归约 store 秩
index_vars 时序	`test_partial_reduction_dynamic_shapes` / `test_std_dynamic_shapes`	`__init__` 阶段属性可用性
双归约输出	`test_var_mean_multi_output`	var_mean 两路 store 秩
constant() scalar	`test_fill_scalar_into_slice` / `test_slice_mutation3` / `test_slice_mutation_chain`	fill/slice mutation rank
fmod 融合	`test_fmod_stay_fused` / `test_fmod_int_stay_fused`	GENERATE_LIST（断言 `libdevice.fmod` 融合、无 `aten.fmod` fallback）

敏感性验证：在父提交（不含本 PR 修复）上跑，fmod 用例如预期失败（fallback 断言命中）；linearize 用例因父提交已含 upstream !45171 部分修复而通过，另 2 个 fmod 用例验证了修复归属本 PR。

回归验证：

ly_all 30 用例（test_pr/run.sh）：与 baseline 失败集完全一致，无回归
全量 1588 用例（inductor_tests + dynamic_shapes）：剔除设备瞬时故障（E39007，已单卡重跑 24/24 通过）后与 baseline 一致，无回归
相对 upstream/master 净修复 4 个用例：test_slice_mutation3 ×2 + test_fmod ×2（pip vs PR 分支 A/B 实测确认）

【资料变更】

请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写"不涉及"

不涉及

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写"不涉及"

不涉及

【功能验证】

说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤> 新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图

测试场景/方法：如上【测试验证】节；两个看护测试文件已随本 PR 提交（test/_inductor/test_triton_experimental_fmod.py、test/_inductor/test_triton_experimental_linearize_regressions.py）
新增/变更内容已新增 UT 看护用例 ✅
验证环境：Ascend 910B2 dev 6，CANN 9.1.0，torch_npu triton_experimental 后端

【CheckList】

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

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

合并受阻

A AACAES成员

27 天前创建了 pull request，commit 8d4aa2bc

ascend-robot成员

27 天前评论：

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

atomgit-bot

27 天前评论：

变更摘要

此 PR 修复了 triton_experimental 后端 codegen 与 launcher 层的三个正确性 bug：标量 Grid1D kernel 的 grid race 导致 in-place 操作结果随机；int64 pointwise kernel 因 downcast 指针缺少 .to(tl.int32) 转换而读到截断/垃圾值；以及 triton_tensor_ndim() 中 _linearize_applied 守卫导致 pre-linearize 阶段 tensor rank 计数偏低引发 rank mismatch。修改涉及 codegen/triton.py（新增 _npu_fixup_in_out_downcast() 后处理、移除 rank 计数守卫）和 npu_triton_heuristics.py（拓宽标量 Grid1D 判定条件、修复 in_out_ptr 临时 buffer 初始化）。

主要改动

拓宽标量 Grid1D 判据（npu_triton_heuristics.py）：将 is_unsplit_scalar_reduction 改为 is_scalar_grid1d，判据从仅覆盖含 R0_BLOCK 的标量归约扩展到所有 npu_num_x_nodes == 0 的 Grid1D kernel，并用 ncfg.codegen_linearize 门控防止非 linearize 模式下将真实 1D pointwise kernel 错误钳位为单 program，从而消除标量 pointwise in-place 操作的 48 路 read-modify-write 竞态
新增 _npu_fixup_in_out_downcast() 后处理（codegen/triton.py）：在 NPUTritonKernel.codegen_kernel() 返回 src 前调用该函数，扫描所有来自 downcast_args 中 *i64/i64 指针且缺少 .to(...) 的 tl.load() 调用，自动追加 .to(tl.int32)，修复 upstream triton compiler 遗漏的 in_ptr 与 in_out_ptr downcast 转换
修复 in_out_ptr 临时 buffer 初始化（npu_triton_heuristics.py）：在 _wrap_launcher_with_downcast() 中，对 in_out_ptr 类别的 downcast 输出参数将 torch.empty_like() 改为 a.to(dst_dtype)，确保 kernel 读取路径获得正确的 downcast 后当前值，而非未初始化垃圾值
移除 _linearize_applied 守卫（codegen/triton.py）：在 triton_tensor_ndim() 中去掉 getattr(self, '_linearize_applied', False) 条件，仅保留 triton_codegen_linearize 门控，使 pre-linearize 阶段也能正确统计 tensor 维度（通过 tree_node_mapping 为空时同一过滤逻辑自然数出全部非 scalar x-node），修复因数 range-tree 数而非 node 数导致的 rank mismatch

不准确？

atomgit-bot

27 天前评论：

代码审查

✅ 未发现问题

不准确？

ascend-robot成员

27 天前添加了label：needs-issue

此处折叠了383条消息查看更多

ascend-robot成员

1 小时前添加了label：docs-ci-pipeline-success

ascend-robot成员

1 小时前评论：

✅ 跳过 docs ci 检查，没有需要检查的文档文件

AtlasAccount成员

1 小时前删除了label：ci-pipeline-running

AtlasAccount成员

1 小时前添加了label：ci-pipeline-passed

AtlasAccount成员

1 小时前评论：

流水线 PR-pipeline_pytorch#66639 [ commitID：23fb9975 ] 已完成

>>>代码风格自动修复执行成功（无修复内容）

阶段	任务名	状态	详情
编译构建	Build_X86	✅ COMPLETED	>>>
	Build_ARM	✅ COMPLETED	>>>
	Build_X86_torchair	⚪ IGNORED	>>>
	Build_ARM_torchair	⚪ IGNORED	>>>
	patch_test	⚪ IGNORED	>>>
	Build_X86_213	✅ COMPLETED	>>>
	Build_ARM_213	✅ COMPLETED	>>>
恶意代码检查	Antipoison	✅ COMPLETED	>>>
编码安全与规范检查	codecheck_pre-commit	✅ COMPLETED	>>>
	check_error	✅ COMPLETED	>>>
	lintrunner	✅ COMPLETED	>>>
开源片段检查	SCA	✅ COMPLETED	>>>
开发者测试	UT_ARM_A3_Part_01	⚪ IGNORED	>>>
	UT_ARM_A3_Part_02	⚪ IGNORED	>>>
	UT_ARM_A2_Part_01	✅ COMPLETED	>>>
	UT_ARM_A2_Part_02	⚪ IGNORED	>>>
	UT_ARM_A2_Part_03	⚪ IGNORED	>>>
	UT_inductor_Part_01	⚪ IGNORED	>>>
	UT_inductor_Part_02	⚪ IGNORED	>>>
	UT_inductor_Part_03	⚪ IGNORED	>>>
	UT_inductor_Part_04	⚪ IGNORED	>>>
	UT_DIST_ARM_Part_01	⚪ IGNORED	>>>
	UT_DIST_ARM_Part_02	⚪ IGNORED	>>>
	UT_DIST_ARM_Part_03	⚪ IGNORED	>>>
	UT_DIST_ARM_Part_04	⚪ IGNORED	>>>
	UT_ARM_A2_Select_Part_01	⚪ IGNORED	>>>
	UT_ARM_A2_Select_Part_02	⚪ IGNORED	>>>
	UT_inductor_Part_213	✅ COMPLETED	>>>
流水线	PR-pipeline_pytorch	✅ COMPLETED	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, huangjingwei, liangsongwei, yashi999, culechan, Dring, wuyouqi1, L1919_snow, qq_52711437, WhiteNight12, nomiz, xiu_21, ffmh, wanglijun55, hss-shuai, husichao, smallsilly, lanshaozuishuai, jimmyisme1, lzy0920232, alpha-junh, Sunshine_Youngster, wei_zhuoyi, zhangyihuiben, zyw-hw, zzzkeke, rmch, yangch0324, LucciC, AACAES, renyujin, wjlflyer, senzhen-town, pengjingyou, qsc97, limuan, yule100, xiaoqi-zhou, kuhn7, chenxingying, hanye02, zichun_ye, anyrenwei, kkjocker, wangzili121, Lu_G, yvjc, puddingfjz, HandsoemLemon, bigprestigee1, huawuyi, zhenyu10, dairenjie, du-jin-hang, zou-jieyu, adelaideliu, TrHan, wanlinan, Windwindzzz, pengqihw, kisnwang, yuheng_wang, honghao_wang, jizewei, zhangguoguang, sunyu-xuan, chenrayray, hbhu_bin, liujunzhu, c_34, LiNuoh, maoyuanpeng1, zzhongmin, zhaoyu65, bellatan, jiabaolin, zhuofanshen, wencaiwen, lu_zhuge, caoshuyang, molly12, lyx324521, LQ1206, gitcode-chenjiao, cai-weiwei1989, CHDong, ogqin, yuanlipingGit, xuqinglin1, lqz2, zouwei1, chaoluoa, paradox325, jackzhang1116, yaoyao, akh, yujiacheng, dengjie0116, Hubert11111, Shine_Ws, wslhj555, longqiand, OYtao666, JiaqingQiang, luyyyy, Kingbelial, zhanghaiyu0101, wenxp1018, yanliu-luoluo, ksun_sekiro, liyong328, wgzheng, tangky, vivi_is_coding, aoiaoisola, weixin_44494597, wangmengmengwang65667, hid57809721, qq_35468730, comeonup, C547032, gcw_m5OQChA4, yao_yao_ling_xian, cnnbwcy, szqfes_12, cora_19, cann_lilin, can, shawnylee233, fanglanyue0916, hhz0, LiNuohang, taohuoquan, Jesse, WSs_321, SCh_zx]评论有效

compile、compile_inductor、compile_torchair : 运行流水线
retry : 重试流水线所有失败子任务
retry <任务名> : 仅重试指定失败子任务
stop : 停止流水线

【合入来源】

【修改方案】

1. 问题现象

2. 问题根因

根因 1 — load mask 不匹配

根因 2 — triton_tensor_ndim() 语义混淆

根因 3 — post_resize 与 store address 的秩不一致

根因 4 — constant() 形状误配

根因 5 — index_vars_per_node 时序问题

根因 6 — aten.fmod 不在 GENERATE_LIST

3. 修改原理

4. 影响范围

5. 测试验证

【资料变更】

【接口变更】

【功能验证】

【CheckList】

合并受阻

变更摘要

代码审查

根因 2 — `triton_tensor_ndim()` 语义混淆

根因 3 — `post_resize` 与 store address 的秩不一致

根因 4 — `constant()` 形状误配

根因 5 — `index_vars_per_node` 时序问题

根因 6 — `aten.fmod` 不在 GENERATE_LIST