已合并
feat(_inductor): add ascendc backend support for aclgraph capture/replay #40263
dingdairong创建于 7月7日
feat(_inductor): add ascendc backend support for aclgraph capture/replay #40263
已合并
共 4 个文件变更+130-14
| @@ -0,0 +1,53 @@ | |||
| 1 | +"""AscendC backend: basic compilation and numeric correctness.""" | ||
| 2 | +import unittest | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from torch.testing._internal.common_utils import ( | ||
| 6 | + instantiate_parametrized_tests, | ||
| 7 | + parametrize, | ||
| 8 | + run_tests, | ||
| 9 | + TestCase, | ||
| 10 | +) | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +class TestAscendcBasic(TestCase): | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + def setUpClass(cls): | ||
| 18 | + super().setUpClass() | ||
| 19 | + cls._ascendc_ok = False | ||
| 20 | + try: | ||
| 21 | + x = torch.randn(4, 4, device="npu") | ||
| 22 | + torch.compile(lambda t: t + 1, backend="inductor", | ||
| 23 | + options={"npu_backend": "ascendc"})(x) | ||
| 24 | + cls._ascendc_ok = True | ||
| 25 | + except Exception: | ||
| 26 | + pass | ||
| 27 | + | ||
| 28 | + def setUp(self): | ||
| 29 | + super().setUp() | ||
| 30 | + if not self._ascendc_ok: | ||
| 31 | + self.skipTest("ascendc backend not available") | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + def test_mul_sub(self, dtype): | ||
| 35 | + """Pointwise pattern: compiled output matches eager.""" | ||
| 36 | + def fn(x, y): | ||
| 37 | + return (x * y - x) | ||
| 38 | + | ||
| 39 | + x = torch.randn(64, 64, dtype=dtype, device="npu") | ||
| 40 | + y = torch.randn(64, 64, dtype=dtype, device="npu") | ||
| 41 | + | ||
| 42 | + eager_out = fn(x, y) | ||
| 43 | + compiled_fn = torch.compile(fn, backend="inductor", | ||
| 44 | + options={"npu_backend": "ascendc"}) | ||
| 45 | + compiled_out = compiled_fn(x, y) | ||
| 46 | + | ||
| 47 | + torch.testing.assert_close(compiled_out, eager_out, rtol=1e-3, atol=1e-3) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +instantiate_parametrized_tests(TestAscendcBasic) | ||
| 51 | + | ||
| 52 | +if __name__ == "__main__": | ||
| 53 | + run_tests() | ||
| @@ -0,0 +1,52 @@ | |||
| 1 | +"""AscendC backend: reduce-overhead mode (aclgraph capture/replay).""" | ||
| 2 | +import unittest | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 6 | + | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestAscendcReduceOverhead(TestCase): | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + def setUpClass(cls): | ||
| 13 | + super().setUpClass() | ||
| 14 | + cls._ascendc_ok = False | ||
| 15 | + try: | ||
| 16 | + x = torch.randn(4, 4, device="npu") | ||
| 17 | + torch.compile(lambda t: t + 1, backend="inductor", | ||
| 18 | + options={"npu_backend": "ascendc"})(x) | ||
| 19 | + cls._ascendc_ok = True | ||
| 20 | + except Exception: | ||
| 21 | + pass | ||
| 22 | + | ||
| 23 | + def setUp(self): | ||
| 24 | + super().setUp() | ||
| 25 | + if not self._ascendc_ok: | ||
| 26 | + self.skipTest("ascendc backend not available") | ||
| 27 | + | ||
| 28 | + def test_capture_replay(self): | ||
| 29 | + """Verify reduce-overhead triggers graph capture and results are correct.""" | ||
| 30 | + def fn(x, y): | ||
| 31 | + return x * y - x | ||
| 32 | + | ||
| 33 | + x = torch.randn(64, 64, device="npu") | ||
| 34 | + y = torch.randn(64, 64, device="npu") | ||
| 35 | + | ||
| 36 | + compiled_fn = torch.compile( | ||
| 37 | + fn, backend="inductor", | ||
| 38 | + options={"npu_backend": "ascendc", "triton.cudagraphs": True}, | ||
| 39 | + ) | ||
| 40 | + | ||
| 41 | + # Warmup (triggers compilation + first capture) | ||
| 42 | + with torch.no_grad(): | ||
| 43 | + for _ in range(3): | ||
| 44 | + out = compiled_fn(x, y) | ||
| 45 | + | ||
| 46 | + # Verify correctness | ||
| 47 | + eager_out = fn(x, y) | ||
| 48 | + torch.testing.assert_close(out, eager_out, rtol=1e-3, atol=1e-3) | ||
| 49 | + | ||
| 50 | + | ||
| 51 | +if __name__ == "__main__": | ||
| 52 | + run_tests() | ||
| @@ -6,18 +6,30 @@ from .graph import patch_codegen_with_cpp_wrapper | |||
| 6 | from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream | 6 | from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, get_current_raw_stream |
| 7 | from ._npu_meta_registration import npu_patch_meta | 7 | from ._npu_meta_registration import npu_patch_meta |
| 8 | 8 | ||
| 9 | +# 顶层 patch：所有 inductor backend（triton / mlir / dvm / ascendc）都需要的 NPU 设备级patch， | ||
| 10 | +# 与 codegen 后端选择无关，在任何 backend loader 之前无条件执行 | ||
| 9 | npu_patch_meta() | 11 | npu_patch_meta() |
| 10 | register_device_op_overrides_npu() | 12 | register_device_op_overrides_npu() |
| 11 | -patch_has_triton() | 13 | + |
| 12 | -patch_is_gpu() | 14 | + |
| 13 | -patch_device_supports_tma() | 15 | +def _apply_common_patches(): |
| 14 | -patch_codegen_with_cpp_wrapper() | 16 | + # triton / mlir 后端共用的 patch |
| 15 | -patch_cache_base_get_system() | 17 | + patch_has_triton() |
| 18 | + patch_is_gpu() | ||
| 19 | + patch_device_supports_tma() | ||
| 20 | + patch_codegen_with_cpp_wrapper() | ||
| 21 | + patch_cache_base_get_system() | ||
| 22 | + | ||
| 16 | 23 | ||
| 17 | def _get_backend() -> str: | 24 | def _get_backend() -> str: |
| 18 | return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default") | 25 | return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default") |
| 19 | 26 | ||
| 27 | + | ||
| 28 | +def _load_ascendc_backend(): | ||
| 29 | + from . import ascendc | ||
 | |||
