Cleanup optimizer

onnxscript/optimizer/__init__.py

@@ -2,160 +2,22 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-import logging
-from typing import Any
-
 import onnx
-import onnx.shape_inference
-
-from onnxscript import ir, rewriter
-from onnxscript.optimizer import _constant_folding, _inliner
-from onnxscript.optimizer.constant_folding import fold_constants
-from onnxscript.optimizer.remove_unused import remove_unused_nodes
-from onnxscript.optimizer.remove_unused_function import remove_unused_functions
-from onnxscript.optimizer.simple_function_folding import (
- inline_functions_with_unused_outputs,
- inline_simple_functions,
-)
-from onnxscript.rewriter import (
- broadcast_to_matmul,
- cast_constant_of_shape,
- gemm_to_matmul_add,
- no_op,
-)
-
-logger = logging.getLogger(__name__)
-
-_DEFAULT_REWRITE_RULES = [
- *no_op.rules.rules, # TODO: merge this rule into constant folding?
- *broadcast_to_matmul.rules.rules,
- gemm_to_matmul_add.rule,
- *cast_constant_of_shape.rules.rules,
-]
-
-
-def optimize(
- model: onnx.ModelProto,
- num_iterations: int = 2,
- *,
- onnx_shape_inference: bool = True,
- stop_if_no_change: bool = True,
- external_data_folder: str = "",
- **kwargs: Any,
-) -> onnx.ModelProto:
- """Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.
-
- Args:
- model (onnx.ModelProto): The model to optimize.
- num_iterations (int, optional): Number of iterations to perform.
- onnx_shape_inference (bool, optional): Whether to perform onnx shape inference on the model.
- Set this to False to turn off onnx shape inference, and rely on model carried shapes and types.
- This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries
- the symbolic shapes recorded from dynamo tracing.
- stop_if_no_change (bool, optional): Whether to stop if no change is detected.
- external_data_folder (str, optional): The folder to store external data.
- **kwargs: Additional keyword arguments. For BC purposes.
- """
- if kwargs.pop("function_aware_folding", None) is not None:
- logger.warning(
- "'function_aware_folding' is deprecated. 'optimize' now supports both fully inlined models and models with functions. "
- "To achieve the same behavior as 'function_aware_folding=True' before, set 'onnx_shape_inference=False'. "
- "This would turn off incremental onnx shape inference and rely on model carried shapes and types. "
- "See 'onnx_shape_inference' for more details."
- )
- for _ in range(num_iterations):
- if onnx_shape_inference:
- if model.ByteSize() < 1024 * 1024 * 1024 * 2:
- # NOTE: strict mode is disabled because it crashes on the models
- # that have different shapes inferred from the model carried shapes.
- # The case can be found in:
- # https://github.com/microsoft/onnxscript/issues/1443
- model = onnx.shape_inference.infer_shapes(
- model, check_type=True, strict_mode=False, data_prop=True
- )
- else:
- logger.warning(
- "The model size is too large for full model shape inference. "
- "Skipping this step."
- )
-
- inline_simple_functions(model)
- modified = fold_constants(
- model, external_data_folder, onnx_shape_inference=onnx_shape_inference
- )
-
- remove_unused_nodes(model)
- inline_simple_functions(model)
- model = remove_unused_functions(model)
- inline_functions_with_unused_outputs(model)
- # NOTE: This is general rewrite rules
- model = rewriter.rewrite(model, pattern_rewrite_rules=_DEFAULT_REWRITE_RULES)
- if stop_if_no_change and not modified:
- logger.debug("Stopping after %d iterations.", _)
- break
-
- for node in model.graph.node:
- logger.debug("Node %s::%s name %s.", node.domain, node.op_type, node.name)
-
- for function in model.functions:
- for node in function.node:
- logger.debug(
- "Function %s::%s node %s::%s name %s.",
- function.domain,
- function.name,
- node.domain,
- node.op_type,
- node.name,
- )
-
- return model
-
-
-_DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT = (
- _constant_folding._DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT
-)
-
-_DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT = (
- _constant_folding._DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT
-)
-
 
