已合并
Fix MLIR full lowering fallback mode #39326
yvjc创建于 6月25日
Fix MLIR full lowering fallback mode #39326
已合并
从已删除 :fusionTest-master合入到Ascend/pytorchmaster
共 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() | ||
| @@ -170,10 +170,11 @@ def _get_compile_mode(): | |||
| 170 | block_dim = 48 | 170 | block_dim = 48 |
| 171 | 171 | ||
| 172 | """ | 172 | """ |
| 173 | -support {"off", "include", "exclude"}, to | 173 | +support {"off", "include", "exclude", "all"}, to |
| 174 | "off": No fallback at all. | 174 | "off": No fallback at all. |
| 175 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. | 175 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. |
| 176 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. | 176 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. |
| 177 | +"all": At compile-time, all Aten IR entering the NPU inductor lowering path will fall back to aten. | ||
| 177 | """ | 178 | """ |
| 178 | fallback_to_aten_mode: str = "exclude" | 179 | fallback_to_aten_mode: str = "exclude" |
| 179 | 180 | ||
| @@ -27,7 +27,7 @@ def _register_npu_inductor_fallbacks(): | |||
| 27 | fallback_set_exclude = OrderedSet() | 27 | fallback_set_exclude = OrderedSet() |
| 28 | env_fallback_list = config.enable_full_lowering_fallback | 28 | env_fallback_list = config.enable_full_lowering_fallback |
| 29 | 29 | ||
| 30 | - if env_fallback_list: | 30 | + if env_fallback_list and config.fallback_to_aten_mode != "all": |
| 31 | for op_name in env_fallback_list.split(","): | 31 | for op_name in env_fallback_list.split(","): |
| 32 | op_name = op_name.strip() | 32 | op_name = op_name.strip() |
| 33 | op = resolve_op_from_name(op_name, logger) | 33 | op = resolve_op_from_name(op_name, logger) |