[Test] Hack to enable optimizer/rewriter integration into dynamo_export

- To test in torchbench.
- Somehow lintrunner changed unrelated files in this commit.

ghstack-source-id: 2f6b038
Pull Request resolved: #1334

Author: BowenBao

onnxscript/function_libs/torch_lib/graph_building.py

@@ -17,11 +17,15 @@
 from typing_extensions import TypeAlias
 
 import onnxscript
+
+# Q: Don't know why vscode only recognizes this import style instead of `from onnxscript import optimizer`.
+import onnxscript.optimizer as optimizer
 from onnxscript import evaluator
 from onnxscript import tensor as onnxscript_tensor
 from onnxscript._internal import param_manipulation, runtime_typing
 from onnxscript.function_libs.torch_lib import _flags
 from onnxscript.function_libs.torch_lib.ops import common as common_ops
+from onnxscript.rewriter import onnxruntime as ort_rewriter
 
 __all__ = [
     "TorchScriptTensor",
@@ -1071,4 +1075,9 @@ def to_model_proto(
                 common_ops.common_opset.domain, common_ops.common_opset.version
             )
         )
+
+        # Not the best integration point. Enables benchmarking the migration.
+        onnx_model = optimizer.optimize(onnx_model)
+        # This also creates contrib op in the model. So definitely not the best integration point.
+        onnx_model = ort_rewriter.rewrite(onnx_model)
         return onnx_model

onnxscript/optimizer/constant_folding.py

@@ -41,9 +41,7 @@ def is_non_deterministic_op(node: onnx.NodeProto) -> bool:
 
 
 def is_constant_op(node: onnx.NodeProto) -> bool:
-    return node.op_type in {"Constant", "ConstantOfShape"} and is_onnx_domain(
-        node.domain
-    )
+    return node.op_type in {"Constant", "ConstantOfShape"} and is_onnx_domain(node.domain)
 
 
 class ConstantFolder(visitor.FunctionCallsiteProtoTransformer):
@@ -119,14 +117,10 @@ def new_constant(self, name, value):
         info.type = onnx.helper.make_tensor_type_proto(
             onnx.helper.np_dtype_to_tensor_dtype(value.dtype), value.shape
         )
-        node = onnx.helper.make_node(
-            "Constant", inputs=[], outputs=[name], value=tensor
-        )
+        node = onnx.helper.make_node("Constant", inputs=[], outputs=[name], value=tensor)
         return [node]
 
-    def convert_attributes(
-        self, attributes: Sequence[onnx.AttributeProto]
-    ) -> dict[str, Any]:
+    def convert_attributes(self, attributes: Sequence[onnx.AttributeProto]) -> dict[str, Any]:
         if self.scopes.current_scope().current_function_scope():
             # Need to resolve ref_attr_name if inside a function.
             attr_dict = {}
@@ -138,9 +132,7 @@ def convert_attributes(
                 )
                 if concrete_attribute is None:
                     continue
-                attr_dict[attribute.name] = onnx.helper.get_attribute_value(
-                    concrete_attribute
-                )
+                attr_dict[attribute.name] = onnx.helper.get_attribute_value(concrete_attribute)
             return attr_dict
         return {attr.name: onnx.helper.get_attribute_value(attr) for attr in attributes}
 
@@ -226,9 +218,7 @@ def process_node(self, node: onnx.NodeProto) -> Sequence[onnx.NodeProto] | None:
             self.add_count(op, outputs.size)
             return replacement
         else:
-            logger.warning(
-                "Skipping constant folding for op %s with multiple outputs.", op
-            )
+            logger.warning("Skipping constant folding for op %s with multiple outputs.", op)
         return None
 
     def process_function_node(
@@ -241,9 +231,7 @@ def process_function_node(
         # Replace function node with Constant if all outputs are constants
         ir_values = [self.lookup(output_name) for output_name in node.output]
         tensors = [
-            self.foldable_value(
-                output_name, ir_value.value if ir_value is not None else None
-            )
+            self.foldable_value(output_name, ir_value.value if ir_value is not None else None)
             for output_name, ir_value in zip(node.output, ir_values)
         ]
         if all(tensor is not None for tensor in tensors):