-def optimize_ir(
- model: ir.Model,
- num_iterations: int = 2,
- *,
- onnx_shape_inference: bool = True,
- stop_if_no_change: bool = True,
- input_size_limit: int = _DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
- output_size_limit: int = _DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
-) -> None:
- """Optimizes a model.
+import onnxscript.optimizer._legacy._optimizer as legacy_optimizer
+from onnxscript import ir
+from onnxscript.optimizer._constant_folding import basic_constant_propagation
+from onnxscript.optimizer._legacy.constant_folding import fold_constants
+from onnxscript.optimizer._optimizer import optimize_ir
+from onnxscript.optimizer._remove_unused import remove_unused_nodes
 
- Args:
- model: The model to be optimized.
- num_iterations: Number of times the optimization loop is repeated.
- onnx_shape_inference: Applies node-level shape-inference as part of optimization
- input_size_limit: Will not apply constant folding to ops with any input of size
- greater than this. Does not apply to special ops like Shape() and Size().
- output_size_limit: Will not rewrite any foldable-op into a Constant op if the size
- of the output tensor is greater than this.
- stop_if_no_change: Not supported currently (has no effect). Meant to stop the
- outer optimization loop if no change is detected in one iteration.
- """
- del stop_if_no_change # Looks like rewriter doesn't support this yet.
- _inliner.inline(model)
- for _ in range(num_iterations):
- _constant_folding.fold_constants(
- model,
- onnx_shape_inference=onnx_shape_inference,
- input_size_limit=input_size_limit,
- output_size_limit=output_size_limit,
- )
- rewriter.rewrite(model, pattern_rewrite_rules=_DEFAULT_REWRITE_RULES)
- remove_unused_nodes(model)
 
+def optimize(model: ir.Model | onnx.ModelProto, *args, **kwargs):
+ if isinstance(model, ir.Model):
+ return optimize_ir(model, *args, **kwargs)
+ else:
+ return legacy_optimizer.optimize(model, *args, **kwargs)
 
-basic_constant_propagation = _constant_folding.basic_constant_propagation
 
 __all__ = [
  "fold_constants",

onnxscript/optimizer/_constant_folding.py

@@ -17,20 +17,32 @@
 
 import onnxscript.ir as ir
 import onnxscript.ir._convenience as _convenience
-import onnxscript.optimizer.constant_folding as constant_folding
 import onnxscript.rewriter.pattern as orp
 import onnxscript.utils.utils as utils
 
+DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT = 1024
+
+DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT = 1024 * 1024
+
 
 def is_control_flow_op(node: ir.Node) -> bool:
  graph_types = {ir.AttributeType.GRAPH, ir.AttributeType.GRAPHS}
  return any(attr.type in graph_types for attr in node.attributes.values())
 
 
+non_deterministic_ops = frozenset(
+ {
+ "RandomUniform",
+ "RandomNormal",
+ "RandomUniformLike",
+ "RandomNormalLike",
+ "Multinomial",
+ }
+)
+
+
 def is_non_deterministic_op(node: ir.Node) -> bool:
- return node.op_type in constant_folding.non_deterministic_ops and utils.is_onnx_domain(
- node.domain
- )
+ return node.op_type in non_deterministic_ops and utils.is_onnx_domain(node.domain)
 
 
 def is_onnx_op(node: ir.Node, op_type: str) -> bool:
@@ -43,10 +55,6 @@ def is_constant_op(node: ir.Node) -> bool:
  )
 
 
-_DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT = 1024
-
-_DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT = constant_folding._DEFAULT_CONSTANT_FOLD_SIZE_LIMIT
-
 logger = logging.getLogger(__name__)
 
 # "Standard" evaluators are used to perform constant-folding.
