chore: import upstream snapshot with attribution

tools/Polygraphy/polygraphy/backend/base/__init__.py

@@ -0,0 +1,2 @@
+from polygraphy.backend.base.runner import *
+from polygraphy.backend.base.loader import *

tools/Polygraphy/polygraphy/backend/base/loader.py

@@ -0,0 +1,40 @@
+#
+# 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 func, mod
+
+
+@mod.export()
+class BaseLoader:
+    """
+    Base class for Polygraphy Loaders.
+    """
+
+    # Implementation for ``__call__``. Derived classes should implement this
+    # method rather than ``__call__``.
+    def call_impl(self, *args, **kwargs):
+        raise NotImplementedError("BaseLoader is an abstract class")
+
+    @func.constantmethod
+    def __call__(self, *args, **kwargs):
+        """
+        Invokes the loader by forwarding arguments to ``call_impl``.
+
+        Note: ``call_impl`` should *not* be called directly - use this function instead.
+        """
+        __doc__ = self.call_impl.__doc__
+        return self.call_impl(*args, **kwargs)

tools/Polygraphy/polygraphy/backend/base/runner.py

@@ -0,0 +1,258 @@
+#
+# 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.
+#
+import copy
+import time
+from collections import defaultdict
+
+from polygraphy import config, func, mod, util
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER, LogMode
+from polygraphy.backend.base import util as base_util
+
+np = mod.lazy_import("numpy")
+
+
+@mod.export()
+class BaseRunner:
+    """
+    Base class for Polygraphy runners. All runners should override the functions and attributes specified here.
+    """
+
+    RUNNER_COUNTS = defaultdict(int)
+
+    def __init__(self, name=None, prefix=None):
+        """
+        Args:
+            name (str):
+                    The name to use for this runner.
+            prefix (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+                    Only used if name is not provided.
+        """
+        prefix = util.default(prefix, "Runner")
+        if name is None:
+            count = BaseRunner.RUNNER_COUNTS[prefix]
+            BaseRunner.RUNNER_COUNTS[prefix] += 1
+            name = f"{prefix}-N{count}-{time.strftime('%x')}-{time.strftime('%X')}"
+        self.name = name
+        self.inference_time = None
+
+        self.is_active = False
+        """bool: Whether this runner has been activated, either via context manager, or by calling ``activate()``."""
+
+    def __enter__(self):
+        """
+        Activate the runner for inference. For example, this may involve allocating CPU or GPU memory.
+        """
+        self.activate()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Deactivate the runner. For example, this may involve freeing CPU or GPU memory.
+        """
+        self.deactivate()
+
+    # Implementation for runner activation. Derived classes should override this function
+    # rather than ``activate()``.
+    def activate_impl(self):
+        pass
+
+    def activate(self):
+        """
+        Activate the runner for inference. For example, this may involve allocating CPU or GPU memory.
+
+        Generally, you should use a context manager instead of manually activating and deactivating.
+        For example:
+        ::
+
+            with RunnerType(...) as runner:
+                runner.infer(...)
+        """
+        if self.is_active:
+            G_LOGGER.warning(
+                f"{self.name:35} | Already active; will not activate again. "
+                "If you really want to activate this runner again, call activate_impl() directly"
+            )
+            return
+
+        if config.INTERNAL_CORRECTNESS_CHECKS:
+            self._pre_activate_runner_state = copy.copy(vars(self))
+
+        self.activate_impl()
+        self.is_active = True
+
+    def get_input_metadata_impl(self):
+        """
+        Implemenation for `get_input_metadata`. Derived classes should override this function
+        rather than `get_input_metadata`.
+
+        Derived classes may return any kind of data type supported by Polygraphy's DataType
+        class (e.g. np.dtype, torch.dtype, etc.)
+        """
+        raise NotImplementedError("BaseRunner is an abstract class")
+
+    @func.constantmethod
+    def get_input_metadata(self, use_numpy_dtypes=None):
+        """
+        Returns information about the inputs of the model.
+        Shapes here may include dynamic dimensions, represented by ``None``.
+        Must be called only after ``activate()`` and before ``deactivate()``.
+
+        Args:
+            use_numpy_dtypes (bool):
+                    [DEPRECATED] Whether to return NumPy data types instead of Polygraphy ``DataType`` s.
+                    This is provided to retain backwards compatibility. In the future,
+                    this parameter will be removed and Polygraphy ``DataType`` s will
+                    always be returned. These can be converted to NumPy data types by calling the `numpy()` method.
+                    Defaults to True.
+
+        Returns:
+            TensorMetadata: Input names, shapes, and data types.
+        """
+        if not self.is_active:
+            G_LOGGER.critical(
+                f"{self.name:35} | Must be activated prior to calling get_input_metadata()"
+            )
+
+        use_numpy_dtypes = util.default(use_numpy_dtypes, True)
+
+        meta = self.get_input_metadata_impl()
+
+        for name, (dtype, _) in meta.items():
+            dtype = DataType.from_dtype(dtype)
+            if use_numpy_dtypes:
+                mod.warn_deprecated(
+                    "Returning NumPy data types instead of Polygraphy `DataType`s from `get_input_metadata()`",
+                    use_instead=None,
+                    remove_in="0.60.0",
+                )
+                meta[name]._dtype = DataType.to_dtype(dtype, "numpy")
+        return meta
+
+    # Implementation for runner inference. Derived classes should override this function
+    # rather than ``infer()``
+    # Derived classes should also set the `inference_time` property so that performance metrics are accurate.
+    def infer_impl(self, feed_dict):
+        raise NotImplementedError("BaseRunner is an abstract class")
+
+    def infer(self, feed_dict, check_inputs=True, *args, **kwargs):
+        """
+        Runs inference using the provided feed_dict.
+
+        Must be called only after ``activate()`` and before ``deactivate()``.
+
+        NOTE: Some runners may accept additional parameters in infer().
+        For details on these, see the documentation for their `infer_impl()` methods.
+
+        Args:
+            feed_dict (OrderedDict[str, numpy.ndarray]):
+                    A mapping of input tensor names to corresponding input NumPy arrays.
+
+            check_inputs (bool):
+                    Whether to check that the provided ``feed_dict`` includes the expected inputs
+                    with the expected data types and shapes.
+                    Disabling this may improve performance.
+                    Defaults to True.
+
+        Attributes:
+            inference_time (float):
+                    The time required to run inference in seconds.
+
+        Returns:
+            OrderedDict[str, numpy.ndarray]:
+                    A mapping of output tensor names to their corresponding NumPy arrays.
+
+                    IMPORTANT: Runners may reuse these output buffers. Thus, if you need to save
+                    outputs from multiple inferences, you should make a copy with ``copy.deepcopy(outputs)``.
+        """
+        if not self.is_active:
+            G_LOGGER.critical(
+                f"{self.name:35} | Must be activated prior to calling infer()"
+            )
+
+        if check_inputs:
+            input_metadata = self.get_input_metadata(use_numpy_dtypes=False)
+            G_LOGGER.verbose(
+                f"{self.name:35} | Input metadata is: {input_metadata}",
+                mode=LogMode.ONCE,
+            )
+            base_util.check_inputs(feed_dict, input_metadata)
+
+        return self.infer_impl(feed_dict, *args, **kwargs)
+
+    @func.constantmethod
+    def last_inference_time(self):
+        """
+        Returns the total inference time in seconds required during the last call to ``infer()``.
+
+        Must be called only after ``activate()`` and before ``deactivate()``.
+
+        Returns:
+            float: The time in seconds, or None if runtime was not measured by the runner.
+        """
+        if self.inference_time is None:
+            msg = f"{self.name:35} | `inference_time` was not set. Inference time will be incorrect! "
+            msg += "To correctly compare runtimes, please set the `inference_time` attribute in `infer_impl()`"
+
+            G_LOGGER.internal_error(msg)
+            G_LOGGER.warning(msg, mode=LogMode.ONCE)
+            return None
+        return self.inference_time
+
+    # Implementation for runner deactivation. Derived classes should override this function
+    # rather than ``deactivate()``.
+    def deactivate_impl(self):
+        pass
+
+    def deactivate(self):
+        """
+        Deactivate the runner. For example, this may involve freeing CPU or GPU memory.
+
+        Generally, you should use a context manager instead of manually activating and deactivating.
+        For example:
+        ::
+
+            with RunnerType(...) as runner:
+                runner.infer(...)
+        """
+        if not self.is_active:
+            G_LOGGER.warning(
+                f"{self.name:35} | Not active; will not deactivate. If you really want to deactivate this runner, call deactivate_impl() directly"
+            )
+            return
+
+        self.inference_time = None
+        self.is_active = None
+
+        self.deactivate_impl()
+        self.is_active = False
+        if config.INTERNAL_CORRECTNESS_CHECKS:
+            old_state = self._pre_activate_runner_state
+            del self._pre_activate_runner_state
+            if old_state != vars(self):
+                G_LOGGER.internal_error(
+                    f"Runner state was not reset after deactivation. Note:\nOld state: {old_state}\nNew state: {vars(self)}"
+                )
+
+    def __del__(self):
+        if self.is_active:
+            # __del__ is not guaranteed to be called, but when it is, this could be a useful warning.
+            print(
+                f"[W] {self.name:35} | Was activated but never deactivated. This could cause a memory leak!"
+            )

tools/Polygraphy/polygraphy/backend/base/util.py

@@ -0,0 +1,53 @@
+#
+# 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 util
+from polygraphy.logger import G_LOGGER
+
+
+def check_inputs(feed_dict, input_metadata):
+    """
+    Checks the provided `feed_dict` against expected input metadata.
+
+    Args:
+        feed_dict (Dict[str, Union[DeviceView, numpy.ndarray, torch.Tensor]]):
+                A mapping of input names to arrays.
+        input_metadata (TensorMetadata):
+                The expected input metadata.
+    """
+    util.check_sequence_contains(
+        feed_dict.keys(), input_metadata.keys(), name="input data", items_name="inputs"
+    )
+
+    for name, inp in feed_dict.items():
+        meta = input_metadata[name]
+
+        # The "buffer" might just be a pointer, in which case we can't do any further checks with it, so we skip it.
+        if isinstance(inp, int):
+            continue
+
+        dtype = util.array.dtype(inp)
+        if dtype != meta.dtype:
+            G_LOGGER.critical(
+                f"Input tensor: {name} | Received unexpected dtype: {dtype}.\nNote: Expected type: {meta.dtype}"
+            )
+
+        shape = util.array.shape(inp)
+        if not util.is_valid_shape_override(shape, meta.shape):
+            G_LOGGER.critical(
+                f"Input tensor: {name} | Received incompatible shape: {shape}.\nNote: Expected a shape compatible with: {meta.shape}"
+            )

tools/Polygraphy/polygraphy/backend/common/__init__.py

@@ -0,0 +1 @@
+from polygraphy.backend.common.loader import *

tools/Polygraphy/polygraphy/backend/common/loader.py

@@ -0,0 +1,101 @@
+#
+# 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.backend.base import BaseLoader
+
+
+@mod.export(funcify=True)
+class BytesFromPath(BaseLoader):
+    """
+    Functor that can load a file in binary mode ('rb').
+    """
+
+    def __init__(self, path):
+        """
+        Loads a file in binary mode ('rb').
+
+        Args:
+            path (str): The file path.
+        """
+        self._path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            bytes: The contents of the file.
+        """
+        return util.load_file(self._path, description="bytes")
+
+
+@mod.export(funcify=True)
+class SaveBytes(BaseLoader):
+    """
+    Functor that can save bytes to a file.
+    """
+
+    def __init__(self, obj, path):
+        """
+        Saves bytes to a file.
+
+        Args:
+            obj (Union[bytes, Callable() -> bytes]):
+                    The bytes to save or a callable that returns them.
+            path (str): The file path.
+        """
+        self._bytes = obj
+        self._path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            bytes: The bytes saved.
+        """
+        obj, _ = util.invoke_if_callable(self._bytes)
+        util.save_file(obj, self._path)
+        return obj
+
+
+@mod.export(funcify=True)
+class InvokeFromScript(BaseLoader):
+    """
+    Functor that invokes a function from a Python script.
+    """
+
+    def __init__(self, path, name):
+        """
+        Invokes the specified function from the specified Python script.
+
+        If you intend to use the function more than once, you should import
+        the function using ``polygraphy.mod.import_from_script`` instead.
+
+        Args:
+            path (str): The path to the Python script. The path must include a '.py' extension.
+            name (str): The name of the function to import and invoke.
+        """
+        self._path = path
+        self._name = name
+
+    @util.check_called_by("__call__")
+    def call_impl(self, *args, **kwargs):
+        """
+        Returns:
+            object:
+                    The return value of the imported function.
+        """
+        return mod.import_from_script(self._path, self._name)(*args, **kwargs)

tools/Polygraphy/polygraphy/backend/onnx/__init__.py

@@ -0,0 +1 @@
+from polygraphy.backend.onnx.loader import *

tools/Polygraphy/polygraphy/backend/onnx/loader.py

@@ -0,0 +1,988 @@
+#
+# 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.
+#
+import copy
+import os
+import sys
+import tempfile
+
+from polygraphy import constants, mod, util
+from polygraphy.backend.base import BaseLoader
+from polygraphy.backend.onnx import util as onnx_util
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER, LogMode
+
+np = mod.lazy_import("numpy")
+onnx = mod.lazy_import("onnx>=1.8.1")
+onnxrt = mod.lazy_import("onnxruntime>=1.10.0")
+onnxmltools = mod.lazy_import(
+    "onnxmltools==1.11.1", requires=["onnxconverter_common>=1.12.2"]
+)
+tf = mod.lazy_import("tensorflow<2.0")
+tf2onnx = mod.lazy_import("tf2onnx")
+tf_util = mod.lazy_import("polygraphy.backend.tf.util", log=False)
+gs = mod.lazy_import("onnx_graphsurgeon>=0.3.27")
+# ONNX-RT's shape inference also requires "sympy", but it is not reported as a dependency,
+# so we work around it by checking for it manually.
+onnxrt_symbolic_shape_inference = mod.lazy_import(
+    "onnxruntime.tools.symbolic_shape_infer>=1.10.0", requires=["sympy"]
+)
+
+LARGE_MODEL_THRESHOLD = 512 << 20  # 512 MiB
+PROTOBUF_THRESHOLD = 2e9
+
+
+class BaseLoadOnnxCopy(BaseLoader):
+    """
+    Abstract base class for loaders that require loading an ONNX model and potentially
+    making a copy.
+    """
+
+    def __init__(self, model, copy=None):
+        """
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+
+            copy (bool): Whether to create a copy of the model first. Defaults to False.
+        """
+        self._model = model
+        self.copy = util.default(copy, False)
+
+    def load(self):
+        model, _ = util.invoke_if_callable(self._model)
+        if self.copy:
+            model = copy.copy(model)
+        return model
+
+
+class _GSGraphManager:
+    """
+    Imports an ONNX-GraphSurgeon graph.
+
+    If the provided model is already a graph, the graph is not
+    exported to ONNX.
+    """
+
+    def __init__(self, model):
+        self._model = model
+
+    def __enter__(self):
+        model, _ = util.invoke_if_callable(self._model)
+        self.USE_GS_GRAPH = isinstance(model, gs.Graph)
+        if self.USE_GS_GRAPH:
+            self.graph = model.copy()
+        else:
+            self.graph = gs_from_onnx(model)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.USE_GS_GRAPH:
+            self.retval = self.graph
+        else:
+            self.retval = gs.export_onnx(self.graph, do_type_check=False)
+
+
+@mod.export(funcify=True)
+class GsFromOnnx(BaseLoader):
+    """
+    Functor that creates an ONNX-GraphSurgeon graph from an ONNX ModelProto.
+    """
+
+    def __init__(self, model):
+        """
+        Creates an ONNX-GraphSurgeon graph from an ONNX ModelProto.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+        """
+        self._model = model
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx_graphsurgeon.Graph: The ONNX-GraphSurgeon representation of the ONNX model
+        """
+        model, _ = util.invoke_if_callable(self._model)
+        return gs.import_onnx(model)
+
+
+@mod.export(funcify=True)
+class OnnxFromPath(BaseLoader):
+    """
+    Functor that loads an ONNX model from a file.
+    """
+
+    def __init__(self, path, external_data_dir=None, ignore_external_data=None):
+        """
+        Loads an ONNX model from a file.
+
+        Args:
+            path (str): The path from which to load the model.
+
+            external_data_dir (str): The directory where external data for the model is stored.
+            ignore_external_data (bool):
+                    Whether to ignore any external data and just load the model structure without any weights.
+                    The model will be usable only for purposes that don't require weights, such as extracting
+                    subgraphs or inspecting model structure.
+                    This can be useful in cases where external data is not available.
+                    Defaults to False.
+        """
+        self.path = path
+        self.external_data_dir = external_data_dir
+        self.ignore_external_data = util.default(ignore_external_data, False)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The ONNX model
+        """
+        G_LOGGER.info(f"Loading model: {self.path}")
+        # If external_data_dir is not None, we'll load external data ourselves
+        auto_load_ext_data = (
+            self.external_data_dir is None and not self.ignore_external_data
+        )
+        try:
+            model = onnx.load(self.path, load_external_data=auto_load_ext_data)
+        except FileNotFoundError:
+            if auto_load_ext_data:
+                G_LOGGER.warning(
+                    "Failed to load model. This could be because external data could not be loaded.\n"
+                    "Hint: If you don't need the model weights, try ignoring external data by setting `ignore_external_data=True` "
+                    "or using the `--ignore-external-data` command-line option."
+                )
+            raise
+
+        if self.external_data_dir is not None:
+            G_LOGGER.verbose(f"Loading external data from: {self.external_data_dir}")
+            onnx.external_data_helper.load_external_data_for_model(
+                model, self.external_data_dir
+            )
+        return model
+
+
+@mod.export(funcify=True)
+class OnnxFromTfGraph(BaseLoader):
+    """
+    Functor that loads a TensorFlow graph and converts it to ONNX using the tf2onnx converter.
+    """
+
+    def __init__(self, graph, opset=None, optimize=None):
+        """
+        Converts a TensorFlow model into ONNX.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            opset (int): The ONNX opset to use during conversion.
+            optimize (bool): Whether to use tf2onnx's graph optimization pass.
+        """
+        self._graph = graph
+        self.opset = util.default(opset, 11)
+        self.optimize = util.default(optimize, True)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The ONNX model.
+        """
+        (graph, output_names), _ = util.invoke_if_callable(self._graph)
+        input_names = list(tf_util.get_input_metadata(graph).keys())
+
+        graphdef = graph.as_graph_def()
+        if self.optimize:
+            graphdef = tf2onnx.tfonnx.tf_optimize(
+                input_names, output_names, graph.as_graph_def()
+            )
+
+        with tf.Graph().as_default() as graph, tf.compat.v1.Session(
+            graph=graph
+        ) as sess:
+            tf.import_graph_def(graphdef, name="")
+
+            onnx_graph = tf2onnx.tfonnx.process_tf_graph(
+                graph,
+                input_names=input_names,
+                output_names=output_names,
+                opset=self.opset,
+            )
+            if self.optimize:
+                onnx_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
+            return onnx_graph.make_model("model")
+
+
+@mod.export(funcify=True)
+class ModifyOutputs(BaseLoadOnnxCopy):
+    """
+    Functor that modifies the outputs of an ONNX model.
+    """
+
+    def __init__(self, model, outputs=None, exclude_outputs=None, copy=None):
+        """
+        Modifies outputs of an ONNX model.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+
+            outputs (Sequence[str]):
+                    Names of tensors to mark as outputs. If provided, this will override the
+                    existing model outputs.
+                    If a value of `constants.MARK_ALL` is used instead of a list, all tensors in the network are marked.
+            exclude_outputs (Sequence[str]):
+                    Names of tensors to exclude as outputs. This can be useful in conjunction with
+                    ``outputs=constants.MARK_ALL`` to omit outputs.
+            copy (bool): Whether to create a copy of the model first. Defaults to False.
+        """
+        super().__init__(model, copy)
+        self.outputs = outputs
+        self.exclude_outputs = exclude_outputs
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The ONNX model with modified outputs.
+        """
+        model = self.load()
+
+        if self.outputs == constants.MARK_ALL:
+            G_LOGGER.verbose("Marking all ONNX tensors as outputs")
+            model = onnx_util.mark_layerwise(model)
+        elif self.outputs is not None:
+            model = onnx_util.mark_outputs(model, self.outputs)
+
+        if self.exclude_outputs is not None:
+            model = onnx_util.unmark_outputs(model, self.exclude_outputs)
+
+        return model
+
+
+@mod.export(funcify=True)
+class ConvertToFp16(BaseLoadOnnxCopy):
+    """
+    Functor that converts all floating point tensors in the model to 16-bit precision.
+    This is *not* needed in order to use TensorRT's fp16 precision, but may be useful for other backends.
+    """
+
+    def __init__(self, model, copy=None):
+        """
+        Converts all floating point tensors in the model to 16-bit precision.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+            copy (bool): Whether to create a copy of the model first. Defaults to False.
+        """
+        super().__init__(model, copy)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The modified ONNX model.
+        """
+        model = self.load()
+
+        G_LOGGER.info("Converting float tensors to float16")
+        model = onnxmltools.utils.float16_converter.convert_float_to_float16(
+            model, keep_io_types=True, disable_shape_infer=True
+        )
+        return model
+
+
+@mod.export(funcify=True)
+class FoldConstants(BaseLoadOnnxCopy):
+    """
+    Functor that folds constants in an ONNX model.
+    """
+
+    def __init__(
+        self,
+        model,
+        num_passes=None,
+        do_shape_inference=None,
+        partitioning=None,
+        fold_shapes=None,
+        copy=None,
+        error_ok=None,
+        size_threshold=None,
+        allow_onnxruntime_shape_inference=None,
+    ):
+        """
+        Fold constants in an ONNX model.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+
+            num_passes (int):
+                    The number of constant folding passes to run.
+                    Sometimes, subgraphs that compute tensor shapes may not be foldable in a single pass.
+                    By default, Polygraphy will automatically determine the number of passes required.
+            do_shape_inference (bool):
+                    Whether to run shape inference in the model between passes.
+                    This enables the loader to fold `Shape` nodes.
+                    Only effective if `fold_shapes` is True.
+                    Defaults to True.
+            partitioning (Union[str, None]):
+                    Whether/How to partition the graph so that errors in folding one
+                    part of a model do not affect other parts. Available modes are:
+
+                    - None: Do not partition the graph. If inference fails, no constants are folded.
+                    - 'basic': Partition the graph. If inference fails in one partition, other partitions will remain unaffected.
+                    - 'recursive': Parition the graph recursively. If inference fails in a partition, the partition will be further partitioned.
+
+                    Defaults to None.
+            fold_shapes (bool):
+                    Whether to fold `Shape` nodes in the graph.
+                    This requires shapes to be inferred in the graph, and can only fold
+                    static shapes.
+                    Defaults to True.
+            copy (bool):
+                    Whether to create a copy of the model first.
+                    Defaults to False.
+            error_ok (bool):
+                    Whether to suppress errors during constant folding.
+                    If this is set to ``False``, errors will be re-raised.
+                    Defaults to True.
+            size_threshold (int):
+                    The maximum size threshold, in bytes, for which to fold constants.
+                    Any tensors larger than this value will not be folded.
+                    Set to ``None`` to disable the size threshold and always fold constants.
+                    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.
+                    Defaults to None.
+            allow_onnxruntime_shape_inference (bool):
+                    Allow ONNX-Runtime's shape inference to be used if available instead of ONNX's
+                    shape inference utilities. The former may provide performance or memory usage benefits.
+                    Has no effect if ``do_shape_inference`` is False.
+                    Defaults to True.
+        """
+        super().__init__(model, copy)
+        self.num_passes = num_passes
+        self.do_shape_inference = util.default(do_shape_inference, True)
+        self.partitioning = partitioning
+        self.fold_shapes = util.default(fold_shapes, True)
+        self.error_ok = util.default(error_ok, True)
+        self.size_threshold = size_threshold
+        self.allow_onnxruntime_shape_inference = allow_onnxruntime_shape_inference
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The new ONNX model with constants folded.
+        """
+
+        def run_const_fold_pass(model):
+            graph = gs_from_onnx(model)
+            del model
+
+            graph.fold_constants(
+                fold_shapes=self.fold_shapes,
+                partitioning=self.partitioning,
+                size_threshold=self.size_threshold,
+            )
+
+            model = gs.export_onnx(graph.cleanup(), do_type_check=False)
+            del graph
+
+            if self.fold_shapes and self.do_shape_inference:
+                model = infer_shapes(
+                    model, allow_onnxruntime=self.allow_onnxruntime_shape_inference
+                )
+            return model
+
+        # Need to manually trigger the autoinstall this since it's used by ONNX-GS, which does not have an autoinstall mechanism.
+        mod.autoinstall(onnxrt)
+        if not onnxrt.is_installed() or not onnxrt.is_importable():
+            G_LOGGER.error(
+                f"ONNX-Runtime is not installed, so constant folding may be suboptimal or not work at all.\n"
+                f"Consider installing ONNX-Runtime: {sys.executable} -m pip install onnxruntime"
+            )
+
+        model = self.load()
+
+        prefold_num_nodes = len(model.graph.node)
+        postfold_num_nodes = -1
+        index = 0
+
+        while (prefold_num_nodes != postfold_num_nodes) and (
+            self.num_passes is None or index < self.num_passes
+        ):
+            prefold_num_nodes = onnx_util.get_num_nodes(model)
+
+            G_LOGGER.start(f"Folding Constants | Pass {index + 1}")
+            try:
+                model = run_const_fold_pass(model)
+            except Exception as err:
+                if not self.error_ok:
+                    raise
+                G_LOGGER.warning(
+                    f"Constant folding pass failed. Skipping subsequent passes.\nNote: Error was:\n{err}"
+                )
+                break
+            else:
+                postfold_num_nodes = onnx_util.get_num_nodes(model)
+                index += 1
+
+                G_LOGGER.finish(
+                    f"{constants.TAB}Total Nodes | Original: {prefold_num_nodes:5}, "
+                    f"After Folding: {postfold_num_nodes:5} | {prefold_num_nodes - postfold_num_nodes:5} Nodes Folded"
+                )
+
+        return model
+
+
+@mod.export(funcify=True)
+class SetUpperBound(BaseLoadOnnxCopy):
+    """
+    Functor that sets upper bounds for tensors with unbounded DDS in an ONNX model.
+
+    Requires that the model has been constant folded and has shapes inferred.
+    """
+
+    def __init__(
+        self,
+        model,
+        upper_bounds,
+        copy=None,
+    ):
+        """
+        Set upper bounds for tensors with unbounded DDS in an ONNX model.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+
+            upper_bounds (Union[int, Dict[str, int]]):
+                    The upper bounds for tensors with unbounded DDS.
+                    If a single integer is provided, it will be used as the default upper bound for all tensors with unbounded DDS.
+                    This can also be provided on a per-tensor basis using a dictionary. In that case, use an empty string ("") as the
+                    key to specify default upper bound for tensors not explicitly listed.
+            copy (bool):
+                    Whether to create a copy of the model first.
+                    Defaults to False.
+        """
+        super().__init__(model, copy)
+        self.upper_bounds = upper_bounds
+
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The new ONNX model.
+        """
+
+        # Set upper bounds for tensors with unbounded DDS in the onnx model.
+        def set_upper_bound(graph, target_tensor_list):
+            applied_bounds = {}
+            for tensor in target_tensor_list:
+                upper_bound = util.value_or_from_dict(self.upper_bounds, tensor.name)
+                if upper_bound is None:
+                    continue
+                # Insert a min operator to set the upper bound for the target tensor.
+                # A target tensor should always be produced from a single node.
+                assert len(tensor.inputs) == 1
+                producer = tensor.inputs[0]
+                producer_idx = producer.outputs.index(tensor)
+                tensor_copy = gs.Variable(
+                    tensor.name + "_copy", dtype=tensor.dtype, shape=tensor.shape
+                )
+                upper_bound_values = np.array(upper_bound)
+                if tensor.shape is not None and len(tensor.shape) > 0:
+                    upper_bound_values = np.array([upper_bound] * len(tensor.shape))
+                tensor_upper_bound = gs.Constant(
+                    tensor.name + "_upper_bound", values=upper_bound_values
+                )
+                min_node = gs.Node(
+                    op="Min", inputs=[tensor_copy, tensor_upper_bound], outputs=[tensor]
+                )
+                producer.outputs[producer_idx] = tensor_copy
+                tensor.inputs = [min_node]
+                graph.nodes.append(min_node)
+                applied_bounds[tensor.name] = upper_bound
+            G_LOGGER.info(f"Set tensor upper bounds: {applied_bounds}")
+            return graph
+
+        model = self.load()
+        graph = gs_from_onnx(model)
+
+        target_tensor_list = onnx_util.get_unbounded_dds_tensors(graph)
+
+        tensor_map = graph.tensors()
+        target_names = {tensor.name for tensor in target_tensor_list}
+        if isinstance(self.upper_bounds, dict):
+            input_names = set(self.upper_bounds.keys()) - {""}
+            # Report error when input tensor name is not in the graph.
+            util.check_sequence_contains(
+                set(tensor_map.keys()),
+                input_names,
+                name="the upper bounds dictionary",
+                items_name="tensors",
+                check_extra=False,
+            )
+            # Report warning when input tensor is not a unbounded DDS tensor.
+            util.check_sequence_contains(
+                set(target_names),
+                input_names,
+                name="the upper bounds dictionary",
+                items_name="tensors",
+                log_func=G_LOGGER.warning,
+                check_extra=False,
+            )
+            # Still set upper bound for input tensors with bounded shapes.
+            target_names.update(input_names)
+        graph = set_upper_bound(graph, [tensor_map[name] for name in target_names])
+        model = gs.export_onnx(graph.cleanup(), do_type_check=False)
+
+        return model
+
+
+@mod.export(funcify=True)
+class InferShapes(BaseLoader):
+    """
+    Functor that runs shape inference on an ONNX model.
+    """
+
+    def __init__(
+        self,
+        model,
+        error_ok=None,
+        external_data_dir=None,
+        save_to_disk_threshold_bytes=None,
+        allow_onnxruntime=None,
+    ):
+        """
+        Run shape inference on an ONNX model.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto, str, Callable() -> str]):
+                    An ONNX model or a callable that returns one, or a path to a model.
+                    Supports models larger than the 2 GB protobuf limit.
+
+            error_ok (bool):
+                    Whether errors during shape inference should be suppressed.
+                    Defaults to True.
+            external_data_dir (str):
+                    The directory where external data for the model is stored.
+                    Only used if the model is provided via a path rather than a loader.
+            save_to_disk_threshold_bytes (int):
+                    The size in bytes above which a ModelProto will be serialized to the disk
+                    before running shape inference.
+                    This can be used to work around the 2 GB protobuf limitation.
+                    Defaults to 2 GB.
+            allow_onnxruntime (bool):
+                    Allow ONNX-Runtime's shape inference to be used if available instead of ONNX's
+                    shape inference utilities. The former may provide performance or memory usage benefits.
+                    Defaults to True.
+        """
+        self._model = model
+        self.error_ok = util.default(error_ok, True)
+        self.external_data_dir = external_data_dir
+        # Subtract a little so we're below the real threshold
+        self.save_to_disk_threshold_bytes = util.default(
+            save_to_disk_threshold_bytes, PROTOBUF_THRESHOLD
+        )
+        self.allow_onnxruntime = util.default(allow_onnxruntime, True)
+
+    def _run_onnx_shape_inference(self, model, external_data_dir):
+        if isinstance(model, onnx.ModelProto):
+            MODEL_SIZE = model.ByteSize()
+            if MODEL_SIZE > LARGE_MODEL_THRESHOLD:
+                G_LOGGER.warning(
+                    f"Attempting to run shape inference on a large model ({MODEL_SIZE // 1024.0 ** 2} MiB). "
+                    "This may require a large amount of memory.\nIf memory consumption becomes too high, "
+                    "the process may be killed. You may want to try disabling shape inference in that case. ",
+                    mode=LogMode.ONCE,
+                )
+
+            if MODEL_SIZE >= self.save_to_disk_threshold_bytes:
+                G_LOGGER.warning(
+                    f"Model size ({MODEL_SIZE / 1024.0 ** 2} MiB) exceeds the in-memory size threshold: "
+                    f"{self.save_to_disk_threshold_bytes / 1024.0 ** 2} MiB.\n"
+                    f"The model will be saved to a temporary file before shape inference is run.",
+                    mode=LogMode.ONCE,
+                )
+                outdir = tempfile.TemporaryDirectory()
+                outpath = os.path.join(outdir.name, "tmp_model.onnx")
+                save_onnx(model, outpath, external_data_path="ext.data")
+                model = outpath
+                external_data_dir = outdir.name
+
+        if isinstance(model, onnx.ModelProto):
+            model = onnx.shape_inference.infer_shapes(model)
+        else:
+            tmp_path = util.NamedTemporaryFile(
+                prefix="tmp_polygraphy_", suffix=".onnx"
+            ).name
+            G_LOGGER.verbose(f"Writing shape-inferred model to: {tmp_path}")
+            onnx.shape_inference.infer_shapes_path(model, tmp_path)
+            # In cases where the original model had external data stored in the same directory,
+            # the external data directory may not be explicitly specified.
+            # In such cases, we need to use the model's directory as the external data path
+            # for the newly generated model.
+            model = onnx_from_path(
+                tmp_path,
+                external_data_dir=util.default(
+                    external_data_dir, os.path.dirname(model) or None
+                ),
+            )
+        return model
+
+    def _run_onnxruntime_shape_inference(self, model, external_data_dir):
+        if not isinstance(model, onnx.ModelProto):
+            model = onnx_from_path(model, external_data_dir=external_data_dir)
+        return onnxrt_symbolic_shape_inference.SymbolicShapeInference.infer_shapes(
+            model, auto_merge=True
+        )
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The new ONNX model with shapes inferred.
+        """
+        model, _ = util.invoke_if_callable(self._model)
+        external_data_dir = self.external_data_dir
+
+        G_LOGGER.verbose("Starting shape inference")
+
+        try:
+            use_onnx_shape_inference = not self.allow_onnxruntime
+            if self.allow_onnxruntime:
+                try:
+                    model = self._run_onnxruntime_shape_inference(
+                        model, external_data_dir
+                    )
+                    G_LOGGER.verbose(
+                        "Inferred shapes in the model with `onnxruntime.tools.symbolic_shape_infer`.\n"
+                        "Note: To force Polygraphy to use `onnx.shape_inference` instead, set `allow_onnxruntime=False` or "
+                        "use the `--no-onnxruntime-shape-inference` command-line option.",
+                        mode=LogMode.ONCE,
+                    )
+                except Exception as err:
+                    use_onnx_shape_inference = True
+                    G_LOGGER.extra_verbose(
+                        f"Error while running `onnxruntime.tools.symbolic_shape_infer`:\n{err}"
+                    )
+                    G_LOGGER.warning(
+                        "Falling back to `onnx.shape_inference` because `onnxruntime.tools.symbolic_shape_infer` either could not be loaded "
+                        "or did not run successfully.\n"
+                        "Note that using ONNX-Runtime for shape inference may be faster and require less memory.\n"
+                        "Consider installing ONNX-Runtime or setting POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment "
+                        "variables to allow Polygraphy to do so automatically.",
+                        mode=LogMode.ONCE,
+                    )
+
+            if use_onnx_shape_inference:
+                model = self._run_onnx_shape_inference(model, external_data_dir)
+        except Exception as err:
+            if not self.error_ok:
+                raise
+            G_LOGGER.warning(f"ONNX shape inference exited with an error:\n{err}")
+            G_LOGGER.internal_error(
+                f"ONNX shape inference exited with an error:\n{err}"
+            )
+
+            if not isinstance(model, onnx.ModelProto):
+                model = onnx_from_path(model, external_data_dir=external_data_dir)
+        else:
+            G_LOGGER.verbose("Shape inference completed successfully")
+
+        return model
+
+
+@mod.export(funcify=True)
+class ExtractSubgraph(BaseLoader):
+    """
+    Functor that extracts a subgraph from an ONNX model.
+    """
+
+    def __init__(
+        self, model, input_metadata=None, output_metadata=None, check_meta=None
+    ):
+        """
+        Extracts a subgraph from an ONNX model.
+
+        Args:
+            model (Union[Union[onnx.ModelProto, onnx_graphsurgeon.Graph], Callable() -> Union[onnx.ModelProto, onnx_graphsurgeon.Graph]]):
+                    An ONNX model or ONNX-GraphSurgeon Graph or a callable that returns one.
+
+            input_metadata (TensorMetadata):
+                    Metadata for the inputs of the subgraph.
+                    Name, shape, and data type are required.
+                    If not provided, the graph outputs are not modified.
+            output_metadata (TensorMetadata):
+                    Metadata for the outputs of the subgraph.
+                    Name and data type are required.
+                    If not provided, the graph outputs are not modified.
+            check_meta (bool):
+                    Whether to check that the provided input and output metadata include
+                    all the expected fields.
+                    Defaults to True.
+        """
+        self._model = model
+        self.input_metadata = input_metadata
+        self.output_metadata = output_metadata
+        self.check_meta = util.default(check_meta, True)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Union[onnx.ModelProto, onnx_graphsurgeon.Graph]:
+                    The new ONNX model or ONNX-GraphSurgeon Graph.
+        """
+        with _GSGraphManager(self._model) as manager:
+            graph = manager.graph
+            TENSOR_MAP = graph.tensors()
+
+            def get_tensor(name):
+                if name not in TENSOR_MAP:
+                    G_LOGGER.critical(f"Tensor: {name} does not exist in the model.")
+                return TENSOR_MAP[name]
+
+            def update_tensor(name, dtype, shape):
+                tensor = get_tensor(name)
+                # No need to update constants
+                if isinstance(tensor, gs.Variable):
+                    tensor.dtype, tensor.shape = (
+                        DataType.to_dtype(DataType.from_dtype(dtype), "onnx")
+                        if dtype is not None
+                        else None
+                    ) or tensor.dtype, shape or tensor.shape
+                return tensor
+
+            def check_meta(name, dtype, shape, meta_type, needs_shape=True):
+                if not self.check_meta:
+                    return
+                if needs_shape and shape is None:
+                    G_LOGGER.warning(
+                        f"{meta_type} metadata should include shape, but no shape was provided for tensor: {name}"
+                    )
+                if dtype is None:
+                    G_LOGGER.warning(
+                        f"{meta_type} metadata should include data type, but no data type was provided for tensor: {name}"
+                    )
+
+            if self.input_metadata is not None:
+                graph.inputs.clear()
+                for name, (dtype, shape) in self.input_metadata.items():
+                    tensor = update_tensor(name, dtype, shape)
+                    check_meta(name, tensor.dtype, tensor.shape, "Input")
+                    tensor.inputs.clear()
+                    graph.inputs.append(tensor)
+
+            if self.output_metadata is not None:
+                graph.outputs.clear()
+                for name, (dtype, shape) in self.output_metadata.items():
+                    tensor = update_tensor(name, dtype, shape)
+                    check_meta(
+                        name, tensor.dtype, tensor.shape, "Output", needs_shape=False
+                    )
+                    graph.outputs.append(tensor)
+
+            graph.cleanup()
+
+            tensor_map = graph.tensors()
+            for tensor in tensor_map.values():
+                if (
+                    isinstance(tensor, gs.Variable)
+                    and not tensor.inputs
+                    and tensor not in graph.inputs
+                ):
+                    consumer_nodes = [
+                        f"Node: '{node.name}' (Op: {node.op})"
+                        for node in tensor.outputs
+                    ]
+                    G_LOGGER.error(
+                        f"Tensor: '{tensor.name}' is a variable tensor consumed by: {consumer_nodes}, "
+                        "but is not produced by a node or marked as a graph input."
+                        f"\nDid you forget to mark a tensor as a graph input? Hint: Try inspecting the resulting model. "
+                        f"\nNote: The resulting model will not be valid!"
+                    )
+
+        return manager.retval
+
+
+@mod.export(funcify=True)
+class SaveOnnx(BaseLoader):
+    """
+    Functor that saves an ONNX model to the specified path.
+    """
+
+    def __init__(
+        self,
+        model,
+        path,
+        external_data_path=None,
+        size_threshold=None,
+        all_tensors_to_one_file=None,
+    ):
+        """
+        Saves an ONNX model to the specified path.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+            path (str): Path at which to write the ONNX model.
+            external_data_path (str):
+                    Path to save external data.
+                    This is always a relative path; external data is always written to the same
+                    directory as the model.
+                    Set to an empty string to use the default path.
+                    Set to None to disable.
+                    Defaults to None if the model is within the protobuf size threshold and an empty string otherwise.
+            size_threshold (int):
+                    Tensor size threshold, in bytes, above which tensor data will be
+                    stored in the external file.
+                    Tensors smaller that this threshold will remain in the ONNX file.
+                    Has no effect if external_data_path is not set.
+                    Defaults to 1024.
+            all_tensors_to_one_file (bool):
+                    Whether to write all tensors to one file when saving external data.
+                    Has no effect if external_data_path is not set.
+                    Defaults to True.
+        """
+        self._model = model
+        self.path = path
+        self.external_data_path = external_data_path
+        self.size_threshold = size_threshold
+        self.all_tensors_to_one_file = all_tensors_to_one_file
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The model, after saving it.
+        """
+        model, _ = util.invoke_if_callable(self._model)
+        G_LOGGER.info(f"Saving ONNX model to: {self.path}")
+
+        model_size = model.ByteSize()
+        if self.external_data_path is None and model_size >= PROTOBUF_THRESHOLD:
+            external_data_path = ""
+            G_LOGGER.warning(
+                f"Model size ({model_size // 1024.0 ** 2} MiB) exceeds protobuf size threshold ({PROTOBUF_THRESHOLD // 1024 ** 2} MiB). "
+                f"Will save weight data to an external file.\n"
+                f"To control the location of this file, use the `external_data_path` parameter or the `--external-data-path` command-line option. "
+            )
+        else:
+            external_data_path = self.external_data_path
+
+        if external_data_path is not None:
+            G_LOGGER.verbose(
+                f"Saving external data for ONNX model to: {external_data_path}"
+            )
+            try:
+                onnx.external_data_helper.convert_model_to_external_data(
+                    model,
+                    location=external_data_path,
+                    all_tensors_to_one_file=util.default(
+                        self.all_tensors_to_one_file, True
+                    ),
+                    size_threshold=util.default(self.size_threshold, 1024),
+                )
+            except TypeError:
+                if self.size_threshold is not None:
+                    G_LOGGER.warning(
+                        "This version of onnx does not support size_threshold in convert_model_to_external_data"
+                    )
+                onnx.external_data_helper.convert_model_to_external_data(
+                    model,
+                    location=external_data_path,
+                    all_tensors_to_one_file=util.default(
+                        self.all_tensors_to_one_file, True
+                    ),
+                )
+        else:
+            if self.size_threshold is not None:
+                G_LOGGER.warning(
+                    "size_threshold is set, but external data path has not been set. "
+                    "No external data will be written."
+                )
+            if self.all_tensors_to_one_file is not None:
+                G_LOGGER.warning(
+                    "all_tensors_to_one_file is set, but external data path has not been set. "
+                    "No external data will be written."
+                )
+
+        util.makedirs(self.path)
+        onnx.save(model, self.path)
+        return model
+
+
+@mod.export(funcify=True)
+class BytesFromOnnx(BaseLoader):
+    """
+    Functor that serializes an ONNX model.
+    """
+
+    def __init__(self, model):
+        """
+        Serializes an ONNX model.
+
+        Args:
+            model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]):
+                    An ONNX model or a callable that returns one.
+        """
+        self._model = model
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            bytes: The serialized model.
+        """
+        model, _ = util.invoke_if_callable(self._model)
+        return model.SerializeToString()
+
+
+@mod.export(funcify=True)
+class OnnxFromBytes(BaseLoader):
+    """
+    Functor that deserializes an ONNX model.
+    """
+
+    def __init__(self, serialized_onnx):
+        """
+        Deserializes an ONNX model.
+
+        Args:
+            serialized_onnx (Union[bytes, Callable() -> bytes]):
+                    A serialized ONNX model or a callable that returns one.
+        """
+        self._serialized_onnx = serialized_onnx
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnx.ModelProto: The ONNX model.
+        """
+        serialized_onnx, _ = util.invoke_if_callable(self._serialized_onnx)
+        model = onnx.ModelProto()
+        model.ParseFromString(serialized_onnx)
+        return model

tools/Polygraphy/polygraphy/backend/onnx/requirements.txt

@@ -0,0 +1 @@
+onnx

tools/Polygraphy/polygraphy/backend/onnx/util.py

@@ -0,0 +1,480 @@
+#
+# 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.
+#
+import copy
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.common import TensorMetadata
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER, LogMode
+
+gs = mod.lazy_import("onnx_graphsurgeon")
+onnx = mod.lazy_import("onnx")
+onnx_numpy_helper = mod.lazy_import("onnx.numpy_helper")
+
+
+def get_num_nodes(model):
+    def _get_num_graph_nodes(graph):
+        num_nodes = len(graph.node)
+        for node in graph.node:
+            for attr in node.attribute:
+                if attr.type == onnx.AttributeProto.GRAPH:
+                    num_nodes += _get_num_graph_nodes(attr.g)
+                elif attr.type == onnx.AttributeProto.GRAPHS:
+                    for subgraph in attr.graphs:
+                        num_nodes += _get_num_graph_nodes(subgraph)
+        return num_nodes
+
+    return _get_num_graph_nodes(model.graph)
+
+
+def all_tensor_names(model, include_inputs=None):
+    include_inputs = util.default(include_inputs, False)
+
+    all_outputs = [
+        output
+        for node in model.graph.node
+        if node.op_type != "Constant"
+        for output in node.output
+    ]
+    if include_inputs:
+        all_outputs += [inp.name for inp in model.graph.input]
+    all_outputs = util.unique_list(all_outputs)
+    return all_outputs
+
+
+def _check_has_tensors(model, outputs):
+    all_outputs = all_tensor_names(model, include_inputs=True)
+    util.check_sequence_contains(
+        all_outputs, outputs, name="the model", items_name="outputs", check_extra=False
+    )
+
+
+def mark_outputs(model, outputs):
+    # Clear the old outputs
+    while model.graph.output:
+        model.graph.output.pop()
+
+    outputs = util.unique_list(outputs)
+    _check_has_tensors(model, outputs)
+
+    value_info_map = {t.name: t for t in model.graph.value_info}
+    out_tensors = []
+    for output in outputs:
+        value_info = value_info_map.get(
+            output, onnx.helper.make_empty_tensor_value_info(output)
+        )
+        out_tensors.append(value_info)
+
+    G_LOGGER.ultra_verbose(f"Marked output tensors in ONNX model: {out_tensors}")
+    model.graph.output.extend(out_tensors)
+    return model
+
+
+def mark_layerwise(model):
+    # Add all non-constant node outputs as graph outputs
+    model = mark_outputs(model, all_tensor_names(model))
+    return model
+
+
+def unmark_outputs(model, outputs):
+    outputs = util.unique_list(outputs)
+    _check_has_tensors(model, outputs)
+
+    cur_outputs = []
+    while model.graph.output:
+        cur_outputs.append(model.graph.output.pop())
+    cur_outputs = list(reversed(cur_outputs))  # Preserve ordering
+
+    for out in cur_outputs:
+        if out.name not in outputs:
+            model.graph.output.extend([out])
+
+    return model
+
+
+def get_shape(tensor):
+    shape = []
+    if isinstance(tensor, onnx.TensorProto):
+        shape = tensor.dims
+    else:
+        for dim in tensor.type.tensor_type.shape.dim:
+            if dim.HasField("dim_param"):
+                shape.append(dim.dim_param)
+            elif dim.HasField("dim_value"):
+                shape.append(dim.dim_value)
+            else:
+                shape.append(-1)
+    return shape
+
+
+def get_dtype(tensor):
+    if isinstance(tensor, onnx.TensorProto):
+        onnx_type = tensor.data_type
+    else:
+        onnx_type = tensor.type.tensor_type.elem_type
+    return DataType.from_dtype(onnx_type, source_module="onnx")
+
+
+def get_values(tensor):
+    try:
+        return onnx_numpy_helper.to_array(tensor)
+    except Exception as err:
+        G_LOGGER.error(
+            f"Failed to load weights.\nNote: Error was: {err}", mode=LogMode.ONCE
+        )
+    return "<error: failed to load weights>"
+
+
+def get_tensor_metadata(tensors):
+    metadata = TensorMetadata()
+    for tensor in tensors:
+        metadata.add(name=tensor.name, dtype=get_dtype(tensor), shape=get_shape(tensor))
+    return metadata
+
+
+def get_input_metadata(graph):
+    # Some "inputs" are actually weights with initalizers, so we need to eliminate those.
+    initializer_names = {tensor.name for tensor in graph.initializer}
+    input_tensors = [
+        tensor for tensor in graph.input if tensor.name not in initializer_names
+    ]
+    return get_tensor_metadata(input_tensors)
+
+
+def get_output_metadata(graph):
+    return get_tensor_metadata(graph.output)
+
+
+def str_from_onnx(model, show_layers=None, show_attrs=None, show_weights=None):
+    """
+    Converts an ONNX Graph to a human-readable representation
+
+    Args:
+        graph (onnx.GraphProto): The onnx graph.
+        show_layers (bool): Whether to display per-layer information.
+        show_attrs (bool): Whether to display per-layer attributes.
+        show_weights (bool): Whether to display the value of weights.
+
+    Returns:
+        str
+    """
+    show_layers = util.default(show_layers, False)
+    show_attrs = util.default(show_attrs, False)
+    show_weights = util.default(show_weights, False)
+
+    def get_opset():
+        default_opset = "Unknown"
+        other_opsets = {}
+        for info in model.opset_import:
+            if not info.domain:
+                default_opset = info.version
+            else:
+                other_opsets[info.domain] = info.version
+        return default_opset, other_opsets
+
+    default_opset, other_opsets = get_opset()
+    onnx_str = ""
+    onnx_str += f"Name: {model.graph.name} | ONNX Opset: {default_opset}"
+    if other_opsets:
+        onnx_str += f" | Other Opsets: {other_opsets}"
+    onnx_str += "\n\n"
+
+    onnx_str += str_from_onnx_graph(
+        model.graph,
+        tensors={},
+        show_layers=show_layers,
+        show_attrs=show_attrs,
+        show_weights=show_weights,
+    )
+    return onnx_str
+
+
+def str_from_onnx_graph(
+    graph, tensors, show_layers, show_attrs, show_weights, indent_level=0
+):
+    input_metadata = get_input_metadata(graph)
+    output_metadata = get_output_metadata(graph)
+    initializer_metadata = get_tensor_metadata(graph.initializer)
+
+    # Subgraph inputs should remain separate from each other, hence copy the tensors map
+    tensors = copy.copy(tensors)
+    tensors.update(get_tensor_metadata(graph.value_info))
+    tensors.update(initializer_metadata)
+    tensors.update(input_metadata)
+    tensors.update(output_metadata)
+
+    graph_type = "Graph" if indent_level == 0 else "Subgraph"
+
+    onnx_str = ""
+    if show_attrs and graph.doc_string:
+        onnx_str += f"---- Docstring ----\n{graph.doc_string}\n\n"
+
+    onnx_str += (
+        f"---- {len(input_metadata)} {graph_type} Input(s) ----\n{input_metadata}\n\n"
+    )
+    onnx_str += f"---- {len(output_metadata)} {graph_type} Output(s) ----\n{output_metadata}\n\n"
+
+    onnx_str += f"---- {len(initializer_metadata)} Initializer(s) ----\n"
+    if show_weights:
+        for init in graph.initializer:
+            onnx_str += f"Initializer | {init.name} [dtype={get_dtype(init)}, shape={get_shape(init)}] | Values:\n{util.indent_block(str(get_values(init)))}\n\n"
+        if not graph.initializer:
+            onnx_str += "{}\n\n"
+    elif show_layers:
+        onnx_str += str(initializer_metadata)
+        onnx_str += "\n\n"
+    else:
+        onnx_str += "\n"
+
+    def get_names_and_meta(names):
+        names_lst = []
+        metadata = TensorMetadata()
+        for name in names:
+            dtype, shape = tensors.get(name, (None, None))
+            if name in initializer_metadata:
+                name = f"Initializer | {name}"
+            names_lst.append(name)
+            metadata.add(name=name, dtype=dtype, shape=shape)
+        return names_lst, metadata
+
+    # Maps values from the AttributeType enum to their string representations, e.g., {1: "FLOAT"}
+    ATTR_TYPE_MAPPING = dict(
+        zip(
+            onnx.AttributeProto.AttributeType.values(),
+            onnx.AttributeProto.AttributeType.keys(),
+        )
+    )
+
+    # Maps an ONNX attribute to the corresponding Python property
+    ONNX_PYTHON_ATTR_MAPPING = {
+        "FLOAT": "f",
+        "INT": "i",
+        "STRING": "s",
+        "TENSOR": "t",
+        "GRAPH": "g",
+        "FLOATS": "floats",
+        "INTS": "ints",
+        "STRINGS": "strings",
+    }
+
+    def attrs_to_dict(attrs):
+        attr_dict = OrderedDict()
+        for attr in attrs:
+
+            def process_attr(attr_str: str):
+                processed = getattr(attr, ONNX_PYTHON_ATTR_MAPPING[attr_str])
+                if attr_str == "STRING":
+                    processed = processed.decode()
+                elif attr_str == "TENSOR":
+                    tensor_str = f"Tensor: [dtype={get_dtype(processed)}, shape={get_shape(processed)}]"
+                    if show_weights:
+                        tensor_str += " | Values:\n" + util.indent_block(
+                            str(get_values(processed))
+                        )
+                    processed = tensor_str
+                elif attr_str == "GRAPH":
+                    processed = "\n" + str_from_onnx_graph(
+                        processed,
+                        tensors,
+                        indent_level=indent_level + 2,
+                        show_layers=show_layers,
+                        show_attrs=show_attrs,
+                        show_weights=show_weights,
+                    )
+                elif attr_str == "FLOATS" or attr_str == "INTS":
+                    # Proto hacky list to normal Python list
+                    processed = [p for p in processed]
+                elif attr_str == "STRINGS":
+                    processed = [p.decode() for p in processed]
+                return processed
+
+            if attr.type in ATTR_TYPE_MAPPING:
+                attr_str = ATTR_TYPE_MAPPING[attr.type]
+                if attr_str in ONNX_PYTHON_ATTR_MAPPING:
+                    attr_dict[attr.name] = process_attr(attr_str)
+                else:
+                    G_LOGGER.warning(
+                        f"Attribute of type {attr_str} is currently unsupported. Skipping attribute."
+                    )
+            else:
+                G_LOGGER.warning(
+                    f"Attribute type: {attr.type} was not recognized. Was the graph generated with a newer IR version than the installed `onnx` package? Skipping attribute."
+                )
+        return attr_dict
+
+    onnx_str += f"---- {len(graph.node)} Node(s) ----\n"
+    if show_layers:
+        for index, node in enumerate(graph.node):
+            input_names, input_meta = get_names_and_meta(node.input)
+            output_names, output_meta = get_names_and_meta(node.output)
+
+            onnx_str += util.str_from_layer(
+                "Node",
+                index,
+                node.name,
+                node.op_type,
+                input_names,
+                input_meta,
+                output_names,
+                output_meta,
+            )
+
+            if show_attrs:
+                attrs = attrs_to_dict(node.attribute)
+                if attrs:
+                    onnx_str += util.indent_block("---- Attributes ----") + "\n"
+                for key, val in attrs.items():
+                    attr_str = ""
+                    if node.name:
+                        attr_str += f"{node.name}."
+                    onnx_str += util.indent_block(f"{attr_str}{key} = {val}") + "\n"
+            onnx_str += "\n"
+
+    return util.indent_block(onnx_str, indent_level)
+
+
+##
+## ONNX-GraphSurgeon utilities
+##
+
+
+def meta_from_gs_tensors(tensors):
+    """Get TensorMetadata from a list of ONNX-GraphSurgeon tensors"""
+    meta = TensorMetadata()
+    for tensor in tensors:
+        meta.add(tensor.name, tensor.dtype, tensor.shape)
+    return meta
+
+
+def set_shapes_from_layerwise_meta(graph, layerwise_meta):
+    """
+    Args:
+        graph (gs.Graph): An ONNX graphsurgeon graph.
+        layerwise_meta (TensorMetadata): Metadata for tensors in the graph.
+    """
+    for tensor in graph.tensors().values():
+        if isinstance(tensor, gs.Variable) and tensor.name in layerwise_meta:
+            tensor.shape = layerwise_meta[tensor.name].shape
+            tensor.dtype = DataType.to_dtype(
+                DataType.from_dtype(layerwise_meta[tensor.name].dtype), "onnx"
+            )
+
+
+def lower_constant_nodes(graph):
+    """Converts the outputs of Constant nodes into constant tensors, removing the nodes"""
+    remove_nodes = set()
+    with graph.node_ids():
+        for node in graph.nodes:
+            if node.op == "Constant" and "value" in node.attrs:
+                node.outputs[0].to_constant(node.attrs["value"].values)
+                remove_nodes.add(node.id)
+        # Iterate from the end so we don't shift the list under us.
+        for node_id in sorted(remove_nodes, reverse=True):
+            del graph.nodes[node_id]
+    return graph
+
+
+def get_unbounded_dds_tensors(graph):
+    graph.toposort()
+    # A dict of operators that might produce a output tensor with unbounded DDS, when the value of the input tensor
+    # at the corresponding index is a runtime value. For example, "Range" => "1" means that if the input 1 of the Range
+    # operator is a runtime value, e.g. not a const tensor or an initializer, then the Range output tensor size is unbounded.
+    dispatcher_dict = {
+        "Range": [1],  # the limit input of the Range operator
+        "Pad": [1],  # the pads input of the Pad operator
+        "Resize": [3],  # the sizes input of the Resize operator
+        "Tile": [1],  # the repeats input of the Tile operator
+        "Expand": [1],  # the shape input of the Expand operator
+    }
+
+    # Check if the given operator produces a output tensor with unbounded DDS.
+    def check_op(node, const_tensor_set):
+        # Check if the operator is inside the dispatcher dict.
+        if node.op in dispatcher_dict:
+            input_idx_list = dispatcher_dict[node.op]
+            for input_idx in input_idx_list:
+                if input_idx < len(node.inputs):
+                    input_tensor = node.inputs[input_idx]
+                    # Check if the corresponding input tensor is a runtime value and its producer is not Min operator.
+                    # If a tensor is produced by a Min operator, its upper bound has already been set.
+                    if (
+                        input_tensor.name not in const_tensor_set
+                        and len(input_tensor.inputs) >= 1
+                        and input_tensor.inputs[0].op != "Min"
+                    ):
+                        return input_tensor
+        return None
+
+    # Find all constant tensors.
+    def get_const_tensors(graph):
+        return {
+            tensor.name
+            for tensor in graph.tensors().values()
+            if isinstance(tensor, gs.Constant)
+        }
+
+    # Find all dynamic shape symbols, customers will set upper bounds for these symbols when building the model in TensorRT.
+    def get_dynamic_shapes(graph):
+        dynamic_shape_set = set()
+        for tensor in graph.inputs:
+            for shape in tensor.shape:
+                if isinstance(shape, str):
+                    dynamic_shape_set.add(shape)
+        return dynamic_shape_set
+
+    # Find all tensors with unbounded DDS.
+    def get_target_tensors(graph):
+        # Find dynamic shapes, these shapes should have upper bounds in TensorRT.
+        dynamic_shape_set = get_dynamic_shapes(graph)
+
+        # Find const tensors. For those operators in the dispatch dict, constant inputs will not introduce outputs with unbounded DDS.
+        const_tensor_set = get_const_tensors(graph)
+
+        # Our target is to find those input tensors that cause its consumer nodes generated unbounded outputs.
+        # If a tensor has named dimensions that appeared before in its symbolic shape, it means that the shape is *not* data dependent,
+        # and so will have an upper bound.
+        target_tensor_names = set()
+        target_tensor_list = []
+        for node in graph.nodes:
+            check_node = False
+            # Check if the node's output contains a new introduced dynamic shape.
+            for tensor in node.outputs:
+                # Always check nodes if tensor.shape is None.
+                # This happens when the symbolic inference does not work correctly due to some restrictions.
+                if tensor.shape is None:
+                    check_node = True
+                else:
+                    for shape in tensor.shape:
+                        # If a shape is a dynamic shape, then it is a str.
+                        # Only check the node that first introduced the dynamic shape.
+                        if isinstance(shape, str) and shape not in dynamic_shape_set:
+                            dynamic_shape_set.add(shape)
+                            check_node = True
+            # Check if the node will generate an unbounded output size.
+            if check_node:
+                target_tensor = check_op(node, const_tensor_set)
+                # Avoid duplication.
+                if (
+                    target_tensor is not None
+                    and target_tensor.name not in target_tensor_names
+                ):
+                    target_tensor_names.add(target_tensor.name)
+                    target_tensor_list.append(target_tensor)
+        return target_tensor_list
+
+    return get_target_tensors(graph)

tools/Polygraphy/polygraphy/backend/onnxrt/__init__.py

@@ -0,0 +1,2 @@
+from polygraphy.backend.onnxrt.loader import *
+from polygraphy.backend.onnxrt.runner import *

tools/Polygraphy/polygraphy/backend/onnxrt/loader.py

@@ -0,0 +1,85 @@
+#
+# 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.backend.base import BaseLoader
+from polygraphy.logger import G_LOGGER
+import os
+
+onnxrt = mod.lazy_import("onnxruntime")
+
+
+@mod.export(funcify=True)
+class SessionFromOnnx(BaseLoader):
+    """
+    Functor that builds an ONNX-Runtime inference session.
+    """
+
+    def __init__(self, model_bytes, providers=None):
+        """
+        Builds an ONNX-Runtime inference session.
+
+        Args:
+            model_bytes (Union[Union[bytes, str], Callable() -> Union[bytes, str]]):
+                    A serialized ONNX model or a path to a model or a callable that returns one of those.
+
+            providers (Sequence[str]):
+                    A sequence of execution providers to use in order of priority.
+                    Each element of the sequence may be either an exact match or a case-insensitive partial match
+                    for the execution providers available in ONNX-Runtime. For example, a value of "cpu" would
+                    match the "CPUExecutionProvider".
+                    Defaults to ``["cpu"]``.
+
+        """
+        self._model_bytes_or_path = model_bytes
+        self.providers = util.default(providers, ["cpu"])
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            onnxruntime.InferenceSession: The inference session.
+        """
+        model_bytes, _ = util.invoke_if_callable(self._model_bytes_or_path)
+
+        available_providers = onnxrt.get_available_providers()
+        providers = []
+        for prov in self.providers:
+            matched_prov_name = util.find_str_in_iterable(prov[0] if isinstance(prov, tuple) else prov, available_providers)
+            matched_prov = (matched_prov_name, prov[1]) if isinstance(prov, tuple) else matched_prov_name
+            if matched_prov is None:
+                G_LOGGER.critical(
+                    f"Could not find specified ONNX-Runtime execution provider.\nNote: Requested provider was: {prov}, but available providers are: {available_providers}"
+                )
+            providers.append(matched_prov)
+
+        G_LOGGER.start(
+            f"Creating ONNX-Runtime Inference Session with providers: {providers}"
+        )
+        # ONNX Runtime tried to bind each thread to a logical CPU, but not all assigned cpu cores are available on some platforms.
+        # Set the number of threads within each operator and between operators the number of usable CPUs to avoid crash in onnxruntime on those platforms.
+        options = onnxrt.SessionOptions()
+        try:
+            # sched_getaffinity is only available on UNIX platforms
+            process_cpu_count = len(os.sched_getaffinity(0))
+        except AttributeError:
+            process_cpu_count = 1
+
+        options.intra_op_num_threads = process_cpu_count
+        options.inter_op_num_threads = process_cpu_count
+        return onnxrt.InferenceSession(
+            model_bytes, providers=providers, sess_options=options
+        )

