已合并
Fix MLIR full lowering fallback mode #39322
yvjc创建于 6月25日
Fix MLIR full lowering fallback mode #39322
已合并
从已删除 :fusionTest-v2.10.0合入到Ascend/pytorchv2.10.0
共 3 个文件变更+44-2
| @@ -0,0 +1,41 @@ | |||
| 1 | +import os | ||
| 2 | + | ||
| 3 | +os.environ["NPU_INDUCTOR_FALLBACK_LIST"] = "allfallback" | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +from torch.testing._internal.common_utils import run_tests | ||
| 7 | +from torch._inductor.utils import run_and_get_code | ||
| 8 | + | ||
| 9 | +from testutils import TestUtils | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class TestAllFallback(TestUtils): | ||
| 13 | + | ||
| 14 | + def add_op(self, x, y): | ||
| 15 | + return x + y + 1 | ||
| 16 | + | ||
| 17 | + def test_all_fallback_detection_mlir(self): | ||
| 18 | + | ||
| 19 | + compiled_add = torch.compile(self.add_op, backend="inductor", options={"npu_backend": "mlir"}) | ||
| 20 | + | ||
| 21 | + x = torch.randn(4, 4, dtype=torch.float32).to("npu") | ||
| 22 | + y = torch.randn(4, 4, dtype=torch.float32).to("npu") | ||
| 23 | + | ||
| 24 | + _, codes = run_and_get_code(compiled_add, x, y) | ||
| 25 | + | ||
| 26 | + self.assertTrue('mlir_fused_' not in codes[0]) | ||
| 27 | + | ||
| 28 | + def test_all_fallback_detection_triton(self): | ||
| 29 | + | ||
| 30 | + compiled_add = torch.compile(self.add_op, backend="inductor") | ||
| 31 | + | ||
| 32 | + x = torch.randn(4, 4, dtype=torch.float32).to("npu") | ||
| 33 | + y = torch.randn(4, 4, dtype=torch.float32).to("npu") | ||
| 34 | + | ||
| 35 | + _, codes = run_and_get_code(compiled_add, x, y) | ||
| 36 | + | ||
| 37 | + self.assertTrue('_fused_' not in codes[0]) | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +if __name__ == "__main__": | ||
| 41 | + run_tests() | ||
| @@ -167,10 +167,11 @@ def _get_compile_mode(): | |||
| 167 | block_dim = 48 | 167 | block_dim = 48 |
| 168 | 168 | ||
| 169 | """ | 169 | """ |
| 170 | -support {"off", "include", "exclude"}, to | 170 | +support {"off", "include", "exclude", "all"}, to |
| 171 | "off": No fallback at all. | 171 | "off": No fallback at all. |
| 172 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. | 172 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. |
| 173 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. | 173 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. |
| 174 | +"all": At compile-time, all Aten IR entering the NPU inductor lowering path will fall back to aten. | ||
| 174 | """ | 175 | """ |
| 175 | fallback_to_aten_mode: str = "exclude" | 176 | fallback_to_aten_mode: str = "exclude" |
| 176 | 177 | ||
| @@ -28,7 +28,7 @@ def _register_npu_inductor_fallbacks(): | |||
| 28 | fallback_set_exclude = OrderedSet() | 28 | fallback_set_exclude = OrderedSet() |
| 29 | env_fallback_list = config.enable_full_lowering_fallback | 29 | env_fallback_list = config.enable_full_lowering_fallback |
| 30 | 30 | ||
| 31 | - if env_fallback_list: | 31 | + if env_fallback_list and config.fallback_to_aten_mode != "all": |
| 32 | for op_name in env_fallback_list.split(","): | 32 | for op_name in env_fallback_list.split(","): |
| 33 | op_name = op_name.strip() | 33 | op_name = op_name.strip() |
| 34 | op = resolve_op_from_name(op_name, logger) | 34 | op = resolve_op_from_name(op_name, logger) |