已合并
test: add MLIR non-fallback kernel regression test #38540
yvjc创建于 6月15日
test: add MLIR non-fallback kernel regression test #38540
已合并
从已删除 :fusionTest-v2.12.0合入到Ascend/pytorchv2.12.0
共 2 个文件变更+90-0
| @@ -0,0 +1,45 @@ | |||
| 1 | +import os | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" | ||
| 5 | +os.environ["TORCHINDUCTOR_USE_AKG"] = "1" | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 9 | +from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.mlir_compiler import AkgCompiler | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class TestMlirKernelNotFallback(TestCase): | ||
| 13 | + def test_basic_op_uses_non_fallback_kernel(self): | ||
| 14 | + def fusion_func(x, y): | ||
| 15 | + return x + y + 1 | ||
| 16 | + | ||
| 17 | + x_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 18 | + y_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 19 | + | ||
| 20 | + expected = fusion_func(x_cpu, y_cpu) | ||
| 21 | + | ||
| 22 | + x_npu = x_cpu.npu() | ||
| 23 | + y_npu = y_cpu.npu() | ||
| 24 | + | ||
| 25 | + has_fallback_kernel = False | ||
| 26 | + original_run = AkgCompiler.run | ||
| 27 | + | ||
| 28 | + def run_and_check_fallback(compiler, *args, **kwargs): | ||
| 29 | + nonlocal has_fallback_kernel | ||
| 30 | + result = original_run(compiler, *args, **kwargs) | ||
| 31 | + launcher_idx = compiler.get_primary_launcher_index() | ||
| 32 | + has_fallback_kernel |= compiler.is_fallback_kernels[launcher_idx] | ||
| 33 | + return result | ||
| 34 | + | ||
| 35 | + compiled = torch.compile(fusion_func, backend="inductor") | ||
| 36 | + | ||
| 37 | + with patch.object(AkgCompiler, "run", run_and_check_fallback): | ||
| 38 | + actual = compiled(x_npu, y_npu) | ||
| 39 | + | ||
| 40 | + self.assertFalse(has_fallback_kernel) | ||
| 41 | + self.assertEqual(expected, actual.cpu(), atol=1e-3, rtol=1e-3) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | +if __name__ == "__main__": | ||
| 45 | + run_tests() | ||
| @@ -0,0 +1,45 @@ | |||
| 1 | +import os | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" | ||
| 5 | + | ||
| 6 | +import torch | ||
| 7 | +import torch_npu | ||
| 8 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 9 | +from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.mlir_compiler import NpuMlirCompiler | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class TestMlirKernelNotFallback(TestCase): | ||
| 13 | + def test_basic_op_uses_non_fallback_kernel(self): | ||
| 14 | + def fusion_func(x, y): | ||
| 15 | + return x + y + 1 | ||
| 16 | + | ||
| 17 | + x_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 18 | + y_cpu = torch.randn((64, 128), dtype=torch.float32) | ||
| 19 | + | ||
| 20 | + expected = fusion_func(x_cpu, y_cpu) | ||
| 21 | + | ||
| 22 | + x_npu = x_cpu.npu() | ||
| 23 | + y_npu = y_cpu.npu() | ||
| 24 | + | ||
| 25 | + has_fallback_kernel = False | ||
| 26 | + original_run = NpuMlirCompiler.run | ||
| 27 | + | ||
| 28 | + def run_and_check_fallback(compiler, *args, **kwargs): | ||
| 29 | + nonlocal has_fallback_kernel | ||
| 30 | + result = original_run(compiler, *args, **kwargs) | ||
| 31 | + launcher_idx = compiler.get_primary_launcher_index() | ||
| 32 | + has_fallback_kernel |= compiler.is_fallback_kernels[launcher_idx] | ||
| 33 | + return result | ||
| 34 | + | ||
| 35 | + compiled = torch.compile(fusion_func, backend="inductor") | ||
| 36 | + | ||
| 37 | + with patch.object(NpuMlirCompiler, "run", run_and_check_fallback): | ||
| 38 | + actual = compiled(x_npu, y_npu) | ||
| 39 | + | ||
| 40 | + self.assertFalse(has_fallback_kernel) | ||
| 41 | + self.assertEqual(expected, actual.cpu(), atol=1e-3, rtol=1e-3) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | +if __name__ == "__main__": | ||
| 45 | + run_tests() | ||