From cfd3816aa4768898d1b7226bc0423ebd969fa035 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 11 Oct 2024 10:19:14 -0700 Subject: [PATCH 1/4] Use input size limits for constant folding --- onnxscript/optimizer/__init__.py | 18 ++++++++++++- onnxscript/optimizer/_constant_folding.py | 33 ++++++++++++++++++----- 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index b35f70a52a..d64a69d1c6 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -111,17 +111,33 @@ def optimize( 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: 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) + _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) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 818fd95e10..3e210ce595 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -43,7 +43,9 @@ def is_constant_op(node: ir.Node) -> bool: ) -_DEFAULT_CONSTANT_FOLD_SIZE_LIMIT = constant_folding._DEFAULT_CONSTANT_FOLD_SIZE_LIMIT +_DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT = 1024 + +_DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT = constant_folding._DEFAULT_CONSTANT_FOLD_SIZE_LIMIT logger = logging.getLogger(__name__) @@ -550,11 +552,16 @@ class ConstantFolder: def __init__( self, + *, external_data_folder: str, - do_shape_inference: bool, + shape_inference: bool, + input_size_limit: int, + output_size_limit: int, ) -> None: self._external_data_folder = external_data_folder - self._do_shape_inference = do_shape_inference + self._shape_inference = shape_inference + self._input_size_limit = input_size_limit + self._output_size_limit = output_size_limit self._init() def _init(self) -> None: @@ -632,7 +639,7 @@ def new_constant(self, irvalue: ir.Value, value): irvalue.const_value = _convenience.tensor(value) - if value.nbytes > _DEFAULT_CONSTANT_FOLD_SIZE_LIMIT: + if value.nbytes > self._output_size_limit: logger.info( "Skip storing constant folded nvalue %s due to large size %s.", irvalue.name, @@ -667,7 +674,7 @@ def process_node(self, node: ir.Node): # TODO(rama): consider merging type/other info from both values # Do incremental shape inference - if self._do_shape_inference and not is_control_flow_op(node): + if self._shape_inference and not is_control_flow_op(node): self._do_inference(node) if node.domain not in self.opset_imports: @@ -696,6 +703,14 @@ def process_node(self, node: ir.Node): if any(x is None for x in input_values): return None + if any(input.size > self._input_size_limit for input in input_values): + if logger.isEnabledFor(logging.DEBUG): + input_sizes = [input.size for input in input_values] + logger.debug( + f"Skipping constant folding for op {node.op_type} due to large input size: {input_sizes}" + ) + return None + # Filter out bfloat16 cases? def convert(av): if av.type == ir.AttributeType.TENSOR: @@ -770,14 +785,18 @@ 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, ) -> bool: """ Applies constant folding optimization to the model. Returns true iff the model was modified. """ folder = ConstantFolder( - external_data_folder, - onnx_shape_inference, + external_data_folder=external_data_folder, + shape_inference=onnx_shape_inference, + input_size_limit=input_size_limit, + output_size_limit=output_size_limit, ) folder.visit_model(model) for op in folder.counts: From ac55fa25b71be88e8c2a6f8711add2e53aa02cd9 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 11 Oct 2024 13:14:54 -0700 Subject: [PATCH 2/4] Address PR feedback --- onnxscript/optimizer/__init__.py | 14 ++++++++++++++ onnxscript/optimizer/_constant_folding.py | 4 +++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index d64a69d1c6..e643762aa7 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -129,6 +129,20 @@ def optimize_ir( input_size_limit: int = _DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT, output_size_limit: int = _DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT, ) -> 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: 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): diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 3e210ce595..5d76379cbf 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -707,7 +707,9 @@ def process_node(self, node: ir.Node): if logger.isEnabledFor(logging.DEBUG): input_sizes = [input.size for input in input_values] logger.debug( - f"Skipping constant folding for op {node.op_type} due to large input size: {input_sizes}" + "Skipping constant folding for op %s due to large input size: %s", + node.op_type, + input_sizes, ) return None From 3c8ad105be1674fa55d55c8132b75533e75903f4 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 11 Oct 2024 13:17:10 -0700 Subject: [PATCH 3/4] Address type warning --- onnxscript/optimizer/_constant_folding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 5d76379cbf..1144f207ab 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -703,9 +703,9 @@ def process_node(self, node: ir.Node): if any(x is None for x in input_values): return None - if any(input.size > self._input_size_limit for input in input_values): + if any(input.size > self._input_size_limit for input in input_values): # type: ignore[union-attr] if logger.isEnabledFor(logging.DEBUG): - input_sizes = [input.size for input in input_values] + input_sizes = [input.size for input in input_values] # type: ignore[union-attr] logger.debug( "Skipping constant folding for op %s due to large input size: %s", node.op_type, From 8b80e16d7d54227bab720190374a0ef683207264 Mon Sep 17 00:00:00 2001 From: "G. Ramalingam" Date: Fri, 11 Oct 2024 21:02:05 -0700 Subject: [PATCH 4/4] Update onnxscript/optimizer/__init__.py Co-authored-by: Justin Chu --- onnxscript/optimizer/__init__.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index e643762aa7..985ac6f109 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -129,12 +129,11 @@ def optimize_ir( input_size_limit: int = _DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT, output_size_limit: int = _DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT, ) -> None: - """ - Optimizes a model. + """Optimizes a model. Args: - model: The model to be optimized - num_iterations: Number of times the optimization loop is repeated + 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().