Merge pull request #187 from fastmachinelearning/feature/transform-subgraph-traveral

Add proper subgraph-traversal for qonnx model_wrapper transform function

Author: Yaman Umuroglu

src/qonnx/core/modelwrapper.py

@@ -1,4 +1,5 @@
-# Copyright (c) 2020 Xilinx, Inc.
+# Copyright (c) 2022 - 2025 Advanced Micro Devices, Inc.
+# Copyright (c) 2020 - 2022 Xilinx, Inc.
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -128,23 +129,58 @@ def analysis(self, analysis_fxn):
         """Runs given anaylsis_fxn on this model and return resulting dict."""
         return analysis_fxn(self)
 
-    def transform(self, transformation, make_deepcopy=True, cleanup=True):
+    def transform_subgraphs(self, transformation, make_deepcopy=True, cleanup=True, apply_to_subgraphs=False, use_preorder_traversal=True):
+        """Applies given Transformation to all subgraphs of this ModelWrapper instance.
+
+        - make_deepcopy : operates on a new (deep)copy of model.
+        - cleanup : execute cleanup transformations before returning
+        - apply_to_subgraphs : if True, transformation is applied to all subgraphs of the model
+        - use_preorder_traversal : if True, uses preorder traversal for subgraph transformation,
+          otherwise postorder traversal is used.
+        """
+        for node in self.model.graph.node:
+                transformed_subgraph_attrs = []
+                for idx, attr in enumerate(node.attribute):
+                    if attr.type == onnx.AttributeProto.GRAPH:
+                        # this is a subgraph, add it to the list
+                        subgraph = self.make_subgraph_modelwrapper(attr.g)
+                        # apply the transformation to the subgraph
+                        subgraph = subgraph.transform(transformation, make_deepcopy, cleanup, apply_to_subgraphs, use_preorder_traversal)
+                        # update the new subgraph in the attrubute
+                        transformed_subgraph_attrs.append((idx, onnx.helper.make_attribute(attr.name, subgraph.model.graph)))
+                # replace the attributes in the node with the transformed subgraph attributes
+                for idx, new_attr in transformed_subgraph_attrs:
+                    # remove the old attribute
+                    node.attribute.pop(idx)
+                    # add the new attribute
+                    node.attribute.insert(idx, new_attr)
+
+    def transform(self, transformation, make_deepcopy=True, cleanup=True, apply_to_subgraphs=False, use_preorder_traversal=True):
         """Applies given Transformation repeatedly until no more changes can be made
         and returns a transformed ModelWrapper instance.
 
         - make_deepcopy : operates on a new (deep)copy of model.
         - cleanup : execute cleanup transformations before returning
+        - apply_to_subgraphs : if True, transformation is applied to all subgraphs of the model
         """
         transformed_model = self
         if make_deepcopy:
             transformed_model = copy.deepcopy(self)
         if self.fix_float64:
             (transformed_model, model_was_changed) = DoubleToSingleFloat().apply(transformed_model)
+
+        if apply_to_subgraphs and use_preorder_traversal == False:
+            transformed_model.transform_subgraphs(transformation, make_deepcopy, cleanup, apply_to_subgraphs, use_preorder_traversal)
+
         model_was_changed = True
         while model_was_changed:
             (transformed_model, model_was_changed) = transformation.apply(transformed_model)
         if cleanup:
             transformed_model.cleanup()
+
+        if apply_to_subgraphs and use_preorder_traversal:
+            transformed_model.transform_subgraphs(transformation, make_deepcopy, cleanup, apply_to_subgraphs, use_preorder_traversal)
+
         return transformed_model
 
     def cleanup(self):
@@ -160,19 +196,8 @@ def cleanup(self):
             transformed_model = transformed_model.transform(trn, cleanup=False, make_deepcopy=False)
         return transformed_model
 
-    def check_compatibility(self):
-        """Checks this model for QONNX compatibility:
-
-        * no embedded subgraphs
-
-        * all tensor shapes are specified, including activations
-
-        * all constants are initializers
-        """
-        # TODO check for no embedded subgraphs
-        # TODO check that all shapes are inferred
-        # TODO check that all constants are initializers
-        return True
+    def make_subgraph_modelwrapper(self, subgraph):
+        return ModelWrapper(util.qonnx_make_model(subgraph, opset_imports=self._model_proto.opset_import))
 
     def get_tensor_datatype(self, tensor_name):
         """Returns the QONNX DataType of tensor with given name."""

tests/core/test_subgraph_traversal.py