tools/Polygraphy/polygraphy/backend/onnxrt/requirements.txt

@@ -0,0 +1 @@
+onnxruntime

tools/Polygraphy/polygraphy/backend/onnxrt/runner.py

@@ -0,0 +1,90 @@
+#
+# 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.
+#
+import time
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.common import TensorMetadata
+from polygraphy.datatype import DataType
+
+
+@mod.export()
+class OnnxrtRunner(BaseRunner):
+    """
+    Runs inference using an ONNX-Runtime inference session.
+    """
+
+    def __init__(self, sess, name=None):
+        """
+        Args:
+            sess (Union[onnxruntime.InferenceSession, Callable() -> onnxruntime.InferenceSession]):
+                    An ONNX-Runtime inference session or a callable that returns one.
+        """
+        super().__init__(name=name, prefix="onnxrt-runner")
+        self._sess = sess
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        self.sess, _ = util.invoke_if_callable(self._sess)
+
+    @util.check_called_by("get_input_metadata")
+    def get_input_metadata_impl(self):
+        meta = TensorMetadata()
+        for node in self.sess.get_inputs():
+            meta.add(
+                node.name,
+                dtype=DataType.from_dtype(node.type, "onnxruntime"),
+                shape=node.shape,
+            )
+        return meta
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict):
+        """
+        Implementation for running inference with ONNX-Runtime.
+        Do not call this method directly - use ``infer()`` instead,
+        which will forward unrecognized arguments to this method.
+
+        Args:
+            feed_dict (OrderedDict[str, Union[numpy.ndarray, torch.Tensor]]):
+                    A mapping of input tensor names to corresponding input NumPy arrays or PyTorch tensors.
+                    If PyTorch tensors are provided in the feed_dict, then this function
+                    will return the outputs also as PyTorch tensors.
+
+        Returns:
+            OrderedDict[str, Union[numpy.ndarray, torch.Tensor]]:
+                    A mapping of output tensor names to corresponding output NumPy arrays
+                    or PyTorch tensors.
+        """
+        use_torch = any(util.array.is_torch(t) for t in feed_dict.values())
+        # `to_numpy()`` and `to_torch()` should be zero-copy whenever possible.
+        feed_dict = {name: util.array.to_numpy(t) for name, t in feed_dict.items()}
+
+        start = time.time()
+        inference_outputs = self.sess.run(None, feed_dict)
+        end = time.time()
+
+        out_dict = OrderedDict()
+        for node, out in zip(self.sess.get_outputs(), inference_outputs):
+            out_dict[node.name] = out if not use_torch else util.array.to_torch(out)
+        self.inference_time = end - start
+        return out_dict
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        del self.sess

