Fix onnxruntime dispatching to incorrect function definition

Signed-off-by: Michael Tuttle <mtuttle@qti.qualcomm.com>

Author: michaelgtuttle

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/fusions/__init__.py

@@ -5,7 +5,7 @@
 
 from .fusion import fuse_supergroups
 from .fusion_registry import AIMET_SUPERGROUP_DOMAIN
-from .ir_utils import inline_all_supergroups
+from .ir_utils import inline_all_supergroups, is_fused_supergroup
 from .layernorm import LayerNormFusion
 from .matmul_add import MatmulAddFusion
 from .rmsnorm import RMSNormFusion

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/fusions/fusion.py

@@ -6,7 +6,7 @@
 from collections import defaultdict
 import onnx_ir
 from onnxscript.rewriter import pattern
-from .fusion_registry import FUSION_PASS_REGISTRY, AIMET_SUPERGROUP_DOMAIN
+from .fusion_registry import FUSION_PASS_REGISTRY
 from . import ir_utils
 
 
@@ -53,6 +53,8 @@ def fuse_supergroups(
         _inline_nested_functions(model)
         onnx_ir.passes.common.RemoveUnusedNodesPass().call(model)
         _rename_supergroup_nodes(model)
+        # Note: ORT does not correctly dispatch using function.overload, workaround by fusing to domain
+        ir_utils.fold_overload_into_op_domain(model)
 
     return model
 
@@ -64,7 +66,7 @@ def _inline_nested_functions(model: onnx_ir.Model):
         node.op_identifier()
         for func in model.functions.values()
         for node in func.graph.all_nodes()
-        if node.domain == AIMET_SUPERGROUP_DOMAIN
+        if ir_utils.is_fused_supergroup(node)
     )
     # Note: To get around name mangling of nested functions by InlinePass, sort functions hierarchically (outermost first)
     ir_utils._sort_functions_hierarchically(model)  # pylint: disable=protected-access
@@ -77,9 +79,7 @@ def _rename_supergroup_nodes(model: onnx_ir.Model):
     """Rename supergroup nodes to have more meaningful names derived from their function body, if possible."""
     node_names = {node.name for node in model.graph.all_nodes()}
     supergroup_nodes = [
-        node
-        for node in model.graph.all_nodes()
-        if node.domain == AIMET_SUPERGROUP_DOMAIN
+        node for node in model.graph.all_nodes() if ir_utils.is_fused_supergroup(node)
     ]
     root_to_nodes = defaultdict(list)
 

TrainingExtensions/onnx/src/python/aimet_onnx/graph_passes/fusions/ir_utils.py

@@ -4,6 +4,7 @@
 """ONNX-ir related utility functions"""
 
 import onnx_ir
+import onnx
 import numpy as np
 
 from .fusion_registry import AIMET_SUPERGROUP_DOMAIN
@@ -64,10 +65,7 @@ def _sort_functions_hierarchically(model: onnx_ir.Model) -> None:
     sorted_funcs = {}
 
     def node_has_impl(node: onnx_ir.Node) -> bool:
-        return (
-            node.domain != AIMET_SUPERGROUP_DOMAIN
-            or node.op_identifier() in sorted_funcs
-        )
+        return not is_fused_supergroup(node) or node.op_identifier() in sorted_funcs
 
     while True:
         runnable_functions = {
@@ -92,15 +90,20 @@ def node_has_impl(node: onnx_ir.Node) -> bool:
 def inline_all_supergroups(model: onnx_ir.Model) -> None:
     """Inline all aimet supergroup functions, restoring original node and value names."""
     supergroup_functions = {
-        func
-        for func in model.functions.values()
-        if func.domain == AIMET_SUPERGROUP_DOMAIN
+        func for func in model.functions.values() if is_fused_supergroup(func)
     }
     if not supergroup_functions:
         return
 
     _sort_functions_hierarchically(model)
     onnx_ir.passes.common.InlinePass(lambda f: f in supergroup_functions).call(model)
+    supergroup_opsets = [
+        opset
+        for opset in model.graph.opset_imports
+        if opset.startswith(AIMET_SUPERGROUP_DOMAIN)
+    ]
+    for name in supergroup_opsets:
+        model.graph.opset_imports.pop(name)
 
 
 def unique_name(base: str, existing: set[str]) -> str:
@@ -120,3 +123,34 @@ def get_upstream_cast_type(value: onnx_ir.Value) -> int | None:
         return None
     to_attr = producer.attributes.get("to")
     return to_attr.as_int() if to_attr is not None else None
+
+
+def fold_overload_into_op_domain(model: onnx_ir.Model):
+    """
+    Works around bug in onnxruntime dispatch to overloaded functions by replacing
+    function.domain with "{function.domain}.{function.overload}" for all supergroup ops
+    """
+    new_functions = {}
+    for function in model.functions.values():
+        if function.domain == AIMET_SUPERGROUP_DOMAIN and function.overload:
+            function.domain = f"{function.domain}.{function.overload}"
+        new_functions[function.identifier()] = function
+
+    for node in model.graph.all_nodes():
+        if node.domain != AIMET_SUPERGROUP_DOMAIN:
+            continue
+        if not node.overload:
+            continue
+        node.domain = f"{node.domain}.{node.overload}"
+
+    model._functions = new_functions  # pylint:disable = protected-access
+    opsets = set(f.domain for f in model.functions.values() if is_fused_supergroup(f))
+    for opset in opsets:
+        model.opset_imports[opset] = 1
+
+
+def is_fused_supergroup(
+    node: onnx_ir.Node | onnx_ir.Function | onnx.NodeProto | onnx.FunctionProto,
+) -> bool:
+    """Return True if ``node`` represents an AIMET supergroup op."""
+    return node.domain.startswith(AIMET_SUPERGROUP_DOMAIN)

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

@@ -109,8 +109,8 @@
     OrtInferenceSession,
 )
 from aimet_onnx.graph_passes.fusions import (
-    AIMET_SUPERGROUP_DOMAIN,
     inline_all_supergroups,
+    is_fused_supergroup,
 )
 from aimet_onnx.batch_norm_fold import _has_unfolded_batchnorms
 import aimet_onnx
