A couple of ort fusion fixes (microsoft#2136)

* Enable the use of SDPA fusions, along with undoing it when it does not
lead to some subsequent final fusion (such as MHA or GQA).
* Fix the use of constants in extracted functions from fusion.
* Fix the use of Gelu instead of FastGelu in the new fusion introduced
earlier today.

---------

Co-authored-by: Justin Chu <[email protected]>

Author: 2 people

noxfile.py

@@ -15,8 +15,7 @@
     "beartype==0.17.2",
     "expecttest==0.1.6",
     "hypothesis",
-    'numpy==1.24.4; python_version<"3.9"',
-    'numpy==1.26.4; python_version>="3.9"',
+    "numpy",
     "packaging",
     "parameterized",
     'psutil; sys_platform != "win32"',

onnxscript/rewriter/ort_fusions/_core.py

@@ -4,7 +4,14 @@
 
 import onnxscript.ir as ir
 from onnxscript.ir.passes.common import shape_inference
-from onnxscript.optimizer import optimize, remove_unused_nodes
+from onnxscript.optimizer import optimize
+from onnxscript.rewriter import rewrite
+from onnxscript.rewriter.ort_fusions import (
+    fused_matmul_rule_sets,
+    # group_normalization_merge_silu,
+    instance_to_group_normalization,
+    softmax,
+)
 from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
 from onnxscript.rewriter.ort_fusions.gelu import fuse_gelu
 from onnxscript.rewriter.ort_fusions.mha import fuse_mha
@@ -21,7 +28,6 @@
 # TODO: There are some potential redundancies below. Can be targeted for optimization
 # once we have robust fusion.
 def _pre_optimize(model: ir.Model) -> ir.Model:
-    optimize(model)
     # TODO: Do we need this dependence on ONNX's partial-data-propagation? There are some
     # extra shape-propagation and partial-data-propagation rules in ONNX that are not yet
     # incorporated in our optimizer.
@@ -30,7 +36,7 @@ def _pre_optimize(model: ir.Model) -> ir.Model:
     return model
 
 
-def fuse_xformers(model: ir.Model) -> None:
+def fuse_xformers(model: ir.Model) -> ir.Model:
     model = _pre_optimize(model)
     fuse_rms_normalization(model)
     fuse_normalization(model)
@@ -40,9 +46,29 @@ def fuse_xformers(model: ir.Model) -> None:
     fuse_sdpa(model)
     fuse_mha(model)
     fuse_gelu(model)
-    remove_unused_nodes(model)
+    # Finally: inline any intermediate fusion functions introduced that were not
+    # consumed by other fusions, and eliminate any remaining unused nodes.
+    optimize(model)
+    return model
+
 
+def optimize_for_ort(model: ir.Model, config_name: str | None = None) -> ir.Model:
+    """
+    Optimize the model for ORT backend.
+
+    TODO: config_name is not used yet. It should be used to select the appropriate
+    optimization configuration (for an EP). Currently, a default implementation is used.
+
+    Args:
+        model: The model to optimize.
+        config_name: The name of the configuration to use for optimization.
+            Typically it identifies the Execution Provider (EP) to optimize for.
+            If None, the default configuration will be used.
+
+    Returns:
+        The optimized model.
+    """
 
-def optimize_for_ort(model: ir.Model) -> None:
-    # TODO(rama): Include the other optimizations
     fuse_xformers(model)
+    rewrite(model, ORT_PATTERN_REWRITE_RULES)
+    return model

onnxscript/rewriter/ort_fusions/_test_utils.py

@@ -32,7 +32,7 @@ def ort_run(model_name: str, model, inputs):
     return ort_outputs
 
 
-def assert_allclose(outputs, expected_outputs, rtol=1e-2, atol=1e-2):
+def assert_allclose(outputs, expected_outputs, rtol=1e-4, atol=1e-4):
     for i, (baseline_output, optimized_output) in enumerate(zip(expected_outputs, outputs)):
         try:
             np.testing.assert_equal(baseline_output.shape, optimized_output.shape)

onnxscript/rewriter/ort_fusions/fuse_xformers_test.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import unittest
+
+import onnxscript.optimizer
+from onnxscript.rewriter.ort_fusions._core import fuse_xformers
+from onnxscript.rewriter.ort_fusions._smollm_1 import smollm_test_1
+from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
+
+
+class TestFuseXformers(unittest.TestCase):
+    def test_fuse_xformers(self):
+        test = smollm_test_1()
+        model = test.get_onnx_model()
+        onnxscript.optimizer.optimize(model)
+        inputs = test.get_ort_inputs()
+        original_outputs = ort_run("original", model, inputs)
+        model = fuse_xformers(model)
+        new_outputs = ort_run("optimized", model, inputs)
+        assert_allclose(new_outputs, original_outputs)
+
+
+if __name__ == "__main__":
+    unittest.main()

onnxscript/rewriter/ort_fusions/gelu.py

@@ -25,7 +25,7 @@ def pattern(self, op, x):
         return result
 
     def rewrite(self, op, x):
-        return op.Gelu(x, _domain="com.microsoft")
+        return op.FastGelu(x, _domain="com.microsoft")
 
 
 _rule = GeluTanhFusion.rule()

onnxscript/rewriter/ort_fusions/gelu_test.py

@@ -47,7 +47,7 @@ def gelu_model(x):
         remove_unused_nodes(model)
 
         self.assertEqual(len(model.graph), 1)
-        self.assertEqual(model.graph.node(0).op_type, "Gelu")
+        self.assertEqual(model.graph.node(0).op_type, "FastGelu")
 
         optimized_output = test_utils.ort_run("Optimized", model, input)
         test_utils.assert_allclose(original_output, optimized_output)

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -71,7 +71,7 @@ def check(self, op, x, scale, epsilon, compute_dtype, target_dtype):
     def rewrite(self, op, x, scale, epsilon, compute_dtype, target_dtype):
         stash_dtype = compute_dtype.value if self._cast_input else x.dtype
         # Note: ORT's SimplifiedLayerNormalization was placed in onnx domain by mistake.
-        # No need to use com.microsoft domain here.
+        # No need to use com.microsoft domain here; but this is a custom op in ORT.
         return op.SimplifiedLayerNormalization(
             x,
             scale,

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -9,8 +9,8 @@
 
 
 class SDPA(pattern.RewriteRuleClassBase):
-    def __init__(self, name: str, *, use_mask: bool, pre_scale: bool):
-        super().__init__(name=name)
+    def __init__(self, name: str, *, use_mask: bool, pre_scale: bool, use_mul: bool):
+        super().__init__(name=name, as_function=True)
         self._use_mask = use_mask
         self._pre_scale = pre_scale
 

onnxscript/rewriter/ort_fusions/sdpa_test.py

@@ -180,3 +180,7 @@ def test_sdpa_fusion(self, name, script_func):
 
         # new_outputs = ort_run("optimized", model, inputs)
         # assert_allclose(new_outputs, original_outputs)
+
+
+if __name__ == "__main__":
+    unittest.main()

onnxscript/rewriter/pattern.py

@@ -1421,6 +1421,7 @@ def replace_pattern(new_pattern):
                 self.remove_nodes,
                 self.graph_pre_visitor,
                 self.graph_post_visitor,
+                self.as_function,
             )
 
         return [replace_pattern(p) for p in self._target_pattern.commute()]
@@ -1502,21 +1503,23 @@ class RewriteRuleClassBase:
     @classmethod
     def rule(cls, *args, **kwargs):
         instance = cls(*args, **kwargs)
-        setup = instance.setup if hasattr(instance, "setup") else None
-        cleanup = instance.cleanup if hasattr(instance, "cleanup") else None
         return RewriteRule(
             instance.pattern,
             instance.rewrite,
             instance.check,
             name=instance.name,
             remove_nodes=instance.remove_nodes,
-            graph_pre_visitor=setup,
-            graph_post_visitor=cleanup,
+            graph_pre_visitor=instance.setup,
+            graph_post_visitor=instance.cleanup,
+            as_function=instance.as_function,
         )
 
-    def __init__(self, name: str | None = None, remove_nodes: bool = True) -> None:
+    def __init__(
+        self, name: str | None = None, remove_nodes: bool = True, as_function: bool = False
+    ) -> None:
         self.name = name or self.__class__.__name__
         self.remove_nodes = remove_nodes
+        self.as_function = as_function
 
     def pattern(self, op, *args, **kwargs):
         raise NotImplementedError("Method 'pattern' must be implemented by derived class.")
@@ -1528,6 +1531,114 @@ def check(self, op, *args, **kwargs):
     def rewrite(self, op, *args, **kwargs):
         raise NotImplementedError("Method 'rewrite' must be implemented by derived class.")
 
+    def setup(self):
+        # Optional setup function that can be overridden by derived classes. Used to do
+        # per model/function initialization.
+        pass
+
+    def cleanup(self):
+        # Optional cleanup function that can be overridden by derived classes. Used to do
+        # per model/function cleanup.
+        pass
+
+
+def _copy_for_function(
+    inputs: Sequence[ir.Value | None], nodes: Sequence[ir.Node], outputs: Sequence[ir.Value]
+):
+    """Utility function to extract a subgraph out as a function."""
+    value_map: dict[ir.Value, ir.Value] = {}
+    function_inputs: list[ir.Value] = []
+    constant_nodes: list[ir.Node] = []
+    for input in inputs:
+        # Create a function input (formal-parameter value) to represent this value:
+        new_value = (
+            ir.Value(
+                name=input.name,
+                shape=input.shape,
+                type=input.type,
+                doc_string=input.doc_string,
+            )
+            if input
+            else ir.Value()  # dummy parameter for a None input
+        )
+        if input is not None:
+            value_map[input] = new_value
+        function_inputs.append(new_value)
+
+    def copy_value(value: ir.Value | None) -> ir.Value | None:
+        if value is None:
+            return None
+        if value not in value_map:
+            const_value = value.const_value
+            if const_value is not None:
+                # create a Constant node to represent the value
+                value_attr = ir.AttrTensor("value", const_value)
+                const_node = ir.Node("", "Constant", [], [value_attr])
+                constant_nodes.append(const_node)
+                value_map[value] = result = const_node.outputs[0]
+                return result
+            raise ValueError(f"Value {value} not found in value_map.")
+        return value_map[value]
+
+    def copy_attr_value(attr: ir.Attr | ir.RefAttr) -> ir.Attr | ir.RefAttr:
+        if not isinstance(attr, ir.Attr):
+            # No need to support this currently, as rewriting inside a function is
+            # not used, as it has several challenges.
+            raise NotImplementedError("RefAttr not supported.")
+        if attr.type in {ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS}:
+            # No need to support this currently, as rewriting control-flow constructs
+            # is not used and has several challenges.
+            raise NotImplementedError("Graph attributes not supported.")
+        # Primitive attributes are immutable by design and can be shared.
+        return attr
+
+    def copy_node(node: ir.Node) -> ir.Node:
+        new_inputs = [copy_value(v) for v in node.inputs]
+        new_attributes = [copy_attr_value(v) for v in node.attributes.values()]
+        new_node = ir.Node(
+            node.domain,
+            node.op_type,
+            new_inputs,
+            new_attributes,
+            overload=node.overload,
+            num_outputs=len(node.outputs),
+            graph=None,
+            name=node.name,
+            doc_string=node.doc_string,  # type: ignore
+            metadata_props=node.metadata_props.copy(),
+        )
+        new_outputs = new_node.outputs
+        for i, output in enumerate(node.outputs):
+            value_map[output] = new_outputs[i]
+            if output.name is not None:
+                new_outputs[i].name = output.name
+        return new_node
+
+    function_nodes = [copy_node(node) for node in nodes]
+    function_outputs = [copy_value(v) for v in outputs]
+    return (function_inputs, constant_nodes + function_nodes, function_outputs)
+
+
+def _get_new_overload(model: ir.Model, domain: str, name: str) -> str:
+    """Get a new overload for the given domain and name.
+
+    Args:
+        model: The model to which the new overload will be added.
+        domain: The domain of the new overload.
+        name: The opname of the new overload.
+
+    Returns:
+        The new overload name.
+    """
+    existing_functions = model.functions
+    # Just a simple implementation for now
+    overload = 1
+    while True:
+        overload_name = str(overload)
+        if (domain, name, overload_name) not in existing_functions:
+            return overload_name
+        overload += 1
+
 
 class RewriteRuleSet:
     def __init__(self, rules: Sequence[RewriteRule], *, commute: bool = False) -> None: