已合并
fix(dvm): guard stride patch behind is_fx_dynamic in fallback kernel. #37883
Margaret_wangrui创建于 6月8日
fix(dvm): guard stride patch behind is_fx_dynamic in fallback kernel. #37883
已合并
共 2 个文件变更+88-4
| @@ -1,4 +1,5 @@ | |||
| 1 | -from unittest import skip | 1 | +from unittest import mock |
| 2 | + | ||
| 2 | import torch | 3 | import torch |
| 3 | 4 | ||
| 4 | from torch.testing._internal.common_utils import TestCase | 5 | from torch.testing._internal.common_utils import TestCase |
| @@ -7,7 +8,12 @@ from torch.testing._internal.common_utils import ( | |||
| 7 | parametrize, | 8 | parametrize, |
| 8 | instantiate_parametrized_tests, | 9 | instantiate_parametrized_tests, |
| 9 | ) | 10 | ) |
| 10 | -from torch_npu._inductor.dvm.graph_fusion import DvmGraphFusionPatch | 11 | +from torch_npu._inductor.dvm.graph_fusion import ( |
| 12 | + DvmGraphFusionPatch, | ||
| 13 | + _dvm_generate_fallback_kernel, | ||
| 14 | + _fused_metas, | ||
| 15 | +) | ||
| 16 | +from torch_npu._inductor.dvm.graph_build import is_fx_dynamic | ||
| 11 | 17 | ||
| 12 | 18 | ||
| 13 | class TestModule(torch.nn.Module): | 19 | class TestModule(torch.nn.Module): |
| @@ -88,5 +94,82 @@ class TestDvmByGraphFusion(TestCase): | |||
| 88 | 94 | ||
| 89 | instantiate_parametrized_tests(TestDvmByGraphFusion) | 95 | instantiate_parametrized_tests(TestDvmByGraphFusion) |
| 90 | 96 | ||
| 97 | + | ||
| 98 | +class _AddModule(torch.nn.Module): | ||
| 99 | + def forward(self, x): | ||
| 100 | + return x + 1 | ||
| 101 | + | ||
| 102 | + | ||
| 103 | +class TestDvmFallbackStridePatchGuard(TestCase): | ||
| 104 | + """Guard adb97b9: skip stride patch in dynamic fused subgraph codegen.""" | ||
| 105 | + | ||
| 106 | + def _make_fallback_kernel(self, fused_id=0): | ||
| 107 | + fallback_kernel = mock.MagicMock() | ||
| 108 | + fallback_kernel.op_overload._name = "dvm::fused_graph_0" | ||
| 109 | + fallback_kernel.codegen_args.return_value = ["buf0", fused_id] | ||
| 110 | + fallback_kernel.codegen_kwargs.return_value = [] | ||
| 111 | + fallback_kernel.get_name.return_value = "buf_out" | ||
| 112 | + return fallback_kernel | ||
| 113 | + | ||
| 114 | + def _make_codegen_wrapper(self): | ||
| 115 | + wrapper = mock.MagicMock() | ||
| 116 | + wrapper.header = mock.MagicMock() | ||
| 117 | + return wrapper | ||
| 118 | + | ||
| 119 | + def _make_fused_meta(self): | ||
| 120 | + meta = mock.MagicMock() | ||
| 121 | + meta.gm = mock.MagicMock() | ||
| 122 | + meta.name = "dvm_graph_fused_0" | ||
| 123 | + codegen = mock.MagicMock() | ||
| 124 | + codegen.cont_flag_input = [True] | ||
| 125 | + codegen.need_trans_input = [False] | ||
| 126 | + meta.codegen.return_value = (codegen, "# dvm kernel\n") | ||
| 127 | + return meta | ||
| 128 | + | ||
| 129 | + | ||
| 130 | + "torch_npu._inductor.dvm.graph_fusion.patch_gm_placeholder_strides_from_codegen_args" | ||
| 131 | + ) | ||
| 132 | + def test_fallback_kernel_stride_patch_guarded_by_is_fx_dynamic(self, mock_patch): | ||
| 133 | + gm_static = torch.fx.symbolic_trace(_AddModule()) | ||
| 134 | + placeholder = next(n for n in gm_static.graph.nodes if n.op == "placeholder") | ||
| 135 | + placeholder.meta["val"] = torch.randn(2, 3) | ||
| 136 | + self.assertFalse(is_fx_dynamic(gm_static)) | ||
| 137 | + | ||
| 138 | + batch = torch.export.Dim("batch", min=1, max=1024) | ||
| 139 | + exported = torch.export.export( | ||
| 140 | + _AddModule(), | ||
| 141 | + (torch.randn(2, 3),), | ||
| 142 | + dynamic_shapes={"x": {0: batch}}, | ||
| 143 | + ) | ||
| 144 | + self.assertTrue(is_fx_dynamic(exported.graph_module)) | ||
| 145 | + | ||
| 146 | + meta = self._make_fused_meta() | ||
| 147 | + try: | ||
| 148 | + _fused_metas[0] = meta | ||
| 149 | + with mock.patch( | ||
| 150 | + "torch_npu._inductor.dvm.graph_fusion.is_fx_dynamic", return_value=True | ||
| 151 | + ): | ||
| 152 | + _dvm_generate_fallback_kernel( | ||
| 153 | + self._make_codegen_wrapper(), | ||
| 154 | + self._make_fallback_kernel(), | ||
| 155 | + ) | ||
| 156 | + mock_patch.assert_not_called() | ||
| 157 | + | ||
| 158 | + mock_patch.reset_mock() | ||
| 159 | + # _dvm_generate_fallback_kernel pops fused_id from _fused_metas; re-seed | ||
| 160 | + # before exercising the static-shape branch in the same test. | ||
| 161 | + _fused_metas[0] = meta | ||
| 162 | + with mock.patch( | ||
| 163 | + "torch_npu._inductor.dvm.graph_fusion.is_fx_dynamic", return_value=False | ||
| 164 | + ): | ||
| 165 | + _dvm_generate_fallback_kernel( | ||
| 166 | + self._make_codegen_wrapper(), | ||
| 167 | + self._make_fallback_kernel(), | ||
| 168 | + ) | ||
| 169 | + mock_patch.assert_called_once_with(meta.gm, ["buf0"]) | ||
| 170 | + finally: | ||
| 171 | + _fused_metas.pop(0, None) | ||
| 172 | + | ||
| 173 | + | ||
| 91 | if __name__ == "__main__": | 174 | if __name__ == "__main__": |
| 92 | run_tests() | 175 | run_tests() |
| @@ -20,7 +20,7 @@ from torch.fx.passes.utils.fuser_utils import ( | |||
| 20 | erase_nodes, | 20 | erase_nodes, |
| 21 | ) | 21 | ) |
| 22 | 22 | ||
| 23 | -from .graph_build import DvmCodegenInterpreter | 23 | +from .graph_build import DvmCodegenInterpreter, is_fx_dynamic |
| 24 | from .util import patch_gm_placeholder_strides_from_codegen_args | 24 | from .util import patch_gm_placeholder_strides_from_codegen_args |
| 25 | from .fx_test import generate_dvm_fx_case | 25 | from .fx_test import generate_dvm_fx_case |
| 26 | from .op_emitter import DVM_OP_REGISTRY | 26 | from .op_emitter import DVM_OP_REGISTRY |
| @@ -376,7 +376,8 @@ def _dvm_generate_fallback_kernel(self, fallback_kernel): | |||
| 376 | meta = _fused_metas.pop(fused_id) | 376 | meta = _fused_metas.pop(fused_id) |
| 377 | 377 | ||
| 378 | args_list = list(args[:-1]) | 378 | args_list = list(args[:-1]) |
| 379 | - patch_gm_placeholder_strides_from_codegen_args(meta.gm, args_list) | 379 | + if not is_fx_dynamic(meta.gm): |
| 380 | + patch_gm_placeholder_strides_from_codegen_args(meta.gm, args_list) | ||
| 380 | cg, code = meta.codegen() | 381 | cg, code = meta.codegen() |
| 381 | self.header.splice(code) | 382 | self.header.splice(code) |
| 382 | 383 | ||