已开启
fix(npu): preserve native spatial ops for matting #44957
fix(npu): preserve native spatial ops for matting #44957
已开启
创建于 16 天前
伦
成员
16 天前

【合入来源】

如有社区issue,请关联issue链接
请勿携带内部流程信息(需求链接、问题单、内部issue等)

【修改方案】

问题现象与根因

Background_Mattingtorch.compile 的 NPU triton_experimental 后端执行时,reflection_pad2dupsample_bilinear2d 的前反向算子会沿用通用 Inductor decomposition,被展开为通用索引/逐点计算,而没有调用已有的 NPU 原生算子。对于该模型中的大尺寸空间特征图,这会引入额外的索引计算、访存和 kernel 调度开销。

移除 decomposition 后,这四个算子又因为没有可供批量扫描发现的上游 lowering,在 implicit_fallbacks=False 严格模式下无法自动注册 fallback,因此还需要显式注册原生 fallback。

修复方案

  1. triton_experimental decomposition 注册阶段移除以下四个算子的通用分解,保留 NPU 原生执行路径:
    • aten.reflection_pad2d.default
    • aten.reflection_pad2d_backward.default
    • aten.upsample_bilinear2d.default
    • aten.upsample_bilinear2d_backward.default
  2. aten.upsample_bilinear2d.vec 注册 NPU 兼容的 dispatcher:根据 output_sizescale_factors 计算目标尺寸,再转发到 aten.upsample_bilinear2d.default,兼容 F.interpolate 的两种常见调用方式。
  3. 新增 EXPLICIT_FALLBACK_LIST,在 _register_npu_inductor_fallbacks() 中显式注册上述四个算子的 native fallback,保证 implicit_fallbacks=False 时同样可编译、执行。
  4. 将改动限定在 triton_experimental 后端,不改变默认 NPU Inductor 后端及其他后端的 decomposition/fallback 行为。
  5. 新增回归测试,同时检查 fallback 注册、前反向数值一致性及生成代码确实调用 NPU 原生算子。

【资料变更】

不涉及。

【接口变更】

不涉及。仅调整 NPU triton_experimental 后端内部的 decomposition 与 fallback 注册逻辑,对外 API 不变。

【功能验证】

验证环境:Ubuntu 22.04.5 LTS(aarch64)、Ascend 910B2、CANN 9.1.0-beta.1、Python 3.11.15、PyTorch 2.13.0、TorchNPU 2.13.0、Triton 3.2.0。

执行 test/_inductor/test_triton_experimental_regressions.py 中本次新增的 4 个定向用例:

  • test_spatial_ops_use_explicit_fallback_list
  • test_reflection_pad2d_native_forward_backward
  • test_bilinear_upsample_size_native_forward_backward
  • test_bilinear_upsample_scale_factor_native_forward_backward

结果:Ran 4 tests in 2.971s - OK

测试覆盖:

  • reflection_pad2d 前向、反向结果与 eager 一致;
  • upsample_bilinear2dsizescale_factor 两种入口及其反向结果与 eager 一致;
  • implicit_fallbacks=False 下编译、执行成功;
  • 生成代码包含上述算子的 torch.ops.aten.*.default 调用,确认走 NPU 原生路径;
  • git diff --check 通过。

【CheckList】

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

likedislike
合并受阻
伦成员
16 天前 创建了 pull request,commit 4ff6aa3c
atomgit-bot
atomgit-bot
16 天前 评论:

变更摘要

本 PR 为 NPU 的 triton_experimental 后端修复 matting(背景抠图)场景下的精度问题,核心思路是保留空间类算子的原生实现并在应用逐点仿射前折叠推理 BatchNorm 参数。改动集中在 torch_npu/_inductor/triton_experimental/overrides.pytorch_npu/_inductor/triton_experimental/__init__.py:新增 prepare_npu_decompositions() 统一清理分解表,并新增两个后端私有改写函数,在 _activate()/apply_npu_overrides() 中接入。

主要改动

  • 新增 prepare_npu_decompositions() 并接入初始化流程: 在 overrides.py 中新增该函数,通过 remove_decompositions 移除 upsample_bilinear2dupsample_bilinear2d_backward 的分解,同时为 upsample_bilinear2d.vec 重新注册自定义 Python 实现(用 upsample_compute_output_size 计算输出尺寸并转发到 default 重载);__init__.py_activate() 中导入并调用该函数,使其在 backend 注册前生效。
  • 新增 _preserve_native_reflection_pad2d(): 调用 remove_decompositions 移除 reflection_pad2dreflection_pad2d_backward 的分解,使反射填充保留为原生 NPU 算子而非索引展开。
  • 新增 _override_inference_batch_norm_decomposition(): 用自定义 inference_batch_norm(先计算 rsqrtscaleshift 折叠推理参数,再按 channel 形状应用仿射)覆盖 torch.ops.aten._native_batch_norm_legit_no_training.default 的分解,并在每次改动后清空 fast_random_decomps 缓存。
  • apply_npu_overrides() 挂载新改写: 在该函数中新增调用 _preserve_native_reflection_pad2d()_override_inference_batch_norm_decomposition(),并更新原有注释,说明分解表相关改写已迁移到共享的 torch_npu._inductor.decomposition 注册器,数值类改写保留在本后端独立维护。
likedislike
atomgit-bot
atomgit-bot
16 天前 评论:

代码审查

✅ 未发现问题

likedislike
ascend-robotascend-robot成员
16 天前 添加了label:ascend-cla/yes
此处折叠了174条消息 查看更多
ascend-robotascend-robot成员
1 天前 添加了label:docs-ci-pipeline-success
AtlasAccountAtlasAccount成员
1 天前 删除了label:ci-pipeline-running
AtlasAccountAtlasAccount成员
1 天前 添加了label:ci-pipeline-passed
AtlasAccount
AtlasAccount成员
1 天前 评论:
流水线 PR-pipeline_pytorch#65758 [ commitID:d4b2b35a ] 已完成
>>>代码风格自动修复执行成功(无修复内容)
阶段 任务名 状态 详情
编译构建 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 ✅ COMPLETED >>>
UT_ARM_A2_Part_03 ✅ COMPLETED >>>
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 ✅ COMPLETED >>>
UT_ARM_A2_Select_Part_02 ✅ COMPLETED >>>
UT_ARM_A2_Part_213 ✅ COMPLETED >>>
UT_inductor_Part_213 ⚪ IGNORED >>>
UT_DIST_ARM_Part_213 ⚪ IGNORED >>>
UT_ARM_A2_Select_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 : 停止流水线
likedislike
rmch成员
1 天前 评论:

/lgtm

likedislike