已合并
Fix MLIR full lowering fallback mode #39298
yvjc创建于 6月25日
Fix MLIR full lowering fallback mode #39298
已合并
从已删除 :bugfix-v2.7.1合入到Ascend/pytorchv2.7.1
共 3 个文件变更+49-7
| @@ -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() | ||
| @@ -160,10 +160,11 @@ def _get_compile_mode(): | |||
| 160 | block_dim = 48 | 160 | block_dim = 48 |
| 161 | 161 | ||
| 162 | """ | 162 | """ |
| 163 | -support {"off", "include", "exclude"}, to | 163 | +support {"off", "include", "exclude", "all"}, to |
| 164 | "off": No fallback at all. | 164 | "off": No fallback at all. |
| 165 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. | 165 | "include": At compile-time, Aten IR included in FALLBACK_LIST will fall back to aten. |
| 166 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. | 166 | "exclude": At compile-time, Aten IR excluded from GENERATE_LIST will fall back to aten. |
| 167 | +"all": At compile-time, all Aten IR entering the NPU inductor lowering path will fall back to aten. | ||
| 167 | """ | 168 | """ |
| 168 | fallback_to_aten_mode: str = "exclude" | 169 | fallback_to_aten_mode: str = "exclude" |
| 169 | 170 | ||
| @@ -311,4 +312,4 @@ decomps_to_exclude_npu = [ | |||
| 311 | aten.reflection_pad2d, | 312 | aten.reflection_pad2d, |
| 312 | aten.grid_sampler_2d, | 313 | aten.grid_sampler_2d, |
| 313 | aten.grid_sampler_2d_backward, | 314 | aten.grid_sampler_2d_backward, |
| 314 | -] | 315 | +] |
| @@ -24,7 +24,10 @@ def _register_npu_inductor_fallbacks(): | |||
| 24 | fallback_set_exclude = set() | 24 | fallback_set_exclude = set() |
| 25 | env_fallback_list = config.enable_full_lowering_fallback | 25 | env_fallback_list = config.enable_full_lowering_fallback |
| 26 | 26 | ||
| 27 | - if env_fallback_list: | 27 | + if config.fallback_to_aten_mode not in {"off", "include", "exclude", "all"}: |
| 28 | + raise AssertionError(f"Error! Unsupported fallback_to_aten_mode: {config.fallback_to_aten_mode} was set!") | ||
| 29 | + | ||
| 30 | + if env_fallback_list and config.fallback_to_aten_mode != 'all': | ||
| 28 | for op_name in env_fallback_list.split(','): | 31 | for op_name in env_fallback_list.split(','): |
| 29 | op_name = op_name.strip() | 32 | op_name = op_name.strip() |
| 30 | op = resolve_op_from_name(op_name, logger) | 33 | op = resolve_op_from_name(op_name, logger) |
| @@ -52,10 +55,7 @@ def _register_npu_inductor_fallbacks(): | |||
| 52 | if isinstance(op, torch._ops.OpOverloadPacket) or \ | 55 | if isinstance(op, torch._ops.OpOverloadPacket) or \ |
| 53 | isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): | 56 | isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): |
| 54 | make_fallback(op) | 57 | make_fallback(op) |
| 55 | - | 58 | + |
| 56 | - if config.fallback_to_aten_mode not in {"off", "include", "exclude"}: | ||
| 57 | - raise AssertionError(f"Error! Unsupported fallback_to_aten_mode: {config.fallback_to_aten_mode} was set!") | ||
| 58 | - | ||
| 59 | if get_anir_mode() == 'O0': | 59 | if get_anir_mode() == 'O0': |
| 60 | fallback_except_gen_set(gen_set=[]) | 60 | fallback_except_gen_set(gen_set=[]) |
| 61 | decomposition.decompositions.clear() | 61 | decomposition.decompositions.clear() |
🔵 Low Priority
新增测试文件
test_inductor_all_fallback_mlir.py中包含两个测试：test_all_fallback_detection_triton（第 17 行）和test_all_fallback_detection_mlir（第 28 行）。文件名以_mlir结尾，暗示仅包含 MLIR 相关测试，但实际上也包含了 triton 后端的测试。现有代码库中已有分离的test_inductor_fallback_list_mlir.py和test_inductor_fallback_list_triton.py，本文件的命名与内容不完全一致，可能让维护者困惑（例如在排查 MLIR 测试失败时发现 triton 测试也在此文件中）。建议：可考虑将文件重命名为
test_inductor_all_fallback.py（去掉_mlir后缀）以准确反映其双后端测试覆盖范围，或将 triton 测试和 MLIR 测试分到各自文件中。此建议纯属可选改进。