已开启
fix(npu): preserve native spatial ops for matting #44957
伦创建于 16 天前
fix(npu): preserve native spatial ops for matting #44957
已开启
合并受阻
atomgit-bot
16 天前 评论:
16 天前 评论:
变更摘要
本 PR 为 NPU 的 triton_experimental 后端修复 matting(背景抠图)场景下的精度问题,核心思路是保留空间类算子的原生实现并在应用逐点仿射前折叠推理 BatchNorm 参数。改动集中在 torch_npu/_inductor/triton_experimental/overrides.py 与 torch_npu/_inductor/triton_experimental/__init__.py:新增 prepare_npu_decompositions() 统一清理分解表,并新增两个后端私有改写函数,在 _activate()/apply_npu_overrides() 中接入。
主要改动
- 新增
prepare_npu_decompositions()并接入初始化流程: 在overrides.py中新增该函数,通过remove_decompositions移除upsample_bilinear2d与upsample_bilinear2d_backward的分解,同时为upsample_bilinear2d.vec重新注册自定义 Python 实现(用upsample_compute_output_size计算输出尺寸并转发到 default 重载);__init__.py的_activate()中导入并调用该函数,使其在 backend 注册前生效。 - 新增
_preserve_native_reflection_pad2d(): 调用remove_decompositions移除reflection_pad2d及reflection_pad2d_backward的分解,使反射填充保留为原生 NPU 算子而非索引展开。 - 新增
_override_inference_batch_norm_decomposition(): 用自定义inference_batch_norm(先计算rsqrt、scale、shift折叠推理参数,再按 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注册器,数值类改写保留在本后端独立维护。


ascend-robot
16 天前 评论:
16 天前 评论:
atomgit-bot
16 天前 评论:
16 天前 评论:
16 天前 添加了label:ascend-cla/yes
此处折叠了174条消息 查看更多
1 天前 添加了label:docs-ci-pipeline-success
1 天前 删除了label:ci-pipeline-running
1 天前 添加了label:ci-pipeline-passed
AtlasAccount
1 天前 评论:
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 | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线


【合入来源】
【修改方案】
问题现象与根因
Background_Matting在torch.compile的 NPUtriton_experimental后端执行时,reflection_pad2d和upsample_bilinear2d的前反向算子会沿用通用 Inductor decomposition,被展开为通用索引/逐点计算,而没有调用已有的 NPU 原生算子。对于该模型中的大尺寸空间特征图,这会引入额外的索引计算、访存和 kernel 调度开销。移除 decomposition 后,这四个算子又因为没有可供批量扫描发现的上游 lowering,在
implicit_fallbacks=False严格模式下无法自动注册 fallback,因此还需要显式注册原生 fallback。修复方案
triton_experimentaldecomposition 注册阶段移除以下四个算子的通用分解,保留 NPU 原生执行路径:aten.reflection_pad2d.defaultaten.reflection_pad2d_backward.defaultaten.upsample_bilinear2d.defaultaten.upsample_bilinear2d_backward.defaultaten.upsample_bilinear2d.vec注册 NPU 兼容的 dispatcher:根据output_size或scale_factors计算目标尺寸,再转发到aten.upsample_bilinear2d.default,兼容F.interpolate的两种常见调用方式。EXPLICIT_FALLBACK_LIST,在_register_npu_inductor_fallbacks()中显式注册上述四个算子的 native fallback,保证implicit_fallbacks=False时同样可编译、执行。triton_experimental后端,不改变默认 NPU Inductor 后端及其他后端的 decomposition/fallback 行为。【资料变更】
不涉及。
【接口变更】
不涉及。仅调整 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_listtest_reflection_pad2d_native_forward_backwardtest_bilinear_upsample_size_native_forward_backwardtest_bilinear_upsample_scale_factor_native_forward_backward结果:
Ran 4 tests in 2.971s - OK。测试覆盖:
reflection_pad2d前向、反向结果与 eager 一致;upsample_bilinear2d的size、scale_factor两种入口及其反向结果与 eager 一致;implicit_fallbacks=False下编译、执行成功;torch.ops.aten.*.default调用,确认走 NPU 原生路径;git diff --check通过。【CheckList】