Allow fuse_xformers to return a count of different fusions applied (microsoft#2159)

Author: shubhambhokare1

onnxscript/rewriter/_fusion_utils.py

@@ -2,9 +2,10 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-from typing import Sequence, Union
+from typing import Callable, Sequence, Union
 
-from onnxscript import ir
+import onnxscript.ir as ir
+from onnxscript.rewriter import pattern
 
 Dim = Union[int, ir.SymbolicDim]
 
@@ -20,3 +21,20 @@ def _check_shape(bindings: dict[str, Dim], val: ir.Value, shape: Sequence[str])
         elif actual != bindings[expected]:
             return False
     return True
+
+
+def apply_fusion_rules(rules: pattern.RewriteRule | pattern.RewriteRuleSet) -> Callable:
+    """
+    Apply the given fusion rules to the model and return the number of fusions applied.
+    If debug is True, enable pattern matching tracer for debugging.
+    """
+
+    def apply_to(model: ir.Model, debug: bool = False) -> int:
+        count = rules.apply_to_model(model)
+        if count == 0 and debug:
+            tracer = pattern.MatchingTracer()
+            rules.apply_to_model(model, tracer=tracer)
+            tracer.report()
+        return count
+
+    return apply_to

onnxscript/rewriter/ort_fusions/_core.py

@@ -14,7 +14,10 @@
     fuse_rotary_embedding,
 )
 from onnxscript.rewriter.ort_fusions.sdpa import fuse_sdpa
-from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_normalization
+from onnxscript.rewriter.ort_fusions.skip_normalization import (
+    fuse_skip_layer_normalization,
+    fuse_skip_rms_normalization,
+)
 
 
 # Preliminary optimizations before applying the transformer fusions.
@@ -29,21 +32,35 @@ def _pre_optimize(model: ir.Model) -> ir.Model:
     return model
 
 
-def fuse_xformers(model: ir.Model) -> ir.Model:
+def fuse_xformers(model: ir.Model) -> tuple[ir.Model, dict[str, int]]:
+    """
+    Apply transformer-specific fusions to the given model.
+
+    Args:
+        model: The input ONNX model represented as an `ir.Model`.
+
+    Returns:
+        A tuple containing:
+        - The optimized `ir.Model` after applying transformer-specific fusions.
+        - A dictionary with a count of each of the fusions applied.
+    """
+    fusion_count = dict()
+
     model = _pre_optimize(model)