| 30 | + | ||
| 20 | def _load_mlir_backend(): | 31 | def _load_mlir_backend(): |
| 32 | + _apply_common_patches() | ||
| 21 | import torch | 33 | import torch |
| 22 | try: | 34 | try: |
| 23 | import torch_mlir | 35 | import torch_mlir |
| @@ -35,6 +47,7 @@ def _load_mlir_backend(): | |||
| 35 | 47 | ||
| 36 | 48 | ||
| 37 | def _load_dvm_backend(): | 49 | def _load_dvm_backend(): |
| 50 | + _apply_common_patches() | ||
| 38 | import torch | 51 | import torch |
| 39 | from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin | 52 | from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin |
| 40 | from .lowering_patch import apply_mlir_inductor_patch | 53 | from .lowering_patch import apply_mlir_inductor_patch |
| @@ -47,6 +60,7 @@ def _load_dvm_backend(): | |||
| 47 | 60 | ||
| 48 | 61 | ||
| 49 | def _load_triton_backend(): | 62 | def _load_triton_backend(): |
| 63 | + _apply_common_patches() | ||
| 50 | import os | 64 | import os |
| 51 | import torch | 65 | import torch |
| 52 | has_triton = torch.utils._triton.has_triton() | 66 | has_triton = torch.utils._triton.has_triton() |
| @@ -169,6 +183,7 @@ def _load_triton_backend(): | |||
| 169 | _BACKEND_LOADERS = { | 183 | _BACKEND_LOADERS = { |
| 170 | "mlir": _load_mlir_backend, | 184 | "mlir": _load_mlir_backend, |
| 171 | "dvm": _load_dvm_backend, | 185 | "dvm": _load_dvm_backend, |
| 186 | + "ascendc": _load_ascendc_backend, | ||
| 172 | "default": _load_triton_backend, | 187 | "default": _load_triton_backend, |
| 173 | } | 188 | } |
| 174 | 189 | ||
| @@ -1,8 +1,6 @@ | |||
| 1 | import logging | 1 | import logging |
| 2 | import os # noqa: C101 | 2 | import os # noqa: C101 |
| 3 | 3 | ||
| 4 | -from triton.runtime.driver import driver | ||
| 5 | - | ||
| 6 | import torch | 4 | import torch |
| 7 | from torch._inductor import config | 5 | from torch._inductor import config |
| 8 | 6 | ||
| @@ -17,13 +15,11 @@ config.trace.enabled = True | |||
| 17 | 15 | ||
| 18 | config.fallback_random = True | 16 | config.fallback_random = True |
| 19 | 17 | ||
| 20 | -# npu hardware params from trion | ||
| 21 | -target = driver.active.get_current_target() | ||
| 22 | -device = driver.active.get_current_device() | ||
| 23 | -prop = driver.active.utils.get_device_properties(device) | ||
| 24 | 18 | ||
| 25 | -num_cube_core = prop["num_aicore"] | 19 | +device = torch.npu.current_device() |
| 26 | -num_vector_core = prop["num_aicore"] | 20 | +prop = torch.npu.get_device_properties(device) |
| 21 | +num_cube_core = prop.cube_core_num | ||
| 22 | +num_vector_core = prop.vector_core_num | ||
| 27 | 23 | ||
| 28 | # unit byte | 24 | # unit byte |
| 29 | npu_block = 32 | 25 | npu_block = 32 |
| @@ -112,7 +108,7 @@ acc_comp_tol = { | |||
| 112 | "default": {"rtol": 1.3e-6, "atol": 1e-5}, | 108 | "default": {"rtol": 1.3e-6, "atol": 1e-5}, |
| 113 | } | 109 | } |
| 114 | 110 | ||
| 115 | -if "Ascend910B" in target.arch: | 111 | +if "Ascend910B" in prop.name: |
| 116 | num_vector_core = num_cube_core * 2 | 112 | num_vector_core = num_cube_core * 2 |
| 117 | 113 | ||
| 118 | log_level_env = os.getenv("INDUCTOR_ASCEND_LOG_LEVEL", "WARNING").upper() | 114 | log_level_env = os.getenv("INDUCTOR_ASCEND_LOG_LEVEL", "WARNING").upper() |
🟠 High Priority
changed line:
from . import ascendc(init.py line 32)该 DIFF 在
_BACKEND_LOADERS中注册了"ascendc": _load_ascendc_backend，但torch_npu/_inductor/目录下不存在ascendcPython 包（既无.py文件也无.so编译产物，经 glob 全仓库扫描确认）。当用户设置
TORCHINDUCTOR_NPU_BACKEND=ascendc时，_load_backend()→_load_ascendc_backend()→from . import ascendc将抛出ImportError，导致 inductor 后端加载失败。如果
ascendc包计划在后续 PR 中添加，建议将_load_ascendc_backend函数及_BACKEND_LOADERS中的注册条目与 ascendc 包一并合入，而不是提前提交。建议：方案一：将
_load_ascendc_backend及_BACKEND_LOADERS["ascendc"]的注册与 ascendc 包的添加放在同一个 PR 中合入。方案二：在_load_ascendc_backend中添加 try/except 并提供清晰的错误提示，例如raise ImportError("ascendc backend is not yet available...")。