Add fusion rules (Whisper optimizations) (#2221)

Add fusion rules to support the optimization of Whisper models.

Fusions added:

- Basic Fusions:
    * additional pattern for erfgelu [moved to #2222]
- SkipLayerNorm:
    * #2259 
    * Fusion patterns where skip_sum is also an output
    * Bias + SkipLayerNorm -> SkipLayerNorm (with bias) [moved to #2222]
- BiasGelu Fusion [moved to #2222]
- SDPA:
    * Support for pattern where only q is pre-scaled
- MHA:
    * Patterns with/without past/present keys/values
    * Patterns with non-rotary embeddings
    * Patterns with/without mask
    * Patterns with cross-attention (only for past key/value patterns)
 - MHA Bias Fusion:
* Bias was offloaded to Attention fusion previously, this fusion fixes
that
 - Attention:
    * Patterns where Q, K and V do not come from slicing

TODO:
- [x]  Fix SDPA singular prescale case, due to lost shape information
- [x] - Enable check conditions when #2210 is merged
- [x] - Improve/Rewrite whisper model test case to be similar to that of
smollm (for eg)
- [x] - Fix failing test cases to account for new patterns
- [x] - Add isolated test cases for new fusions like BiasGelu,
SkipLayerNorm etc

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/_core.py

@@ -16,6 +16,7 @@
 from onnxscript.rewriter.ort_fusions.bias_gelu import fuse_bias_gelu
 from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
 from onnxscript.rewriter.ort_fusions.erfgelu import fuse_erfgelu
+from onnxscript.rewriter.ort_fusions.fuse_mha_bias import fuse_mha_bias
 from onnxscript.rewriter.ort_fusions.fuse_packed_qkv_gqa import fuse_qkv_gqa
 from onnxscript.rewriter.ort_fusions.gelu import fuse_gelu
 from onnxscript.rewriter.ort_fusions.gqa import fuse_gqa
@@ -79,6 +80,8 @@ def fuse(func, apply_shape_inference: bool = False):
     fusion_count["rotary_embedding"] = fuse(fuse_rotary_embedding)
     fusion_count["partial_rotary_embedding"] = fuse(fuse_partial_rotary_embedding)
     fusion_count["cos_sin_cache"] = fuse(fuse_cos_sin_cache)
+    # We apply shape inference after the SDPA fusion as new nodes are added
+    # in the rewrite rule for certain patterns of SDPA.
     fusion_count["sdpa"] = fuse(fuse_sdpa, apply_shape_inference=True)
     # Optimize to avoid trying multiple attention-based fusions
     fusion_count["mha"] = fuse(fuse_mha)
@@ -87,8 +90,10 @@ def fuse(func, apply_shape_inference: bool = False):
         # and avoid trying the attention fusion.
         fusion_count["gqa"] = fuse(fuse_gqa)
         fusion_count["packed_qkv_for_gqa"] = fuse(fuse_qkv_gqa)
+        fusion_count["mha_bias"] = 0
         fusion_count["attention"] = 0
     else:
+        fusion_count["mha_bias"] = fuse(fuse_mha_bias)
         fusion_count["attention"] = fuse(fuse_attention)
         fusion_count["gqa"] = 0
     fusion_count["gelu"] = fuse(fuse_gelu)