modify to add callable

Author: shubhambhokare1

onnxscript/rewriter/_fusion_utils.py

@@ -2,9 +2,11 @@
 # 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,20 +22,20 @@ def _check_shape(bindings: dict[str, Dim], val: ir.Value, shape: Sequence[str])
         elif actual != bindings[expected]:
             return False
     return True
-import onnxscript.ir as ir
-from onnxscript.rewriter import pattern
 
 
-def apply_fusion_rules(
-    rules: pattern.RewriteRule | pattern.RewriteRuleSet, model: ir.Model, debug: bool = False
-) -> int:
+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.
     """
-    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
+
+    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/cos_sin_cache.py

@@ -170,7 +170,8 @@ def rewrite(
 
 
 def fuse_cos_sin_cache(model: ir.Model, debug: bool = False) -> int:
-    count = _fusion_utils.apply_fusion_rules(cos_sin_cache_rules, model, debug=debug)
+    fuse_cos_sin_cache = _fusion_utils.apply_fusion_rules(cos_sin_cache_rules)
+    count = fuse_cos_sin_cache(model, debug=debug)
     if count != 0:
         remove_unused_nodes(model)
     return count

onnxscript/rewriter/ort_fusions/gelu.py

@@ -34,4 +34,5 @@ def rewrite(self, op, x):
 
 
 def fuse_gelu(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(gelu_rules, model, debug=debug)
+    fuse_gelu = _fusion_utils.apply_fusion_rules(gelu_rules)
+    return fuse_gelu(model, debug=debug)

onnxscript/rewriter/ort_fusions/gqa.py

@@ -151,6 +151,7 @@ def rewrite(
 
 
 def fuse_gqa(model: ir.Model, debug: bool = False) -> int:
-    count = _fusion_utils.apply_fusion_rules(gqa_rules, model, debug=debug)
+    fuse_gqa = _fusion_utils.apply_fusion_rules(gqa_rules)
+    count = fuse_gqa(model, debug=debug)
     remove_unused_nodes(model)
     return count

onnxscript/rewriter/ort_fusions/mha.py

@@ -261,4 +261,5 @@ def rewrite(
 
 
 def fuse_mha(model: ir.Model, *, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(mha_rules, model, debug=debug)
+    fuse_mha = _fusion_utils.apply_fusion_rules(mha_rules)
+    return fuse_mha(model, debug=debug)

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -92,4 +92,5 @@ def rewrite(self, op, x, scale, epsilon, compute_dtype, target_dtype):
 
 
 def fuse_rms_normalization(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(rms_normalization_ruleset, model, debug=debug)
+    fuse_rms_normalization = _fusion_utils.apply_fusion_rules(rms_normalization_ruleset)
+    return fuse_rms_normalization(model, debug=debug)

onnxscript/rewriter/ort_fusions/rotary_embedding.py

@@ -121,8 +121,10 @@ def rewrite(self, op, x, end1, x_part_1_rope, **_):
 
 
 def fuse_rotary_embedding(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(rotary_embedding_rules, model, debug=debug)
+    fuse_rotary_embedding = _fusion_utils.apply_fusion_rules(rotary_embedding_rules)
+    return fuse_rotary_embedding(model, debug=debug)
 
 
 def fuse_partial_rotary_embedding(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(partial_embedding_rules, model, debug=debug)
+    fuse_partial_rotary_embedding = _fusion_utils.apply_fusion_rules(partial_embedding_rules)
+    return fuse_partial_rotary_embedding(model, debug=debug)

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -131,4 +131,5 @@ def rewrite(self, op, query, key_transposed, value, mask, **_):
 
 
 def fuse_sdpa(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(sdpa_rules, model, debug=debug)
+    fuse_sdpa = _fusion_utils.apply_fusion_rules(sdpa_rules)
+    return fuse_sdpa(model, debug=debug)

onnxscript/rewriter/ort_fusions/skip_normalization.py

@@ -6,7 +6,7 @@
 from onnxscript.rewriter import _fusion_utils, pattern
 
 
-def _skip_rmsnorm_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,
@@ -33,14 +33,14 @@ def _skip_rms_normalization(op, input, skip, gamma, epsilon, stash_type):
 
 
 _skip_rms_rule = pattern.RewriteRule(
-    _skip_rmsnorm_pattern, _skip_rms_normalization, matcher=pattern.SimplePatternMatcher
+    _skip_rms_norm_pattern, _skip_rms_normalization, matcher=pattern.SimplePatternMatcher
 )
 
 skip_rms_normalization_rules = [_skip_rms_rule]
 skip_rms_normalization_ruleset = pattern.RewriteRuleSet(skip_rms_normalization_rules)
 
 
-def _skip_layernorm_pattern(op, input, skip, gamma, beta, epsilon, stash_type):
+def _skip_layer_norm_pattern(op, input, skip, gamma, beta, epsilon, stash_type):
     skip_sum = op.Add(input, skip)
     normalized = op.LayerNormalization(
         skip_sum,
@@ -69,18 +69,22 @@ def _skip_layer_normalization(op, input, skip, gamma, beta, epsilon, stash_type)
 
 
 _skip_layer_rule = pattern.RewriteRule(
-    _skip_layernorm_pattern, _skip_layer_normalization, matcher=pattern.SimplePatternMatcher
+    _skip_layer_norm_pattern, _skip_layer_normalization, matcher=pattern.SimplePatternMatcher
 )
 
 skip_layer_normalization_rules = [_skip_layer_rule]
 skip_layer_normalization_ruleset = pattern.RewriteRuleSet(skip_layer_normalization_rules)
 
 
 def fuse_skip_rms_normalization(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(skip_rms_normalization_ruleset, model, debug=debug)
+    fuse_skip_rms_normalization = _fusion_utils.apply_fusion_rules(
+        skip_rms_normalization_ruleset
+    )
+    return fuse_skip_rms_normalization(model, debug=debug)
 
 
 def fuse_skip_layer_normalization(model: ir.Model, debug: bool = False) -> int:
-    return _fusion_utils.apply_fusion_rules(
-        skip_layer_normalization_ruleset, model, debug=debug
+    fuse_skip_layer_normalization = _fusion_utils.apply_fusion_rules(
+        skip_layer_normalization_ruleset
     )
+    return fuse_skip_layer_normalization(model, debug=debug)