# # SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from polygraphy import mod, util from polygraphy.logger import G_LOGGER from polygraphy.tools import util as tools_util from polygraphy.tools.args import ( DataLoaderArgs, ModelArgs, OnnxInferShapesArgs, OnnxLoadArgs, OnnxSaveArgs, ) from polygraphy.tools.args import util as args_util from polygraphy.tools.args.base import BaseArgs from polygraphy.tools.script import make_invocable from polygraphy.tools.surgeon.subtool.base import BaseSurgeonSubtool onnx_backend = mod.lazy_import("polygraphy.backend.onnx") onnx_util = mod.lazy_import("polygraphy.backend.onnx.util") gs = mod.lazy_import("onnx_graphsurgeon") class ConstFoldArgs(BaseArgs): """ Constant Folding: folding constants Depends on: OnnxInferShapesArgs """ def add_parser_args_impl(self): self.group.add_argument( "--fold-constants", help="Fold constants in the graph by computing subgraphs whose values " "are not dependent on runtime inputs.", action="store_true", default=None, ) self.group.add_argument( "--num-passes", "--num-const-fold-passes", help="The number of constant folding passes to run. " "Sometimes, subgraphs that compute tensor shapes may not be foldable in a single pass. " "If not specified, Polygraphy will automatically determine the number of passes required. ", type=int, default=None, dest="num_const_fold_passes", ) self.group.add_argument( "--partitioning", help="Controls how to partition the graph during constant folding: {{" "'basic': Partition the graph so failures in one part do not affect other parts, " "'recursive': In addition to partitioning the graph, partition partitions where needed}} ", choices=["basic", "recursive"], default=None, ) self.group.add_argument( "--no-fold-shapes", help="Disable folding Shape nodes and subgraphs that operate on shapes", dest="fold_shapes", default=None, action="store_false", ) self.group.add_argument( "--no-per-pass-shape-inference", help="Disable shape inference between passes of constant folding", dest="per_pass_shape_inference", default=None, action="store_false", ) self.group.add_argument( "--fold-size-threshold", help=""" The maximum per-tensor size threshold, in bytes, for which to apply constant folding. Any nodes generating tensors larger than this size will not be folded away. For example, some models may apply ops like `Tile` or `Expand` to constants, which can result in very large tensors. Rather than pre-computing those constants and bloating the model size, it may be desirable to skip folding them and allow them to be computed at runtime. Optionally, use a `K`, `M`, or `G` suffix to indicate KiB, MiB, or GiB respectively. For example, `--fold-size-threshold 16M` is equivalent to `--fold-size-threshold 16777216`. """, default=None, ) def parse_impl(self, args): """ Parses command-line arguments and populates the following attributes: Attributes: fold_constants (bool): Whether to apply constant folding. num_passes (int): The number of constant folding passes to apply. partitioning (str): The partitioning mode to use when constant folding. fold_shapes (bool): Whether to allow shape folding. per_pass_shape_inference (bool): Whether to apply shape inference between constant folding passes. size_threshold (int): The threshold in bytes over which to avoid folding constants. """ self.fold_constants = args_util.get(args, "fold_constants") self.num_passes = args_util.get(args, "num_const_fold_passes") self.partitioning = args_util.get(args, "partitioning") self.fold_shapes = args_util.get(args, "fold_shapes") self.per_pass_shape_inference = args_util.get(args, "per_pass_shape_inference") self.size_threshold = args_util.parse_num_bytes( args_util.get(args, "fold_size_threshold") ) if not self.fold_constants: for arg in [ "num_const_fold_passes", "partitioning", "fold_shapes", "per_pass_shape_inference", "fold_size_threshold", ]: val = args_util.get(args, arg) if val is not None: G_LOGGER.warning( f"Argument: '--{arg.replace('_', '-')}' will be ignored since constant folding is not enabled.\n" "This argument is only valid when the `--fold-constants` option is provided." ) def add_to_script_impl(self, script, loader_name): if not self.fold_constants: return loader_name script.add_import(imports=["FoldConstants"], frm="polygraphy.backend.onnx") loader_name = script.add_loader( make_invocable( "FoldConstants", loader_name, num_passes=self.num_passes, do_shape_inference=( self.arg_groups[OnnxInferShapesArgs].do_shape_inference if self.per_pass_shape_inference is not False # since `None` indicates default value else False ), fold_shapes=self.fold_shapes, partitioning=self.partitioning, size_threshold=self.size_threshold, allow_onnxruntime_shape_inference=self.arg_groups[ OnnxInferShapesArgs ].allow_onnxruntime, ), "fold_constants", ) return loader_name def fold(self, model): loader = args_util.run_script(self.add_to_script, model) return util.invoke_if_callable(loader)[0] class Sanitize(BaseSurgeonSubtool): """ Clean up, optimize, and/or change input shapes in an ONNX model. """ def __init__(self): super().__init__("sanitize") def get_subscriptions_impl(self): return [ ModelArgs( model_opt_required=True, input_shapes_opt_name="override-inputs", required_model_type="onnx", input_shapes_opt_doc="Override input shapes in the model for the given inputs", ), DataLoaderArgs(), OnnxInferShapesArgs(default=True, allow_force_fallback=True), OnnxLoadArgs(outputs_opt_prefix="", allow_setting_upper_bounds=True), OnnxSaveArgs(allow_shape_inference=True, output_opt_required=True), ConstFoldArgs(), ] def add_parser_args_impl(self, parser): parser.add_argument( "--cleanup", help="Run dead layer removal on the graph. This is generally not required if other options are set. ", action="store_true", default=False, ) parser.add_argument( "--toposort", help="Topologically sort nodes in the graph. ", action="store_true", default=False, ) def show_start_end_logging_impl(self, args): return True def run_impl_surgeon(self, args): # First do all processing that requires an ONNX-GraphSurgeon graph, then do everything # that operates on the ONNX model. This lets us avoid ONNX-GraphSurgeon import if we don't # need it. def do_graph_processing(model): graph = None rerun_shape_inference = False def get_graph(): nonlocal graph if graph is None: graph = onnx_backend.gs_from_onnx(model) return graph user_input_metadata = self.arg_groups[ModelArgs].input_shapes if user_input_metadata: graph = get_graph() graph = tools_util.override_input_shapes(graph, user_input_metadata) rerun_shape_inference = True if self.arg_groups[OnnxInferShapesArgs].force_fallback: _, layerwise_meta = self.arg_groups[ OnnxInferShapesArgs ].fallback_inference(model) graph = get_graph() onnx_util.set_shapes_from_layerwise_meta(graph, layerwise_meta) if args.cleanup: graph = get_graph() graph.cleanup() if args.toposort: graph = get_graph() graph.toposort() if graph is not None: model = gs.export_onnx(graph) return model, rerun_shape_inference def do_model_processing(model): model = self.arg_groups[ConstFoldArgs].fold(model) return model model = super().load_model() model, rerun_shape_inference = do_graph_processing(model) if rerun_shape_inference: model = self.arg_groups[OnnxInferShapesArgs].infer_shapes(model) model = do_model_processing(model) super().save_model(model)