tools/Polygraphy/polygraphy/backend/pluginref/__init__.py

@@ -0,0 +1,17 @@
+#
+# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.backend.pluginref.runner import *

tools/Polygraphy/polygraphy/backend/pluginref/references.py

@@ -0,0 +1,102 @@
+#
+# 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
+from polygraphy.logger import G_LOGGER
+
+np = mod.lazy_import("numpy")
+gs = mod.lazy_import("onnx_graphsurgeon")
+
+OP_REGISTRY = {}  # Dict[str, Callable]: Maps op names to reference implementations
+
+
+def register(op):
+    """
+    Registers a function as the reference implementation for a given op.
+
+    Args:
+         op (str): The name of the op for which to register this function.
+    """
+
+    def register_impl(func):
+        def wrapped_func(node, intermediate_tensors):
+            inputs = []
+            for inp in node.inputs:
+                if inp.is_empty():  # Optional input
+                    inputs.append(None)
+                elif isinstance(inp, gs.Constant):
+                    inputs.append(inp.values)
+                elif inp.name in intermediate_tensors:
+                    inputs.append(intermediate_tensors[inp.name])
+                else:
+                    G_LOGGER.internal_error(
+                        f"Input: {inp.name} was not found in intermediate tensors and is not a constant.\nNote: Intermediate tensors include: {list(intermediate_tensors.keys())}"
+                    )
+
+            outputs = func(node.attrs, *inputs)
+            if len(outputs) != len(node.outputs):
+                G_LOGGER.internal_error(
+                    f"{op} reference implementation returned the wrong number of outputs.\nNote: Expected {len(node.outputs)} but recevied {len(outputs)}"
+                )
+
+            return {
+                out_tensor.name: out for out_tensor, out in zip(node.outputs, outputs)
+            }
+
+        OP_REGISTRY[op] = wrapped_func
+        return wrapped_func
+
+    return register_impl
+
+
+@register("Identity")
+def run_identity(attrs, x):
+    return [x]
+
+
+@register("InstanceNormalization")
+def run_instancenorm(attrs, x, weights, bias):
+    epsilon = attrs.get("epsilon", 1.0e-5)
+
+    rank = len(x.shape)
+    axis = tuple(range(2, rank))
+    mean = np.mean(x, axis=axis, keepdims=True)
+    var = np.var(x, axis=axis, keepdims=True)
+
+    # Weights and bias needs to be broadcasted to shape of X. C dimension should be a wildcard.
+    broadcast_shape = [-1] + [1] * (rank - 2)
+    weights = weights.reshape(broadcast_shape)
+    bias = bias.reshape(broadcast_shape)
+
+    res = weights * (x - mean) / np.sqrt(var + epsilon) + bias
+    return [res]
+
+
+@register("MeanVarianceNormalization")
+def run_meanvarnorm(attrs, x):
+    epsilon = 1.0e-9
+    axes = attrs.get("axes", [0, 2, 3])
+    axes = tuple(axes)
+
+    data_mean = np.mean(x, axis=axes, keepdims=True)
+    data_mean_squared = np.power(data_mean, 2)
+    data_squared = np.power(x, 2)
+    data_squared_mean = np.mean(data_squared, axis=axes, keepdims=True)
+    std = np.sqrt(data_squared_mean - data_mean_squared)
+    res = (x - data_mean) / (std + epsilon)
+
+    return [res]

