|
|
|
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
import logging |
|
|
|
|
|
import onnx_ir as ir |
|
|
import onnx_ir.passes.common as common_passes |
|
|
|
|
|
from onnxscript import rewriter |
|
|
from onnxscript.optimizer import _constant_folding |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
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 = _constant_folding.DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT, |
|
|
output_size_limit: int = _constant_folding.DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT, |
|
|
inline: bool = True, |
|
|
) -> None: |
|
|
"""Optimizes a model. |
|
|
|
|
|
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: Stop the optimization loop if no change is detected in an iteration. |
|
|
inline: If True, inlines all functions in the model. |
|
|
""" |
|
|
passes = [ |
|
|
ir.passes.PassManager( |
|
|
[ |
|
|
_constant_folding.FoldConstantsPass( |
|
|
shape_inference=onnx_shape_inference, |
|
|
input_size_limit=input_size_limit, |
|
|
output_size_limit=output_size_limit, |
|
|
), |
|
|
rewriter.RewritePass(rewriter._DEFAULT_REWRITE_RULES), |
|
|
common_passes.RemoveUnusedNodesPass(), |
|
|
common_passes.RemoveUnusedFunctionsPass(), |
|
|
common_passes.RemoveUnusedOpsetsPass(), |
|
|
], |
|
|
steps=num_iterations, |
|
|
early_stop=stop_if_no_change, |
|
|
), |
|
|
common_passes.RemoveUnusedNodesPass(), |
|
|
common_passes.LiftConstantsToInitializersPass(lift_all_constants=True, size_limit=0), |
|
|
common_passes.LiftSubgraphInitializersToMainGraphPass(), |
|
|
common_passes.DeduplicateInitializersPass(), |
|
|
common_passes.CommonSubexpressionEliminationPass(), |
|
|
] |
|
|
if inline: |
|
|
|
|
|
passes = [common_passes.InlinePass(), *passes] |
|
|
optimizer_pass = ir.passes.Sequential(*passes) |
|
|
assert optimizer_pass.in_place |
|
|
result = optimizer_pass(model) |
|
|
assert result.model is model |
|
|
|
