Set priority among supergroups

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
Co-authored-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

Author: 2 people

TrainingExtensions/common/src/python/aimet_common/graph_searcher.py

@@ -36,7 +36,9 @@
 # =============================================================================
 """Main class for pattern match based graph searcher"""
 
-from typing import Callable, Optional
+from typing import Callable, Optional, List, Set
+from aimet_common.connected_graph.connectedgraph import ConnectedGraph
+from aimet_common.graph_pattern_matcher import PatternType
 from aimet_common.utils import AimetLogger
 from aimet_common.connected_graph.operation import Op
 
@@ -110,7 +112,9 @@ class GraphSearcher:
     It uses SlidingWindow to maintain the search window and PatternMatcher to match sub graph patterns.
     """
 
-    def __init__(self, conn_graph, patterns_with_callback):
+    def __init__(
+        self, conn_graph: ConnectedGraph, patterns_with_callback: List[PatternType]
+    ):
         """
         initializes params required for pattern matching
         :param patterns_with_callback: patterns with corresponding call back functions
@@ -124,18 +128,22 @@ def __init__(self, conn_graph, patterns_with_callback):
             else:
                 self.type_to_op_dict[op.type] = [op]
 
+        self._already_matched: Set[Op] = set()
+
     # pylint: disable=too-many-nested-blocks
     def find_all_patterns_in_graph_apply_actions(
         self,
         ignore: Optional[Op] = None,
         op_pattern_to_reject: Callable[[Op], bool] = None,
+        disjoint: bool = False,
     ):
         """
         Find corresponding op sequences and apply actions.
         :param ignore: List of operations to ignore during searching
         :param op_pattern_to_reject: Callable to perform additional checks on Op to reject pattern match.
             This is useful to express intent on patterns that should not be matched.
             Since GraphSearcher performs high level pattern match, this enables to provide override for aggressive rejection for a given op config.
+        :param disjoint: If True, ensures that matched patterns do not share any ops.
         """
 
         if ignore is None:
@@ -151,19 +159,29 @@ def find_all_patterns_in_graph_apply_actions(
                     matched_ops = self._match_pattern(
                         op, pattern_type.pattern, ignore, op_pattern_to_reject
                     )
-                    if matched_ops:
-                        for matched_ops_list in matched_ops:
-                            pattern_type.action(pattern_type, matched_ops_list)
-                            logger.debug("found match: %s", matched_ops_list)
+                    if not matched_ops:
+                        continue
+                    for matched_ops_list in matched_ops:
+                        if disjoint and any(
+                            matched_op in self._already_matched
+                            for matched_op in matched_ops_list
+                        ):
+                            # This pattern has already been matched as part of a longer pattern
+                            continue
+                        else:
+                            self._already_matched |= set(matched_ops_list)
+
+                        pattern_type.action(pattern_type, matched_ops_list)
+                        logger.debug("found match: %s", matched_ops_list)
 
     # pylint: disable=too-many-branches, too-many-return-statements
     def _match_pattern(
         self,
-        op,
-        pattern,
-        ignored_ops,
-        op_pattern_to_reject: Callable[[Op], bool] = None,
-    ):
+        op: Op,
+        pattern: List[str],
+        ignored_ops: List[Op],
+        op_pattern_to_reject: Optional[Callable[[Op], bool]] = None,
+    ) -> Optional[List[List[Op]]]:
         if not pattern:
             return []
 

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/graph_pass.py