tools/Polygraphy/polygraphy/backend/pluginref/requirements.txt

@@ -0,0 +1,2 @@
+numpy
+onnx_graphsurgeon

tools/Polygraphy/polygraphy/backend/pluginref/runner.py

@@ -0,0 +1,83 @@
+#
+# 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.
+#
+
+import copy
+import time
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.backend.pluginref.references import OP_REGISTRY
+from polygraphy.logger import G_LOGGER
+
+np = mod.lazy_import("numpy")
+onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
+
+
+@mod.export()
+class PluginRefRunner(BaseRunner):
+    """
+    Runs inference using custom CPU reference implementations
+    """
+
+    def __init__(self, graph, name=None):
+        """
+        Args:
+            graph (Union[onnx_graphsurgeon.Graph, Callable() -> onnx_graphsurgeon.Graph]):
+                    An ONNX-GraphSurgeon graph or a callable that returns one.
+            name (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+        """
+        super().__init__(name=name, prefix="pluginref-runner")
+        self._graph = graph
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        self.graph, _ = util.invoke_if_callable(self._graph)
+
+    @util.check_called_by("get_input_metadata")
+    def get_input_metadata_impl(self):
+        return onnx_util.meta_from_gs_tensors(self.graph.inputs)
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict):
+        start = time.time()
+
+        intermediate_tensors = copy.copy(feed_dict)
+        for node in self.graph.nodes:
+            if node.op not in OP_REGISTRY:
+                G_LOGGER.critical(
+                    f"Op: {node.op} does not have a reference implementation registered!"
+                )
+
+            intermediate_tensors.update(
+                OP_REGISTRY[node.op](node, intermediate_tensors)
+            )
+
+        outputs = OrderedDict()
+        for out in self.graph.outputs:
+            outputs[out.name] = intermediate_tensors[out.name]
+
+        end = time.time()
+
+        self.inference_time = end - start
+        return outputs
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        del self.graph

tools/Polygraphy/polygraphy/backend/pyt/__init__.py

@@ -0,0 +1 @@
+from polygraphy.backend.pyt.runner import *

tools/Polygraphy/polygraphy/backend/pyt/requirements.txt

@@ -0,0 +1 @@
+torch

tools/Polygraphy/polygraphy/backend/pyt/runner.py

@@ -0,0 +1,81 @@
+#
+# 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.
+#
+import time
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+
+torch = mod.lazy_import("torch>=1.13.0")
+
+
+@mod.export()
+class PytRunner(BaseRunner):
+    """
+    Runs inference using PyTorch.
+    """
+
+    def __init__(self, model, input_metadata, output_names, name=None):
+        """
+        Args:
+            model (Union[torch.nn.Module, Callable() -> torch.nn.Module]):
+                    A torch.nn.Module or subclass or a callable that returns one.
+            input_metadata (TensorMetadata): Mapping of input names to their data types and shapes.
+            output_names (List[str]):
+                    A list of output names of the model. This information is used by the
+                    Comparator to determine which outputs to compare.
+
+
+            name (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+        """
+        super().__init__(name=name, prefix="pytorch-runner")
+        self._model = model
+        self.input_metadata = input_metadata
+        self.output_names = output_names
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        self.model, _ = util.invoke_if_callable(self._model)
+        self.model.eval()
+
+    @util.check_called_by("get_input_metadata")
+    def get_input_metadata_impl(self):
+        return self.input_metadata
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict):
+        with torch.no_grad():
+            inputs = [
+                torch.from_numpy(val.astype(dtype)).cuda()
+                for (val, (dtype, _)) in zip(
+                    feed_dict.values(), self.input_metadata.values()
+                )
+            ]
+            start = time.time()
+            outputs = self.model(*inputs)
+            end = time.time()
+
+        out_dict = OrderedDict()
+        for name, output in zip(self.output_names, outputs):
+            out_dict[name] = output.cpu().numpy()
+        return out_dict, end - start
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        del self.model

tools/Polygraphy/polygraphy/backend/tensorrt_rtx/__init__.py

@@ -0,0 +1 @@
+from polygraphy.backend.tensorrt_rtx.config import *

tools/Polygraphy/polygraphy/backend/tensorrt_rtx/config.py

@@ -0,0 +1,117 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 config as polygraphy_config, mod, util
+from polygraphy.backend.trt.config import _CreateConfigCommon
+from polygraphy.backend.trt.util import inherit_and_extend_docstring
+from polygraphy.logger import G_LOGGER
+from polygraphy.mod.trt_importer import lazy_import_trt
+
+trt = lazy_import_trt()
+
+
+@mod.export(funcify=True, func_name="create_config_rtx")
+class CreateConfigRTX(_CreateConfigCommon):
+    """
+    Functor that creates an IBuilderConfig with TensorRT-RTX specific features.
+    """
+
+    @inherit_and_extend_docstring(_CreateConfigCommon.__init__)
+    def __init__(
+        self,
+        use_gpu=None,
+        compute_capabilities=None,
+        **kwargs
+    ):
+        """
+        Creates an IBuilderConfig with TensorRT-RTX specific features.
+
+        Args:
+            use_gpu (bool):
+                    Whether to use the current GPU device as target for engine compilation.
+                    Equivalent to setting ComputeCapability.CURRENT. This is mutually exclusive with compute_capabilities.
+                    Defaults to False.
+            compute_capabilities (List[Tuple[int, int]]):
+                    List of (major, minor) compute capability tuples to target for engine compilation.
+                    This is mutually exclusive with use_gpu. When specified, the engine can only run on devices
+                    with the specified compute capabilities.
+                    Defaults to None.
+        """
+        super().__init__(**kwargs)
+        self.use_gpu = util.default(use_gpu, False)
+        self.compute_capabilities = compute_capabilities
+
+        if self.use_gpu and self.compute_capabilities:
+            G_LOGGER.critical("use_gpu and compute_capabilities are mutually exclusive.")
+
+        self._validator()
+
+    def _validator(self):
+        """
+        Validates initialization parameters for TensorRT-RTX specific features.
+        """
+        if self.use_gpu or self.compute_capabilities is not None:
+            if not polygraphy_config.USE_TENSORRT_RTX:
+                G_LOGGER.critical("--compute-capabilities and --use-gpu settings are only supported with USE_TENSORRT_RTX=1.")
+            
+            # Validate compute capabilities format and availability
+            if self.compute_capabilities:
+                for major, minor in self.compute_capabilities:
+                    cap_name = f"SM{major}{minor}"
+                    if not hasattr(trt.ComputeCapability, cap_name):
+                        G_LOGGER.critical(f"Compute capability {major}.{minor} ({cap_name})"
+                                           " not supported by this TensorRT-RTX version.")
+
+    def _configure_flags(self, builder, network, config):
+        """
+        Validates and configures TensorRT-RTX-specific features.
+
+        Args:
+            builder (trt.Builder): The TensorRT builder
+            network (trt.INetworkDefinition): The TensorRT network
+            config (trt.IBuilderConfig): The TensorRT builder config to modify
+        """
+        # Set compute capabilities if specified
+        if self.use_gpu or self.compute_capabilities is not None:
+            try:
+                if self.use_gpu:
+                    # Use current GPU device
+                    config.num_compute_capabilities = 1
+                    config.set_compute_capability(trt.ComputeCapability.CURRENT, 0)
+                    G_LOGGER.info("Using current GPU device for engine compilation (ComputeCapability.CURRENT)")
+                elif self.compute_capabilities:
+                    # Set specific compute capabilities
+                    config.num_compute_capabilities = len(self.compute_capabilities)
+                    G_LOGGER.info(f"Setting {len(self.compute_capabilities)} target compute capabilities: {self.compute_capabilities}")
+                    for i, (major, minor) in enumerate(self.compute_capabilities):
+                        cap_name = f"SM{major}{minor}"
+                        compute_cap = getattr(trt.ComputeCapability, cap_name)
+                        config.set_compute_capability(compute_cap, i)
+            except Exception as e:
+                G_LOGGER.critical(f"Failed to set compute capabilities: {e}. You are likely not using a TensorRT-RTX build.")
+
+    @util.check_called_by("__call__")
+    def call_impl(self, builder, network):
+        """
+        Callable implementation that creates and configures the IBuilderConfig with TensorRT-RTX features.
+        """
+        # Enable all common config options
+        config = super().call_impl(builder, network)
+
+        self._configure_flags(builder, network, config)
+
+        return config

tools/Polygraphy/polygraphy/backend/tf/README.md

@@ -0,0 +1 @@
+NOTE: The `tf` backend currently only supports TensorFlow 1.X.

tools/Polygraphy/polygraphy/backend/tf/__init__.py

@@ -0,0 +1,37 @@
+from polygraphy.backend.tf.loader import *
+from polygraphy.backend.tf.runner import *
+
+
+def register_logger_callback():
+    from polygraphy.logger import G_LOGGER
+
+    def set_tf_logging_level(severity_trie):
+        import os
+        from polygraphy import mod
+
+        tf = mod.lazy_import("tensorflow<2.0")
+
+        if not tf.is_installed() or not tf.is_importable():
+            return
+
+        sev = severity_trie.get(G_LOGGER.module_path(os.path.dirname(__file__)))
+        if sev > G_LOGGER.WARNING:
+            tf_sev = tf.compat.v1.logging.ERROR
+            tf_logging_level = "3"
+        elif sev > G_LOGGER.INFO:
+            tf_sev = tf.compat.v1.logging.WARN
+            tf_logging_level = "2"
+        elif sev > G_LOGGER.VERBOSE:
+            tf_sev = tf.compat.v1.logging.INFO
+            tf_logging_level = "1"
+        else:
+            tf_sev = tf.compat.v1.logging.DEBUG
+            tf_logging_level = "0"
+
+        tf.compat.v1.logging.set_verbosity(tf_sev)
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = tf_logging_level
+
+    G_LOGGER.register_callback(set_tf_logging_level)  # Will be registered when this backend is imported.
+
+
+register_logger_callback()

tools/Polygraphy/polygraphy/backend/tf/loader.py

@@ -0,0 +1,496 @@
+#
+# 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.
+#
+# Sets up everything needed to perform inference in TensorFlow.
+import os
+
+from polygraphy import constants, mod, util
+from polygraphy.backend.base import BaseLoader
+from polygraphy.backend.tf import util as tf_util
+from polygraphy.logger import G_LOGGER
+
+tf = mod.lazy_import("tensorflow<2.0")
+
+
+@mod.export(funcify=True)
+class OptimizeGraph(BaseLoader):
+    """
+    Functor that freezes a TensorFlow graph, and folds constants.
+    """
+
+    def __init__(self, graph):
+        """
+        Freezes a TensorFlow graph and folds constants.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+        """
+        self._graph = graph
+
+    def constfold(self, graphdef, output_names):
+        from tensorflow.core.protobuf import (
+            config_pb2,
+            meta_graph_pb2,
+            rewriter_config_pb2,
+        )
+        from tensorflow.python.framework import importer, ops
+        from tensorflow.python.grappler import tf_optimizer
+        from tensorflow.python.training import saver
+
+        graph = ops.Graph()
+        with graph.as_default():
+            output_collection = meta_graph_pb2.CollectionDef()
+            output_list = output_collection.node_list.value
+            for output in output_names:
+                output_list.append(output.encode("utf-8"))
+
+            importer.import_graph_def(graphdef, name="")
+            metagraph = saver.export_meta_graph(
+                graph_def=graph.as_graph_def(add_shapes=True), graph=graph
+            )
+            metagraph.collection_def["train_op"].CopyFrom(output_collection)
+
+        rewriter_config = rewriter_config_pb2.RewriterConfig()
+        rewriter_config.optimizers.extend(["constfold"])
+        rewriter_config.meta_optimizer_iterations = (
+            rewriter_config_pb2.RewriterConfig.ONE
+        )
+
+        session_config = config_pb2.ConfigProto()
+        session_config.graph_options.resave_options.CopyFrom(rewriter_config)
+        return tf_optimizer.OptimizeGraph(session_config, metagraph, graph_id=b"graph")
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, output_names), _ = util.invoke_if_callable(self._graph)
+        with tf.Session(graph=graph) as sess:
+            sess.run(tf.initializers.global_variables())
+            sess.run(tf.initializers.local_variables())
+
+            graphdef = sess.graph.as_graph_def()
+            removed = tf.graph_util.remove_training_nodes(graphdef)
+            G_LOGGER.ultra_verbose(f"Removed nodes: {removed}")
+
+            for node in graphdef.node:
+                if node.op == "RefSwitch":
+                    node.op = "Switch"
+                    for index in range(len(node.input)):
+                        if "moving_" in node.input[index]:
+                            node.input[index] = node.input[index] + "/read"
+                elif node.op == "AssignSub":
+                    node.op = "Sub"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                elif node.op == "AssignAdd":
+                    node.op = "Add"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                elif node.op == "Assign":
+                    node.op = "Identity"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                    if "validate_shape" in node.attr:
+                        del node.attr["validate_shape"]
+                    if len(node.input) == 2:
+                        # input0: ref: Should be from a Variable node. May be uninitialized.
+                        # input1: value: The value to be assigned to the variable.
+                        node.input[0] = node.input[1]
+                        del node.input[1]
+
+            # Strip port information from outputs
+            output_names = [name.split(":")[0] for name in output_names]
+            output_graph_def = tf.graph_util.convert_variables_to_constants(
+                sess, graphdef, output_names
+            )
+            output_graph_def = self.constfold(output_graph_def, output_names)
+            return graph_from_frozen(output_graph_def)
+
+
+@mod.export(funcify=True)
+class GraphFromKeras(BaseLoader):
+    """
+    Functor that loads a TensorFlow model from Keras.
+    """
+
+    def __init__(self, path):
+        """
+        Loads a TensorFlow model from Keras.
+
+        Args:
+            path (Union[str, h5py.File]): A path to the saved model, or the file object.
+        """
+        self.path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+
+        from tensorflow.python import keras
+        from tensorflow.python.keras import backend
+
+        model = keras.models.load_model(self.path)
+        graph = backend.get_session().graph
+        return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class GraphFromFrozen(BaseLoader):
+    """
+    Functor that loads a TensorFlow frozen model.
+    """
+
+    def __init__(self, path):
+        """
+        Loads a TensorFlow frozen model.
+
+        Args:
+            path (Union[str, tf.Graph, tf.GraphDef]):
+                    A path to the frozen model, or a frozen TensorFlow graph or graphdef.
+        """
+        self.path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        graph = tf_util.load_graph(self.path)
+        return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class GraphFromCkpt(BaseLoader):
+    """
+    Functor that loads a TensorFlow model from a checkpoint. Note that in order to use checkpoints,
+    you must NOT use subprocesses in the Comparator.
+    """
+
+    def __init__(self, dir, name=None):
+        """
+        Loads a TensorFlow model from a checkpoint.
+
+        Args:
+            dir (str): Path to a directory containing checkpoints.
+
+
+            name (str):
+                    The name of the checkpoint to load, not including the file extension.
+                    For example, to load `model.meta`, the argument would be `model`.
+        """
+        self.dir = dir
+        self.name = name
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        # If `name` is not provided, this expects that the directory contains a `checkpoint` file with the contents:
+        #
+        # model_checkpoint_path: "model"
+        # all_model_checkpoint_paths: "model"
+        #
+        # where "model" is the checkpoint name
+        if not os.path.isdir(self.dir):
+            G_LOGGER.warning(
+                f"Specified checkpoint directory: {self.dir} does not look like a directory."
+            )
+
+        if self.name is None:
+            G_LOGGER.verbose(
+                "Checkpoint name was not explicitly provided, searching for `checkpoint` file"
+            )
+            checkpoint = tf.train.get_checkpoint_state(self.dir)
+            if checkpoint is None:
+                ckpt_file_contents = '\nmodel_checkpoint_path: "model"\nall_model_checkpoint_paths: "model"\n'
+                G_LOGGER.critical(
+                    f"Checkpoint directory: {self.dir} does not contain a `checkpoint` file, and the checkpoint name was not provided. Please either create a checkpoint file with the contents:\n{ckpt_file_contents} \nWhere `model` is the name of the checkpoint, or explicitly provide the name with --ckpt, not including file extensions"
+                )
+            input_checkpoint = checkpoint.model_checkpoint_path
+        else:
+            input_checkpoint = os.path.join(self.dir, self.name)
+
+        meta_file = input_checkpoint + ".meta"
+        with tf.Graph().as_default() as graph, tf.compat.v1.Session(
+            graph=graph
+        ).as_default() as sess:
+            saver = tf.compat.v1.train.import_meta_graph(meta_file, clear_devices=True)
+            saver.restore(sess, input_checkpoint)
+            return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class UseTfTrt(BaseLoader):
+    """
+    [UNTESTED] Functor that optimizes a TensorFlow model using TF-TRT.
+    """
+
+    def __init__(
+        self,
+        graph,
+        max_workspace_size=None,
+        fp16=None,
+        int8=None,
+        max_batch_size=None,
+        is_dynamic_op=False,
+        minimum_segment_size=None,
+    ):
+        """
+        Optimizes a TensorFlow model using TF-TRT.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+            max_workspace_size (int): The maximum workspace size.
+            fp16 (bool): Whether to run in FP16 mode.
+            max_batch_size (int): The maximum batch size.
+        """
+        self._graph = graph
+        self.max_workspace_size = util.default(max_workspace_size, 1 << 24)
+        self.fp16 = util.default(fp16, False)
+        self.fp8 = util.default(fp8, False)
+        self.int8 = util.default(int8, False)
+        self.max_batch_size = util.default(max_batch_size, 1)
+        self.is_dynamic_op = is_dynamic_op
+        self.minimum_segment_size = util.default(minimum_segment_size, 3)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        from tensorflow.contrib import tensorrt as tf_trt
+
+        (graph, output_names), _ = util.invoke_if_callable(self._graph)
+
+        precision_mode = "FP16" if self.fp16 else "FP32"
+        precision_mode = "INT8" if self.int8 else precision_mode
+        precision_mode = "FP8" if self.fp8 else precision_mode
+
+        G_LOGGER.info(
+            f"For TF-TRT, using outputs={output_names}, max_workspace_size_bytes={self.max_workspace_size}, max_batch_size={self.max_batch_size}, minimum_segment_size={self.minimum_segment_size}, is_dynamic_op={self.is_dynamic_op}, precision_mode={precision_mode}"
+        )
+
+        graphdef = tf_trt.create_inference_graph(
+            graph.as_graph_def(),
+            outputs=output_names,
+            max_workspace_size_bytes=self.max_workspace_size,
+            max_batch_size=self.max_batch_size,
+            minimum_segment_size=self.minimum_segment_size,
+            is_dynamic_op=self.is_dynamic_op,
+            precision_mode=precision_mode,
+        )
+
+        segment_number = 0
+        for node in graphdef.node:
+            if node.op == "TRTEngineOp":
+                engine = node.attr["serialized_segment"].s
+                segment_number += 1
+        G_LOGGER.info(f"Found {segment_number} engines in TFTRT graph")
+
+        with tf.Graph().as_default() as graph:
+            tf.import_graph_def(graphdef, name="")
+            return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class ModifyGraphOutputs(BaseLoader):
+    """
+    Functor that modifies outputs of a TensorFlow graph.
+    """
+
+    def __init__(self, graph, outputs=None):
+        """
+        Modifies outputs of a TensorFlow graph.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            outputs (List[str]):
+                    Names of output tensors. If provided, this will override the outputs
+                    determined by the loader.
+                    If a value of `constants.MARK_ALL` is used instead of a list, all tensors in the network are marked.
+        """
+        self._graph = graph
+        self.outputs = outputs
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, outputs), _ = util.invoke_if_callable(self._graph)
+
+        if self.outputs == constants.MARK_ALL:
+            outputs = list(tf_util.get_output_metadata(graph, layerwise=True).keys())
+        elif self.outputs is not None:
+            outputs = self.outputs
+
+        return graph, outputs
+
+
+@mod.export(funcify=True)
+class SaveGraph(BaseLoader):
+    """
+    Functor that writes out artifacts from a TensorFlow graph.
+    """
+
+    def __init__(self, graph, path=None, tensorboard_dir=None, engine_dir=None):
+        """
+        Writes out artifacts from a TensorFlow Graph.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            path (str): Path at which to save the frozen graphdef.
+            tensorboard_dir (str): The directory in which to write TensorBoard visualizations.
+            engine_dir (str): The directory in which to save TF-TRT engines,
+        """
+        self._graph = graph
+        self.path = path
+        self.tensorboard_dir = tensorboard_dir
+        self.engine_dir = engine_dir
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, outputs), _ = util.invoke_if_callable(self._graph)
+
+        if self.path:
+            util.save_file(graph.as_graph_def().SerializeToString(), dest=self.path)
+        if self.tensorboard_dir:
+            G_LOGGER.info(f"Writing tensorboard events to {self.tensorboard_dir}")
+            train_writer = tf.compat.v1.summary.FileWriter(self.tensorboard_dir)
+            train_writer.add_graph(graph)
+
+        if self.engine_dir is not None:
+            graphdef = graph.as_graph_def()
+            segment_number = 0
+            for node in graphdef.node:
+                if node.op == "TRTEngineOp":
+                    engine = node.attr["serialized_segment"].s
+                    if self.engine_dir is not None:
+                        util.save_file(
+                            contents=engine,
+                            dest=os.path.join(
+                                self.engine_dir, f"segment-{segment_number}"
+                            ),
+                        )
+                    segment_number += 1
+
+        return graph, outputs
+
+
+@mod.export(funcify=True)
+class CreateConfig(BaseLoader):
+    """
+    Functor that creates a TensorFlow config.
+    """
+
+    def __init__(self, gpu_memory_fraction=None, allow_growth=None, use_xla=None):
+        """
+        Creates a TensorFlow config.
+
+        Args:
+            gpu_memory_fraction (float):
+                The fraction of GPU memory that will be made available to TensorFlow.
+                This should be a value between 0.0 and 1.0.
+            allow_growth (bool): Whether to allow GPU memory allocated by TensorFlow to grow.
+            use_xla (bool): Whether to attempt to enable XLA.
+        """
+        self.gpu_memory_fraction = util.default(gpu_memory_fraction, 0.9)
+        self.allow_growth = util.default(allow_growth, False)
+        self.use_xla = util.default(use_xla, False)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            tf.ConfigProto: The TensorFlow config.
+        """
+
+        # Session configuration
+        gpu_options = tf.compat.v1.GPUOptions(
+            per_process_gpu_memory_fraction=self.gpu_memory_fraction,
+            allow_growth=self.allow_growth,
+        )
+        config = tf.compat.v1.ConfigProto(gpu_options=gpu_options)
+        if self.use_xla:
+            config.graph_options.optimizer_options.global_jit_level = (
+                tf.OptimizerOptions.ON_1
+            )
+        G_LOGGER.verbose(
+            f"Using gpu memory fraction: {self.gpu_memory_fraction}, XLA: {self.use_xla}"
+        )
+        return config
+
+
+@mod.export(funcify=True)
+class SessionFromGraph(BaseLoader):
+    """
+    Functor that creates a TensorFlow session that can be used for inference.
+    """
+
+    def __init__(self, graph, config=None):
+        """
+        Creates a TensorFlow session.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            config (Union[tf.ConfigProto, Callable() -> tf.ConfigProto]):
+                    A TensorFlow ConfigProto or a callable that returns one.
+        """
+        self.graph = graph
+        self.config = util.default(config, CreateConfig())
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            tf.Session: The TensorFlow session.
+        """
+        config, _ = util.invoke_if_callable(self.config)
+        (graph, output_names), _ = util.invoke_if_callable(self.graph)
+
+        with graph.as_default() as graph, tf.compat.v1.Session(
+            graph=graph, config=config
+        ).as_default() as sess:
+            G_LOGGER.verbose(f"Using TensorFlow outputs: {output_names}")
+            G_LOGGER.extra_verbose("Initializing variables in TensorFlow Graph")
+            sess.run(tf.compat.v1.initializers.global_variables())
+            return sess, output_names

