已合并

refactor(inductor): isolate NPU integrations by device #42564

Xuan Peng创建于 7月23日从HinPeng/pytorch： br/inductor-refactor-part1-0723 合入到 Ascend/pytorch： master

refactor(inductor): isolate NPU integrations by device #42564

已合并

Xuan Peng创建于 7月23日

从HinPeng/pytorch br/inductor-refactor-part1-0723合入到Ascend/pytorch master

Xuan Peng

7月23日

Issue: https://gitcode.com/Ascend/pytorch/issues/3329

Summary

PyTorch Inductor keeps lowerings in a global registry. Previously, importing
torch_npu replaced community lowerings globally, so CPU nodes in a mixed
CPU/NPU FX graph could enter NPU-specific lowering and autotuning paths.

This change installs a device-aware lowering dispatcher: NPU nodes use the
NPU lowering, while non-NPU nodes retain the original upstream lowering.

Changes

Consolidate the NPU autotune runtime and simplify NPU scheduling setup needed
by the device-specific lowering path.
Isolate FX patches by device to avoid applying NPU behavior to non-NPU graphs.
Add _make_device_lowering_dispatcher and install it around overridden
Inductor lowerings.
Keep the original upstream lowering available for CPU and other non-NPU
devices.
Update mm, bmm, flex_attention, and grouped MM call sites for the
PyTorch 2.13 (node, choice) algorithm-selector ABI.
Make the NPU algorithm selector delegate non-NPU layouts to the upstream
selector with the full PyTorch 2.13 call contract
Add coverage for device dispatch, selector compatibility, CPU GEMM fallback,
and mixed-device graph partitioning.

【合入来源】

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

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

【修改方案】

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

【资料变更】

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

【接口变更】

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

【功能验证】

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

【CheckList】

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

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

Pull Request已成功合入, 合并人@ascend-robot

（感谢 Xuan Peng 的贡献）

X Xuan Peng

7月23日创建了 pull request，commit 9c336316

atomgit-bot

7月23日评论：

变更摘要

本次 PR 对 torch_npu 的 Inductor 后端进行了大规模重构，核心目标是简化代码结构、消除与上游 PyTorch 的耦合冲突、并建立更清晰的设备分发机制。主要涉及以下几个方面：引入基于 lowering_patch 的设备 lowering 分发器，使 NPU 算子注册不再覆盖社区 handler；重构 AlgorithmSelectorCache 的 monkey-patch 机制，使其返回 (node, choice) 元组并按设备类型正确派发；移除 NPUTritonKernelWithLoop 与 inductor_ascend_linear_mode 配置项，统一走 index-first fallback 路径；重构 NPUCachingAutotuner 的预编译流程，采用 plugin 模式并统一串行/并行预编译逻辑；将 fasta_autotune.py 和 tile_generator.py 从 codegen 迁移至 runtime，明确运行时与代码生成层的边界。

主要改动

新增 lowering_patch 设备分发机制: 在 torch_npu/_inductor/lowering_patch.py 中新增 install_device_lowering_dispatch 函数，通过 _make_device_lowering_dispatcher 创建包装器，当 layout.device.type == "npu" 时调用 NPU handler，否则委托给上游社区 handler；在 __init__.py 中通过 LOWERING_OVERRIDE_OP 列表调用该函数，替代了原先修改 TritonScheduling 的 patch_triton_scheduling 调用。
AlgorithmSelectorCache 接口重构: select_algorithm.py 中的 __call__ 和 make_benchmark_fn 现在通过 functools.wraps 包装原始社区函数，按 layout.device.type 判断派发；NPU 路径的 __call__ 返回值从单个 output_node() 改为 (selected_node, selected_key) 元组；相应地，mm.py、addmm、bmm.py、flex_attention.py、mm_grouped.py 等调用方全部解包该元组。
移除 NPUTritonKernelWithLoop 与 inductor_ascend_linear_mode: 从 triton.py 中删除约 500 行的 NPUTritonKernelWithLoop 类及其 codegen_kernel/codegen_body 等复杂方法；从 config.py 删除 inductor_ascend_linear_mode = "linear" 配置项；NPUCombinedScheduling.codegen_node 简化为始终先尝试 _triton_scheduling（index 模式），异常时回退到 _nolinear_triton_scheduling；block remap 逻辑浓缩为 _remap_fallback_block_subs 函数，由 inductor_meta["requires_no_linear_block_remap"] 标志控制。
NPUCachingAutotuner 预编译流程重构: 引入 _NPUSkipPrecompilePlugin 插件机制替代硬编码的 skip_precompile 检查；将 precompile_parallel、benchmark_all_configs、_benchmark_all_configs 等独立函数整合进类方法 _precompile_worker_serial/_precompile_worker_parallel；提取 _create_launcher_grid 函数统一 Grid 创建逻辑；_precompile_configs_with_vf_retry 统一了串行与并行路径的 vf_fusion 重试逻辑。
post_grad 与 joint_graph 的 NPU 感知改造: post_grad.py 中 patch_pattern_mm_plus_mm 不再删除 mm_plus_mm 的 pattern entry，而是通过 _npu_aware_extra_check 包裹 extra_check，使 NPU 场景跳过该融合但不影响其他设备；joint_graph.py 中 constant_fold_uniform_value 增加 _is_npu_graph 判断，仅对 NPU 图执行 eliminate_dead_code；autotune_process.py 移除 patch_tuning_process_pool，get_device_list 回归社区实现。
模块重组与清理: fasta_autotune.py 和 tile_generator.py 从 torch_npu/_inductor/codegen/ 重命名移动至 torch_npu/_inductor/runtime/，tile_generator 的导入路径相应更新；NPUNoLinearTritonScheduling.kernel_type 改为类级别属性（= NPUTritonKernel）；删除了 codegen/triton.py 中的 patch_triton_scheduling 和 select_index_dtype 静态方法。

