Pass: op.cond all test

Author: dayo09

tico/experimental/controlflow/passes/map_subgraph.py

@@ -20,15 +20,14 @@
 import torch
 from torch.export import ExportedProgram
 
-from tico.serialize.quant_param import QPARAM_KEY, QuantParam
 from tico.utils import logging
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import get_quant_dtype
 from tico.utils.validate_args_kwargs import CondArgs
 from tico.utils.graph import create_node
-from tico.utils.subgraph import get_gm_map
+from tico.utils.subgraph import get_all_graph_modules, freeze_subgraphs
 import operator
+from torch.utils import _pytree as pytree
 
 @trace_graph_diff_on_pass
 class MapSubgraph(PassBase):
@@ -53,27 +52,50 @@ def call(self, exported_program: ExportedProgram, _) -> PassResult:
                 continue
 
             cond_args = CondArgs(*node.args, **node.kwargs)
+            true_graph = cond_args.true_graph
+            false_graph = cond_args.false_graph
+            graph_args = cond_args.cond_args
 
-            true_graph_idx = None
-            false_graph_idx = None
-            for gm_info in get_gm_map(exported_program):
-                if gm_info["name"] == cond_args.true_graph.name:
-                    true_graph_idx = gm_info["index"]
-                    continue
-                if gm_info["name"] == cond_args.false_graph.name:
-                    false_graph_idx = gm_info["index"]
-                    continue
-            assert true_graph_idx is not None
-            assert false_graph_idx is not None
+            def _set_meta_val(graph_node, graph_module, graph_args):
+                def _get_meta_val(node):
+                    assert hasattr(node, 'meta'), f"'node' has no attribute named 'meta' (node: {node})"
+                    assert "val" in node.meta, f"val key not in node.meta (node: {node}, meta: {node.meta})"
+                    return node.meta["val"]
+
+                args, kwargs = pytree.tree_map_only(
+                    torch.fx.Node,
+                    _get_meta_val,
+                    (graph_args, {}),
+                )
+                
+                new_val = graph_module(*args, **kwargs)  # type: ignore[operator]
+                graph_node.meta["val"] = new_val
+    
+            for graph_module, name in get_all_graph_modules(exported_program, subgraph_only=True):
+                if true_graph.name == name:
+                    _set_meta_val(true_graph, graph_module, graph_args)
+                if false_graph.name == name:
+                    _set_meta_val(false_graph, graph_module, graph_args)
+            
+            assert "val" in true_graph.meta, f"{true_graph} has no node.meta['val']"
+            assert "val" in false_graph.meta, f"{false_graph} has no node.meta['val']"
 
+            freeze_subgraphs(exported_program)
             with graph.inserting_before(node):
                 circle_if = create_node(
                     graph,
                     torch.ops.circle_custom.if_,
-                    args=(cond_args.condition, true_graph_idx, false_graph_idx, cond_args.cond_args),
+                    args=(cond_args.condition, cond_args.true_graph, cond_args.false_graph, cond_args.cond_args),
                     kwargs={},
                     origin=node,
                 )
+                
+            for t, f in zip(true_graph.meta['val'], false_graph.meta['val']):
+                assert type(t) == type(f)
+                assert t.shape == f.shape, f"{t.shape} != {f.shape}"
+                assert t.dtype == f.dtype, f"{t.dtype} != {f.dtype}"
+                
+            circle_if.meta["val"] = true_graph.meta['val'][0]
 
             # FIX ME UNLESS torch.ops.higher_order.cond generates this pattern
             assert len(node.users) == 1

tico/serialize/circle_graph.py