tools/Polygraphy/polygraphy/backend/tf/requirements.txt

@@ -0,0 +1 @@
+tensorflow<2.0

tools/Polygraphy/polygraphy/backend/tf/runner.py

@@ -0,0 +1,107 @@
+#
+# 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.
+#
+# Sets up everything needed to perform inference in TensorFlow.
+import os
+import time
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.backend.tf import util as tf_util
+from polygraphy.logger import G_LOGGER
+
+tf = mod.lazy_import("tensorflow<2.0")
+
+
+@mod.export()
+class TfRunner(BaseRunner):
+    """
+    Runs inference using a TensorFlow session.
+    """
+
+    def __init__(self, sess, timeline_dir=None, name=None):
+        """
+        Args:
+            sess (Union[Tuple[tf.Session, Sequence[str]], Callable() -> Tuple[tf.Session, Sequence[str]]]):
+                    A tuple containing a TensorFlow session and output names or a callable that returns one.
+
+
+            timeline_dir (str):
+                    Path to write a TensorFlow timeline.
+                    Note that profiling may affect execution time.
+            name (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+        """
+        super().__init__(name=name, prefix="tf-runner")
+
+        self._sess = sess
+
+        self.timeline_dir = timeline_dir
+        self.num_inferences = 0
+        self.run_options = None
+        self.run_metadata = None
+        if self.timeline_dir is not None:
+            # Enable profiling
+            G_LOGGER.warning("Profiling is enabled. This will impact performance")
+            self.run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
+            self.run_metadata = tf.RunMetadata()
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        (self.sess, self.output_names), _ = util.invoke_if_callable(self._sess)
+
+    @util.check_called_by("get_input_metadata")
+    def get_input_metadata_impl(self):
+        return tf_util.get_input_metadata(self.sess.graph)
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict):
+        G_LOGGER.extra_verbose(f"Received feed_dict: {feed_dict}")
+        start = time.time()
+        inference_outputs = self.sess.run(
+            self.output_names,
+            feed_dict=feed_dict,
+            options=self.run_options,
+            run_metadata=self.run_metadata,
+        )
+        end = time.time()
+
+        out_dict = OrderedDict()
+        for name, out in zip(self.output_names, inference_outputs):
+            out_dict[name] = out
+        self.inference_time = end - start
+
+        if self.timeline_dir is not None:
+            from tensorflow.python.client import timeline
+
+            t1 = timeline.Timeline(self.run_metadata.step_stats)
+
+            util.save_file(
+                contents=t1.generate_chrome_trace_format(),
+                dest=os.path.join(self.timeline_dir, f"run-{self.num_inferences}"),
+                mode="w",
+            )
+        self.num_inferences += 1
+
+        return out_dict
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        self.sess.close()
+        del (self.sess, self.output_names)
+        self.num_inferences = 0

tools/Polygraphy/polygraphy/backend/tf/util.py

@@ -0,0 +1,204 @@
+#
+# 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 collections import defaultdict
+
+from polygraphy import mod, util
+from polygraphy.common import TensorMetadata
+from polygraphy.logger import G_LOGGER
+
+tf = mod.lazy_import("tensorflow<2.0")
+
+
+def load_graph(path):
+    """
+    Loads a TensorFlow frozen model.
+
+    Args:
+        path (Union[str, tf.Graph, tf.GraphDef]):
+                A path to the frozen model, or a frozen TensorFlow graph or graphdef.
+
+    Returns:
+        tf.Graph: The TensorFlow graph
+    """
+    if isinstance(path, tf.Graph):
+        return path
+
+    if isinstance(path, str):
+        graphdef = tf.compat.v1.GraphDef()
+
+        import google
+
+        try:
+            graphdef.ParseFromString(util.load_file(path, description="GraphDef"))
+        except google.protobuf.message.DecodeError:
+            G_LOGGER.backtrace()
+            G_LOGGER.critical(
+                f"Could not import TensorFlow GraphDef from: {path}. Is this a valid TensorFlow model?"
+            )
+    elif isinstance(path, tf.compat.v1.GraphDef):
+        graphdef = path
+
+    with tf.Graph().as_default() as graph:
+        tf.import_graph_def(graphdef, name="")
+        return graph
+
+
+def find_nodes_by_ops(graphdef, ops):
+    ops = set(ops)
+    return [node for node in graphdef.node if any([op in node.op for op in ops])]
+
+
+def map_node_outputs(graphdef):
+    def sanitize_input_name(input_name):
+        # Strip port information and control symbol
+        split_input = input_name.split(":")
+        if len(split_input) > 1:
+            split_input.pop(-1)
+        return ":".join(split_input).replace("^", "")
+
+    node_outputs = defaultdict(list)
+    for node in graphdef.node:
+        for input_name in node.input:
+            node_outputs[sanitize_input_name(input_name)].append(node)
+    return node_outputs
+
+
+def get_tensor_metadata(tensors):
+    metadata = TensorMetadata()
+    for tensor in tensors:
+        try:
+            shape = [
+                elem.value if hasattr(elem, "value") else elem for elem in tensor.shape
+            ]
+        except ValueError:
+            # Happens when rank is unknown
+            shape = None
+        metadata.add(tensor.name, dtype=tensor.dtype.as_numpy_dtype, shape=shape)
+    return metadata
+
+
+def get_input_metadata(graph):
+    input_tensors = []
+    input_nodes = find_nodes_by_ops(graph.as_graph_def(), ["Placeholder", "FIFOQueue"])
+    G_LOGGER.verbose(
+        f"Found input tensors: {[f'{n.name}: {n.op}' for n in input_nodes]}"
+    )
+    for node in input_nodes:
+        input_tensors.append(graph.get_tensor_by_name(node.name + ":0"))
+
+    G_LOGGER.verbose(f"Retrieved TensorFlow input_tensors: {input_tensors}")
+    return get_tensor_metadata(input_tensors)
+
+
+def get_output_metadata(graph, layerwise=False):
+    graphdef = graph.as_graph_def()
+
+    node_output_map = map_node_outputs(graphdef)
+
+    def is_output_node(node):
+        # Make sure that we're not using hanging nodes as outputs - must have at least one input.
+        if len(node_output_map[node.name]) != 0 or len(node.input) == 0:
+            return False
+
+        # Tensors with no shape cannot be outputs and TensorFlow doesn't like certain ops as outputs.
+        EXCLUDE_OPS = [
+            "Switch",
+            "FusedBatchNorm",
+            "Assert",
+            "NextIteration",
+            "Enter",
+            "LoopCond",
+            "Exit",
+            "Print",
+            "Assign",
+            "NoOp",
+            "ReadVariableOp",
+            "VarIsInitializedOp",
+            "Const",
+        ]
+
+        # Additionally, we sometimes need to exclude entire namespaces e.g. while loops.
+        EXCLUDE_NAMESPACES = ["while", "Assert"]
+
+        if any([ex_op in node.op for ex_op in EXCLUDE_OPS]) or any(
+            [ns in node.name for ns in EXCLUDE_NAMESPACES]
+        ):
+            G_LOGGER.extra_verbose(
+                f"Excluding {node.name}, op {node.op} is not a valid output op or is part of an excluded namespace (Note: excluded namespaces: {EXCLUDE_NAMESPACES})"
+            )
+            return False
+
+        return True
+
+    # For layerwise mode, every layer becomes an output.
+    if layerwise:
+        output_nodes = list(graphdef.node)
+        G_LOGGER.verbose(
+            f"Running in layerwise mode. Marking {len(output_nodes)} layers as potential outputs"
+        )
+    else:
+        output_nodes = [node for node in graphdef.node if is_output_node(node)]
+    G_LOGGER.extra_verbose(f"Found likely output nodes: {output_nodes}")
+
+    output_tensors = []
+    for node in output_nodes:
+        tensor_name = node.name + ":0"
+        try:
+            tensor = graph.get_tensor_by_name(tensor_name)
+            output_tensors.append(tensor)
+        except KeyError:
+            G_LOGGER.warning(f"Could not import: {tensor_name}. Skipping.")
+    if len(output_tensors) != len(output_nodes):
+        G_LOGGER.warning(
+            f"Excluded {len(output_nodes) - len(output_tensors)} ops that don't seem like outputs. Use -vv/--super-verbose, or set logging verbosity to EXTRA_VERBOSE to view them."
+        )
+
+    G_LOGGER.extra_verbose(
+        f"Found output op types in graph: {set(tensor.op.type for tensor in output_tensors)}"
+    )
+    G_LOGGER.verbose(f"Retrieved TensorFlow output_tensors: {output_tensors}")
+    return get_tensor_metadata(output_tensors)
+
+
+def get_graph_output_names(graph):
+    return list(get_output_metadata(graph).keys())
+
+
+def str_from_graph(graph, show_layers=None, show_attrs=None, show_weights=None):
+    show_layers = util.default(show_layers, False)
+    show_attrs = util.default(show_attrs, False)
+    show_weights = util.default(show_weights, False)
+
+    graph_str = ""
+    input_metadata = get_input_metadata(graph)
+    output_metadata = get_output_metadata(graph)
+
+    graph_str += f"---- {len(input_metadata)} Graph Inputs ----\n{input_metadata}\n\n"
+    graph_str += (
+        f"---- {len(output_metadata)} Graph Outputs ----\n{output_metadata}\n\n"
+    )
+    graph_str += f"---- {len(graph.as_graph_def().node)} Nodes ----\n"
+    if show_layers:
+        G_LOGGER.warning(
+            "Displaying layer information is unsupported for TensorFlow graphs. "
+            "Please use --show layers attrs weights if you would like to see the raw nodes"
+        )
+        if show_attrs or show_weights:
+            for node in graph.as_graph_def().node:
+                graph_str += str(node) + "\n"
+        graph_str += "\n"
+    return util.indent_block(graph_str, level=0)

tools/Polygraphy/polygraphy/backend/trt/__init__.py

@@ -0,0 +1,9 @@
+from polygraphy.backend.trt.algorithm_selector import *
+from polygraphy.backend.trt.calibrator import *
+from polygraphy.backend.trt.config import *
+from polygraphy.backend.trt.file_reader import *
+from polygraphy.backend.trt.loader import *
+from polygraphy.backend.trt.profile import *
+from polygraphy.mod.trt_importer import *
+from polygraphy.backend.trt.runner import *
+from polygraphy.backend.trt.util import *

tools/Polygraphy/polygraphy/backend/trt/algorithm_selector.py

@@ -0,0 +1,422 @@
+#
+# 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 func, mod, util, constants
+from polygraphy.backend.trt import util as trt_util
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.common.interface import TypedDict
+from polygraphy.json import Decoder, Encoder, add_json_methods
+from polygraphy.logger import G_LOGGER, LogMode
+
+from typing import Sequence
+
+trt = lazy_import_trt()
+
+
+##
+## Data Structures
+##
+
+#
+# NOTE: Modifying the structure of the data classes below will break backwards compatiblity
+#
+
+
+def check_is_instance(obj, cls, name):
+    if not isinstance(obj, cls):
+        G_LOGGER.critical(
+            f"'{name}' must be an instance of {cls.__name__}, but is: {obj}."
+        )
+
+
+@mod.export()
+class TensorInfo:
+    """
+    Tracks information about a tensor, such as format and data type.
+    """
+
+    @staticmethod
+    def from_trt(io_info):
+        """
+        Creates a Polygraphy ``TensorInfo`` instance from a TensorRT ``IAlgorithmIOInfo``.
+
+        Args:
+            io_info (trt.IAlgorithmIOInfo): The algorithm I/O information.
+
+        Returns:
+            TensorInfo
+        """
+        return TensorInfo(
+            io_info.dtype,
+            tuple(io_info.strides),
+            # These fields were added in 8.6
+            util.try_getattr(io_info, "vectorized_dim"),
+            util.try_getattr(io_info, "components_per_element"),
+        )
+
+    def __init__(self, dtype, strides, vectorized_dim, components_per_element):
+        """
+        Args:
+            dtype (trt.DataType): The data type.
+            strides (Sequence[int]): The strides.
+            vectorized_dim (int): The index of the vectorized dimensions.
+            components_per_element (int): The number of components per element.
+        """
+        check_is_instance(dtype, trt.DataType, "dtype")
+        check_is_instance(strides, Sequence, "strides")
+        if vectorized_dim is not None:
+            check_is_instance(vectorized_dim, int, "vectorized_dim")
+        if components_per_element is not None:
+            check_is_instance(components_per_element, int, "components_per_element")
+
+        self.dtype = dtype
+        self.strides = tuple(strides)
+        self.vectorized_dim = vectorized_dim
+        self.components_per_element = components_per_element
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        return f"TensorInfo({str(self.dtype)}, {self.strides}, {self.vectorized_dim}, {self.components_per_element})"
+
+    def __hash__(self):
+        return hash(
+            (self.dtype, self.strides, self.vectorized_dim, self.components_per_element)
+        )
+
+
+@Encoder.register(TensorInfo)
+def encode(tensor_info):
+    return {
+        "dtype": str(tensor_info.dtype),
+        "strides": tensor_info.strides,
+        "vectorized_dim": tensor_info.vectorized_dim,
+        "components_per_element": tensor_info.components_per_element,
+    }
+
+
+@Decoder.register(TensorInfo)
+def decode(dct):
+    return TensorInfo(
+        util.getattr_nested(trt, dct["dtype"]),
+        dct["strides"],
+        dct["vectorized_dim"],
+        dct["components_per_element"],
+    )
+
+
+@mod.export()
+class Algorithm:
+    """
+    Represents a TensorRT algorithm variant, which can be uniquely represented
+    by an implementation ID, tactic ID, and I/O tensor information.
+    """
+
+    @staticmethod
+    def from_trt(context, algorithm):
+        """
+        Creates a Polygraphy ``Algorithm`` instance from a TensorRT
+        ``IAlgorithmContext`` and ``IAlgorithm``.
+
+        Args:
+            context (trt.IAlgorithmContext):
+                    The algorithm context corresponding to the layer.
+            algorithm (trt.IAlgorithm):
+                    The algorithm variant provided by TensorRT.
+
+        Returns:
+            Algorithm
+        """
+
+        implementation = algorithm.algorithm_variant.implementation
+        tactic = algorithm.algorithm_variant.tactic
+        inputs = tuple(
+            TensorInfo.from_trt(algorithm.get_algorithm_io_info(i))
+            for i in range(context.num_inputs)
+        )
+        outputs = tuple(
+            TensorInfo.from_trt(algorithm.get_algorithm_io_info(i))
+            for i in range(context.num_inputs, context.num_inputs + context.num_outputs)
+        )
+        return Algorithm(implementation, tactic, inputs, outputs)
+
+    def __init__(self, implementation, tactic, inputs, outputs):
+        """
+        Args:
+            implementation (int):
+                    The implementation for this Algorithm.
+            tactic (int):
+                    The tactic for this Algorithm.
+            inputs (Sequence[TensorInfo]):
+                    A sequence of TensorInfos for each input.
+            outputs (Sequence[TensorInfo]):
+                    A sequence of TensorInfos for each output.
+        """
+        self.implementation = implementation
+        self.tactic = tactic
+
+        def check_io(lst, name):
+            for index, io in enumerate(lst):
+                check_is_instance(io, TensorInfo, f"{name}[{index}]")
+
+        check_io(inputs, "inputs")
+        check_io(outputs, "outputs")
+
+        # Use tuples here so the class is hashable.
+        self.inputs = tuple(inputs)
+        self.outputs = tuple(outputs)
+
+    def __str__(self):
+        return f"(Implementation: {self.implementation}, Tactic: {self.tactic}) | Inputs: {self.inputs} | Outputs: {self.outputs}"
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __hash__(self):
+        return hash((self.implementation, self.tactic, self.inputs, self.outputs))
+
+
+@Encoder.register(Algorithm)
+def encode(algo):
+    return {
+        "implementation": algo.implementation,
+        "tactic": algo.tactic,
+        "inputs": algo.inputs,
+        "outputs": algo.outputs,
+    }
+
+
+@Decoder.register(Algorithm)
+def decode(dct):
+    return Algorithm(
+        implementation=dct["implementation"],
+        tactic=dct["tactic"],
+        inputs=dct["inputs"],
+        outputs=dct["outputs"],
+    )
+
+
+@mod.export()
+@add_json_methods("tactic replay file")
+class TacticReplayData(TypedDict(lambda: str, lambda: Algorithm)):
+    """
+    Maps layer names to corresponding tactics.
+    More specifically, it is an ``OrderedDict[str, Algorithm]``.
+    """
+
+    def add(self, name, algorithm):
+        """
+        Add an entry into the tactic replay data.
+
+        Args:
+            name (str): The name of the layer
+            algorithm (Algorithm): The algorithm to use for the layer.
+
+        Returns:
+            TacticReplayData: self, to allow for method chaining.
+        """
+        self[name] = algorithm
+        return self
+
+    def __str__(self):
+        return "\n".join(
+            [
+                f"Layer: {name}\n{constants.TAB}Algorithm: {algorithm}"
+                for (name, algorithm) in self.items()
+            ]
+        )
+
+
+@Encoder.register(TacticReplayData)
+def encode(replay):
+    return {"replay": replay.dct}
+
+
+@Decoder.register(TacticReplayData)
+def decode(dct):
+    return TacticReplayData(dct["replay"])
+
+
+##
+## Algorithm Selectors
+##
+
+
+# Everything is encapsulated in functions so that we don't create a dependency on TensorRT
+# when objects from this file are imported.
+def get_base_selector_type():
+    class BaseSelector(trt.IAlgorithmSelector):
+        def __init__(self, data):
+            # Must explicitly initialize parent for any trampoline class! Will mysteriously segfault without this.
+            trt.IAlgorithmSelector.__init__(self)
+
+            self.path = None
+            self.data = TacticReplayData()
+            if isinstance(data, TacticReplayData):
+                self.data = data
+            else:
+                self.path = data
+
+        def select_algorithms(self, context, choices):
+            return list(range(len(choices)))
+
+    return BaseSelector
+
+
+@mod.deprecate(remove_in="0.50.0", use_instead=None)
+@mod.export()
+def TacticRecorder(record):
+    """
+    A TensorRT algorithm selector that can record tactics selected by TensorRT.
+
+    The generated tactic replay file is specific to network and builder configuration.
+    Changing either of these may render the tactic replay file unusable.
+
+    Args:
+        record (Union[path, file-like, TacticReplayData]):
+                A path or file-like object or an empty ``TacticReplayData`` instance.
+                Tactics will be recorded and stored here.
+    """
+
+    class TacticRecorderClass(get_base_selector_type()):
+        def __init__(self):
+            super().__init__(record)
+            # The function that constructed this instance
+            self.make_func = TacticRecorder
+
+        @G_LOGGER.log_exception
+        def report_algorithms(self, contexts, choices):
+            """
+            Records algorithms selected by TensorRT into the provided path or
+            ``TacticReplayData`` instance.
+
+            Args:
+                contexts (List[trt.IAlgorithmContext]):
+                        The list of TensorRT algorithm contexts. Generally, there is one per layer.
+                choices (List[trt.IAlgorithm]):
+                        A list of selected algorithms for each context.
+
+            Returns:
+                None
+            """
+            for context, choice in zip(contexts, choices):
+                self.data.add(context.name, Algorithm.from_trt(context, choice))
+
+            if self.path is not None:
+                self.data.save(self.path)
+
+    return TacticRecorderClass()
+
+@mod.deprecate(remove_in="0.50.0", use_instead=None)
+@mod.export()
+def TacticReplayer(replay):
+    """
+    A TensorRT algorithm selector that can replay tactics according to a tactic replay file.
+
+    Args:
+        replay (Union[path, file-like, TacticReplayData]):
+                A path or file-like object containing a JSON-ified ``TacticReplayData`` instance,
+                or a ``TacticReplayData`` instance.
+    """
+
+    class TacticReplayerClass(get_base_selector_type()):
+        def __init__(self):
+            super().__init__(replay)
+
+            if self.path is not None:
+                self.data = TacticReplayData.load(self.path)
+
+            # The function that constructed this instance
+            self.make_func = TacticReplayer
+
+        @G_LOGGER.log_exception
+        @func.constantmethod
+        def select_algorithms(self, context, choices):
+            """
+            Selects an algorithm based on ``self.data`` if possible. Otherwise, returns
+            default tactics.
+
+            Args:
+                context (trt.IAlgorithmContext):
+                        The TensorRT algorithm context.
+                choices (List[trt.IAlgorithm]):
+                        A list of TensorRT algorithm choices.
+
+            Returns:
+                List[int]:
+                        The indices of selected tactics. If ``self.data`` includes the layer and
+                        TensorRT provides a matching tactic, this will always be of length 1.
+
+            Raises:
+                PolygraphyException:
+                        If a tactic is set for a layer in ``self.data`` but is not provided by
+                        TensorRT as a choice for that layer.
+            """
+            default_choices = super().select_algorithms(context, choices)
+
+            if not self.data:  # No replay data, we are in recording mode.
+                return default_choices
+
+            if context.name not in self.data:
+                G_LOGGER.warning(
+                    f"Layer: {context.name} was not found in the tactic replay. Falling back to default tactics."
+                )
+                sep = f"\n{constants.TAB}"
+                G_LOGGER.warning(
+                    "Has the network changed since the tactic replay file was generated?\n"
+                    f"Note: Layers in the tactic replay are:{sep}{sep.join(self.data.keys())}",
+                    mode=LogMode.ONCE,
+                )
+                return default_choices
+
+            # Need to find the index of the tactic we want.
+            to_select = self.data[context.name]
+            tactic_choices = [Algorithm.from_trt(context, algo) for algo in choices]
+
+            if to_select not in tactic_choices:
+                sep = f"\n{constants.TAB}"
+                G_LOGGER.critical(
+                    f"Layer: {context.name} | Tactic in replay was not provided by TensorRT as a choice for this layer.\n"
+                    f"Has the network or builder configuration changed since the replay file was generated?\n"
+                    f"Note: Tactic in replay was:{sep}{to_select}\nProvided choices were:{sep}{sep.join(map(str, tactic_choices))}"
+                )
+
+            return [tactic_choices.index(to_select)]
+
+        @G_LOGGER.log_exception
+        @func.constantmethod
+        def report_algorithms(self, contexts, choices):
+            """
+            Checks if the tactics specified in ``self.data`` were selected and raises an exception
+            if not.
+
+            Raises:
+                PolygraphyException:
+                        If a tactic specified in ``self.data`` was not selected for a layer.
+            """
+            for context, choice in zip(contexts, choices):
+                if context.name in self.data:
+                    to_select = self.data[context.name]
+                    selected = Algorithm.from_trt(context, choice)
+                    if to_select != selected:
+                        G_LOGGER.critical(
+                            f"Layer: {context.name} | TensorRT selected a tactic different than the one specified in the tactic replay."
+                            f"\nNote: Tactic in replay was:\n{constants.TAB}{to_select}, but TensorRT selected:\n{constants.TAB}{selected}"
+                        )
+
+    return TacticReplayerClass()

