rebase modifs

Author: shubhambhokare1

onnxscript/rewriter/__init__.py

@@ -18,7 +18,6 @@
     broadcast_to_matmul,
     cast_constant_of_shape,
     collapse_slices,
-    erfgelu,
     gemm_to_matmul_add,
     llama_rule_sets,
     no_op,
@@ -29,7 +28,6 @@
 _DEFAULT_REWRITE_RULES: tuple[pattern.RewriteRule, ...] = (
     *no_op.rules.rules,  # TODO: merge this rule into constant folding?
     *broadcast_to_matmul.rules.rules,
-    *erfgelu.rules.rules,  # type: ignore[has-type]
     gemm_to_matmul_add.rule,  # type: ignore[has-type]
     *cast_constant_of_shape.rules.rules,
     *collapse_slices.rules.rules,

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -2,7 +2,9 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-from onnxscript.rewriter import _fusion_utils, pattern
+import math
+
+from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
 
 class SDPA(pattern.RewriteRuleClassBase):
@@ -166,34 +168,23 @@ def rewrite(
         return op.SDPA(*sdpa_args, scale=self._scale, _domain="ai.onnxruntime.fusion")
 
 
-# Rules for SDPA without mask
-unmasked_pre_div_sdpa_rule = SDPA.rule(
-    "unmasked_pre_div_sdpa", use_mask=False, pre_scale=True, use_mul=False
-)
-unmasked_pre_mul_sdpa_rule = SDPA.rule(
-    "unmasked_pre_mul_sdpa", use_mask=False, pre_scale=True, use_mul=True
-)
-unmasked_post_div_sdpa_rule = SDPA.rule(
-    "unmasked_post_div_sdpa", use_mask=False, pre_scale=False, use_mul=False
-)
-unmasked_post_mul_sdpa_rule = SDPA.rule(
-    "unmasked_post_mul_sdpa", use_mask=False, pre_scale=False, use_mul=True
-)
-
-# Rules for SDPA with mask
-masked_pre_div_sdpa_rule = SDPA.rule(
-    "masked_pre_div_sdpa", use_mask=True, pre_scale=True, use_mul=False
-)
-masked_pre_mul_sdpa_rule = SDPA.rule(
-    "masked_pre_mul_sdpa", use_mask=True, pre_scale=True, use_mul=True
-)
-masked_post_div_sdpa_rule = SDPA.rule(
-    "masked_post_div_sdpa", use_mask=True, pre_scale=False, use_mul=False
-)
-masked_post_mul_sdpa_rule = SDPA.rule(
-    "masked_post_mul_sdpa", use_mask=True, pre_scale=False, use_mul=True
-)
-
+parameter_combinations = [
+    {
+        "name": f"sdpa_{'masked_' if use_mask else 'unmasked_'}{'pre_' if pre_scale else 'post_'}{'only_q_' if pre_scale_q else ''}{'mul' if use_mul else 'div'}{'_3d' if has_3d_inputs else ''}",
+        "use_mask": use_mask,
+        "pre_scale": pre_scale,
+        "pre_scale_q": pre_scale_q,
+        "use_mul": use_mul,
+        "has_3d_inputs": has_3d_inputs,
+    }
+    for use_mask in [False, True]
+    for pre_scale in [False, True]
+    for pre_scale_q in [False, True]
+    for use_mul in [False, True]
+    for has_3d_inputs in [False, True]
+]
+
+# Dynamically create the rules
 sdpa_rules = pattern.RewriteRuleSet(
     [
         SDPA.rule(