@@ -324,5 +324,4 @@ def get_tid(
             return self.name_to_tid[node_name]
 
         # Unreachable
-        breakpoint()
         raise RuntimeError("fx Node was not converted to tensor.")

tico/serialize/circle_serializer.py

@@ -28,11 +28,14 @@
 from tico.serialize.operators.node_visitor import get_node_visitors
 from tico.utils import logging
 from tico.utils.serialize import finalise_tensor_names, validate_tensor_shapes
+from tico.utils.subgraph import get_all_graph_modules
 
 
 multiple_output_ops = [
     torch.ops.aten.split_with_sizes.default,
     torch.ops.aten.max.dim,
+    # torch.ops.circle_custom.if_.default,
+    # torch.ops.circle_custom.if_,
 ]
 
 def _initialize_model() -> tuple[CircleModel, CircleSubgraph]:
@@ -46,7 +49,6 @@ def _initialize_model() -> tuple[CircleModel, CircleSubgraph]:
     graph = CircleSubgraph(model)
     return model, graph
 
-from tico.utils.subgraph import get_gm_map
 
 def build_circle(
     ep: ExportedProgram, config: CompileConfigBase = get_default_config()
@@ -64,21 +66,17 @@ def build_circle(
     model = CircleModel()
 
     op_codes: Dict[OpCode, int] = {}
-    
-    for gm_info in get_gm_map(ep):
-        if gm_info["name"]: #non-root subgraph
-            graph_module = getattr(ep.graph_module, gm_info["name"])
-        else:
-            graph_module = ep.graph_module
+    for graph_module, name in get_all_graph_modules(ep):
         ep_graph = graph_module.graph
-        
         graph = CircleSubgraph(model)
+        
         # Export tensors
-        if gm_info["name"]: #non-root subgraph
-            _export_tensors_for_subgraph(graph, ep_graph, ep)
-        else:
+        if name == '': # root graph
             _export_tensors(graph, ep_graph, ep)
-        if gm_info["index"] == 0: # Root graph
+        else:
+            _export_tensors_for_subgraph(graph, ep_graph, ep)
+
+        if name == '': # root graph
             # Register inputs
             logger.debug("---------------Register inputs--------------")
             for in_spec in ep.graph_signature.input_specs:
@@ -108,7 +106,6 @@ def build_circle(
         # Export operators
         logger.debug("---------------Export operators--------------")
         visitors = get_node_visitors(op_codes, graph)
-        ep_graph.print_tabular()
         for node in ep_graph.nodes:
             if node.op != "call_function":
                 continue
@@ -161,8 +158,6 @@ def _export_tensors(graph: CircleSubgraph, ep_graph, ep: ExportedProgram) -> Non
             if node.target in multiple_output_ops:
                 continue
             node_val = node.meta["val"]
-            if node.name == 'cond':
-                continue
             if node_val.layout != torch.strided:
                 raise RuntimeError(
                     f"Only support dense tensors (node layout: {node_val.layout})"
@@ -179,7 +174,7 @@ def _export_tensors(graph: CircleSubgraph, ep_graph, ep: ExportedProgram) -> Non
         elif node.op == "output":
             for output in node.args[0]:
                 if isinstance(output, torch.fx.Node):
-                    assert graph.has_tensor(output.name)
+                    assert graph.has_tensor(output.name), f"{output}"
             continue
 
         elif node.op == "call_method":

tico/serialize/operators/op_circle_if.py

@@ -26,11 +26,11 @@
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
 from tico.utils.validate_args_kwargs import CircleIfArgs
 from tico.utils.errors import NotYetSupportedError
-
+from tico.utils.subgraph import get_frozen_subgraphs
 
 @register_node_visitor
 class CircleIfVisitor(NodeVisitor):
-    target: List[torch._ops.OpOverload] = [torch.ops.circle_custom.if_]
+    target: List[torch._ops.OpOverload] = [torch.ops.circle_custom.if_, torch.ops.circle_custom.if_.default]
 
     def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
         super().__init__(op_codes, graph)
@@ -45,17 +45,28 @@ def define_node(
         if_args = CircleIfArgs(*node.args, **node.kwargs)
 
         pred = if_args.pred
-        then_idx = if_args.then_graph_idx
-        else_idx = if_args.else_graph_idx
+        then_graph = if_args.then_graph
+        else_graph = if_args.else_graph
         arguments = if_args.if_args
 
+        then_graph_idx = None
+        else_graph_idx = None
+        for frozen_subgraph in get_frozen_subgraphs():
+            if frozen_subgraph.name == then_graph.name:
+                then_graph_idx = frozen_subgraph.idx
+            if frozen_subgraph.name == else_graph.name:
+                else_graph_idx = frozen_subgraph.idx
+        assert then_graph_idx is not None
+        assert else_graph_idx is not None
+            
+        
         inputs = [pred, *arguments]
         outputs = [node]
-
+        # outputs = [i for i in node.users.keys()]
         operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
         operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.IfOptions
         operator.builtinOptions = circle.IfOptions.IfOptionsT()
-        operator.builtinOptions.thenSubgraphIndex = then_idx
-        operator.builtinOptions.elseSubgraphIndex = else_idx
+        operator.builtinOptions.thenSubgraphIndex = then_graph_idx
+        operator.builtinOptions.elseSubgraphIndex = else_graph_idx
 
         return operator

tico/utils/convert.py

@@ -87,7 +87,7 @@
     trace_graph_diff_on_func,
 )
 from tico.utils.utils import has_quantization_ops, SuppressWarning
-from tico.utils.subgraph import get_gm_map
+from tico.utils.subgraph import get_all_graph_modules
 
 
 @trace_const_diff_on_func
@@ -199,11 +199,7 @@ def convert_exported_module_to_circle(
 
     assert isinstance(config, CompileConfigBase)
 
-    for gm_info in get_gm_map(exported_program):
-        if gm_info["name"]: #non-root subgraph
-            graph_module = getattr(exported_program.graph_module, gm_info["name"])
-        else:
-            graph_module = exported_program.graph_module
+    for graph_module, _ in get_all_graph_modules(exported_program):
         logger = logging.getLogger(__name__)
         logger.debug("Input ExportedProgram (must be core aten)")
         logger.debug(exported_program)
@@ -232,13 +228,9 @@ def convert_exported_module_to_circle(
         #     CompositeImplicitAutograd and have functional schema are safe to not decompose.
         exported_program = traced_run_decompositions(exported_program)
 
-    for gm_info in get_gm_map(exported_program):
-        if gm_info["name"]: #non-root subgraph
-            graph_module = getattr(exported_program.graph_module, gm_info["name"])
-        else:
-            graph_module = exported_program.graph_module
+    for graph_module, _ in get_all_graph_modules(exported_program):
         graph = graph_module.graph
-        
+
         reinterpret_pass = PassManager(
             passes=[
                 MapSubgraph(),

tico/utils/register_custom_op.py

@@ -19,22 +19,10 @@
 from torch.library import custom_op, register_fake
 
 from tico.utils.mx.mx_ops import _quantize_mx
-from tico.utils.subgraph import get_gm_map
 
-# Note that an operator assumes input tensor has NHWC format.
 def CircleIf():
     @custom_op("circle_custom::if_", mutates_args=())
-    def if_(pred: torch.Tensor, true_graph_idx: int, false_graph_idx: int, if_args: List[torch.Tensor]) -> torch.Tensor:
-        true_graph = None
-        false_graph = None
-        for gm_info in get_gm_map():
-            if gm_info["index"] == true_graph_idx:
-                true_graph = gm_info["gm"]
-                continue
-            if gm_info["index"] == false_graph_idx:
-                false_graph = gm_info["gm"]
-                continue
-            
+    def if_(pred: torch.Tensor, true_graph: torch.Tensor, false_graph: torch.Tensor, if_args: List[torch.Tensor]) -> torch.Tensor:
         if pred:
             result = true_graph(*if_args)
             assert len(result) == 1 # TODO: Support tuple of result
@@ -45,23 +33,12 @@ def if_(pred: torch.Tensor, true_graph_idx: int, false_graph_idx: int, if_args:
             return result[0]
 
     @register_fake("circle_custom::if_")
-    def _(pred: torch.Tensor, true_graph_idx: int, false_graph_idx: int, if_args: List):
-        true_graph = None
-        false_graph = None
-        for gm_info in get_gm_map():
-            if gm_info["index"] == true_graph_idx:
-                true_graph = gm_info["gm"]
-                continue
-            if gm_info["index"] == false_graph_idx:
-                false_graph = gm_info["gm"]
-                continue
-            
+    def _(pred: torch.Tensor, true_graph: torch.Tensor, false_graph: torch.Tensor, if_args: List[torch.Tensor]):
         result = true_graph(*if_args)
         assert len(result) == 1 # TODO: Support tuple of result
 
         return result[0]
 
-
 # Note that an operator assumes input tensor has NHWC format.
 def CircleResizeNearestNeighbor():
     @custom_op("circle_custom::resize_nearest_neighbor", mutates_args=())

tico/utils/subgraph.py

@@ -1,43 +1,42 @@
 import torch
 from torch.export import ExportedProgram
-from typing import Optional
-import functools
+from copy import deepcopy
+from typing import Iterator, List, Iterator
+from dataclasses import dataclass
+@dataclass
+class FrozenSubgraph:
+    idx: int
+    name: str # model-wise, unique name
+    frozen_graph_module: torch.fx.GraphModule # copied subgraph
 
-_gm_map = None
-def get_gm_map(ep: Optional[ExportedProgram] = None):
+_frozen_subgraphs: List[FrozenSubgraph] = []
+
+def freeze_subgraphs(ep: ExportedProgram):
     """
-    Returns [{"index":0, "name": "true_graph_0", "getter": lambda ep: ep.graph_module}, ...}]
+    Freeze subgraphs to provide shape inference logic of FakeTensor.
     """
-    # Build _gm_map only once while compiler running
-    global _gm_map
-    if _gm_map is None:
-        assert ep is not None
-        _gm_map = _build_gm_map(ep)
-    return _gm_map
+    for idx, (graph_module, name) in enumerate(get_all_graph_modules(ep, subgraph_only=True), start = 1):
+        global _frozen_subgraphs
+        _frozen_subgraphs += [FrozenSubgraph(idx = idx, name = name, frozen_graph_module = deepcopy(graph_module))]
 
-def _build_gm_map(ep: ExportedProgram):
-    ret = []
-    
-    # root GraphModule 추가
-    ret.append({
-        "index": len(ret),
-        "name": "",
-        "gm": ep.graph_module,
-    })
+def get_frozen_subgraphs() -> List[FrozenSubgraph]:
+    global _frozen_subgraphs
+    return _frozen_subgraphs
+
+
+def get_all_graph_modules(ep: ExportedProgram, subgraph_only: bool = False) -> Iterator[tuple[torch.fx.GraphModule, str]]:
+    """
+    Get all graph modules and its name
+    """
+    if not subgraph_only:
+        yield ep.graph_module, "" # root has no name
 
-    # Inspect non-root subgraphs
+    # yield subgraphs
     for node in ep.graph.nodes:
         if node.op == "get_attr":
-            attr = getattr(node.graph.owning_module, node.target)
+            graph_module = getattr(node.graph.owning_module, node.target)
 
             # TODO: Enable recursion (n-depth)
-            if isinstance(attr, torch.fx.graph_module.GraphModule):
-                assert hasattr(node, 'name')
-                assert getattr(node, 'name') != ret[0]["name"]
-                graph_name = getattr(node, 'name')
-                ret.append({
-                    "index": len(ret),
-                    "name": graph_name,
-                    "gm": attr,
-                })
-    return ret
+            if isinstance(graph_module, torch.fx.graph_module.GraphModule):
+                assert hasattr(graph_module, 'meta')
+                yield graph_module, getattr(node, 'name')