tools/Polygraphy/polygraphy/backend/trt/calibrator.py

@@ -0,0 +1,295 @@
+#
+# 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.
+#
+import contextlib
+import copy
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import util as base_util
+from polygraphy.backend.trt import util as trt_util
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER, LogMode
+
+trt = lazy_import_trt()
+np = mod.lazy_import("numpy")
+
+
+@mod.export()
+def Calibrator(
+    data_loader,
+    cache=None,
+    BaseClass=None,
+    batch_size=None,
+    quantile=None,
+    regression_cutoff=None,
+    algo=None,
+):
+    """
+    Supplies calibration data to TensorRT to calibrate the network for INT8 inference.
+
+    Args:
+        data_loader (Sequence[OrderedDict[str, Union[numpy.ndarray, DeviceView, torch.Tensor, int]]]):
+            A generator or iterable that yields a dictionary that maps input names to NumPy
+            arrays, Polygraphy DeviceViews, PyTorch tensors, or GPU pointers. If NumPy arrays,
+            DeviceViews, or PyTorch tensors are provided, the calibrator will check the data types
+            and shapes if possible to ensure that they match those expected by the model.
+
+            In case you don't know details about the inputs ahead of time, you can access the
+            `input_metadata` property in your data loader, which will be set to a ``TensorMetadata``
+            instance by Polygraphy APIs like ``CreateConfig`` and ``EngineFromNetwork``.
+            Note that this does not work for generators or lists.
+
+            The number of calibration batches is controlled by the number of items supplied
+            by the data loader.
+
+
+        cache (Union[str, file-like]):
+                Path or file-like object to save/load the calibration cache.
+                By default, the calibration cache is not saved.
+        BaseClass (type):
+                The type of calibrator to inherit from.
+                Defaults to ``trt.IInt8EntropyCalibrator2``.
+        batch_size (int):
+                [DEPRECATED] The size of each batch provided by the data loader.
+        quantile (float):
+                The quantile to use for ``trt.IInt8LegacyCalibrator``.
+                Has no effect for other calibrator types.
+                Defaults to 0.5.
+        regression_cutoff (float):
+                The regression cutoff to use for ``trt.IInt8LegacyCalibrator``.
+                Has no effect for other calibrator types.
+                Defaults to 0.5.
+        algo (trt.CalibrationAlgoType):
+                Calibration algorithm to use for ``trt.IInt8Calibrator``.
+                Has no effect for other calibrator types.
+                Defaults to ``trt.CalibrationAlgoType.MINMAX_CALIBRATION``.
+    """
+    BaseClass = util.default(BaseClass, trt.IInt8EntropyCalibrator2)
+
+    class CalibratorClass(BaseClass):
+        """
+        Calibrator that supplies calibration data to TensorRT to calibrate the network for INT8 inference.
+        """
+
+        def __init__(self):
+            # Must explicitly initialize parent for any trampoline class! Will mysteriously segfault without this.
+            BaseClass.__init__(self)  # type: ignore
+
+            self.data_loader = data_loader
+            self._cache = cache
+            self.device_buffers = OrderedDict()
+            self.input_metadata = None
+            self.reset()
+            G_LOGGER.verbose(f"Created calibrator [cache={self._cache}]")
+
+            self.batch_size = util.default(batch_size, 1)
+
+            self.is_polygraphy_calibrator = True
+            # The function that constructed this instance
+            self.make_func = Calibrator
+
+        def set_input_metadata(self, input_metadata):
+            """
+            Sets the input metadata for the calibrator.
+
+            This is passed along to the data loader and is also used for
+            input data type and shape checks.
+
+            NOTE: This generally does not need to be called manually if the calibrator is being used
+            with Polygraphy's loaders, like ``CreateConfig`` or ``EngineFromNetwork``.
+
+            Args:
+                input_metadata (TensorMetadata):
+                        Mapping of input names to their data types and shapes.
+                        Passed along to the data loader if provided. This is required if
+                        using Polygraphy's included `DataLoader` to provide calibration data,
+                        or if data type and shape checking is desired.
+            """
+            calibration_metadata = copy.copy(input_metadata)
+            for name, meta_tuple in calibration_metadata.items():
+                if meta_tuple.dtype not in {
+                    DataType.FLOAT32,
+                    DataType.INT32,
+                    DataType.INT64,
+                    DataType.BOOL,
+                }:
+                    G_LOGGER.warning(
+                        f"TensorRT requires non-index calibration inputs to be provided in float32. "
+                        f"Input: {name} has datatype: {meta_tuple.dtype}, so will override to float32 in the calibrator's metadata. "
+                        f"If you are using a custom data loader with the calibrator, please ensure that you return a float32 tensor for this input."
+                    )
+                    meta_tuple.dtype = DataType.FLOAT32
+
+            self.input_metadata = calibration_metadata
+            if calibration_metadata is not None:
+                with contextlib.suppress(AttributeError):
+                    self.data_loader.input_metadata = calibration_metadata
+
+        def reset(self):
+            """
+            Reset this calibrator for reuse.
+
+            The calibrator will clear any dynamic ranges cached from previous calibration runs, and will
+            attempt to rewind the data loader (note that generators cannot be rewound).
+
+            Typically, this is only required if the same calibrator is used for multiple different networks.
+            """
+            # Attempt to reset data loader
+            self.data_loader_iter = iter(self.data_loader)
+            self.num_batches = 0
+
+            # Make sure calibrator will check the cache again when reset.
+            self.cache_contents = None
+
+        def get_batch_size(self):
+            return self.batch_size
+
+        def _get_batch_impl(self, names):
+            try:
+                buffers = next(self.data_loader_iter)
+            except StopIteration:
+                if not self.num_batches:
+                    G_LOGGER.critical(
+                        "Calibrator data loader provided no data.\nPossible reasons for this include:\n(1) data loader "
+                        "has no data to provide\n(2) data loader was a generator, and the calibrator is being "
+                        "used multiple times (generators cannot be rewound)"
+                    )
+                return None
+
+            self.num_batches += 1
+
+            if self.input_metadata is not None:
+                base_util.check_inputs(buffers, self.input_metadata)
+
+            ptrs = []
+            for name in names:
+                buf = buffers[name]
+
+                if isinstance(buf, int):
+                    ptrs.append(buf)
+                else:
+                    ptrs.append(
+                        trt_util._get_array_on_gpu(buf, name, self.device_buffers)
+                    )
+
+            return ptrs
+
+        def get_batch(self, names):
+            ptrs = None
+            try:
+                ptrs = self._get_batch_impl(names)
+            except PolygraphyException:
+                pass
+            if ptrs is None:
+                self.free()
+            return ptrs
+
+        def read_calibration_cache(self):
+            def load_from_cache():
+                if self._cache is None or not util.get_file_size(self._cache):
+                    return None
+
+                try:
+                    return util.load_file(self._cache, description="calibration cache")
+                except Exception as err:
+                    G_LOGGER.error(
+                        f"Could not read from calibration cache: {self._cache}\nNote: Error was: {err}"
+                    )
+                    return None
+
+            if self.cache_contents is not None:
+                return self.cache_contents
+
+            self.cache_contents = load_from_cache()
+
+            if not self.cache_contents:
+                if self.cache_contents is not None:
+                    G_LOGGER.warning(
+                        "Calibration cache was provided, but is empty. "
+                        "Will regenerate scales by running calibration.",
+                        mode=LogMode.ONCE,
+                    )
+                self.cache_contents = None
+
+            return self.cache_contents
+
+        def write_calibration_cache(self, cache):
+            self.cache_contents = cache.tobytes()
+
+            if self._cache is None:
+                return
+
+            try:
+                util.save_file(
+                    contents=self.cache_contents,
+                    dest=self._cache,
+                    description="calibration cache",
+                )
+            except Exception as err:
+                G_LOGGER.error(
+                    f"Could not write to calibration cache: {self._cache}.\nNote: Error was: {err}"
+                )
+
+        def free(self):
+            """
+            Frees all device buffers associated with this calibrator
+            """
+            for device_buffer in self.device_buffers.values():
+                device_buffer.free()
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.free()
+
+        # IInt8LegacyCalibrator methods
+        if BaseClass == trt.IInt8LegacyCalibrator:
+
+            def get_quantile(self):
+                return util.default(quantile, 0.5)
+
+            def get_regression_cutoff(self):
+                return util.default(regression_cutoff, 0.5)
+
+            def read_histogram_cache(self, length):
+                pass
+
+            def write_histogram_cache(self, ptr, length):
+                pass
+
+        # IInt8Calibrator methods
+        if BaseClass == trt.IInt8Calibrator:
+
+            def get_algorithm(self):
+                return util.default(algo, trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2)
+
+        def __repr__(self):
+            return util.make_repr(
+                "Calibrator",
+                data_loader,
+                cache=cache,
+                BaseClass=BaseClass,
+                batch_size=batch_size,
+                quantile=quantile,
+                regression_cutoff=regression_cutoff,
+                algo=algo,
+            )[0]
+
+    return CalibratorClass()

tools/Polygraphy/polygraphy/backend/trt/config.py