@@ -51,7 +51,7 @@ class GraphPass:
     """
 
     @abstractmethod
-    def match_pattern(self, op: Op, model: ModelProto):
+    def match_pattern(self, op: Op, model: ModelProto) -> List[Op]:
         """
         Pattern match and collect ops starting from given Op.
         """
@@ -126,7 +126,7 @@ def __init__(self):
         self.disable_quantizers: List[str] = []
 
     @abstractmethod
-    def match_pattern(self, op: Op, model: ModelProto):
+    def match_pattern(self, op: Op, model: ModelProto) -> List[Op]:
         """
         Pattern match and collect ops starting from given Op.
         """

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/pass_registry.py

@@ -41,6 +41,7 @@
 from aimet_onnx.meta.connectedgraph import ConnectedGraph
 from aimet_onnx.qc_quantize_op import QcQuantizeOp
 from aimet_onnx.utils import ModelProto
+from ..meta.operations import Op
 
 
 class PassRegistry:
@@ -126,3 +127,33 @@ def apply_graph_passes(
             PASS_REGISTRY[p](model, connected_graph, op_to_quantizers)
         else:
             raise ValueError(f"Graph pass requested but not found: {p}")
+
+
+def find_all_matches(
+    model: ModelProto,
+    connected_graph: ConnectedGraph,
+    passes_to_run: List[str],
+) -> List[List[Op]]:
+    """
+    Runs list of graph passes on input ConnectedGraph
+
+    Args:
+        connected_graph (ConnectedGraph): Input graph to run graph passes on
+        passes_to_run (List[str]): List of graph passes to run.
+
+    Raises:
+        ValueError: If requested GraphPass does not exists.
+    """
+    matches = []
+
+    for p in passes_to_run:
+        for op in connected_graph.ordered_ops:
+            if p in PASS_REGISTRY:
+                graph_pass = PASS_REGISTRY[p]
+                match = graph_pass.match_pattern(op, model)
+                if match:
+                    matches.append(match)
+            else:
+                raise ValueError(f"Graph pass requested but not found: {p}")
+
+    return matches

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/common_patterns.py

@@ -3,6 +3,7 @@
 
 # pylint: disable=missing-docstring
 
+from typing import List
 from aimet_common.connected_graph.operation import Op
 from aimet_onnx.utils import ModelProto
 
@@ -13,7 +14,7 @@
 )
 
 
-def match_rms_norm_pattern(op: Op, model: ModelProto):
+def match_rms_norm_pattern(op: Op, model: ModelProto) -> List[Op]:
     """Common pattern for RMSNormalization which can be re-used"""
     # Match Mul(x, x) or Pow(x, 2)
     match = match_pow_2_pattern(op, model)

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/decoder_block.py

@@ -38,15 +38,15 @@ def apply_on_op(self, op: Op, model: ModelProto, _: Dict[str, QcQuantizeOp]):
             self.pattern_last_op = None
 
     # pylint: disable=too-many-branches, too-many-return-statements
-    def match_pattern(self, op: Op, model: ModelProto):
+    def match_pattern(self, op: Op, model: ModelProto) -> List[Op]:
         """
         Match RMSNorm pattern and collect ops to disable output quantizers
         """
         all_ops = match_rms_norm_pattern(op, model)
         if not all_ops:
-            return False
+            return []
 
         # Check if weights are present
         self.pattern_last_op = all_ops[0]
 
-        return True
+        return all_ops

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/layernorm.py

@@ -36,6 +36,7 @@
 # =============================================================================
 # pylint: disable=missing-module-docstring
 
+from typing import List
 from aimet_common.connected_graph.operation import Op
 from aimet_onnx.graph_passes.graph_pass import SupergroupGraphPass
 from aimet_onnx.graph_passes.pass_registry import register_pass
@@ -80,7 +81,7 @@ class LayerNormalization(SupergroupGraphPass):
     """
 
     # pylint: disable=too-many-branches, too-many-return-statements
-    def match_pattern(self, op: Op, _: ModelProto):
+    def match_pattern(self, op: Op, _: ModelProto) -> List[Op]:
         """
         Match LayerNormalization pattern and collect ops to disable output quantizers
         """
@@ -94,26 +95,26 @@ def match_pattern(self, op: Op, _: ModelProto):
             or len(sub_1.output_ops) != 2
             or sub_1.inputs[0] != op.inputs[0]
         ):
-            return False
+            return []
 
         pow_1 = get_op_from_outputs(sub_1, "Pow")
         div_1 = get_op_from_outputs(sub_1, "Div")
         if pow_1 is None or div_1 is None:
-            return False
+            return []
 
         # Sqrt(Var(x) + ε)
         match, denominator_ops = check_consecutive_ops(
             pow_1, ["Pow", "ReduceMean", "Add", "Sqrt"]
         )
         if not match:
-            return False
+            return []
 
         # (x - E(x)) / Sqrt(Var(x) + ε)
         if (
             div_1.inputs[0].producer != sub_1
             or div_1.inputs[1].producer != denominator_ops[-1]
         ):
-            return False
+            return []
 
         # Collect quantizers to disable.
         all_ops = [op, sub_1] + denominator_ops + [div_1]
@@ -124,8 +125,10 @@ def match_pattern(self, op: Op, _: ModelProto):
         # Check if affine_transform is set.
         # (x - E(x)) / Sqrt(Var(x) + ε) * γ
         match, div_mul_ops = check_consecutive_ops(div_1, ["Div", "Mul"])
-        if not match:
-            return True
+        if match:
+            all_ops += div_mul_ops[1:]
+        else:
+            return all_ops
 
         # NOTE: keep weights quantized
         self.disable_output_quantizers(op_list=[div_1])
@@ -135,10 +138,12 @@ def match_pattern(self, op: Op, _: ModelProto):
         match, mul_add_ops = check_consecutive_ops(
             div_mul_ops[-1], ["Mul", "Add"], validate_last_op_consumers=False
         )
-        if not match:
-            return True
+        if match:
+            all_ops += mul_add_ops[1:]
+        else:
+            return all_ops
 
         # NOTE: skip bias quantization
         self.disable_output_quantizers(op_list=mul_add_ops[:1])
         self.disable_const_quantizers(op_list=mul_add_ops[-1:])