不准确？

ascend-robot成员

7月23日评论：

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

atomgit-bot

7月23日评论：

代码审查

审查总结

本次 PR 共审查 27 个文件，发现 5 个问题：

优先级	数量	说明
P1	1	`test_cpu_compile.py` 缺少 `import torch_npu._inductor`，导致 lowering dispatch 测试必然失败
P2	2	`post_grad.py` 中 `_npu_aware_extra_check` 对 `None` extra_check 无防护；`select_algorithm.py` 中 `__call__` 返回值从单值变为二元组的 breaking change
P3	2	`npu_combined_scheduling.py` 中裸 `except Exception` 过于宽泛；`test_inductor_graph_partition.py` 中测试对上游实现有侵入性假设

整体风险评估：中低风险。

重构范围大但逻辑合理——主要是将 NPU Inductor 的实现从 monkey-patch 模式迁移到更规范的设备分发（device dispatch）和插件（plugin）模式，与 PyTorch v2.13 的接口契约对齐。所有调用方均已适配新的二元组返回值。

各文件审查结果

文件	结果
`test/_inductor/test_algorithm_selector_dispatch.py`	无问题
`test/_inductor/test_autotune_process.py`	无问题
`test/_inductor/test_cpu_compile.py`	P1：缺少 `import torch_npu._inductor`
`test/_inductor/test_exceptions.py`	无问题
`test/_inductor/test_inductor_graph_partition.py`	P3（架构注释）
`test/_inductor/test_lowering_device_dispatch.py`	无问题
`test/_inductor/test_scheduling_contract.py`	无问题
`test/npu/test_public_bindings.py`	无问题
`torch_npu/_inductor/__init__.py`	无问题
`torch_npu/_inductor/autotune_process.py`	无问题
`torch_npu/_inductor/codegen/npu_combined_scheduling.py`	P3：裸 `except Exception`
`torch_npu/_inductor/codegen/scheduling.py`	无问题
`torch_npu/_inductor/codegen/triton.py`	无问题
`torch_npu/_inductor/config.py`	无问题
`torch_npu/_inductor/fx_passes/joint_graph.py`	无问题
`torch_npu/_inductor/fx_passes/post_grad.py`	P2：`None` extra_check 崩溃风险
`torch_npu/_inductor/kernel/bmm.py`	无问题
`torch_npu/_inductor/kernel/flex_attention.py`	无问题
`torch_npu/_inductor/kernel/mm.py`	无问题
`torch_npu/_inductor/kernel/mm_grouped.py`	无问题
`torch_npu/_inductor/lowering_patch.py`	无问题
`torch_npu/_inductor/profiler.py`	无问题
`torch_npu/_inductor/runtime/fasta_autotune.py`	无问题
`torch_npu/_inductor/runtime/tile_generator.py`	无问题
`torch_npu/_inductor/runtime/triton_heuristics.py`	无问题
`torch_npu/_inductor/select_algorithm.py`	P2：返回值 ABI breaking change
`torch_npu/_inductor/utils.py`	无问题

类型	数量
🔴 阻塞	1
🟡 建议	1

⛔ 需要修改

不准确？

ascend-robot成员

7月23日添加了label：ascend-cla/yes

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

ascend-robot成员

7月25日评论：

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
zhucehw

ascend-robot成员

7月25日添加了label：approved

ascend-robot成员

7月25日合入了pull request

ascend-robot成员

7月25日评论：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

ascend-robot成员

7月25日评论：

流水线 pytorch_gitcode_PR_multiVersion#13158 [ commitID：954b7a52 ] 已完成