@@ -0,0 +1,657 @@
+#
+# 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.
+#
+import contextlib
+import copy
+import re
+
+from polygraphy import config as polygraphy_config, mod, util
+from polygraphy.backend.base import BaseLoader
+from polygraphy.backend.trt import util as trt_util
+from polygraphy.backend.trt.profile import Profile
+from polygraphy.backend.trt.util import inherit_and_extend_docstring
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.logger import G_LOGGER
+
+trt = lazy_import_trt()
+
+
+class _CreateConfigCommon(BaseLoader):
+    """
+    Generic TensorRT IBuilderConfig.
+    """
+
+    def __init__(
+        self,
+        profiles=None,
+        precision_constraints=None,
+        load_timing_cache=None,
+        algorithm_selector=None,
+        sparse_weights=None,
+        tactic_sources=None,
+        restricted=None,
+        profiling_verbosity=None,
+        memory_pool_limits=None,
+        refittable=None,
+        strip_plan=None,
+        preview_features=None,
+        engine_capability=None,
+        direct_io=None,
+        builder_optimization_level=None,
+        hardware_compatibility_level=None,
+        max_aux_streams=None,
+        version_compatible=None,
+        exclude_lean_runtime=None,
+        quantization_flags=None,
+        error_on_timing_cache_miss=None,
+        disable_compilation_cache=None,
+        progress_monitor=None,
+        weight_streaming=None,
+        runtime_platform=None,
+        tiling_optimization_level=None,
+    ):
+        """
+        Creates an IBuilderConfig that can be used by EngineFromNetwork.
+
+        Args:
+            profiles (List[Profile]):
+                    A list of optimization profiles to add to the configuration. Only needed for
+                    networks with dynamic input shapes. If this is omitted for a network with
+                    dynamic shapes, a default profile is created, where dynamic dimensions are
+                    replaced with Polygraphy's DEFAULT_SHAPE_VALUE (defined in constants.py).
+                    A partially populated profile will be automatically filled using values from ``Profile.fill_defaults()``
+                    See ``Profile`` for details.
+            precision_constraints (Optional[str]):
+                    If set to "obey", require that layers execute in specified precisions.
+                    If set to "prefer", prefer that layers execute in specified precisions but allow TRT to fall back to
+                    other precisions if no implementation exists for the requested precision.
+                    Otherwise, precision constraints are ignored.
+                    Defaults to None.
+            load_timing_cache (Union[str, file-like]):
+                    A path or file-like object from which to load a tactic timing cache.
+                    Providing a tactic timing cache can speed up the engine building process.
+                    Caches can be generated while building an engine with, for example, EngineFromNetwork.
+                    If a path is provided, the file will be locked for exclusive access so that other processes
+                    cannot update the cache while it is being read.
+                    If the file specified by the path does not exist, CreateConfig will emit a warning and fall back
+                    to using an empty timing cache.
+            algorithm_selector (trt.IAlgorithmSelector):
+                    An algorithm selector. Allows the user to control how tactics are selected
+                    instead of letting TensorRT select them automatically.
+            sparse_weights (bool):
+                    Whether to enable optimizations for sparse weights.
+                    Defaults to False.
+            tactic_sources (List[trt.TacticSource]):
+                    The tactic sources to enable. This controls which libraries (e.g. cudnn, cublas, etc.)
+                    TensorRT is allowed to load tactics from.
+                    Use an empty list to disable all tactic sources.
+                    Defaults to TensorRT's default tactic sources.
+            restricted (bool):
+                    Whether to enable safety scope checking in the builder. This will check if the network
+                    and builder configuration are compatible with safety scope.
+                    Defaults to False.
+            profiling_verbosity (trt.ProfilingVerbosity):
+                    The verbosity of NVTX annotations in the generated engine.
+                    Higher verbosity allows you to determine more information about the engine.
+                    Defaults to ``trt.ProfilingVerbosity.VERBOSE``.
+            memory_pool_limits (Dict[trt.MemoryPoolType, int]):
+                    Limits for different memory pools.
+                    This should be a mapping of pool types to their respective limits in bytes.
+            refittable (bool):
+                    Enables the engine to be refitted with new weights after it is built.
+                    Defaults to False.
+            strip_plan (bool):
+                    Strips the refittable weights from the engine plan file.
+                    Defaults to False.
+            preview_features (List[trt.PreviewFeature]):
+                    The preview features to enable.
+                    Use an empty list to disable all preview features.
+                    Defaults to TensorRT's default preview features.
+            engine_capability (trt.EngineCapability):
+                    The engine capability to build for.
+                    Defaults to the default TensorRT engine capability.
+            direct_io (bool):
+                    Whether to disallow reformatting layers at network input/output tensors with
+                    user-specified formats.
+                    Defaults to False.
+            builder_optimization_level (int):
+                    The builder optimization level. A higher optimization level allows the optimizer to spend more time
+                    searching for optimization opportunities. The resulting engine may have better performance compared
+                    to an engine built with a lower optimization level.
+                    Refer to the TensorRT API documentation for details.
+                    Defaults to TensorRT's default optimization level.
+            hardware_compatibility_level (trt.HardwareCompatibilityLevel):
+                    The hardware compatibility level. This allows engines built on one GPU architecture to work on GPUs
+                    of other architectures.
+                    Defaults to TensorRT's default hardware compatibility level.
+            max_aux_streams (int):
+                    The maximum number of auxiliary streams that TensorRT is allowed to use. If the network contains
+                    operators that can run in parallel, TRT can execute them using auxiliary streams in addition to the
+                    one provided to the IExecutionContext::enqueueV3() call.
+                    The default maximum number of auxiliary streams is determined by the heuristics in TensorRT on
+                    whether enabling multi-stream would improve the performance.
+            version_compatible (bool):
+                    Whether to build an engine that is version compatible.
+            exclude_lean_runtime (bool):
+                    Whether to exclude the lean runtime in version compatible engines.
+                    Requires that version compatibility is enabled.
+            quantization_flags (List[trt.QuantizationFlag]):
+                    The quantization flags to enable.
+                    Use an empty list to disable all quantization flags.
+                    Defaults to TensorRT's default quantization flags.
+            error_on_timing_cache_miss (bool):
+                    Emit error when a tactic being timed is not present in the timing cache.
+                    This flag has an effect only when IBuilderConfig has an associated ITimingCache.
+                    Defaults to False.
+            disable_compilation_cache (bool):
+                    Whether to disable caching JIT-compiled code.
+                    Defaults to False.
+            progress_monitor (trt.IProgressMonitor):
+                    A progress monitor. Allow users to view engine building progress through CLI.
+            weight_streaming (bool):
+                    TWhether to enable weight streaming for the TensorRT Engine.
+            runtime_platform (trt.RuntimePlatform):
+                    Describes the intended runtime platform (operating system and CPU architecture) for the execution of the TensorRT engine. 
+                    TensorRT provides support for cross-platform engine compatibility when the target runtime platform is different from the build platform.
+                    Defaults to TensorRT's default runtime platform.
+            tiling_optimization_level (trt.TilingOptimizationLevel):
+                    The tiling optimization level. Setting a higher optimization level allows TensorRT to spend more building time for more tiling strategies.
+                    Defaults to TensorRT's default tiling optimization level. Refer to the TensorRT API documentation for details.
+        """
+        self.profiles = util.default(profiles, [Profile()])
+        self.precision_constraints = precision_constraints
+        self.restricted = util.default(restricted, False)
+        self.refittable = util.default(refittable, False)
+        self.strip_plan = util.default(strip_plan, False)
+        self.timing_cache_path = load_timing_cache
+        self.algorithm_selector = algorithm_selector
+        self.sparse_weights = util.default(sparse_weights, False)
+        self.tactic_sources = tactic_sources
+        self.profiling_verbosity = profiling_verbosity
+        self.memory_pool_limits = memory_pool_limits
+        self.preview_features = preview_features
+        self.engine_capability = engine_capability
+        self.direct_io = util.default(direct_io, False)
+        self.builder_optimization_level = builder_optimization_level
+        self.hardware_compatibility_level = hardware_compatibility_level
+        self.max_aux_streams = max_aux_streams
+        self.version_compatible = version_compatible
+        self.exclude_lean_runtime = exclude_lean_runtime
+        self.quantization_flags = quantization_flags
+        self.error_on_timing_cache_miss = util.default(
+            error_on_timing_cache_miss, False
+        )
+        self.disable_compilation_cache = util.default(disable_compilation_cache, False)
+        self.progress_monitor = progress_monitor
+        self.weight_streaming = weight_streaming
+        self.runtime_platform = runtime_platform
+        self.tiling_optimization_level = tiling_optimization_level
+
+    @util.check_called_by("__call__")
+    def call_impl(self, builder, network):
+        """
+        Args:
+            builder (trt.Builder):
+                    The TensorRT builder to use to create the configuration.
+            network (trt.INetworkDefinition):
+                    The TensorRT network for which to create the config. The network is used to
+                    automatically create a default optimization profile if none are provided.
+
+        Returns:
+            trt.IBuilderConfig: The TensorRT builder configuration.
+        """
+        config = builder.create_builder_config()
+
+        def try_run(func, name):
+            try:
+                return func()
+            except AttributeError:
+                trt_util.fail_unavailable(f"{name} in {self.__class__.__name__}")
+
+        def try_set_flag(flag_name):
+            return try_run(
+                lambda: config.set_flag(getattr(trt.BuilderFlag, flag_name)),
+                flag_name.lower(),
+            )
+
+        if self.preview_features is not None:
+            for preview_feature in trt.PreviewFeature.__members__.values():
+                try_run(
+                    lambda: config.set_preview_feature(
+                        preview_feature, preview_feature in self.preview_features
+                    ),
+                    "preview_features",
+                )
+
+        G_LOGGER.verbose("Setting TensorRT Optimization Profiles")
+        profiles = copy.deepcopy(self.profiles)
+        for profile in profiles:
+            # Last profile is used for set_calibration_profile.
+            calib_profile = profile.fill_defaults(network)
+            config.add_optimization_profile(calib_profile.to_trt(builder, network))
+        newline = "\n"
+        sep = ",\n"
+        G_LOGGER.info(
+            f"Configuring with profiles:[\n"
+            f"{util.indent_block(sep.join([f'Profile {index}:{newline}{util.indent_block(profile)}' for index, profile in enumerate(profiles)]))}\n]"
+        )
+
+        layer_with_precisions = {
+            layer.name: layer.precision.name
+            for layer in network
+            if layer.precision_is_set and not layer.type == trt.LayerType.SHAPE
+        }
+        if self.precision_constraints == "obey":
+            try_set_flag("OBEY_PRECISION_CONSTRAINTS")
+        elif self.precision_constraints == "prefer":
+            try_set_flag("PREFER_PRECISION_CONSTRAINTS")
+        elif layer_with_precisions:
+            G_LOGGER.warning(
+                "It looks like some layers in the network have compute precision set, but precision constraints were not enabled. "
+                "\nPrecision constraints must be set to 'prefer' or 'obey' for layer compute precision to take effect. "
+                f"\nNote: Layers and their requested precisions were: {layer_with_precisions}"
+            )
+        if self.restricted:
+            try_set_flag("SAFETY_SCOPE")
+
+        if self.refittable:
+            try_set_flag("REFIT")
+
+        if self.strip_plan:
+            try_set_flag("STRIP_PLAN")
+
+        if self.direct_io:
+            try_set_flag("DIRECT_IO")
+
+        if self.sparse_weights:
+            try_set_flag("SPARSE_WEIGHTS")
+
+        if self.profiling_verbosity is not None:
+
+            def set_profiling_verbosity():
+                config.profiling_verbosity = self.profiling_verbosity
+
+            try_run(set_profiling_verbosity, name="profiling_verbosity")
+        else:
+            try:
+                config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+            except AttributeError:
+                pass
+
+        if self.memory_pool_limits is not None:
+            for pool_type, pool_size in self.memory_pool_limits.items():
+                try_run(
+                    lambda: config.set_memory_pool_limit(pool_type, pool_size),
+                    name="memory_pool_limits",
+                )
+
+        if self.tactic_sources is not None:
+            tactic_sources_flag = 0
+            for source in self.tactic_sources:
+                tactic_sources_flag |= 1 << int(source)
+            try_run(
+                lambda: config.set_tactic_sources(tactic_sources_flag),
+                name="tactic_sources",
+            )
+
+        try:
+            cache = None
+            if self.timing_cache_path:
+                try:
+                    with util.LockFile(self.timing_cache_path):
+                        timing_cache_data = util.load_file(
+                            self.timing_cache_path, description="tactic timing cache"
+                        )
+                        cache = config.create_timing_cache(timing_cache_data)
+                except FileNotFoundError:
+                    G_LOGGER.warning(
+                        "Timing cache file {} not found, falling back to empty timing cache.".format(
+                            self.timing_cache_path
+                        )
+                    )
+            if cache is None:
+                # Create an empty timing cache by default so it will be populated during engine build.
+                # This way, consumers of CreateConfig have the option to use the cache later.
+                cache = config.create_timing_cache(b"")
+        except AttributeError:
+            if self.timing_cache_path:
+                trt_util.fail_unavailable(f"load_timing_cache in {self.__class__.__name__}")
+        else:
+            config.set_timing_cache(cache, ignore_mismatch=False)
+
+        if self.algorithm_selector is not None:
+
+            def set_algo_selector():
+                config.algorithm_selector = self.algorithm_selector
+
+            try_run(set_algo_selector, name="algorithm_selector")
+
+            if not self.timing_cache_path:
+                G_LOGGER.warning(
+                    "Disabling tactic timing cache because algorithm selector is enabled."
+                )
+                try_set_flag("DISABLE_TIMING_CACHE")
+
+        if self.engine_capability is not None:
+
+            def set_engine_cap():
+                config.engine_capability = self.engine_capability
+
+            try_run(set_engine_cap, "engine_capability")
+
+        if self.builder_optimization_level is not None:
+
+            def set_builder_optimization_level():
+                config.builder_optimization_level = self.builder_optimization_level
+
+            try_run(set_builder_optimization_level, "builder_optimization_level")
+
+        if self.hardware_compatibility_level is not None:
+
+            def set_hardware_compatibility_level():
+                config.hardware_compatibility_level = self.hardware_compatibility_level
+
+            try_run(set_hardware_compatibility_level, "hardware_compatibility_level")
+
+        if self.version_compatible:
+            try_set_flag("VERSION_COMPATIBLE")
+
+        if self.exclude_lean_runtime:
+            if not self.version_compatible:
+                G_LOGGER.critical(
+                    f"Cannot set EXCLUDE_LEAN_RUNTIME if version compatibility is not enabled. "
+                )
+            try_set_flag("EXCLUDE_LEAN_RUNTIME")
+
+        if self.hardware_compatibility_level is not None or self.version_compatible:
+            G_LOGGER.info(
+                "Version or hardware compatibility was enabled. "
+                "If you are using an ONNX model, please set the NATIVE_INSTANCENORM ONNX parser flag, e.g. `--onnx-flags NATIVE_INSTANCENORM`"
+            )
+
+        if self.max_aux_streams is not None:
+
+            def set_max_aux_streams():
+                config.max_aux_streams = self.max_aux_streams
+
+            try_run(set_max_aux_streams, "max_aux_streams")
+
+        if self.quantization_flags is not None:
+            for quantization_flag in trt.QuantizationFlag.__members__.values():
+                if quantization_flag in self.quantization_flags:
+                    try_run(
+                        lambda: config.set_quantization_flag(quantization_flag),
+                        "quantization_flag",
+                    )
+                else:
+                    try_run(
+                        lambda: config.clear_quantization_flag(quantization_flag),
+                        "quantization_flag",
+                    )
+
+        if self.error_on_timing_cache_miss:
+            try_set_flag("ERROR_ON_TIMING_CACHE_MISS")
+
+        if self.disable_compilation_cache:
+            try_set_flag("DISABLE_COMPILATION_CACHE")
+
+        if self.progress_monitor is not None:
+
+            def set_progress_monitor():
+                config.progress_monitor = self.progress_monitor
+
+            try_run(set_progress_monitor, name="progress_monitor")
+
+        if self.weight_streaming:
+            try_set_flag("WEIGHT_STREAMING")
+        
+        if self.runtime_platform is not None:
+
+            def set_runtime_platform():
+                config.runtime_platform = self.runtime_platform
+
+            try_run(set_runtime_platform, "runtime_platform")
+
+        if self.tiling_optimization_level is not None:
+
+            def set_tiling_optimization_level():
+                config.tiling_optimization_level = self.tiling_optimization_level
+
+            try_run(set_tiling_optimization_level, "tiling_optimization_level")
+
+        return config
+
+
+@mod.export(funcify=True)
+class CreateConfig(_CreateConfigCommon):
+    """
+    Functor that creates an IBuilderConfig with TensorRT features.
+    """
+
+    @inherit_and_extend_docstring(_CreateConfigCommon.__init__)
+    def __init__(
+        self,
+        tf32=None,
+        fp16=None,
+        int8=None,
+        fp8=None,
+        bf16=None,
+        calibrator=None,
+        use_dla=None,
+        allow_gpu_fallback=None,
+        **kwargs
+    ):
+        """
+        Creates an IBuilderConfig with TensorRT-specific features.
+
+        Args:
+            tf32 (bool):
+                    Whether to enable TF32 precision. Defaults to False.
+            fp16 (bool):
+                    Whether to enable FP16 precision. Defaults to False.
+            int8 (bool):
+                    Whether to enable INT8 precision. Defaults to False.
+            fp8 (bool):
+                    Whether to enable FP8 precision. Defaults to False.
+            bf16 (bool):
+                    Whether to enable BF16 precision. Defaults to False.
+            calibrator (trt.IInt8Calibrator):
+                    An int8 calibrator. Only required in int8 mode when
+                    the network does not have explicit precision. For networks with
+                    dynamic shapes, the last profile provided (or default profile if
+                    no profiles are provided) is used during calibration.
+            use_dla (bool):
+                    [EXPERIMENTAL] Whether to enable DLA as the default device type.
+                    Defaults to False.
+            allow_gpu_fallback (bool):
+                    [EXPERIMENTAL] When DLA is enabled, whether to allow layers to fall back to GPU if they cannot be run on DLA.
+                    Has no effect if DLA is not enabled.
+                    Defaults to False.
+            **kwargs: All other arguments from _CreateConfigCommon.
+        """
+        super().__init__(**kwargs)
+        self.tf32 = util.default(tf32, False)
+        self.fp16 = util.default(fp16, False)
+        self.bf16 = util.default(bf16, False)
+        self.int8 = util.default(int8, False)
+        self.fp8 = util.default(fp8, False)
+        self.calibrator = calibrator
+        self.use_dla = util.default(use_dla, False)
+        self.allow_gpu_fallback = util.default(allow_gpu_fallback, False)
+
+        if self.calibrator is not None and not self.int8:
+            G_LOGGER.warning(
+                "A calibrator was provided to `CreateConfig`, but int8 mode was not enabled. "
+                "Did you mean to set `int8=True` to enable building with int8 precision?"
+            )
+
+        # Print a message to tell users that TF32 can be enabled to improve perf with minor accuracy differences.
+        if not self.tf32:
+            G_LOGGER.info(
+                "TF32 is disabled by default. Turn on TF32 for better performance with minor accuracy differences."
+            )
+
+        self._validator()
+
+    def _validator(self):
+        """
+        Validates initialization parameters for TensorRT-specific features.
+        """
+        # Validate that TensorRT-RTX specific flags are not used in regular TensorRT mode
+        if polygraphy_config.USE_TENSORRT_RTX:
+            if self.fp16 or self.int8 or self.bf16 or self.fp8:
+                G_LOGGER.critical("Precision flags (fp16, int8, bf16, fp8) are not supported with USE_TENSORRT_RTX=1.")
+            if self.use_dla:
+                G_LOGGER.critical("DLA is not supported with USE_TENSORRT_RTX=1.")
+            if self.calibrator is not None:
+                G_LOGGER.critical("Custom calibrator is not supported with USE_TENSORRT_RTX=1.")
+
+    def _configure_flags(self, builder, network, config):
+        """
+        Validates and configures TensorRT-specific features.
+
+        Args:
+            builder (trt.Builder): The TensorRT builder
+            network (trt.INetworkDefinition): The TensorRT network
+            config (trt.IBuilderConfig): The TensorRT builder config to modify
+        """
+        def try_run(func, name):
+            try:
+                return func()
+            except AttributeError:
+                trt_util.fail_unavailable(f"{name} in CreateConfig")
+
+        def try_set_flag(flag_name):
+            return try_run(
+                lambda: config.set_flag(getattr(trt.BuilderFlag, flag_name)),
+                flag_name.lower(),
+            )
+
+        # Add precision-related logic
+        if self.tf32:
+            try_set_flag("TF32")
+        else:  # TF32 is on by default
+            with contextlib.suppress(AttributeError):
+                config.clear_flag(trt.BuilderFlag.TF32)
+
+        if self.fp16:
+            try_set_flag("FP16")
+
+        if self.bf16:
+            try_set_flag("BF16")
+
+        if self.fp8:
+            try_set_flag("FP8")
+
+        if self.int8:
+            try_set_flag("INT8")
+
+        if self.int8:
+            # No Q/DQ layers means that we will need to calibrate.
+            if not any(
+                layer.type in [trt.LayerType.QUANTIZE, trt.LayerType.DEQUANTIZE]
+                for layer in network
+            ):
+                if self.calibrator is not None:
+                    config.int8_calibrator = self.calibrator
+                    try:
+                        profiles = copy.deepcopy(self.profiles)
+                        calib_profile = profiles[-1].fill_defaults(network)
+                        config.set_calibration_profile(
+                            calib_profile.to_trt(builder, network)
+                        )
+                        G_LOGGER.info(f"Using calibration profile: {calib_profile}")
+                    except AttributeError:
+                        G_LOGGER.extra_verbose(
+                            "Cannot set calibration profile on TensorRT 7.0 and older."
+                        )
+
+                    trt_util.try_setup_polygraphy_calibrator(
+                        config,
+                        network,
+                        calib_profile=calib_profile.to_trt(builder, network),
+                    )
+                else:
+                    G_LOGGER.warning(
+                        "Network does not have explicit precision and no calibrator was provided. Please ensure "
+                        "that tensors in the network have dynamic ranges set, or provide a calibrator in order to use int8 mode."
+                    )
+
+        if self.use_dla:
+            config.default_device_type = trt.DeviceType.DLA
+            config.DLA_core = 0
+
+        if self.allow_gpu_fallback:
+            try_set_flag("GPU_FALLBACK")
+
+    @util.check_called_by("__call__")
+    def call_impl(self, builder, network):
+        """
+        Callable implementation that creates and configures the IBuilderConfig with TensorRT features.
+        """
+        config = super().call_impl(builder, network)
+
+        self._configure_flags(builder, network, config)
+
+        return config
+
+
+@mod.export(funcify=True)
+class PostprocessConfig(BaseLoader):
+    """
+    [EXPERIMENTAL] Functor that applies a given post-processing function to a TensorRT ``IBuilderConfig``.
+    """
+
+    def __init__(self, config, func):
+        """
+        Applies a given post-processing function to a TensorRT ``IBuilderConfig``.
+
+        Args:
+            config (Union[trt.IBuilderConfig, Callable[[trt.Builder, trt.INetworkDefinition], trt.IBuilderConfig]):
+                    A TensorRT IBuilderConfig or a callable that accepts a TensorRT builder and network and returns a config.
+            func (Callable[[trt.Builder, trt.INetworkDefinition, trt.IBuilderConfig], None])
+                    A callable which takes a builder, network, and config parameter and modifies the config in place.
+        """
+
+        self._config = config
+
+        # Sanity-check that the function passed in is callable
+        if not callable(func):
+            G_LOGGER.critical(
+                f"Object {func} (of type {type(func)}) is not a callable."
+            )
+
+        self._func = func
+
+    @util.check_called_by("__call__")
+    def call_impl(self, builder, network):
+        """
+        Args:
+            builder (trt.Builder):
+                    The TensorRT builder to use to create the configuration.
+            network (trt.INetworkDefinition):
+                    The TensorRT network for which to create the config. The network is used to
+                    automatically create a default optimization profile if none are provided.
+
+        Returns:
+            trt.IBuilderConfig:
+                    The modified builder configuration.
+        """
+        config, _ = util.invoke_if_callable(self._config, builder, network)
+
+        self._func(builder, network, config)
+
+        return config

tools/Polygraphy/polygraphy/backend/trt/file_reader.py

@@ -0,0 +1,94 @@
+#
+# 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 pathlib import Path
+
+from polygraphy import mod, util
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.logger import G_LOGGER
+
+trt = lazy_import_trt()
+
+@mod.export()
+def FileReader(
+    filepath,
+    BaseClass=None,
+):
+    """
+    Class that supplies data to TensorRT from a stream. This may help reduce memory usage during deserialization.
+
+    Args:
+        filepath (str): 
+                The path to the serialized file.
+
+    """
+    BaseClass = util.default(BaseClass, trt.IStreamReader)
+
+    class FileReaderClass(BaseClass):
+        """
+        Class that supplies data to TensorRT from a stream. This may help reduce memory usage during deserialization. 
+        """
+
+        def __init__(self):
+            # Must explicitly initialize parent for any trampoline class! Will mysteriously segfault without this.
+            BaseClass.__init__(self)  # type: ignore
+
+            self.filepath = filepath
+
+            if not Path(self.filepath).exists():
+                G_LOGGER.error(f"File at {self.filepath} does not exist!")
+
+            self.mode = 'rb'
+            self.file = open(self.filepath, self.mode)
+            if not self.file:
+                G_LOGGER.error(f"Failed to open file at {self.filepath}!")
+
+            self.make_func = FileReader
+
+        def read(self, size: int) -> bytes:
+            return self.file.read(size)
+
+        def seek(self, offset: int, whence: int = 0) -> int:
+            """
+            Seek to a position in the stream. Required for IStreamReaderV2.
+
+            Args:
+                offset: The offset to seek to
+                whence: How to interpret the offset (0=absolute, 1=relative to current, 2=relative to end)
+
+            Returns:
+                The new absolute position
+            """
+            return self.file.seek(offset, whence)
+
+        def free(self):
+            if self.file:
+                self.file.close()
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.free()
+
+        def __repr__(self):
+            return util.make_repr(
+                "FileReader",
+                self.filepath,
+                BaseClass=BaseClass,
+            )[0]
+
+    return FileReaderClass()

tools/Polygraphy/polygraphy/backend/trt/profile.py

@@ -0,0 +1,215 @@
+#
+# 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 constants, mod, util
+from polygraphy.backend.trt import util as trt_util
+from polygraphy.common.interface import TypedDict
+from polygraphy.logger import G_LOGGER, LogMode
+from fnmatch import fnmatch
+
+
+@mod.export()
+class ShapeTuple:
+    """
+    Represents a set of shapes for a single binding in a profile.
+    """
+
+    def __init__(self, min, opt, max):
+        """
+        Args:
+            min (Tuple[int]): The minimum shape that the profile will support.
+            opt (Tuple[int]): The shape for which TensorRT will optimize the engine.
+            max (Tuple[int]): The maximum shape that the profile will support.
+        """
+        self.min = min
+        self.opt = opt
+        self.max = max
+
+    def __str__(self):
+        return f"(min={self.min}, opt={self.opt}, max={self.max})"
+
+    def __repr__(self):
+        return type(self).__name__ + self.__str__()
+
+    def __iter__(self):
+        yield from [self.min, self.opt, self.max]
+
+
+@mod.export()
+class Profile(TypedDict(lambda: str, lambda: ShapeTuple)):
+    """
+    An ordered dictionary that represents a single optimization profile that
+    can be used to build an engine.
+
+    More specifically, it is an ``OrderedDict[str, ShapeTuple]`` which maps binding
+    names to a set of min/opt/max shapes.
+    """
+
+    def add(self, name, min, opt, max):
+        """
+        A convenience function to add shapes for a single binding.
+
+        Args:
+            name (str): The name of the binding.
+            min (Tuple[int]): The minimum shape that the profile will support.
+            opt (Tuple[int]): The shape for which TensorRT will optimize the engine.
+            max (Tuple[int]): The maximum shape that the profile will support.
+
+        Returns:
+            Profile:
+                self, which allows this function to be easily chained to add multiple bindings,
+                e.g., Profile().add(...).add(...)
+        """
+        self[name] = ShapeTuple(min, opt, max)
+        return self
+
+    def __getitem__(self, key):
+        """
+        Retrieves the shapes registered for a given input name.
+
+        Returns:
+            ShapeTuple:
+                    A named tuple including ``min``, ``opt``, and ``max`` members for the shapes
+                    corresponding to the input.
+        """
+        if key not in self:
+            G_LOGGER.critical(
+                f"Binding: {key} does not have shapes set in this profile"
+            )
+        return super().__getitem__(key)
+
+    def fill_defaults(self, network, default_shape_value=None):
+        """
+        Fill this profile with sane default values for any bindings whose
+        shapes have not been set explicitly.
+
+        Args:
+            network (trt.INetworkDefinition):
+                    The TensorRT network this profile is meant for.
+                    This will be used to determine model inputs and their shapes.
+            default_shape_value (int):
+                    The value to use to override dynamic dimensions.
+
+        Returns:
+            Profile: Self
+        """
+        default_shape_value = util.default(
+            default_shape_value, constants.DEFAULT_SHAPE_VALUE
+        )
+
+        for idx in range(network.num_inputs):
+            inp = network.get_input(idx)
+
+            if any(fnmatch(inp.name, wc) for wc in self):
+                continue
+
+            with G_LOGGER.verbosity(G_LOGGER.CRITICAL):  # WAR for spam from TRT
+                is_shape_tensor = inp.is_shape_tensor
+            if is_shape_tensor:
+                rank = inp.shape[0] if len(inp.shape) > 0 else 1
+                shape = (default_shape_value,) * rank
+                G_LOGGER.warning(
+                    f"{trt_util.str_from_tensor(inp, is_shape_tensor)} | No values provided; "
+                    f"Will use input values: {shape} for min/opt/max in profile.\n",
+                    mode=LogMode.ONCE,
+                )
+                G_LOGGER.warning(
+                    "This will cause the shape-tensor to have static values. If this is incorrect, please "
+                    "set the range of values for this input shape-tensor.",
+                    mode=LogMode.ONCE,
+                )
+            else:
+                shape = util.override_dynamic_shape(inp.shape, default_shape_value)
+                if shape != inp.shape:
+                    G_LOGGER.warning(
+                        f"{trt_util.str_from_tensor(inp, is_shape_tensor)} | No shapes provided; Will use shape: {shape} for min/opt/max in profile.\n",
+                        mode=LogMode.ONCE,
+                    )
+                    G_LOGGER.warning(
+                        "This will cause the tensor to have a static shape. If this is incorrect, please "
+                        "set the range of shapes for this input tensor.",
+                        mode=LogMode.ONCE,
+                    )
+
+            self.add(inp.name, shape, shape, shape)
+        return self
+
+    def to_trt(self, builder, network):
+        """
+        Creates a TensorRT IOptimizationProfile based on the values set in this Profile.
+
+        Args:
+            builder (trt.Builder):
+                    A TensorRT builder. This will be used to construct the IOptimizationProfile.
+            network (trt.INetworkDefinition):
+                    The TensorRT network the profile applies to.
+
+        Returns:
+            trt.IOptimizationProfile: A TensorRT optimization profile.
+        """
+        trt_profile = builder.create_optimization_profile()
+        unused_keys = set(self.keys())
+        inp_names = [network.get_input(idx).name for idx in range(network.num_inputs)]
+        name_to_key, unmatched_inps = util.match_keys(unused_keys, inp_names)
+        if unmatched_inps:
+            G_LOGGER.critical(
+                f"Invalid inputs were provided to the optimization profile: {set(unmatched_inps)}\n"
+                f"Note: Inputs available in the TensorRT network are: {set(inp_names)}"
+            )
+
+        for idx in range(network.num_inputs):
+            inp = network.get_input(idx)
+            key = name_to_key[inp.name] if inp.name in name_to_key else None
+
+            with G_LOGGER.verbosity():  # WAR for spam from TRT
+                is_shape_tensor = inp.is_shape_tensor
+
+            if is_shape_tensor:
+                if key:
+                    shapes = self[key]
+                    trt_profile.set_shape_input(
+                        inp.name, shapes.min, shapes.opt, shapes.max
+                    )
+                    G_LOGGER.verbose(
+                        f"{trt_util.str_from_tensor(inp, is_shape_tensor)} | Setting input shape-tensor value range to: {shapes}"
+                    )
+                else:
+                    G_LOGGER.warning(
+                        f"{trt_util.str_from_tensor(inp, is_shape_tensor)} | No values provided. Assuming this is not a dynamic shape-tensor.",
+                        mode=LogMode.ONCE,
+                    )
+            else:
+                shapes = self[key if key else inp.name]
+                trt_profile.set_shape(inp.name, shapes.min, shapes.opt, shapes.max)
+                G_LOGGER.verbose(
+                    f"{trt_util.str_from_tensor(inp, is_shape_tensor)} | Setting input tensor shapes to: {shapes}"
+                )
+
+        return trt_util.check_profile(trt_profile)
+
+    def __repr__(self):
+        ret = "Profile()"
+        for name, (min, opt, max) in self.items():
+            ret += f".add('{name}', min={min}, opt={opt}, max={max})"
+        return ret
+
+    def __str__(self):
+        elems = []
+        for name, (min, opt, max) in self.items():
+            elems.append(f"{name} [min={min}, opt={opt}, max={max}]")
+
+        sep = ",\n "
+        return "{" + sep.join(elems) + "}"