@@ -786,8 +786,8 @@ def _is_tensor_quantizable(self, name: str) -> bool:
                     return False
 
             if all(
-                (consumer.domain, consumer.op_type)
-                == (AIMET_SUPERGROUP_DOMAIN, "MaskedSoftmax")
+                consumer.op_type == "MaskedSoftmax"
+                and is_fused_supergroup(consumer)
                 and name == consumer.input[2]
                 for consumer in consumer_nodes
             ):
@@ -1954,8 +1954,7 @@ def export(
             with self._remove_quantization_nodes():
                 ir_model = onnx_ir.from_proto(self.model.model)
                 if any(
-                    node.domain == AIMET_SUPERGROUP_DOMAIN
-                    for node in ir_model.graph.all_nodes()
+                    is_fused_supergroup(node) for node in ir_model.graph.all_nodes()
                 ):
                     inline_all_supergroups(ir_model)
 
@@ -2565,10 +2564,7 @@ def _to_onnx_qdq(
 
         # Note: Must inline supergroups before version conversion, since version_converter does not handle functions
         ir_model = onnx_ir.from_proto(model_copy)
-        if any(
-            node.domain == AIMET_SUPERGROUP_DOMAIN
-            for node in ir_model.graph.all_nodes()
-        ):
+        if any(is_fused_supergroup(node) for node in ir_model.graph.all_nodes()):
             inline_all_supergroups(ir_model)
             model_copy = onnx_ir.to_proto(ir_model)
 

TrainingExtensions/onnx/test/python/models/models_for_tests.py

@@ -5432,3 +5432,37 @@ def matmul_add_with_shared_root_naming():
     )
     onnx.checker.check_model(model)
     return model
+
+
+def model_with_transposed_and_non_transposed_gemm():
+    class Model(torch.nn.Module):
+        """
+        W -> Transpose -----v        W2 -----v
+        x --------------> MatMul -> Add -> MatMul -> Add
+        """
+
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(in_features=10, out_features=5)
+            self.weight_2 = torch.nn.Parameter(torch.ones(5, 3))
+            self.bias_2 = torch.nn.Parameter(torch.ones(3))
+
+        def forward(self, x):
+            return self.linear(x) @ self.weight_2 + self.bias_2
+
+    model = Model()
+    x = torch.randn(1, 1, 10)
+    buffer = io.BytesIO()
+    torch.onnx.export(
+        model,
+        (x,),
+        buffer,
+        input_names=["input"],
+        output_names=["output"],
+        opset_version=20,
+        optimize=False,  # prevent folding transB
+    )
+    buffer.seek(0)
+    model = load_model(buffer)
+    onnx.checker.check_model(model)
+    return model

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

@@ -16,8 +16,11 @@
 from aimet_onnx.utils import make_dummy_input, get_node_attribute
 from aimet_onnx import QuantizationSimModel
 
-from aimet_onnx.graph_passes.fusions import fuse_supergroups, inline_all_supergroups
-from aimet_onnx.graph_passes.fusions.fusion_registry import AIMET_SUPERGROUP_DOMAIN
+from aimet_onnx.graph_passes.fusions import (
+    fuse_supergroups,
+    inline_all_supergroups,
+    is_fused_supergroup,
+)
 from ..models import models_for_tests, test_models
 from ..models.test_models import rmsnorm_model
 
@@ -365,6 +368,10 @@ class TestMatmulAddFusion:
     @pytest.mark.parametrize(
         "model_factory, expected_matches",
         [
+            (
+                lambda path: models_for_tests.model_with_transposed_and_non_transposed_gemm(),
+                2,
+            ),
             (lambda path: models_for_tests.matmul_bias_add_model(bias_first=True), 1),
             (lambda path: models_for_tests.matmul_bias_add_model(bias_first=False), 1),
             (lambda path: models_for_tests.matmul_add_model(), 0),  # Not a bias add
@@ -416,7 +423,7 @@ def test_fuses_matmul_add(self, tmp_path, model_factory, expected_matches):
         supergroups = [
             node
             for node in model_proto.graph.node
-            if node.op_type == "Gemm" and node.domain == "aimet.supergroup"
+            if node.op_type == "Gemm" and is_fused_supergroup(node)
         ]
         assert len(supergroups) == expected_matches
 
@@ -445,7 +452,7 @@ def test_transb_attribute_with_transpose(self):
         gemm_nodes = [
             node
             for node in model_proto.graph.node
-            if node.op_type == "Gemm" and node.domain == "aimet.supergroup"
+            if node.op_type == "Gemm" and is_fused_supergroup(node)
         ]
         assert len(gemm_nodes) == 1
 
@@ -468,7 +475,7 @@ def test_transb_attribute_without_transpose(self):
         gemm_nodes = [
             node
             for node in model_proto.graph.node
-            if node.op_type == "Gemm" and node.domain == "aimet.supergroup"
+            if node.op_type == "Gemm" and is_fused_supergroup(node)
         ]
         assert len(gemm_nodes) == 1
 
@@ -790,14 +797,10 @@ def test_fuse_then_inline_restores_original_names(self, tmp_path, model_factory)
 
         # No supergroup functions remain
         assert not any(
-            func.domain == AIMET_SUPERGROUP_DOMAIN
-            for func in ir_model.functions.values()
+            is_fused_supergroup(func) for func in ir_model.functions.values()
         )
 
-        assert not any(
-            node.domain == AIMET_SUPERGROUP_DOMAIN
-            for node in ir_model.graph.all_nodes()
-        )
+        assert not any(is_fused_supergroup(node) for node in ir_model.graph.all_nodes())
 
         # Original non-Constant names are fully restored
         final_node_names = [
@@ -856,9 +859,7 @@ def test_conventional_naming_derives_module_name(
         model = fuse_supergroups(model, patterns=ALL_PATTERNS)
 
         supergroup_names = {
-            node.name
-            for node in model.graph.all_nodes()
-            if node.domain == AIMET_SUPERGROUP_DOMAIN
+            node.name for node in model.graph.all_nodes() if is_fused_supergroup(node)
         }
         assert supergroup_names == expected_names
 
@@ -881,7 +882,7 @@ def test_non_conventional_naming_keeps_default_with_unique_names(
         fuse_supergroups(model, patterns=ALL_PATTERNS)
 
         for node in model.graph.all_nodes():
-            if node.domain == AIMET_SUPERGROUP_DOMAIN:
+            if is_fused_supergroup(node):
                 assert node.name
 
         all_names = [node.name for node in model.graph.all_nodes()]

TrainingExtensions/onnx/test/python/test_quantsim.py

@@ -102,7 +102,7 @@
     identity_tree,
     back_to_back_qdq_pairs,
 )
-from aimet_onnx.graph_passes.fusions import fuse_supergroups, AIMET_SUPERGROUP_DOMAIN
+from aimet_onnx.graph_passes.fusions import fuse_supergroups, is_fused_supergroup
 
 CPU_PROVIDERS = ["CPUExecutionProvider"]
 CUDA_PROVIDERS = ["CUDAExecutionProvider", "CPUExecutionProvider"]
@@ -155,7 +155,7 @@ def _get_tensor_dtypes(model: onnx.ModelProto):
 
 
 def _has_aimet_supergroups(model: onnx.ModelProto):
-    return any(node.domain == AIMET_SUPERGROUP_DOMAIN for node in model.graph.node)
+    return any(is_fused_supergroup(node) for node in model.graph.node)
 
 
 def _fuse_all_supergroups(model: onnx.ModelProto):
@@ -7749,6 +7749,7 @@ def forward(self, x, y):
 @pytest.mark.parametrize(
     "model_factory",
     [
+        lambda _: models_for_tests.model_with_transposed_and_non_transposed_gemm(),
         models_for_tests.llama_rmsnorm_model,
         models_for_tests.rmsnorm_model,
         partial(models_for_tests.rmsnorm_model, elementwise_affine=False),
@@ -7777,7 +7778,7 @@ def test_e2e_quantsim_with_fused_supergroups(model_factory, tmp_path, providers)
     supergroup_bias = set()
     supergroup_nodes = []
     for node in model.graph.node:
-        if node.domain != AIMET_SUPERGROUP_DOMAIN:
+        if not is_fused_supergroup(node):
             continue
         supergroup_nodes.append(node)
         if node.op_type in ("Gemm", "LayerNormalization"):