@@ -787,8 +795,8 @@ def fold_constants(
  external_data_folder: str = "",
  *,
  onnx_shape_inference: bool = False,
- input_size_limit: int = _DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
- output_size_limit: int = _DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
+ input_size_limit: int = DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
+ output_size_limit: int = DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
 ) -> bool:
  """
  Applies constant folding optimization to the model.

onnxscript/optimizer/constant_folding_test.py → onnxscript/optimizer/_constant_folding_test.py

@@ -8,7 +8,8 @@
 
 import onnxscript.optimizer as optimizer
 from onnxscript.ir import serde
-from onnxscript.optimizer import _constant_folding, constant_folding
+from onnxscript.optimizer import _constant_folding
+from onnxscript.optimizer._legacy import constant_folding
 
 
 @parameterized.parameterized_class(("using_ir",), [(False,), (True,)])

onnxscript/optimizer/_legacy/_optimizer.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import onnx
+import onnx.shape_inference
+
+import onnxscript.optimizer
+from onnxscript import rewriter
+from onnxscript.optimizer._optimizer import _DEFAULT_REWRITE_RULES
+from onnxscript.optimizer._legacy.constant_folding import fold_constants
+from onnxscript.optimizer._remove_unused import remove_unused_nodes
+from onnxscript.optimizer._remove_unused_function import remove_unused_functions
+from onnxscript.optimizer._legacy._simple_function_folding import (
+ inline_functions_with_unused_outputs,
+ inline_simple_functions,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def optimize(
+ model: onnx.ModelProto,
+ num_iterations: int = 2,
+ *,
+ onnx_shape_inference: bool = True,
+ stop_if_no_change: bool = True,
+ external_data_folder: str = "",
+ **kwargs: Any,
+) -> onnx.ModelProto:
+ """Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.
+
+ Args:
+ model (onnx.ModelProto): The model to optimize.
+ num_iterations (int, optional): Number of iterations to perform.
+ onnx_shape_inference (bool, optional): Whether to perform onnx shape inference on the model.
+ Set this to False to turn off onnx shape inference, and rely on model carried shapes and types.
+ This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries
+ the symbolic shapes recorded from dynamo tracing.
+ stop_if_no_change (bool, optional): Whether to stop if no change is detected.
+ external_data_folder (str, optional): The folder to store external data.
+ **kwargs: Additional keyword arguments. For BC purposes.
+ """
+ if kwargs.pop("function_aware_folding", None) is not None:
+ logger.warning(
+ "'function_aware_folding' is deprecated. 'optimize' now supports both fully inlined models and models with functions. "
+ "To achieve the same behavior as 'function_aware_folding=True' before, set 'onnx_shape_inference=False'. "
+ "This would turn off incremental onnx shape inference and rely on model carried shapes and types. "
+ "See 'onnx_shape_inference' for more details."
+ )
+ for _ in range(num_iterations):
+ if onnx_shape_inference:
+ if model.ByteSize() < 1024 * 1024 * 1024 * 2:
+ # NOTE: strict mode is disabled because it crashes on the models
+ # that have different shapes inferred from the model carried shapes.
+ # The case can be found in:
+ # https://github.com/microsoft/onnxscript/issues/1443
+ model = onnx.shape_inference.infer_shapes(
+ model, check_type=True, strict_mode=False, data_prop=True
+ )
+ else:
+ logger.warning(
+ "The model size is too large for full model shape inference. "
+ "Skipping this step."
+ )
+
+ inline_simple_functions(model)
+ modified = fold_constants(
+ model, external_data_folder, onnx_shape_inference=onnx_shape_inference
+ )
+
+ remove_unused_nodes(model)
+ inline_simple_functions(model)
+ model = remove_unused_functions(model)
+ inline_functions_with_unused_outputs(model)
+ # NOTE: This is general rewrite rules
+ model = rewriter.rewrite(model, pattern_rewrite_rules=_DEFAULT_REWRITE_RULES)
+ if stop_if_no_change and not modified:
+ logger.debug("Stopping after %d iterations.", _)
+ break
+
+ for node in model.graph.node:
+ logger.debug("Node %s::%s name %s.", node.domain, node.op_type, node.name)
+
+ for function in model.functions:
+ for node in function.node:
+ logger.debug(
+ "Function %s::%s node %s::%s name %s.",
+ function.domain,
+ function.name,
+ node.domain,
+ node.op_type,
+ node.name,
+ )
+
+ return model

onnxscript/optimizer/simple_function_folding.py → onnxscript/optimizer/_legacy/_simple_function_folding.py

@@ -11,7 +11,7 @@
 
 import onnxscript._legacy_ir as ir
 from onnxscript._legacy_ir import visitor
-from onnxscript.optimizer import remove_unused_proto
+from onnxscript.optimizer._legacy import _remove_unused_proto
 
 logger = logging.getLogger(__name__)
 
@@ -168,7 +168,7 @@ def _find_nodes_with_any_unused_output(
  # All unused output means the node is not used at all.
  # Hence do not update used_values with the node's inputs.
  continue
- used_values |= remove_unused_proto.compute_used_in_node(node)
+ used_values |= _remove_unused_proto.compute_used_in_node(node)
  return target_nodes
 
  def visit_model(self, model: onnx.ModelProto) -> None: