已合并
fix(mfusion): preserve output container kind in contract restore #38664
shaoshengqi创建于 6月16日
fix(mfusion): preserve output container kind in contract restore #38664
已合并
共 2 个文件变更+83-7
| @@ -32,12 +32,12 @@ try: | |||
| 32 | ) | 32 | ) |
| 33 | from torch_npu._inductor.mfusion.subgraph_registry import Payload | 33 | from torch_npu._inductor.mfusion.subgraph_registry import Payload |
| 34 | except ImportError: | 34 | except ImportError: |
| 35 | - _emit_mfusion_dvm_codegen = None # type: ignore[misc, assignment] | 35 | + _emit_mfusion_dvm_codegen = None |
| 36 | - _find_output_node = None # type: ignore[misc, assignment] | 36 | + _find_output_node = None |
| 37 | - _restore_output_contract = None # type: ignore[misc, assignment] | 37 | + _restore_output_contract = None |
| 38 | - _snapshot_output_contract = None # type: ignore[misc, assignment] | 38 | + _snapshot_output_contract = None |
| 39 | - MFusionPatch = None # type: ignore[misc, assignment] | 39 | + MFusionPatch = None |
| 40 | - Payload = None # type: ignore[misc, assignment] | 40 | + Payload = None |
| 41 | 41 | ||
| 42 | # Gate omits AKG mfusion; keep module importable even if torch_npu mfusion tree is absent. | 42 | # Gate omits AKG mfusion; keep module importable even if torch_npu mfusion tree is absent. |
| 43 | HAS_MFUSION_STACK = HAS_AKG_MFUSION and MFusionPatch is not None | 43 | HAS_MFUSION_STACK = HAS_AKG_MFUSION and MFusionPatch is not None |
| @@ -96,6 +96,56 @@ def _make_add_mul_subgraph(): | |||
| 96 | "torch_npu._inductor.mfusion.graph_fusion not available", | 96 | "torch_npu._inductor.mfusion.graph_fusion not available", |
| 97 | ) | 97 | ) |
| 98 | class TestMfusionOutputContract(TestCase): | 98 | class TestMfusionOutputContract(TestCase): |
| 99 | + def test_restore_output_contract_preserves_single_output(self): | ||
| 100 | + graph = Graph() | ||
| 101 | + a = graph.placeholder("a") | ||
| 102 | + b = graph.placeholder("b") | ||
| 103 | + add = graph.call_function(torch.ops.aten.add.Tensor, (a, b)) | ||
| 104 | + graph.output(add) | ||
| 105 | + gm = GraphModule(torch.nn.Module(), graph) | ||
| 106 | + | ||
| 107 | + snapshot = _snapshot_output_contract(_find_output_node(gm.graph)) | ||
| 108 | + | ||
| 109 | + roundtrip_graph = Graph() | ||
| 110 | + rt_a = roundtrip_graph.placeholder("a") | ||
| 111 | + rt_b = roundtrip_graph.placeholder("b") | ||
| 112 | + rt_add = roundtrip_graph.call_function(torch.ops.aten.add.Tensor, (rt_a, rt_b)) | ||
| 113 | + roundtrip_graph.output(rt_add) | ||
| 114 | + rt_gm = GraphModule(torch.nn.Module(), roundtrip_graph) | ||
| 115 | + | ||
| 116 | + _restore_output_contract(_find_output_node(rt_gm.graph), snapshot) | ||
| 117 | + | ||
| 118 | + restored_output = _find_output_node(rt_gm.graph) | ||
| 119 | + self.assertIsNotNone(restored_output) | ||
| 120 | + self.assertIs(restored_output.args[0], rt_add) | ||
| 121 | + | ||
| 122 | + def test_restore_output_contract_preserves_list_output(self): | ||
| 123 | + graph = Graph() | ||
| 124 | + a = graph.placeholder("a") | ||
| 125 | + b = graph.placeholder("b") | ||
| 126 | + add = graph.call_function(torch.ops.aten.add.Tensor, (a, b)) | ||
| 127 | + graph.output([add, b]) | ||
| 128 | + gm = GraphModule(torch.nn.Module(), graph) | ||
| 129 | + | ||
| 130 | + snapshot = _snapshot_output_contract(_find_output_node(gm.graph)) | ||
| 131 | + | ||
| 132 | + roundtrip_graph = Graph() | ||
| 133 | + rt_a = roundtrip_graph.placeholder("a") | ||
| 134 | + rt_b = roundtrip_graph.placeholder("b") | ||
| 135 | + rt_add = roundtrip_graph.call_function(torch.ops.aten.add.Tensor, (rt_a, rt_b)) | ||
| 136 | + roundtrip_graph.output((rt_add, rt_b)) | ||
| 137 | + rt_gm = GraphModule(torch.nn.Module(), roundtrip_graph) | ||
| 138 | + | ||
| 139 | + _restore_output_contract(_find_output_node(rt_gm.graph), snapshot) | ||
| 140 | + | ||
| 141 | + restored_output = _find_output_node(rt_gm.graph) | ||
| 142 | + self.assertIsNotNone(restored_output) | ||
| 143 | + restored_args = restored_output.args[0] | ||
| 144 | + self.assertIsInstance(restored_args, list) | ||
| 145 | + self.assertEqual(len(restored_args), 2) | ||
| 146 | + self.assertIs(restored_args[0], rt_add) | ||
| 147 | + self.assertIs(restored_args[1], rt_b) | ||
| 148 | + | ||
| 99 | def test_restore_output_contract_preserves_none_slots_and_stride_meta(self): | 149 | def test_restore_output_contract_preserves_none_slots_and_stride_meta(self): |
| 100 | graph = Graph() | 150 | graph = Graph() |
| 101 | a = graph.placeholder("a") | 151 | a = graph.placeholder("a") |
| @@ -388,6 +388,15 @@ def _output_arg_list(output_node: Node) -> list[Any]: | |||
| 388 | return list(args) if isinstance(args, (list, tuple)) else [args] | 388 | return list(args) if isinstance(args, (list, tuple)) else [args] |
| 389 | 389 | ||
| 390 | 390 | ||
| 391 | +def _output_container_kind(output_node: Node) -> str: | ||
| 392 | + args = output_node.args[0] | ||
| 393 | + if isinstance(args, tuple): | ||
| 394 | + return "tuple" | ||
| 395 | + if isinstance(args, list): | ||
| 396 | + return "list" | ||
| 397 | + return "single" | ||
| 398 | + | ||
| 399 | + | ||
| 391 | def _copy_output_meta_value(value: Any) -> Any: | 400 | def _copy_output_meta_value(value: Any) -> Any: |
| 392 | if isinstance(value, list): | 401 | if isinstance(value, list): |
| 393 | return list(value) | 402 | return list(value) |
| @@ -404,6 +413,7 @@ def _snapshot_output_contract(output_node: Node | None) -> dict[str, Any]: | |||
| 404 | 413 | ||
| 405 | output_args = _output_arg_list(output_node) | 414 | output_args = _output_arg_list(output_node) |
| 406 | return { | 415 | return { |
| 416 | + "output_container": _output_container_kind(output_node), | ||
| 407 | "output_len": len(output_args), | 417 | "output_len": len(output_args), |
| 408 | "none_positions": [i for i, arg in enumerate(output_args) if arg is None], | 418 | "none_positions": [i for i, arg in enumerate(output_args) if arg is None], |
| 409 | "valid_positions": [i for i, arg in enumerate(output_args) if arg is not None], | 419 | "valid_positions": [i for i, arg in enumerate(output_args) if arg is not None], |
| @@ -470,7 +480,23 @@ def _restore_output_contract(output_node: Node | None, snapshot: dict[str, Any]) | |||
| 470 | original_output_strides | 480 | original_output_strides |
| 471 | ) | 481 | ) |
| 472 | 482 | ||
| 473 | - output_node.args = (tuple(restored),) | 483 | + output_container = snapshot.get("output_container", "tuple") |
| 484 | + if output_container == "tuple": | ||
| 485 | + output_node.args = (tuple(restored),) | ||
| 486 | + elif output_container == "list": | ||
| 487 | + output_node.args = (list(restored),) | ||
| 488 | + elif output_container == "single": | ||
| 489 | + if len(restored) != 1: | ||
| 490 | + raise RuntimeError( | ||
| 491 | + "mfusion output ABI mismatch while restoring single output: " | ||
| 492 | + f"restored_len={len(restored)}, original_len={output_len}" | ||
| 493 | + ) | ||
| 494 | + output_node.args = (restored[0],) | ||
| 495 | + else: | ||
| 496 | + raise RuntimeError( | ||
| 497 | + "mfusion output ABI has unsupported output container: " | ||
| 498 | + f"{output_container!r}" | ||
| 499 | + ) | ||
| 474 | 500 | ||
| 475 | 501 | ||
| 476 | def _layout_get_size(layout): | 502 | def _layout_get_size(layout): |