已合并
[fix]mlir_enable #31966
cuiduo创建于 3月18日
[fix]mlir_enable #31966
已合并
共 3 个文件变更+54-0
| @@ -0,0 +1,43 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 3 | +from testutils import TestUtils | ||
| 4 | +from torch._inductor.utils import run_and_get_code | ||
| 5 | +import torch_npu | ||
| 6 | +import torch_npu._inductor | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestAdd(TestUtils): | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + def op_calc(self, first_element, second_element): | ||
| 13 | + result = first_element + second_element | ||
| 14 | + return result | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + def test_options_environ_cases(self, shape, dtype): | ||
| 19 | + x = self._generate_tensor(shape, dtype) | ||
| 20 | + y = self._generate_tensor(shape, dtype) | ||
| 21 | + std_out = self.op_calc(x, y) | ||
| 22 | + compile_func = torch.compile(self.op_calc, options={"npu_backend": "mlir"}) | ||
| 23 | + compile_out, codes = run_and_get_code(compile_func, x, y) | ||
| 24 | + self.assertEqual(std_out, compile_out) | ||
| 25 | + self.assertTrue("mlir_fused" in codes[0]) | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + def test_config_environ_cases(self, shape, dtype): | ||
| 30 | + torch._inductor.config.npu_backend = "mlir" | ||
| 31 | + x = self._generate_tensor(shape, dtype) | ||
| 32 | + y = self._generate_tensor(shape, dtype) | ||
| 33 | + std_out = self.op_calc(x, y) | ||
| 34 | + compile_func = torch.compile(self.op_calc) | ||
| 35 | + compile_out, codes = run_and_get_code(compile_func, x, y) | ||
| 36 | + self.assertEqual(std_out, compile_out) | ||
| 37 | + self.assertTrue("mlir_fused" in codes[0]) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +instantiate_parametrized_tests(TestAdd) | ||
| 41 | + | ||
| 42 | +if __name__ == "__main__": | ||
| 43 | + run_tests() | ||
| @@ -21,6 +21,7 @@ else: | |||
| 21 | import torch | 21 | import torch |
| 22 | from torch._dynamo.device_interface import register_interface_for_device, get_interface_for_device | 22 | from torch._dynamo.device_interface import register_interface_for_device, get_interface_for_device |
| 23 | from torch._inductor import lowering as inductor_lowering | 23 | from torch._inductor import lowering as inductor_lowering |
| 24 | + from torch._inductor.lowering import make_fallback as ori_make_fallback | ||
| 24 | from torch._inductor.choices import InductorChoices | 25 | from torch._inductor.choices import InductorChoices |
| 25 | from torch._inductor.codegen.common import register_backend_for_device, register_device_op_overrides | 26 | from torch._inductor.codegen.common import register_backend_for_device, register_device_op_overrides |
| 26 | from torch._inductor.runtime import autotune_cache | 27 | from torch._inductor.runtime import autotune_cache |
| @@ -182,6 +182,16 @@ def patch_inductor_wrapper(): | |||
| 182 | os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir' | 182 | os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir' |
| 183 | device_id = torch_npu.npu.current_device() | 183 | device_id = torch_npu.npu.current_device() |
| 184 | torch_npu._C._recovery_all_npu_stream(device_id) | 184 | torch_npu._C._recovery_all_npu_stream(device_id) |
| 185 | + try: | ||
| 186 | + import torch_mlir | ||
| 187 | + from torch_mlir import ir | ||
| 188 | + except ImportError as e: | ||
| 189 | + raise ImportError("torch_mlir is not installed, install it first.") from e | ||
| 190 | + from torch_npu._inductor import ori_make_fallback | ||
| 191 | + torch._inductor.lowering.make_fallback = ori_make_fallback | ||
| 192 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin | ||
| 193 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu import torch_mlir_patch | ||
| 194 | + | ||
| 185 | elif self.config.get("npu_backend") == "dvm" or torch._inductor.config.npu_backend == "dvm": | 195 | elif self.config.get("npu_backend") == "dvm" or torch._inductor.config.npu_backend == "dvm": |
| 186 | os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'dvm' | 196 | os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'dvm' |
| 187 | _TorchCompileInductorWrapper.__call__ = new_call | 197 | _TorchCompileInductorWrapper.__call__ = new_call |