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

@@ -21,7 +21,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 +29,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,7 +39,10 @@ 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) -> None:

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_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,16 @@ 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
+
 
 class RewriteRuleSet:
     def __init__(self, rules: Sequence[RewriteRule], *, commute: bool = False) -> None: