Extend optimize_for_ort to cover passes

onnxscript/rewriter/ort_fusions/_core.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import onnxscript.ir as ir
-from onnxscript.ir.passes.common import shape_inference
+import onnxscript.ir.passes.common as common_passes
 from onnxscript.optimizer import optimize
 from onnxscript.rewriter import rewrite
 from onnxscript.rewriter.ort_fusions import (
@@ -48,7 +48,7 @@
     # 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.
-    shape_inference.infer_shapes(model)
+    common_passes.ShapeInferencePass()(model)
     optimize(model)
     return model
 
@@ -135,4 +135,18 @@
     )
     # Apply the ORT pattern rewrite rules.
     rewrite(model, ORT_PATTERN_REWRITE_RULES)
-    return model, fusion_count
+
+    passes = [
+        # TODO(exporter team): Fold transpose into initializers
+        # Apply the ORT optimization passes.
+        # https://github.com/microsoft/onnxruntime/blob/74dcf7e296639095dfa55d31336998b6f719ed76/onnxruntime/python/tools/transformers/dynamo_onnx_helper.py#L172
+        common_passes.ClearMetadataAndDocStringPass(),
+        # https://github.com/microsoft/onnxruntime/blob/74dcf7e296639095dfa55d31336998b6f719ed76/onnxruntime/python/tools/transformers/dynamo_onnx_helper.py#L139
+        common_passes.LiftConstantsToInitializersPass(lift_all_constants=False, size_limit=1),
+        common_passes.RemoveInitializersFromInputsPass(),
+        common_passes.ShapeInferencePass(),
+        common_passes.CheckerPass(),
+    ]
+    optimize_for_ort_passes = ir.passes.Sequential(*passes)
+    result = optimize_for_ort_passes(model)
+    return result.model, fusion_count