Cleanup ort transformer fusions

onnxscript/rewriter/ort_fusions/_core.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import onnxscript.ir as ir
+from onnxscript.ir.passes.common import shape_inference
 from onnxscript.optimizer import optimize, remove_unused_nodes
 from onnxscript.rewriter import rewrite
 from onnxscript.rewriter.ort_fusions import (
@@ -12,9 +13,13 @@
     softmax,
 )
 from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
+from onnxscript.rewriter.ort_fusions.gelu import fuse_gelu
 from onnxscript.rewriter.ort_fusions.mha import fuse_mha
 from onnxscript.rewriter.ort_fusions.rms_normalization import fuse_rms_normalization
-from onnxscript.rewriter.ort_fusions.rotary_embedding import fuse_rotary_embedding
+from onnxscript.rewriter.ort_fusions.rotary_embedding import (
+    fuse_partial_rotary_embedding,
+    fuse_rotary_embedding,
+)
 from onnxscript.rewriter.ort_fusions.sdpa import fuse_sdpa
 from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_normalization
 
@@ -27,14 +32,29 @@
 ]
 
 
-def fuse_xformers(model: ir.Model) -> None:
+# Preliminary optimizations before applying the transformer fusions.
+# 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.
+    model = shape_inference.infer_shapes(model)
     optimize(model)
+    return model
+
+
+def fuse_xformers(model: ir.Model) -> None:
+    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_gelu(model)
     remove_unused_nodes(model)