Minor cleanup

onnxscript/rewriter/ort_fusions/__init__.py

@@ -7,18 +7,5 @@
     "ORT_PATTERN_REWRITE_RULES",
 ]
 
-from onnxscript.rewriter.ort_fusions import (
-    fused_matmul_rule_sets,
-    # group_normalization_merge_silu,
-    instance_to_group_normalization,
-    softmax,
-)
-from onnxscript.rewriter.ort_fusions._core import optimize_for_ort
 
-ORT_PATTERN_REWRITE_RULES = [
-    *softmax.rules.rules,
-    *instance_to_group_normalization.rules.rules,
-    # NOTE: group normalization merge silu should be applied after instance to group normalization
-    # *group_normalization_merge_silu.rules.rules,
-    *fused_matmul_rule_sets.fused_matmul_rule_sets(),
-]
+from onnxscript.rewriter.ort_fusions._core import ORT_PATTERN_REWRITE_RULES, optimize_for_ort

onnxscript/rewriter/ort_fusions/_core.py

@@ -4,13 +4,28 @@
 
 import onnxscript.ir as ir
 from onnxscript.optimizer import optimize, remove_unused_nodes
+from onnxscript.rewriter import rewrite
+from onnxscript.rewriter.ort_fusions import (
+    fused_matmul_rule_sets,
+    # group_normalization_merge_silu,
+    instance_to_group_normalization,
+    softmax,
+)
 from onnxscript.rewriter.ort_fusions.cos_sin_cache import fuse_cos_sin_cache
 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.sdpa import fuse_sdpa
 from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_normalization
 
+ORT_PATTERN_REWRITE_RULES = [
+    *softmax.rules.rules,
+    *instance_to_group_normalization.rules.rules,
+    # NOTE: group normalization merge silu should be applied after instance to group normalization
+    # *group_normalization_merge_silu.rules.rules,
+    *fused_matmul_rule_sets.fused_matmul_rule_sets(),
+]
+
 
 def fuse_xformers(model: ir.Model) -> None:
     optimize(model)
@@ -24,5 +39,5 @@
 
 
 def optimize_for_ort(model: ir.Model) -> None:
-    # TODO(rama): Include the other optimizations
+    rewrite(model, ORT_PATTERN_REWRITE_RULES)
     fuse_xformers(model)