-    fuse_rms_normalization(model)
-    fuse_normalization(model)
-    fuse_rotary_embedding(model)
-    fuse_partial_rotary_embedding(model)
-    fuse_cos_sin_cache(model)
-    fuse_sdpa(model)
-    fuse_mha(model)
-    fuse_attention(model)
-    fuse_gelu(model)
+    fusion_count["rms_normalization"] = fuse_rms_normalization(model)
+    fusion_count["skip_layer_normalization"] = fuse_skip_layer_normalization(model)
+    fusion_count["skip_rms_normalization"] = fuse_skip_rms_normalization(model)
+    fusion_count["rotary_embedding"] = fuse_rotary_embedding(model)
+    fusion_count["partial_rotary_embedding"] = fuse_partial_rotary_embedding(model)
+    fusion_count["cos_sin_cache"] = fuse_cos_sin_cache(model)
+    fusion_count["sdpa"] = fuse_sdpa(model)
+    fusion_count["mha"] = fuse_mha(model)
+    fusion_count["attention"] = fuse_attention(model)
+    fusion_count["gelu"] = fuse_gelu(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
+    return model, fusion_count
 
 
 def optimize_for_ort(model: ir.Model) -> None:

onnxscript/rewriter/ort_fusions/attention.py

@@ -268,10 +268,4 @@ def rewrite(
 )
 
 
-def fuse_attention(model: ir.Model, *, debug: bool = False) -> int:
-    count = attention_rules.apply_to_model(model)
-    if debug and count == 0:
-        tracer = pattern.MatchingTracer()
-        attention_rules.apply_to_model(model, tracer=tracer)
-        tracer.report()
-    return count
+fuse_attention = _fusion_utils.apply_fusion_rules(attention_rules)

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -5,8 +5,7 @@
 import numpy as np
 
 import onnxscript.ir as ir
-from onnxscript.optimizer import remove_unused_nodes
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
 # Rewrite the computation of cos/sin cache into the form expected by ORT's custom ops.
 

onnxscript/rewriter/ort_fusions/fuse_xformers_test.py

@@ -17,7 +17,20 @@ def test_fuse_xformers(self):
         onnxscript.optimizer.optimize(model)
         inputs = test.get_ort_inputs()
         original_outputs = ort_run("original", model, inputs)
-        model = fuse_xformers(model)
+        model, fusion_count = fuse_xformers(model)
+
+        # Check if the number of fusions applied for each fusion is correct
+        self.assertEqual(fusion_count["rms_normalization"], 3)
+        self.assertEqual(fusion_count["skip_layer_normalization"], 0)
+        self.assertEqual(fusion_count["skip_rms_normalization"], 2)
+        self.assertEqual(fusion_count["rotary_embedding"], 2)
+        self.assertEqual(fusion_count["partial_rotary_embedding"], 0)
+        self.assertEqual(fusion_count["cos_sin_cache"], 2)
+        self.assertEqual(fusion_count["sdpa"], 1)
+        self.assertEqual(fusion_count["mha"], 0)
+        self.assertEqual(fusion_count["attention"], 0)
+        self.assertEqual(fusion_count["gelu"], 0)
+
         new_outputs = ort_run("optimized", model, inputs)
         assert_allclose(new_outputs, original_outputs)
 

onnxscript/rewriter/ort_fusions/gelu.py

@@ -4,8 +4,7 @@
 
 import math
 
-from onnxscript import ir
-from onnxscript.rewriter import pattern
+from onnxscript.rewriter import _fusion_utils, pattern
 
 _sqrt_two_over_pi = math.sqrt(2.0 / math.pi)
 
@@ -33,5 +32,4 @@ def rewrite(self, op, x):
 gelu_rules = pattern.RewriteRuleSet([_rule])
 
 
-def fuse_gelu(model: ir.Model) -> None:
-    gelu_rules.apply_to_model(model)
+fuse_gelu = _fusion_utils.apply_fusion_rules(gelu_rules)

onnxscript/rewriter/ort_fusions/gqa.py

@@ -2,9 +2,7 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-import onnxscript.ir as ir
-from onnxscript.optimizer import remove_unused_nodes
-from onnxscript.rewriter import pattern
+from onnxscript.rewriter import _fusion_utils, pattern
 
 
 class GroupQueryAttention(pattern.RewriteRuleClassBase):
@@ -150,8 +148,4 @@ def rewrite(
 gqa_rules = pattern.RewriteRuleSet([_rule1])
 
 
-def fuse_gqa(model: ir.Model) -> int:
-    count = gqa_rules.apply_to_model(model)
-    print(f"GQA count: {count}")
-    remove_unused_nodes(model)
-    return count
+fuse_gqa = _fusion_utils.apply_fusion_rules(gqa_rules)

onnxscript/rewriter/ort_fusions/mha_test.py

@@ -17,7 +17,7 @@ def test_smollm(self):
         model = smollm_test.get_onnx_model()
         onnxscript.optimizer.optimize(model)
         xformers.fuse_rms_normalization(model)
-        xformers.fuse_normalization(model)
+        xformers.fuse_skip_rms_normalization(model)
         xformers.fuse_rotary_embedding(model)
         xformers.fuse_cos_sin_cache(model)
 

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import onnxscript.ir as ir
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
 """
 RMS Normalization: This is referred to as SimplifiedLayerNormalization in the ORT codebase.
@@ -91,6 +91,4 @@ def rewrite(self, op, x, scale, epsilon, compute_dtype, target_dtype):
 rms_normalization_ruleset = pattern.RewriteRuleSet(rms_normalization_rules)
 
 
-def fuse_rms_normalization(model: ir.Model) -> None:
-    count = rms_normalization_ruleset.apply_to_model(model)
-    print(f"RMS Normalization count: {count}")
+fuse_rms_normalization = _fusion_utils.apply_fusion_rules(rms_normalization_ruleset)

onnxscript/rewriter/ort_fusions/rotary_embedding.py

@@ -2,8 +2,7 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-import onnxscript.ir as ir
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
 # Add first version of the RotaryEmbeddingFusion rule. This considers only one simple pattern
 # for full rotation without interleaving.

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -4,8 +4,7 @@
 
 import math
 
-import onnxscript.ir as ir
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
 
 class SDPA(pattern.RewriteRuleClassBase):

onnxscript/rewriter/ort_fusions/skip_normalization.py

@@ -2,11 +2,10 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-from onnxscript.rewriter import pattern
-from onnxscript.rewriter.ort_fusions.rms_normalization import rms_normalization_rules
+from onnxscript.rewriter import _fusion_utils, pattern
 
 
-def _skip_norm_pattern(op, input, skip, gamma, epsilon, stash_type):
+def _skip_rms_norm_pattern(op, input, skip, gamma, epsilon, stash_type):
     skip_sum = op.Add(input, skip)
     normalized = op.SimplifiedLayerNormalization(
         skip_sum,
@@ -18,7 +17,7 @@ def _skip_norm_pattern(op, input, skip, gamma, epsilon, stash_type):
     return normalized, skip_sum
 
 
-def _skip_normalization(op, input, skip, gamma, epsilon, stash_type):
+def _skip_rms_normalization(op, input, skip, gamma, epsilon, stash_type):
     if stash_type.value != 1:  # FLOAT type
         return None
     normalized, _mean, _inv_std_var, skip_sum = op.SkipSimplifiedLayerNormalization(
@@ -32,15 +31,49 @@ def _skip_normalization(op, input, skip, gamma, epsilon, stash_type):
     return normalized, skip_sum
 
 
-_rule = pattern.RewriteRule(
-    _skip_norm_pattern, _skip_normalization, matcher=pattern.SimplePatternMatcher
-)
+_skip_rms_rule = pattern.RewriteRule(_skip_rms_norm_pattern, _skip_rms_normalization)
+
+skip_rms_normalization_rules = [_skip_rms_rule]
+skip_rms_normalization_ruleset = pattern.RewriteRuleSet(skip_rms_normalization_rules)
+
+
+def _skip_layer_norm_pattern(op, input, skip, gamma, beta, epsilon, stash_type):
+    skip_sum = op.Add(input, skip)
+    normalized = op.LayerNormalization(
+        skip_sum,
+        gamma,
+        beta,
+        axis=-1,
+        epsilon=epsilon,
+        stash_type=stash_type,
+    )
+    return normalized
+
+
+def _skip_layer_normalization(op, input, skip, gamma, beta, epsilon, stash_type):
+    if stash_type.value != 1:  # FLOAT type
+        return None
+    normalized, _mean, _inv_std_var = op.SkipLayerNormalization(
+        input,
+        skip,
+        gamma,
+        beta,
+        epsilon=epsilon,
+        _outputs=3,
+        _domain="com.microsoft",
+    )
+    return normalized
+
+
+_skip_layer_rule = pattern.RewriteRule(_skip_layer_norm_pattern, _skip_layer_normalization)
 
-skip_normalization_rules = [_rule]
-normalization_rules = rms_normalization_rules + skip_normalization_rules
-normalization_ruleset = pattern.RewriteRuleSet(normalization_rules)
+skip_layer_normalization_rules = [_skip_layer_rule]
+skip_layer_normalization_ruleset = pattern.RewriteRuleSet(skip_layer_normalization_rules)
 
 
-def fuse_normalization(model):
-    count = normalization_ruleset.apply_to_model(model)
-    print(f"Normalization count: {count}")
+fuse_skip_rms_normalization = _fusion_utils.apply_fusion_rules(skip_rms_normalization_ruleset)
+
+
+fuse_skip_layer_normalization = _fusion_utils.apply_fusion_rules(
+    skip_layer_normalization_ruleset
+)

onnxscript/rewriter/ort_fusions/skip_normalization_test.py

@@ -7,7 +7,8 @@
 import onnxscript.optimizer
 from onnxscript.rewriter.ort_fusions._smollm_1 import TestData
 from onnxscript.rewriter.ort_fusions._test_utils import assert_allclose, ort_run
-from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_normalization
+from onnxscript.rewriter.ort_fusions.rms_normalization import fuse_rms_normalization
+from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_skip_rms_normalization
 
 
 class TestSkipNormalization(unittest.TestCase):
@@ -17,7 +18,8 @@ def test_smollm(self):
         onnxscript.optimizer.optimize(model)
         inputs = smollm_test.get_ort_inputs()
         original_outputs = ort_run("original", model, inputs)
-        fuse_normalization(model)
+        fuse_rms_normalization(model)
+        fuse_skip_rms_normalization(model)
         op_types = [n.op_type for n in model.graph]
         self.assertIn("SkipSimplifiedLayerNormalization", op_types)
         new_outputs = ort_run("optimized", model, inputs)