-        return True
+        return all_ops

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/matmul_add.py

@@ -30,12 +30,15 @@ def match_pattern(self, op: Op, model: ModelProto):
         Match RMSNormalization pattern and collect ops to disable output quantizers
         """
         if op.type != "MatMul":
-            return False
+            return []
 
         if _get_matmul_add_bias_idx(op, model) is None:
-            return False
+            return []
 
-        return True
+        matmul_op: Op = op
+        add_op: Op = op.output_ops[0]
+
+        return [matmul_op, add_op]
 
     def apply_on_op(
         self, op: Op, model: ModelProto, op_quantizers: Dict[str, QcQuantizeOp]

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/passes/rmsnorm.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # pylint: disable=missing-module-docstring
 
+from typing import List
 from aimet_common.connected_graph.operation import Op
 from aimet_onnx.graph_passes.graph_pass import SupergroupGraphPass
 from aimet_onnx.graph_passes.pass_registry import register_pass
@@ -55,13 +56,13 @@ class RMSNormalization(SupergroupGraphPass):
     """
 
     # pylint: disable=too-many-branches, too-many-return-statements
-    def match_pattern(self, op: Op, model: ModelProto):
+    def match_pattern(self, op: Op, model: ModelProto) -> List[Op]:
         """
         Match RMSNormalization pattern and collect ops to disable output quantizers
         """
         all_ops = match_rms_norm_pattern(op, model)
         if not all_ops:
-            return False
+            return []
 
         # Check if weights are present
         elementwise_affine = False
@@ -74,4 +75,4 @@ def match_pattern(self, op: Op, model: ModelProto):
         self.disable_output_quantizers(all_ops[:-1])
         # Disable all constant quantizers except weights
         self.disable_const_quantizers(all_ops[:-1] if elementwise_affine else all_ops)
-        return True
+        return all_ops

TrainingExtensions/onnx/src/python/aimet_onnx/quantsim_config/quantsim_config.py

@@ -72,7 +72,7 @@
 from aimet_onnx.meta.connectedgraph import ConnectedGraph, CONSTANT_TYPE
 from aimet_onnx.utils import get_product_name_from_quantized_name
 from aimet_onnx.qc_quantize_op import OpMode, QcQuantizeOp
-from aimet_onnx.graph_passes.pass_registry import apply_graph_passes
+from aimet_onnx.graph_passes.pass_registry import apply_graph_passes, find_all_matches
 
 # pylint: disable=no-name-in-module, ungrouped-imports
 if version.parse(onnx.__version__) >= version.parse("1.14.0"):
@@ -337,6 +337,15 @@ def _set_supergroup_configs(self, supergroups_configs: List[SupergroupType]):
         Set supergroup specific configurations (fourth level of specificity in configuration file)
         :param supergroups_configs: Configurations for supergroups
         """
+        matched_by_graph_pass = find_all_matches(
+            self._model.model,
+            self._conn_graph,
+            self._get_supergroup_pass_list(),
+        )
+        matched_by_graph_pass = set(
+            node_name for match in matched_by_graph_pass for node_name in match
+        )
+
         patterns_with_callbacks = []
         for supergroup_config in supergroups_configs:
             callback = SupergroupConfigCallback(self._model, self._op_to_quantizers)
@@ -347,10 +356,19 @@ def _set_supergroup_configs(self, supergroups_configs: List[SupergroupType]):
             for pattern in patterns:
                 patterns_with_callbacks.append(pattern)
 
+        def exclude_from_supergroup(op: NodeProto) -> bool:
+            return (
+                _check_if_conv3d_or_depthwise_conv(op)
+                # Don't apply "supergroups" config if the op will be handled by graph passes
+                # since "supergroup_pass_list" is a higher priority than "supergroups"
+                or op in matched_by_graph_pass
+            )
+
         if patterns_with_callbacks:
             graph_searcher = GraphSearcher(self._conn_graph, patterns_with_callbacks)
             graph_searcher.find_all_patterns_in_graph_apply_actions(
-                op_pattern_to_reject=_check_if_conv3d_or_depthwise_conv
+                op_pattern_to_reject=exclude_from_supergroup,
+                disjoint=True,
             )
 
     def _set_model_input_configs(self, model_input_configs: ConfigType):

TrainingExtensions/onnx/test/python/test_graph_passes/test_graph_registry.py

@@ -83,7 +83,7 @@ def match_pattern(self, op: Op, _: ModelProto):
         self.disable_quantizers = get_const_input_names(
             op_list=[op]
         ) + get_output_names(op_list=[op])
-        return True
+        return [op]
 
 
 def test_register_and_apply_graph_pass():