已合并
Fix mfusion FX-to-torch-mlir conversion for torch.ops.higher_order.run_and_save_rng_state #38582
shuqian0创建于 6月15日
Fix mfusion FX-to-torch-mlir conversion for torch.ops.higher_order.run_and_save_rng_state #38582
已合并
共 4 个文件变更+181-8
| @@ -15,6 +15,7 @@ | |||
| 15 | 15 | ||
| 16 | from unittest import skipUnless | 16 | from unittest import skipUnless |
| 17 | 17 | ||
| 18 | +import operator | ||
| 18 | import torch | 19 | import torch |
| 19 | import torch.nn as nn | 20 | import torch.nn as nn |
| 20 | from torch.func import functionalize | 21 | from torch.func import functionalize |
| @@ -57,6 +58,10 @@ class TestFusedModule(nn.Module): | |||
| 57 | return b | 58 | return b |
| 58 | 59 | ||
| 59 | 60 | ||
| 61 | +def opaque_nested_tuple_target(x): | ||
| 62 | + raise RuntimeError("test-only opaque target should not execute") | ||
| 63 | + | ||
| 64 | + | ||
| 60 | class WrapperModule(torch.nn.Module): | 65 | class WrapperModule(torch.nn.Module): |
| 61 | def __init__(self, gm): | 66 | def __init__(self, gm): |
| 62 | super().__init__() | 67 | super().__init__() |
| @@ -203,6 +208,70 @@ class TestFxRoundtrip(TestCase): | |||
| 203 | self.assertIn("func.func @main", mlir_text) | 208 | self.assertIn("func.func @main", mlir_text) |
| 204 | self.assertIn("torch.aten.mul", mlir_text) | 209 | self.assertIn("torch.aten.mul", mlir_text) |
| 205 | 210 | ||
| 211 | + def test_opaque_run_and_save_rng_state_roundtrip(self): | ||
| 212 | + from torch._prims.rng_prims import run_and_save_rng_state | ||
| 213 | + | ||
| 214 | + def f(x): | ||
| 215 | + return run_and_save_rng_state(torch.ops.aten.rand_like.default, x) | ||
| 216 | + | ||
| 217 | + gm = make_fx(f)(torch.empty(2, 3)) | ||
| 218 | + mlir_module = import_mlir_module_from_fx(gm) | ||
| 219 | + self.assertIn("torch.mfusion_opaque", str(mlir_module)) | ||
| 220 | + | ||
| 221 | + new_gm = export_mlir_module_to_fx(mlir_module) | ||
| 222 | + self.assertTrue( | ||
| 223 | + any( | ||
| 224 | + node.op == "call_function" | ||
| 225 | + and node.target is run_and_save_rng_state | ||
| 226 | + for node in new_gm.graph.nodes | ||
| 227 | + ) | ||
| 228 | + ) | ||
| 229 | + | ||
| 230 | + def test_opaque_nested_tuple_result_roundtrip(self): | ||
| 231 | + from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 232 | + | ||
| 233 | + mode = FakeTensorMode() | ||
| 234 | + with mode: | ||
| 235 | + fake = mode.from_tensor(torch.empty(2, 3)) | ||
| 236 | + | ||
| 237 | + graph = torch.fx.Graph() | ||
| 238 | + x = graph.placeholder("x") | ||
| 239 | + x.meta["val"] = fake | ||
| 240 | + root = graph.call_function(opaque_nested_tuple_target, (x,)) | ||
| 241 | + root.meta["val"] = (fake, (fake, fake)) | ||
| 242 | + first = graph.call_function(operator.getitem, (root, 0)) | ||
| 243 | + first.meta["val"] = fake | ||
| 244 | + second = graph.call_function(operator.getitem, (root, 1)) | ||
| 245 | + second.meta["val"] = (fake, fake) | ||
| 246 | + nested0 = graph.call_function(operator.getitem, (second, 0)) | ||
| 247 | + nested0.meta["val"] = fake | ||
| 248 | + nested1 = graph.call_function(operator.getitem, (second, 1)) | ||
| 249 | + nested1.meta["val"] = fake | ||
| 250 | + graph.output((first, nested0, nested1)) | ||
| 251 | + gm = torch.fx.GraphModule(torch.nn.Module(), graph) | ||
| 252 | + | ||
| 253 | + mlir_module = import_mlir_module_from_fx(gm) | ||
| 254 | + self.assertIn("torch.mfusion_opaque", str(mlir_module)) | ||
| 255 | + | ||
| 256 | + new_gm = export_mlir_module_to_fx(mlir_module) | ||
| 257 | + self.assertTrue( | ||
| 258 | + any( | ||
| 259 | + node.op == "call_function" | ||
| 260 | + and node.target is opaque_nested_tuple_target | ||
| 261 | + for node in new_gm.graph.nodes | ||
| 262 | + ) | ||
| 263 | + ) | ||
| 264 | + self.assertTrue( | ||
| 265 | + any( | ||
| 266 | + node.op == "call_function" | ||
| 267 | + and node.target == operator.getitem | ||
| 268 | + and isinstance(node.args[0], torch.fx.Node) | ||
| 269 | + and node.args[0].op == "call_function" | ||
| 270 | + and node.args[0].target == operator.getitem | ||
| 271 | + for node in new_gm.graph.nodes | ||
| 272 | + ) | ||
| 273 | + ) | ||
| 274 | + | ||
| 206 | def test_fused_op_roundtrip(self): | 275 | def test_fused_op_roundtrip(self): |
| 207 | def fused_backend(gm: torch.fx.GraphModule, example_inputs): | 276 | def fused_backend(gm: torch.fx.GraphModule, example_inputs): |
| 208 | # Clone inputs to avoid mutating them during tracing | 277 | # Clone inputs to avoid mutating them during tracing |
| @@ -138,6 +138,12 @@ def _make_opaque_fake_forward(output_meta: Any): | |||
| 138 | return _fn | 138 | return _fn |
| 139 | 139 | ||
| 140 | 140 | ||
| 141 | +def _meta_value_at_path(val: Any, path: tuple[int, ...]) -> Any: | ||
| 142 | + for index in path: | ||
| 143 | + val = val[index] | ||
| 144 | + return val | ||
| 145 | + | ||
| 146 | + | ||
| 141 | def patch_opaque_roundtrip_targets_for_fake_eval( | 147 | def patch_opaque_roundtrip_targets_for_fake_eval( |
| 142 | gm: torch.fx.GraphModule, | 148 | gm: torch.fx.GraphModule, |
| 143 | ) -> list[tuple[torch.fx.Node, Any]]: | 149 | ) -> list[tuple[torch.fx.Node, Any]]: |
| @@ -501,7 +507,29 @@ class FxExporter: | |||
| 501 | node.meta = dict(payload.meta) | 507 | node.meta = dict(payload.meta) |
| 502 | node.meta["mfusion_opaque_roundtrip"] = True | 508 | node.meta["mfusion_opaque_roundtrip"] = True |
| 503 | 509 | ||
| 504 | - if len(op.results) == 1: | 510 | + result_paths = payload.result_paths |
| 511 | + if result_paths: | ||
| 512 | + if len(result_paths) != len(op.results): | ||
| 513 | + raise RuntimeError( | ||
| 514 | + "opaque result path count does not match MLIR results: " | ||
| 515 | + f"{len(result_paths)} != {len(op.results)}" | ||
| 516 | + ) | ||
| 517 | + path_nodes: dict[tuple[int, ...], torch.fx.Node] = {(): node} | ||
| 518 | + for path, res in zip(result_paths, op.results): | ||
| 519 | + prefix = () | ||
| 520 | + for index in path: | ||
| 521 | + child_path = prefix + (index,) | ||
| 522 | + if child_path not in path_nodes: | ||
| 523 | + result_node = self.graph.call_function( | ||
| 524 | + operator.getitem, (path_nodes[prefix], index) | ||
| 525 | + ) | ||
| 526 | + result_node.meta["val"] = _meta_value_at_path( | ||
| 527 | + payload.meta["val"], child_path | ||
| 528 | + ) | ||
| 529 | + path_nodes[child_path] = result_node | ||
| 530 | + prefix = child_path | ||
| 531 | + self.value_map[res] = path_nodes[path] | ||
| 532 | + elif len(op.results) == 1: | ||
| 505 | self.value_map[op.results[0]] = node | 533 | self.value_map[op.results[0]] = node |
| 506 | elif len(op.results) > 1: | 534 | elif len(op.results) > 1: |
| 507 | for i, res in enumerate(op.results): | 535 | for i, res in enumerate(op.results): |
| @@ -51,6 +51,7 @@ __all__ = ["import_mlir_module_from_fx"] | |||
| 51 | 51 | ||
| 52 | _INT64_MIN = -(1 << 63) | 52 | _INT64_MIN = -(1 << 63) |
| 53 | _INT64_MAX = (1 << 63) - 1 | 53 | _INT64_MAX = (1 << 63) - 1 |
| 54 | +_TORCH_MLIR_SUPPORTED_HOPS = frozenset({"auto_functionalized"}) | ||
| 54 | 55 | ||
| 55 | 56 | ||
| 56 | def _normalize_var_to_range(var_to_range): | 57 | def _normalize_var_to_range(var_to_range): |
| @@ -76,7 +77,7 @@ def _is_supported_call_function_target(node: torch.fx.Node) -> bool: | |||
| 76 | if isinstance(target, TorchOpOverload): | 77 | if isinstance(target, TorchOpOverload): |
| 77 | return True | 78 | return True |
| 78 | if isinstance(target, HigherOrderOperator): | 79 | if isinstance(target, HigherOrderOperator): |
| 79 | - return True | 80 | + return target.name() in _TORCH_MLIR_SUPPORTED_HOPS |
| 80 | return False | 81 | return False |
| 81 | 82 | ||
| 82 | 83 | ||
| @@ -121,6 +122,17 @@ def _schema_type_from_meta_value(val: Any) -> str: | |||
| 121 | raise RuntimeError(f"unsupported opaque schema value type: {type(val).__name__}") | 122 | raise RuntimeError(f"unsupported opaque schema value type: {type(val).__name__}") |
| 122 | 123 | ||
| 123 | 124 | ||
| 125 | +def _flatten_result_meta_tree( | ||
| 126 | + val: Any, path: tuple[int, ...] = () | ||
| 127 | +) -> list[tuple[tuple[int, ...], Any]]: | ||
| 128 | + if isinstance(val, tuple): | ||
| 129 | + leaves: list[tuple[tuple[int, ...], Any]] = [] | ||
| 130 | + for i, item in enumerate(val): | ||
| 131 | + leaves.extend(_flatten_result_meta_tree(item, path + (i,))) | ||
| 132 | + return leaves | ||
| 133 | + return [(path, val)] | ||
| 134 | + | ||
| 135 | + | ||
| 124 | def _encode_fx_arg_tree(arg: Any, flat_nodes: list[torch.fx.Node]) -> Any: | 136 | def _encode_fx_arg_tree(arg: Any, flat_nodes: list[torch.fx.Node]) -> Any: |
| 125 | if isinstance(arg, torch.fx.Node): | 137 | if isinstance(arg, torch.fx.Node): |
| 126 | index = len(flat_nodes) | 138 | index = len(flat_nodes) |
| @@ -139,8 +151,10 @@ def _encode_fx_arg_tree(arg: Any, flat_nodes: list[torch.fx.Node]) -> Any: | |||
| 139 | 151 | ||
| 140 | 152 | ||
| 141 | def _register_opaque_custom_op( | 153 | def _register_opaque_custom_op( |
| 142 | - node: torch.fx.Node, flat_nodes: list[torch.fx.Node] | 154 | + node: torch.fx.Node, |
| 143 | -) -> Any: | 155 | + flat_nodes: list[torch.fx.Node], |
| 156 | + flat_result_metas: list[Any], | ||
| 157 | +) -> tuple[Any, tuple[tuple[int, ...], ...]]: | ||
| 144 | target_name = opaque_registry.new_target_name() | 158 | target_name = opaque_registry.new_target_name() |
| 145 | _, lib_name, op_name = target_name.split(".", 2) | 159 | _, lib_name, op_name = target_name.split(".", 2) |
| 146 | full_op_name = f"{lib_name}::{op_name.replace('.', '_')}" | 160 | full_op_name = f"{lib_name}::{op_name.replace('.', '_')}" |
| @@ -154,7 +168,14 @@ def _register_opaque_custom_op( | |||
| 154 | f"{_schema_type_from_meta_value(arg.meta['val'])} arg{i}" | 168 | f"{_schema_type_from_meta_value(arg.meta['val'])} arg{i}" |
| 155 | for i, arg in enumerate(flat_nodes) | 169 | for i, arg in enumerate(flat_nodes) |
| 156 | ) | 170 | ) |
| 157 | - ret_schema = _schema_type_from_meta_value(node.meta["val"]) | 171 | + flat_results = _flatten_result_meta_tree(node.meta["val"]) |
| 172 | + result_paths = tuple(path for path, _ in flat_results) | ||
| 173 | + flat_result_metas[:] = [meta for _, meta in flat_results] | ||
| 174 | + ret_schemas = [_schema_type_from_meta_value(meta) for meta in flat_result_metas] | ||
| 175 | + if len(ret_schemas) == 1: | ||
| 176 | + ret_schema = ret_schemas[0] | ||
| 177 | + else: | ||
| 178 | + ret_schema = f"({', '.join(ret_schemas)})" | ||
| 158 | schema = f"({arg_schema}) -> {ret_schema}" | 179 | schema = f"({arg_schema}) -> {ret_schema}" |
| 159 | 180 | ||
| 160 | 181 | ||
| @@ -168,9 +189,42 @@ def _register_opaque_custom_op( | |||
| 168 | args_spec=args_spec, | 189 | args_spec=args_spec, |
| 169 | kwargs_spec=kwargs_spec, | 190 | kwargs_spec=kwargs_spec, |
| 170 | meta=dict(node.meta), | 191 | meta=dict(node.meta), |
| 192 | + result_paths=result_paths, | ||
| 171 | ), | 193 | ), |
| 172 | ) | 194 | ) |
| 173 | - return getattr(getattr(torch.ops, lib_name), op_name).default | 195 | + return getattr(getattr(torch.ops, lib_name), op_name).default, result_paths |
| 196 | + | ||
| 197 | + | ||
| 198 | +def _rewrite_opaque_result_users( | ||
| 199 | + gm: torch.fx.GraphModule, | ||
| 200 | + source: torch.fx.Node, | ||
| 201 | + leaf_nodes: dict[tuple[int, ...], torch.fx.Node], | ||
| 202 | + base_path: tuple[int, ...] = (), | ||
| 203 | +) -> None: | ||
| 204 | + for user in list(source.users): | ||
| 205 | + if ( | ||
| 206 | + user.op == "call_function" | ||
| 207 | + and user.target == operator.getitem | ||
| 208 | + and len(user.args) >= 2 | ||
| 209 | + and isinstance(user.args[1], int) | ||
| 210 | + ): | ||
| 211 | + path = base_path + (user.args[1],) | ||
| 212 | + if path in leaf_nodes: | ||
| 213 | + user.replace_all_uses_with(leaf_nodes[path]) | ||
| 214 | + gm.graph.erase_node(user) | ||
| 215 | + continue | ||
| 216 | + if any(leaf_path[: len(path)] == path for leaf_path in leaf_nodes): | ||
| 217 | + _rewrite_opaque_result_users(gm, user, leaf_nodes, path) | ||
| 218 | + if len(user.users) == 0: | ||
| 219 | + gm.graph.erase_node(user) | ||
| 220 | + continue | ||
| 221 | + | ||
| 222 | + raise RuntimeError(f"opaque result getitem path is out of range: {path}") | ||
| 223 | + | ||
| 224 | + raise RuntimeError( | ||
| 225 | + "unsupported opaque result use: " | ||
| 226 | + f"name={user.name}, target={user.target!r}, path={base_path}" | ||
| 227 | + ) | ||
| 174 | 228 | ||
| 175 | 229 | ||
| 176 | def _wrap_unsupported_call_function_targets(gm: torch.fx.GraphModule) -> None: | 230 | def _wrap_unsupported_call_function_targets(gm: torch.fx.GraphModule) -> None: |
| @@ -179,11 +233,32 @@ def _wrap_unsupported_call_function_targets(gm: torch.fx.GraphModule) -> None: | |||
| 179 | if node.op != "call_function" or _is_supported_call_function_target(node): | 233 | if node.op != "call_function" or _is_supported_call_function_target(node): |
| 180 | continue | 234 | continue |
| 181 | flat_nodes: list[torch.fx.Node] = [] | 235 | flat_nodes: list[torch.fx.Node] = [] |
| 182 | - opaque_target = _register_opaque_custom_op(node, flat_nodes) | 236 | + flat_result_metas: list[Any] = [] |
| 237 | + opaque_target, result_paths = _register_opaque_custom_op( | ||
| 238 | + node, flat_nodes, flat_result_metas | ||
| 239 | + ) | ||
| 240 | + if len(flat_result_metas) == 1 and result_paths[0] != (): | ||
| 241 | + raise RuntimeError( | ||
| 242 | + "unsupported opaque single-leaf tuple result: " | ||
| 243 | + f"name={node.name}, target={node.target!r}, path={result_paths[0]}" | ||
| 244 | + ) | ||
| 183 | with gm.graph.inserting_before(node): | 245 | with gm.graph.inserting_before(node): |
| 184 | replacement = gm.graph.call_function(opaque_target, tuple(flat_nodes)) | 246 | replacement = gm.graph.call_function(opaque_target, tuple(flat_nodes)) |
| 185 | replacement.meta = dict(node.meta) | 247 | replacement.meta = dict(node.meta) |
| 186 | - node.replace_all_uses_with(replacement) | 248 | + if len(flat_result_metas) > 1: |
| 249 | + replacement.meta["val"] = tuple(flat_result_metas) | ||
| 250 | + leaf_nodes: dict[tuple[int, ...], torch.fx.Node] = {} | ||
| 251 | + for i, (path, meta) in enumerate(zip(result_paths, flat_result_metas)): | ||
| 252 | + with gm.graph.inserting_before(node): | ||
| 253 | + leaf_node = gm.graph.call_function( | ||
| 254 | + operator.getitem, (replacement, i) | ||
| 255 | + ) | ||
| 256 | + leaf_node.meta["val"] = meta | ||
| 257 | + leaf_nodes[path] = leaf_node | ||
| 258 | + _rewrite_opaque_result_users(gm, node, leaf_nodes) | ||
| 259 | + else: | ||
| 260 | + replacement.meta["val"] = flat_result_metas[0] | ||
| 261 | + node.replace_all_uses_with(replacement) | ||
| 187 | gm.graph.erase_node(node) | 262 | gm.graph.erase_node(node) |
| 188 | changed = True | 263 | changed = True |
| 189 | 264 | ||
| @@ -10,6 +10,7 @@ class Payload: | |||
| 10 | args_spec: Any | 10 | args_spec: Any |
| 11 | kwargs_spec: Any | 11 | kwargs_spec: Any |
| 12 | meta: dict[str, Any] | 12 | meta: dict[str, Any] |
| 13 | + result_paths: tuple[tuple[int, ...], ...] = () | ||
| 13 | 14 | ||
| 14 | 15 | ||
| 15 | _registry: dict[str, Payload] = {} | 16 | _registry: dict[str, Payload] = {} |