已合并
test: add MLIR non-fallback kernel regression test #38520
yvjc创建于 6月15日
test: add MLIR non-fallback kernel regression test #38520
已合并
从已删除 :fusionTest合入到Ascend/pytorchv2.7.1
共 5 个文件变更+99-3
| @@ -59,7 +59,6 @@ class TestInductorBackendColdStart(TestCase): | |||
| 59 | 59 | ||
| 60 | import torch_npu._inductor | 60 | import torch_npu._inductor |
| 61 | 61 | ||
| 62 | - assert torch._inductor.config.trace.enabled | ||
| 63 | assert torch_npu._inductor._get_backend() == backend_env | 62 | assert torch_npu._inductor._get_backend() == backend_env |
| 64 | 63 | ||
| 65 | def fn(x, y): | 64 | def fn(x, y): |
| @@ -0,0 +1,47 @@ | |||
| 1 | +import os | ||
| 2 | +from unittest import skip | ||
| 3 | +from unittest.mock import patch | ||
| 4 | + | ||
| 5 | +os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" | ||
| 6 | +os.environ["TORCHINDUCTOR_USE_AKG"] = "1" | ||
| 7 | + | ||
| 8 | +import torch | ||
| 9 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 10 | +from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.mlir_compiler import AkgCompiler | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class TestMlirKernelNotFallback(TestCase): | ||
| 14 | + | ||
| 15 | + def test_basic_op_uses_non_fallback_kernel(self): | ||
| 16 | + def fusion_func(x, y): | ||
| 17 | + return x + y + 1 | ||
| 18 | + | ||
| 19 | + x_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 20 | + y_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 21 | + | ||
| 22 | + expected = fusion_func(x_cpu, y_cpu) | ||
| 23 | + | ||
| 24 | + x_npu = x_cpu.npu() | ||
| 25 | + y_npu = y_cpu.npu() | ||
| 26 | + | ||
| 27 | + has_fallback_kernel = False | ||
| 28 | + original_run = AkgCompiler.run | ||
| 29 | + | ||
| 30 | + def run_and_check_fallback(compiler, *args, **kwargs): | ||
| 31 | + nonlocal has_fallback_kernel | ||
| 32 | + result = original_run(compiler, *args, **kwargs) | ||
| 33 | + launcher_idx = compiler.get_primary_launcher_index() | ||
| 34 | + has_fallback_kernel |= compiler.is_fallback_kernels[launcher_idx] | ||
| 35 | + return result | ||
| 36 | + | ||
| 37 | + compiled = torch.compile(fusion_func, backend="inductor") | ||
| 38 | + | ||
| 39 | + with patch.object(AkgCompiler, "run", run_and_check_fallback): | ||
| 40 | + actual = compiled(x_npu, y_npu) | ||
| 41 | + | ||
| 42 | + self.assertFalse(has_fallback_kernel) | ||
| 43 | + self.assertEqual(expected, actual.cpu(), atol=1e-3, rtol=1e-3) | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +if __name__ == "__main__": | ||
| 47 | + run_tests() | ||
| @@ -0,0 +1,47 @@ | |||
| 1 | +import os | ||
| 2 | +from unittest import skip | ||
| 3 | +from unittest.mock import patch | ||
| 4 | + | ||
| 5 | +os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +import torch_npu | ||
| 9 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 10 | +from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.mlir_compiler import NpuMlirCompiler | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class TestMlirKernelNotFallback(TestCase): | ||
| 14 | + | ||
| 15 | + def test_basic_op_uses_non_fallback_kernel(self): | ||
| 16 | + def fusion_func(x, y): | ||
| 17 | + return x + y + 1 | ||
| 18 | + | ||
| 19 | + x_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 20 | + y_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 21 | + | ||
| 22 | + expected = fusion_func(x_cpu, y_cpu) | ||
| 23 | + | ||
| 24 | + x_npu = x_cpu.npu() | ||
| 25 | + y_npu = y_cpu.npu() | ||
| 26 | + | ||
| 27 | + has_fallback_kernel = False | ||
| 28 | + original_run = NpuMlirCompiler.run | ||
| 29 | + | ||
| 30 | + def run_and_check_fallback(compiler, *args, **kwargs): | ||
| 31 | + nonlocal has_fallback_kernel | ||
| 32 | + result = original_run(compiler, *args, **kwargs) | ||
| 33 | + launcher_idx = compiler.get_primary_launcher_index() | ||
| 34 | + has_fallback_kernel |= compiler.is_fallback_kernels[launcher_idx] | ||
| 35 | + return result | ||
| 36 | + | ||
| 37 | + compiled = torch.compile(fusion_func, backend="inductor") | ||
| 38 | + | ||
| 39 | + with patch.object(NpuMlirCompiler, "run", run_and_check_fallback): | ||
| 40 | + actual = compiled(x_npu, y_npu) | ||
| 41 | + | ||
| 42 | + self.assertFalse(has_fallback_kernel) | ||
| 43 | + self.assertEqual(expected, actual.cpu(), atol=1e-3, rtol=1e-3) | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +if __name__ == "__main__": | ||
| 47 | + run_tests() | ||
| @@ -1,10 +1,13 @@ | |||
| 1 | from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides | 1 | from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides |
| 2 | import torch_npu | 2 | import torch_npu |
| 3 | -from torch_npu._inductor.codegen.catlass.catlass_utils import try_import_catlass | ||
| 4 | 3 | ||
| 5 | 4 | ||
| 6 | class NewNPUDeviceOpOverrides(DeviceOpOverrides): | 5 | class NewNPUDeviceOpOverrides(DeviceOpOverrides): |
| 7 | def import_get_raw_stream_as(self, name): | 6 | def import_get_raw_stream_as(self, name): |
| 7 | + # Importing CATLASS loads the NPU config, which initializes NPU state. | ||
| 8 | + # Keep it lazy so forked compile workers can import device overrides. | ||
| 9 | + from torch_npu._inductor.codegen.catlass.catlass_utils import try_import_catlass | ||
| 10 | + | ||
| 8 | enabled_catlass = try_import_catlass() | 11 | enabled_catlass = try_import_catlass() |
| 9 | if not enabled_catlass and hasattr(torch_npu._C, "_npu_getCurrentRawStreamNoWait"): | 12 | if not enabled_catlass and hasattr(torch_npu._C, "_npu_getCurrentRawStreamNoWait"): |
| 10 | return f"from torch_npu._C import _npu_getCurrentRawStreamNoWait as {name}" | 13 | return f"from torch_npu._C import _npu_getCurrentRawStreamNoWait as {name}" |
| @@ -9,7 +9,7 @@ from torch._C import DispatchKey | |||
| 9 | from torch._decomp import remove_decompositions | 9 | from torch._decomp import remove_decompositions |
| 10 | from torch._prims_common.wrappers import out_wrapper | 10 | from torch._prims_common.wrappers import out_wrapper |
| 11 | import torch.nn.functional as F | 11 | import torch.nn.functional as F |
| 12 | -from .config import is_ascend950 | 12 | + |
| 13 | from .lowering_common import add_overload | 13 | from .lowering_common import add_overload |
| 14 | from .ascend_npu_ir.ascend_npu_ir import config as anir_config | 14 | from .ascend_npu_ir.ascend_npu_ir import config as anir_config |
| 15 | from .ascend_npu_ir.ascend_npu_ir.npu.utils import run_once | 15 | from .ascend_npu_ir.ascend_npu_ir.npu.utils import run_once |