@@ -0,0 +1,243 @@
+import pytest
+from collections import Counter
+
+from qonnx.core.modelwrapper import ModelWrapper
+from qonnx.transformation.base import Transformation
+
+from qonnx.util.basic import qonnx_make_model, get_by_name
+import onnx
+from onnx import helper
+
+# Helper to recursively build a graph with subgraphs attached to nodes
+def make_graph(tree):
+    """
+    Recursively build a ModelWrapper tree from a nested tuple/list structure.
+    Each graph will have one node per subgraph, with the subgraph attached as a node attribute.
+    Example input:
+        ("top", [("sub1", []), ("sub2", [("sub2_1", [])])])
+    Returns the top-level ModelWrapper.
+    """
+    name, subtrees = tree
+    # Create subgraphs recursively
+    subgraph_nodes = []
+    inputs = []
+    outputs = []
+    for subtree in subtrees:
+        subgraph = make_graph(subtree)
+        sg_name_in = f"{subgraph.name}_in"
+        sg_name_out = f"{subgraph.name}_out"
+        inputs.append(onnx.helper.make_tensor_value_info(sg_name_in, onnx.TensorProto.FLOAT, [4, 4]))
+        outputs.append(onnx.helper.make_tensor_value_info(sg_name_out, onnx.TensorProto.FLOAT, [4, 4]))
+        # Attach subgraph as attribute to node
+        node = helper.make_node(
+            op_type="SubgraphNode",  # dummy op_type
+            inputs=[sg_name_in],
+            outputs=[sg_name_out],
+            name=f"{subgraph.name}_node",
+        )
+        # ONNX expects subgraphs as AttributeProto, so we set it below
+        attr = onnx.helper.make_attribute("body", subgraph)
+        node.attribute.append(attr)
+        subgraph_nodes.append(node)
+    # Create the graph for this level
+    graph = helper.make_graph(
+        nodes=subgraph_nodes,
+        name=name,
+        inputs=inputs,
+        outputs=outputs,
+    )
+
+    return graph
+
+def make_subgraph_model(tree):
+    """
+    Build a ModelWrapper with a graph structure based on the provided tree.
+    The tree is a nested tuple/list structure where each node can have subgraphs.
+    """
+    return ModelWrapper(qonnx_make_model(make_graph(tree), opset_imports=[helper.make_opsetid("", 10)]))
+
+
+class DummyTransform(Transformation):
+    def __init__(self):
+        self.visited = list()
+
+    def apply(self, model_wrapper):
+        # get the name of the graph being transformed
+        graph_name = model_wrapper.model.graph.name
+        # set a metadata property to test whether metadata is preserved
+        model_wrapper.set_metadata_prop(graph_name, "visited")
+        model_wrapper.set_metadata_prop("opset_id", str(model_wrapper.model.opset_import[0].version))
+        # add a dummy node to the graph to simulate a transformation
+        # to see if the subgraph transformation is presered
+
+        dummy_name_in = f"{graph_name}_dummy_in"
+        dummy_name_out = f"{graph_name}_dummy_out"
+        model_wrapper.model.graph.input.append(helper.make_tensor_value_info(dummy_name_in, onnx.TensorProto.FLOAT, [4, 4]))
+        model_wrapper.model.graph.output.append(helper.make_tensor_value_info(dummy_name_out, onnx.TensorProto.FLOAT, [4, 4]))
+        model_wrapper.model.graph.node.append(
+            helper.make_node(
+                "DummyNode",  # dummy op_type
+                inputs=[dummy_name_in],
+                outputs=[dummy_name_out],
+                name=f"{graph_name}_dummy_node",
+            )
+        )
+
+        # collect the name of the graph being transformed to check how many times each graph was visited
+        self.visited.append(graph_name)
+        #import pdb; pdb.set_trace()
+        return model_wrapper, False
+
+class NestedTransform(Transformation):
+    def __init__(self):
+        self.dummy_transform = DummyTransform()
+    def apply(self, model_wrapper):
+        return model_wrapper.transform(self.dummy_transform), False
+
+def get_subgraph_names(tree):
+    """
+    Recursively collect the names of all subgraphs in the tree structure.
+    """
+    names = set()
+
+    def traverse(tree):
+        name = tree[0]
+        subgraphs = tree[1]
+        names.add(name)
+        for subgraph in subgraphs:
+            traverse(subgraph)
+
+    traverse(tree)
+    return names
+
+
+def check_all_visted_once(tree, transform):
+    """
+    Check that all subgraphs in the tree structure were visited exactly once.
+    """
+    visited  = transform.visited
+    expected = get_subgraph_names(tree)
+    assert Counter(visited) == Counter(expected), f"Visited: {visited}, Expected: {expected}"
+
+def check_visit_order(tree, transform, order):
+    """
+    Check that the order of visited subgraphs matches the expected preorder or postorder traversal.
+    """
+    visited = transform.visited
+    expected = order(tree)
+    assert visited == expected, f"Visited: {visited}, Expected: {expected}"
+
+def check_all_subgraphs_transformed(graph):
+    """
+    Check that all subgraphs in the tree structure have been transformed.
+    """
+
+    # look for the optype "DummyNode" in the model graph
+    dummynode_found = False
+    for node in graph.node:
+        if node.op_type == "DummyNode":
+            dummynode_found = True
+            break
+    if not dummynode_found:
+        raise AssertionError(f"DummyNode not found in the transformed model graph {graph.name}")
+
+    # check that metadata is set for all subgraphs
+    def get_metadata_props(graph, key):
+        metadata_prop = get_by_name(graph.metadata_props, key, "key")
+        if metadata_prop is None:
+            return None
+        else:
+            return metadata_prop.value
+
+    assert(get_metadata_props(graph, graph.name) == "visited"), f"Metadata for {graph.name} not set correctly"
+    assert(get_metadata_props(graph, "opset_id") == "10"), "Metadata for opset_id not set correctly"
+    # recursively check all subgraphs
+    for node in graph.node:
+         for attr in node.attribute:
+            if attr.type == onnx.AttributeProto.GRAPH:
+                check_all_subgraphs_transformed(attr.g)
+
+@pytest.mark.parametrize("cleanup", [False, True])
+@pytest.mark.parametrize("make_deepcopy", [False, True])
+@pytest.mark.parametrize("tree, apply_to_subgraphs",
+                         [(("top", []), True),
+                          (("top", []), False),
+                          (("top", [("sub1", [])]), False)])
+def test_no_traversal(tree, cleanup, make_deepcopy, apply_to_subgraphs):
+    # Check that the top-level model is transformed exactly once when there are no subgraphs.
+    # Check that the top-level model is transformed exactly once when there are subgraphs, but apply_to_subgraphs is False.
+    # This should always be done correctly regardless of cleanup and make_deepcopy.
+
+    model = make_subgraph_model(tree)
+    transform = DummyTransform()
+    t_model = model.transform(transform, cleanup, make_deepcopy, apply_to_subgraphs)
+
+    assert transform.visited == ["top"]
+    assert t_model.get_metadata_prop("top") == "visited"
+
+def build_preorder_traversal(tree):
+    """
+    Build a preorder traversal of the tree structure.
+    """
+    traversal = []
+
+    def traverse(node):
+        name, subtrees = node
+        traversal.append(name)
+        for subtree in subtrees:
+            traverse(subtree)
+
+    traverse(tree)
+    return traversal
+
+def build_postorder_traversal(tree):
+    """
+    Build a postorder traversal of the tree structure.
+    """
+    traversal = []
+
+    def traverse(node):
+        name, subtrees = node
+        for subtree in subtrees:
+            traverse(subtree)
+        traversal.append(name)
+
+    traverse(tree)
+    return traversal
+
+@pytest.mark.parametrize("cleanup", [False, True])
+@pytest.mark.parametrize("make_deepcopy", [False, True])
+@pytest.mark.parametrize("tree", [("top", [("sub1", []), ("sub2", [])]),
+                                  ("top", [("sub1", [("sub1_1", []), ("sub1_2",[])]), ("sub2", [("sub2_1", [])])])])
+@pytest.mark.parametrize("use_preorder_traversal", [True, False])
+def test_traversal(tree, cleanup, make_deepcopy, use_preorder_traversal):
+    # Check that the top-level model and all subgraphs are transformed when apply_to_subgraphs is True.
+    # This should always be done correctly regardless of cleanup and make_deepcopy.
+    print(f"Testing tree: {tree}, cleanup: {cleanup}, make_deepcopy: {make_deepcopy}")
+    model = make_subgraph_model(tree)
+    transform = DummyTransform()
+    t_model = model.transform(transform, cleanup, make_deepcopy, apply_to_subgraphs=True, use_preorder_traversal=use_preorder_traversal)
+
+    check_all_visted_once(tree, transform)
+    check_all_subgraphs_transformed(t_model.model.graph)
+
+    if use_preorder_traversal:
+        traversal_order = build_preorder_traversal
+    else:
+        traversal_order = build_postorder_traversal
+    check_visit_order(tree, transform, traversal_order)
+
+
+@pytest.mark.parametrize("cleanup", [False, True])
+@pytest.mark.parametrize("make_deepcopy", [False, True])
+@pytest.mark.parametrize("tree", [("top", [("sub1", []), ("sub2", [])]),
+                                  ("top", [("sub1", [("sub1_1", []), ("sub1_2",[])]), ("sub2", [("sub2_1", [])])])])
+def test_traversal_nested(tree, cleanup, make_deepcopy):
+    # Check that the top-level model and all subgraphs are transformed when apply_to_subgraphs is True.
+    # This should always be done correctly regardless of cleanup and make_deepcopy.
+    model = make_subgraph_model(tree)
+    transform = NestedTransform()
+    t_model = model.transform(transform, cleanup, make_deepcopy, apply_to_subgraphs=True)
+
+    check_all_visted_once(tree, transform.dummy_transform)
+    check_all_subgraphs_transformed(t_model.model.graph)