Merge branch 'feature/sort_commutative_inputs' of https://github.com/iksnagreb/qonnx into iksnagreb-feature/sort_commutative_inputs

Author: maltanar

src/qonnx/core/modelwrapper.py

@@ -38,7 +38,12 @@
 import qonnx.util.onnx as onnxutil
 from qonnx.core.datatype import DataType
 from qonnx.transformation.double_to_single_float import DoubleToSingleFloat
-from qonnx.transformation.general import RemoveStaticGraphInputs, RemoveUnusedTensors, SortGraph
+from qonnx.transformation.general import (
+    RemoveStaticGraphInputs,
+    RemoveUnusedTensors,
+    SortCommutativeInputsInitializerLast,
+    SortGraph,
+)
 
 
 class ModelWrapper:
@@ -149,6 +154,7 @@ def cleanup(self):
             RemoveUnusedTensors(),
             RemoveStaticGraphInputs(),
             SortGraph(),
+            SortCommutativeInputsInitializerLast(),
         ]
         for trn in cleanup_transforms:
             transformed_model = transformed_model.transform(trn, cleanup=False, make_deepcopy=False)

src/qonnx/transformation/general.py

@@ -29,6 +29,9 @@
 import json
 import numpy as np
 import warnings
+
+# Protobuf onnx graph node type
+from onnx import NodeProto  # noqa
 from onnx import mapping
 from toposort import toposort_flatten
 
@@ -359,3 +362,57 @@ def apply(self, model):
 
         # one iteration is enough
         return (model, False)
+
+
+# Groups inputs by categories, i.e., groups dynamic inputs first, followed by
+# initializers. Keeps order of inputs in each category.
+def group_inputs_by_category(node: NodeProto, model):  # noqa
+    # Select all dynamic inputs, which are those without initializer tensor
+    dynamics = [i for i in node.input if model.get_initializer(i) is None]
+    # Select all input which are initializers, which, by exclusion, are all
+    # those not among the dynamic inputs
+    initializers = [i for i in node.input if i not in dynamics]
+    # Return lists of dynamic anc initializer inputs
+    return dynamics, initializers
+
+
+# Tidy-Up transformation sorting the inputs to all commutative operations to
+# have initializer inputs last
+class SortCommutativeInputsInitializerLast(Transformation):
+    """
+    Sorts inputs of nodes describing commutative operations to have initializer
+    inputs last. This order of inputs is assumed by many other transformations.
+    """
+
+    # Set of supported commutative operations
+    #   TODO: There might be more valid operations
+    SUPPORTED_COMMUTATIVE_OPS = {"Add", "Mul", "And", "Or", "Xor", "Sum"}
+
+    # Applies the transform to a whole model graph
+    def apply(self, model):  # noqa
+        # Get the model graph out of the model wrapper object
+        graph = model.graph
+        # Keep track of whether the graph has been modified
+        graph_modified = False
+        # Iterate all nodes in the graph keeping track of the index
+        for index, node in enumerate(graph.node):
+            # Check whether this node is among the supported
+            if node.op_type in self.SUPPORTED_COMMUTATIVE_OPS:
+                # Group node inputs by category
+                dynamics, initializers = group_inputs_by_category(node, model)
+                # Flatten the grouped input list
+                inputs = [*dynamics, *initializers]
+                # Length of sorted and original input list must match
+                assert len(inputs) == len(node.input)
+                # Reassigned inputs from sorted categories
+                #   Note: ONNX does not allow direct assignment to node.input
+                for i, name in enumerate(inputs):
+                    # The graph has been modified if any input is reordered
+                    if node.input[i] != name:
+                        # Note: This is never reset back to False
+                        graph_modified = True
+                    # Reassign input name at the new index
+                    node.input[i] = name
+        # Return the transformed model and indicate whether the graph actually
+        # has been transformed
+        return model, graph_modified