tools/Polygraphy/polygraphy/backend/trt/requirements.txt

@@ -0,0 +1,2 @@
+-i https://pypi.ngc.nvidia.com/
+nvidia-tensorrt

tools/Polygraphy/polygraphy/backend/trt/runner.py

@@ -0,0 +1,544 @@
+#
+# 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.
+#
+import math
+import time
+import ctypes
+from collections import OrderedDict
+
+from polygraphy import config, cuda, mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.backend.trt import util as trt_util
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.common import FormattedArray
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER
+
+np = mod.lazy_import("numpy")
+torch = mod.lazy_import("torch>=1.13.0")
+trt = lazy_import_trt()
+
+
+def _make_debug_listener():
+    class DebugTensorWriter(trt.IDebugListener):
+        def __init__(self):
+            trt.IDebugListener.__init__(self)
+            self.debug_tensor_outputs = {}
+
+        def process_debug_tensor(self, addr, location, type, shape, name, stream):
+            if type in [util.try_getattr(trt, "fp8"), util.try_getattr(trt, "int4"), util.try_getattr(trt, "fp4"), util.try_getattr(trt, "bfloat16")]:
+                G_LOGGER.warning(f"Not supported datatype for debug tensor in polygraphy: {type}")
+                return
+
+            cuda.wrapper().stream_synchronize(stream)
+            datatype = DataType.from_dtype(type)
+            size = util.volume(shape)
+            buffer = np.zeros(shape, dtype=DataType.to_dtype(datatype, "numpy"))
+            buffer = util.array.resize_or_reallocate(buffer, size)
+            if location == trt.TensorLocation.HOST:
+                ctypes.memmove(util.array.data_ptr(buffer), addr, size * datatype.itemsize)
+            else:
+                cuda.wrapper().memcpy(
+                    dst=util.array.data_ptr(buffer),
+                    src=addr,
+                    nbytes=size * datatype.itemsize,
+                    kind=cuda.MemcpyKind.DeviceToHost,
+                    stream_ptr=stream,
+                )
+                cuda.wrapper().stream_synchronize(stream)
+            self.debug_tensor_outputs[name] = util.array.resize_or_reallocate(buffer, shape)
+
+    return DebugTensorWriter()
+
+
+def _make_output_allocator():
+
+    class OutputAllocator(trt.IOutputAllocator):
+        def __init__(self):
+            trt.IOutputAllocator.__init__(self)
+            self.buffers = {}
+            self.shapes = {}
+            self.use_torch = False
+
+        def reallocate_output(self, tensor_name, memory, size, alignment):
+            shape = (size,)
+            if tensor_name not in self.buffers:
+                self.buffers[tensor_name] = (
+                    cuda.DeviceArray.raw(shape)
+                    if not self.use_torch
+                    else torch.empty(shape, dtype=torch.uint8, device="cuda")
+                )
+            else:
+                self.buffers[tensor_name] = util.array.resize_or_reallocate(self.buffers[tensor_name], shape)
+            G_LOGGER.extra_verbose(f"Reallocated output tensor: {tensor_name} to: {self.buffers[tensor_name]}")
+            return util.array.data_ptr(self.buffers[tensor_name])
+
+        def notify_shape(self, tensor_name, shape):
+            self.shapes[tensor_name] = tuple(shape)
+
+        def set_use_torch(self, use_torch):
+            self.use_torch = use_torch
+
+    return OutputAllocator()
+
+
+def _get_array_on_cpu(arr, name, host_buffers, stream, nbytes, use_torch):
+    """
+    Copies the provided array to CPU memory and returns it.
+    If sufficient CPU memory has not been allocated for the array in
+    ``host_bufffers``, this function will allocate new memory.
+
+    If the input is a `torch.Tensor`, then a `torch.Tensor` is returned.
+    Otherwise, if the input is a `DeviceView`, a `NumPy` array is returned.
+
+    Args:
+        arr (Union[DeviceView, torch.Tensor]): The array.
+        name (str): The name of the array.
+        host_buffers (Dict[str, Union[numpy.ndarray, torch.Tensor]]):
+                A mapping of names to host buffers.
+        stream (cuda.Stream): The CUDA stream to use.
+        nbytes (int): The number of bytes to copy. This may be smaller than the size of the GPU memory.
+        use_torch (bool): Whether to use PyTorch tensors instead of NumPy arrays.
+
+    Returns:
+        Union[numpy.ndarray, torch.Tensor]: The host buffer as a flat array of bytes.
+    """
+    if not util.array.is_on_gpu(arr):
+        G_LOGGER.internal_error(f"_get_array_on_cpu() should only be called with input arrays on the GPU!")
+
+    # The host buffer will always be a "raw" array, i.e. a flat array of bytes.
+    shape = (nbytes,)
+    dtype = DataType.UINT8
+    # If we switch between torch tensors and DeviceViews between inferences, we need to reallocate the host buffer.
+    if name not in host_buffers or util.array.is_torch(host_buffers[name]) != use_torch:
+        host_buffers[name] = (
+            np.empty(shape, dtype=DataType.to_dtype(dtype, "numpy"))
+            if not use_torch
+            else torch.empty(shape, dtype=DataType.to_dtype(dtype, "torch"), device="cpu")
+        )
+
+    host_buffers[name] = util.array.resize_or_reallocate(host_buffers[name], shape)
+    cuda.wrapper().memcpy(
+        dst=util.array.data_ptr(host_buffers[name]),
+        src=util.array.data_ptr(arr),
+        nbytes=nbytes,
+        kind=cuda.MemcpyKind.DeviceToHost,
+        stream_ptr=stream.ptr,
+    )
+    return host_buffers[name]
+
+
+@mod.export()
+class TrtRunner(BaseRunner):
+    """
+    Runs inference using TensorRT.
+
+    Note that runners are not designed for production deployment and should generally
+    be used only for prototyping, testing, and debugging.
+    """
+
+    def __init__(
+        self,
+        engine,
+        name: str = None,
+        optimization_profile: int = None,
+        allocation_strategy: str = None,
+        weight_streaming_budget: int = None,
+        weight_streaming_percent: float = None,
+    ):
+        """
+        Args:
+            engine (Union[Union[trt.ICudaEngine, trt.IExecutionContext], Callable() -> Union[trt.ICudaEngine, trt.IExecutionContext]]):
+                    A TensorRT engine or execution context or a callable that returns one.
+                    If an engine is provided, the runner will create a context automatically.
+
+            name (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+            optimization_profile (int):
+                    The index of the optimization profile to set each time this runner is activated.
+                    When this is not provided, the profile is not set explicitly and will default to the 0th profile.
+                    You can also change the profile after the runner is active using the ``set_profile()`` method.
+            allocation_strategy (str):
+                    The way device memory (internal activation and scratch memory) is allocated for the execution context. The value of this argument can be:
+                        - "static": The default value. The execution context will pre-allocate a block of memory that is sufficient for any possible input size across all profiles.
+                        - "profile": Allocate device memory enough for the current profile based on profile max shapes.
+                        - "runtime": Allocate device meomry enough for the current input shapes.
+            weight_streaming_budget (int):
+                    The amount of GPU memory that TensorRT can use for weights at runtime. It can take on the following values:
+                        None or -2: Disables weight streaming at runtime.
+                        -1: TensorRT will decide the streaming budget automatically.
+                        >= 0: The maximum amount of GPU memory TensorRT is allowed to use for weights in bytes.
+            weight_streaming_percent (float):
+                    The percentage of weights that TRT will keep on the GPU. It can take on the following values:
+                        None or 100%: Disables weight streaming at runtime.
+                        [0 to 100]: The percentage of weights TRT will stream. 0 will stream the maximum number of weights.
+        """
+        super().__init__(name=name, prefix="trt-runner")
+        self._engine_or_context = engine
+        self.optimization_profile = optimization_profile
+        self.allocation_strategy = allocation_strategy
+        self.weight_streaming_budget = weight_streaming_budget
+        self.weight_streaming_percent = weight_streaming_percent
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        engine_or_context, _ = util.invoke_if_callable(self._engine_or_context)
+
+        if isinstance(engine_or_context, trt.ICudaEngine):
+            self.engine = engine_or_context
+            self._set_weight_streaming_budget()
+            allocation_strategy = util.default(self.allocation_strategy, "static")
+            if allocation_strategy == "static":
+                self.context = self.engine.create_execution_context()
+            elif allocation_strategy in ["profile", "runtime"]:
+                # Device memory will be managed by polygraphy
+                self.context = self.engine.create_execution_context(trt.ExecutionContextAllocationStrategy.USER_MANAGED)
+            else:
+                G_LOGGER.critical("Invalid allocation strategy specified.")
+            if not self.context:
+                G_LOGGER.critical("Invalid Context. See error log for details.")
+        elif isinstance(engine_or_context, trt.IExecutionContext):
+            self.context = engine_or_context
+            self.engine = self.context.engine
+            if self.allocation_strategy is not None:
+                G_LOGGER.warning(
+                    "An allocation strategy was specified. Please ensure the provided execution context uses the same strategy."
+                )
+
+        else:
+            G_LOGGER.critical(
+                "Invalid Engine or Context. Please ensure the engine was built correctly. See error log for details."
+            )
+
+        self.device_input_buffers = OrderedDict()
+        self.host_output_buffers = OrderedDict()
+        self.stream = cuda.Stream()
+        self.context_memory_buffer = None
+        self.output_allocator = _make_output_allocator()
+
+        if self.optimization_profile is not None:
+            self.set_profile(self.optimization_profile)
+
+    def set_profile(self, index: int):
+        """
+        Sets the active optimization profile for this runner.
+        The runner must already be active (see ``__enter__()`` or ``activate()``).
+
+        This only applies if your engine was built with multiple
+        optimization profiles.
+
+        In TensorRT 8.0 and newer, the profile will be set asynchronously
+        using this runner's CUDA stream (``runner.stream``).
+
+        By default, the runner uses the first profile (profile 0).
+
+        Args:
+            index (int):
+                    The index of the optimization profile to use.
+        """
+        if not hasattr(self, "context") or self.context is None:
+            G_LOGGER.critical(f"{self.name:35} | Must be activated prior to calling set_profile()")
+
+        try:
+            self.context.set_optimization_profile_async
+        except AttributeError:
+            self.context.active_optimization_profile = index
+        else:
+            if not self.context.set_optimization_profile_async(index, self.stream.ptr):
+                G_LOGGER.critical(f"Failed to set optimization profile to: {index}")
+
+    @util.check_called_by("get_input_metadata")
+    def get_input_metadata_impl(self):
+        return trt_util.get_metadata_from_engine(self.engine, self.context, mode=trt.TensorIOMode.INPUT)
+
+    def _infer_impl(self, feed_dict, copy_outputs_to_host, return_raw_buffers):
+        def get_io(mode):
+            for idx in range(self.engine.num_io_tensors):
+                name = self.engine.get_tensor_name(idx)
+
+                if self.engine.get_tensor_mode(name) == mode:
+                    yield name
+
+        use_torch = False
+
+        for name in get_io(trt.TensorIOMode.INPUT):
+            # Set up input tensor shapes and copy from host memory if needed
+            array = feed_dict[name]
+            if not isinstance(array, FormattedArray):
+                array = FormattedArray(array, shape=util.array.shape(array))
+
+            underlying_array = array.array
+            use_torch = use_torch or util.array.is_torch(underlying_array)
+
+            ptr = None
+            if self.engine.is_shape_inference_io(name):
+                if not util.array.is_on_cpu(underlying_array):
+                    G_LOGGER.critical(
+                        f"A {type(underlying_array).__name__} was provided for input: {name}, but since this is a shape tensor, "
+                        "it must reside in host memory. "
+                    )
+
+                ptr = util.array.data_ptr(underlying_array)
+            else:
+                ptr = trt_util._get_array_on_gpu(underlying_array, name, self.device_input_buffers, self.stream)
+
+            # If the format is HWC, make sure array.shape is considered after transposing back to CHW
+            if trt_util.get_tensor_format(self.engine, self.context, name) == trt.TensorFormat.HWC:
+                array_shape = trt_util.get_chw_shape_from_hwc(array.shape, self.context.get_tensor_strides(name))
+            else:
+                array_shape = array.shape
+
+            # Only update the input shape/address if something has changed. Otherwise, we'd be
+            # doing extra work unnecessarily.
+            # We retrieve the semantic shape from the FormattedArray, *not* the underlying array.
+            if self.context.get_tensor_shape(name) != array_shape:
+                G_LOGGER.ultra_verbose(f"Setting {name} input shape to: {array_shape}")
+                if not self.context.set_input_shape(name, array_shape):
+                    G_LOGGER.critical(f"For input: {name}, failed to set shape to: {array_shape}")
+
+            if self.context.get_tensor_address(name) != ptr:
+                if not self.context.set_tensor_address(name, ptr):
+                    G_LOGGER.critical(f"For input: {name}, failed to set tensor address to: {ptr}")
+
+        try:
+            self.context.set_all_tensors_debug_state
+        except AttributeError:
+            pass
+        else:
+            # Set up the debug listener before running inference.
+            debug_listener = _make_debug_listener()
+            self.context.set_all_tensors_debug_state(True)
+            if not self.context.set_debug_listener(debug_listener):
+                G_LOGGER.critical(f"Failed to set debug listener.")
+
+        # Set up the output allocator before running inference.
+        self.output_allocator.set_use_torch(use_torch and torch.cuda.is_available())
+        for name in get_io(trt.TensorIOMode.OUTPUT):
+            if not self.context.set_output_allocator(name, self.output_allocator):
+                G_LOGGER.critical(f"For output: {name}, failed to set output allocator")
+
+        if self.allocation_strategy in ["profile", "runtime"]:
+            if self.allocation_strategy == "profile":
+                # Perform per-profile allocation.
+                size_to_allocate = 0
+                if config.USE_TENSORRT_RTX or mod.version(trt.__version__) >= mod.version("10.1"):
+                    size_to_allocate = self.engine.get_device_memory_size_for_profile_v2(
+                        self.context.active_optimization_profile
+                    )
+                else:
+                    size_to_allocate = self.engine.get_device_memory_size_for_profile(
+                        self.context.active_optimization_profile
+                    )
+            elif self.allocation_strategy == "runtime":
+                # Perform runtime allocation.
+                size_to_allocate = self.context.update_device_memory_size_for_shapes()
+
+            if self.context_memory_buffer is None:
+                self.context_memory_buffer = cuda.DeviceArray.raw((size_to_allocate,))
+
+            self.context_memory_buffer.resize((size_to_allocate,))
+            if config.USE_TENSORRT_RTX or mod.version(trt.__version__) >= mod.version("10.1"):
+                self.context.set_device_memory(self.context_memory_buffer.ptr, self.context_memory_buffer.allocated_nbytes)
+            else:
+                self.context.device_memory = self.context_memory_buffer.ptr
+
+        if not self.context.execute_async_v3(self.stream.ptr):
+            G_LOGGER.critical("`execute_async_v3()` failed. Please see the logging output above for details.")
+
+        output_buffers = OrderedDict()
+        for name in get_io(trt.TensorIOMode.OUTPUT):
+            # If we're dealing with vectorized formats, we need to return a FormattedArray.
+            # Otherwise, we create a view instead with the correct shape/dtype.
+            raw_array = self.output_allocator.buffers[name]
+
+            shape = self.output_allocator.shapes[name]
+            # If the format is HWC, make sure the result is shaped accordingly
+            tensor_format = trt_util.get_tensor_format(self.engine, self.context, name)
+            if tensor_format == trt.TensorFormat.HWC:
+                shape = trt_util.get_hwc_shape_from_chw(shape, self.context.get_tensor_strides(name))
+            using_vectorized_format = tensor_format != trt.TensorFormat.LINEAR and tensor_format != trt.TensorFormat.HWC
+            should_use_formatted_array = return_raw_buffers or using_vectorized_format
+
+            dtype = DataType.from_dtype(self.engine.get_tensor_dtype(name), source_module="tensorrt")
+
+            # The memory allocated by the output allocator may be larger than actually required.
+            # If we're using a vectorized format, then we need to copy the whole thing.
+            # Otherwise, we can determine how much we actually need.
+            nbytes = (
+                util.array.nbytes(raw_array)
+                if using_vectorized_format
+                # Some data types have fractional sizes, in which case we round up to the nearest byte.
+                else int(math.ceil(util.volume(shape) * dtype.itemsize))
+            )
+
+            if copy_outputs_to_host:
+                raw_array = _get_array_on_cpu(
+                    raw_array,
+                    name,
+                    self.host_output_buffers,
+                    self.stream,
+                    nbytes,
+                    use_torch=use_torch,
+                )
+
+            if should_use_formatted_array:
+                array = FormattedArray(raw_array, shape=shape)
+            else:
+                array = util.array.view(raw_array, dtype, shape)
+            output_buffers[name] = array
+
+        self.stream.synchronize()
+
+        try:
+            self.context.set_all_tensors_debug_state
+        except AttributeError:
+            pass
+        else:
+            if debug_listener.debug_tensor_outputs:
+                output_buffers.update(debug_listener.debug_tensor_outputs)
+
+        return output_buffers
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict, copy_outputs_to_host=None, return_raw_buffers=None):
+        """
+        Implementation for running inference with TensorRT.
+        Do not call this method directly - use ``infer()`` instead,
+        which will forward unrecognized arguments to this method.
+
+        Args:
+            feed_dict (OrderedDict[str, Union[numpy.ndarray, DeviceView, torch.Tensor]]):
+                    A mapping of input tensor names to corresponding input NumPy arrays,
+                    Polygraphy DeviceViews, or PyTorch tensors.
+                    If PyTorch tensors are provided in the feed_dict, then this function
+                    will return the outputs also as PyTorch tensors.
+                    If the provided inputs already reside in GPU memory, no additional copies are made.
+
+            copy_outputs_to_host (bool):
+                    Whether to copy inference outputs back to host memory.
+                    If this is False, PyTorch GPU tensors or Polygraphy DeviceViews
+                    are returned instead of PyTorch CPU tensors or NumPy arrays respectively.
+                    Defaults to True.
+
+        Returns:
+            OrderedDict[str, Union[numpy.ndarray, DeviceView, torch.Tensor]]:
+                    A mapping of output tensor names to corresponding output NumPy arrays,
+                    Polygraphy DeviceViews, or PyTorch tensors.
+        """
+        copy_outputs_to_host = util.default(copy_outputs_to_host, True)
+        return_raw_buffers = util.default(return_raw_buffers, False)
+
+        start = time.time()
+        output_buffers = self._infer_impl(feed_dict, copy_outputs_to_host, return_raw_buffers)
+        end = time.time()
+        self.inference_time = end - start
+
+        return output_buffers
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        [buf.free() for buf in self.device_input_buffers.values()]
+        if self.context_memory_buffer is not None:
+            self.context_memory_buffer.free()
+        self.stream.free()
+
+        del (
+            self.engine,
+            self.context,
+            self.device_input_buffers,
+            self.host_output_buffers,
+            self.stream,
+            self.context_memory_buffer,
+            self.output_allocator,
+        )
+
+    def _set_weight_streaming_budget(self):
+        # Setup weight streaming if applicable
+        if self.weight_streaming_budget != None and self.weight_streaming_percent != None:
+            G_LOGGER.warning(f"Cannot specify the weight streaming budget both in bytes and percentage. Prioritizing the bytes value.")
+
+        if self.weight_streaming_budget is not None:
+            assert self.weight_streaming_budget == -2 or self.weight_streaming_budget == -1 or self.weight_streaming_budget >= 0
+
+        if config.USE_TENSORRT_RTX or mod.version(trt.__version__) >= mod.version("10.1"):
+            self._set_weight_streaming_budget_v2()
+        else:
+            self._set_weight_streaming_budget_v1()
+
+    def _set_weight_streaming_budget_v1(self):
+        budget_bytes = None
+        if self.weight_streaming_budget is not None:
+            if self.weight_streaming_budget == -2:
+                budget_bytes = 0
+            else:
+                budget_bytes = self.weight_streaming_budget
+
+        elif self.weight_streaming_percent is not None:
+            assert 0 <= self.weight_streaming_percent <= 100
+            if self.weight_streaming_percent == 0:
+                budget_bytes = 0  # Disable weight streaming
+            else:
+                try:
+                    min_budget = self.engine.minimum_weight_streaming_budget
+                except AttributeError:
+                    # minimum_weight_streaming_budget is deprecated in TensorRT 10.1 and removed in
+                    # TensorRT RTX 1.0. For the new / V2 path, the minimum budget is 0.
+                    min_budget = 0
+                max_budget = self.engine.streamable_weights_size
+                budget_bytes = (1 - self.weight_streaming_percent / 100.0) * (max_budget - min_budget) + min_budget
+
+        if budget_bytes is not None:
+            budget_bytes = int(budget_bytes)
+            self.engine.weight_streaming_budget = budget_bytes
+            if self.engine.weight_streaming_budget != budget_bytes:
+                G_LOGGER.critical(f"Failed to set weight streaming budget to {budget_bytes}!")
+            if budget_bytes == 0:
+                G_LOGGER.info(f"Weight streaming is disabled.")
+            elif budget_bytes == -1:
+                G_LOGGER.info(f"Weight streaming is enabled with TensorRT automatically determiing the budget.")
+            else:
+                G_LOGGER.info(f"Weight streaming is enabled with a memory budget of {budget_bytes} bytes.")
+
+
+    def _set_weight_streaming_budget_v2(self):
+        budget_bytes = None
+        if self.weight_streaming_budget is not None:
+            # use V2 path
+            assert self.weight_streaming_budget == -2 or self.weight_streaming_budget == -1 or self.weight_streaming_budget >= 0
+            if self.weight_streaming_budget == -2:
+                budget_bytes = self.engine.streamable_weights_size
+            elif self.weight_streaming_budget == -1:
+                budget_bytes = self.engine.get_weight_streaming_automatic_budget()
+            else:
+                budget_bytes = self.weight_streaming_budget
+
+        elif self.weight_streaming_percent is not None:
+            assert 0 <= self.weight_streaming_percent <= 100
+            if self.weight_streaming_percent == 100:
+                budget_bytes = self.engine.streamable_weights_size
+            else:
+                budget_bytes = self.weight_streaming_percent / 100.0 * (self.engine.streamable_weights_size)
+
+        if budget_bytes is not None:
+            budget_bytes = int(budget_bytes)
+            self.engine.weight_streaming_budget_v2 = budget_bytes
+            if self.engine.weight_streaming_budget_v2 != budget_bytes:
+                G_LOGGER.critical(f"Failed to set weight streaming budget to {budget_bytes}!")
+            if budget_bytes == self.engine.streamable_weights_size:
+                G_LOGGER.info(f"Weight streaming is disabled.")
+            else:
+                G_LOGGER.info(f"Weight streaming is enabled with a memory budget of {budget_bytes} bytes.")