A couple of ort fusion fixes (#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: gramalingam

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,7 @@
 
 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,
@@ -36,7 +36,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.
@@ -45,7 +44,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)
@@ -55,9 +54,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:
     fuse_xformers(model)
     rewrite(model, ORT_PATTERN_REWRITE_RULES)
+    return model

onnxscript/rewriter/ort_fusions/_test_utils.py

@@ -33,7 +33,7 @@ def ort_run(model_name: str, model, inputs):
     return session.run(None, inputs)
 
 
-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

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

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

@@ -1428,6 +1428,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()]
@@ -1509,21 +1510,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.")
@@ -1535,30 +1538,52 @@ 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:
-        if input is None:
-            raise NotImplementedError("None inputs not supported.")
-        new_value = ir.Value(
-            name=input.name,
-            shape=input.shape,
-            type=input.type,
-            doc_string=input.doc_string,
+        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
         )
-        value_map[input] = new_value
+        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]
 
@@ -1598,7 +1623,7 @@ def copy_node(node: ir.Node) -> ir.Node:
 
     function_nodes = [copy_node(node) for node in nodes]
     function_outputs = [copy_value(v) for v in outputs]
-    return (function_inputs, function_nodes, function_outputs)
+    return (function_inputs, constant_nodes + function_nodes, function_outputs)
 
 
 def _get_new_overload(model: ir.Model, domain: str, name: str) -> str: