chore: import upstream snapshot with attribution

tools/Polygraphy/polygraphy/README.md

@@ -0,0 +1,257 @@
+# Polygraphy Python API
+
+## Table of Contents
+
+- [Introduction](#introduction)
+- [Backends](#backends)
+    - [Loaders](#loaders)
+    - [Runners](#runners)
+        - [Writing A Custom Runner](#writing-a-custom-runner)
+- [Comparator](#comparator)
+    - [Data Loaders](#data-loaders)
+- [Logger](#logger)
+- [Putting It All Together](#putting-it-all-together)
+- [Enabling PyTorch](#enabling-pytorch)
+- [Examples](#examples)
+- [Python API Reference Documentation](#python-api-reference-documentation)
+    - [Building Python API Documentation Locally](#building-python-api-documentation-locally)
+- [Deprecation Policy](#deprecation-policy)
+
+## Introduction
+
+The Polygraphy API consists broadly of two major components:
+[`Backend`s](#backends) and the [`Comparator`](#comparator).
+
+**NOTE:** To help you get started with the API, you can use the [`run`](./tools/run/) tool
+with the `--gen-script` option to auto-generate template scripts that use the Polygraphy API.
+
+> :warning: Any APIs not documented in the [API reference documentation](#python-api-reference-documentation)
+    should be considered internal only and do not adhere to the [deprecation policy](#deprecation-policy)
+    for public APIs. Thus, they may be modified or removed at any time without warning.
+    Avoid using undocumented APIs!
+
+
+## Backends
+
+A Polygraphy backend provides an interface for a deep learning framework.
+Backends are comprised of two components: Loaders and Runners.
+
+
+### Loaders
+
+A `Loader` is a functor or callable that loads or operates on models in some way.
+
+Existing `Loader`s can be composed for more advanced behaviors.
+For example, we can implement a conversion like `ONNX -> TensorRT Network -> TensorRT Engine`:
+
+```python
+from polygraphy.backend.trt import EngineFromNetwork, NetworkFromOnnxPath
+
+build_engine = EngineFromNetwork(NetworkFromOnnxPath("/path/to/model.onnx"))
+```
+
+`build_engine` is a callable that will build a TensorRT engine.
+
+Polygraphy also provides immediately evaluated functional variants of each Loader.
+These use the same names, except `snake_case` instead of `PascalCase`, and expose the same APIs.
+Using the functional loaders, the conversion above would be:
+
+```python
+from polygraphy.backend.trt import engine_from_network, network_from_onnx_path
+
+engine = engine_from_network(network_from_onnx_path("/path/to/model.onnx"))
+```
+
+`engine` is a TensorRT engine as opposed to a callable.
+
+
+### Runners
+
+A `Runner` uses a loader to load a model and can then run inference (see [`BaseRunner`](./backend/base/runner.py)).
+
+**IMPORTANT:** Runners may reuse their output buffers. Thus, if you need to save outputs from multiple inferences, you should
+make a copy of the outputs with `copy.deepcopy(outputs)`.
+
+To use a runner, you just need to activate it, and then call `infer()`.
+Note that activating a runner can be very expensive, so you should minimize the
+number of times you activate a runner - ideally do not do this more than once.
+
+It is recommended to use a context manager to activate and deactivate the
+runner rather than calling the functions manually:
+
+```python
+from polygraphy.backend.trt import TrtRunner
+
+with TrtRunner(build_engine) as runner:
+    outputs = runner.infer(feed_dict={"input0": input_data})
+```
+
+#### Writing A Custom Runner
+
+Generally, you do not need to write custom runners unless you want to support a new backend.
+
+In case you do, in the simplest case, you only need to implement two functions:
+- `infer_impl`: Accepts a dictionary of numpy buffers, runs inference, and finally returns a dictionary containing the outputs.
+- `get_input_metadata_impl`: Returns a [`TensorMetadata`](./common/struct.py) mapping input names to their shapes and data types.
+    You may use `None`, negative numbers, or strings to indicate dynamic dimensions.
+
+For more advanced runners, where some setup is required, you may also need to implement the `activate_impl()` and `deactivate_impl()` functions.
+
+For example, in the `TrtRunner`, engines are created in `activate_impl()` and destroyed in `deactivate_impl()`.
+Importantly, the GPU is *not used at all* until these functions have been called (notice, for example,
+    that in the `TrtRunner`, the CUDA runtime library is only loaded in the `activate_impl()` function).
+This allows the `Comparator` to optionally provide each runner with exclusive access to the GPU, to prevent any interference between runners.
+
+
+## Comparator
+
+The `Comparator` is used to run inference for runners, and then compare accuracy (see [Comparator.py](./comparator/comparator.py)).
+This process is divided into two phases:
+
+1. Running inference:
+
+    ```python
+    run_results = Comparator.run(runners)
+    ```
+
+    This function accepts a list of runners and returns a `RunResults` object (see [Comparator.py](./comparator/comparator.py))
+    containing the inference outputs of each run.
+    It also accepts an optional `data_loader` argument to control the input data. If not provided, it will use the
+    default data loader. `Comparator.run()` continues until inputs from the data loader are exhausted.
+
+2. Comparing results:
+
+    ```python
+    Comparator.compare_accuracy(run_results)
+    ```
+
+    This function accepts the results returned by `Comparator.run` and compares them between runners.
+
+**IMPORTANT:** The Comparator is designed for scenarios where you need to compare a small number
+of inputs across multiple runners. It is **not** a good idea to use it
+to validate a model with an entire dataset! Instead, runners should be used directly for such
+cases (see the [example](../examples/api/02_validating_on_a_dataset)).
+
+
+### Data Loaders
+
+A data loader is used by the `Comparator` to load input data to feed to each runner
+(for example, see Polygraphy's [default data loader](./comparator/data_loader.py)).
+
+A data loader can be any generator or iterable that yields
+a dictionary of input buffers. In the simplest case, this can just be a `list` of `dict`s.
+
+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 an `TensorMetadata` instance by the Comparator.
+
+**NOTE:** Polygraphy provides a default `DataLoader` class that uses numpy to generate random input buffers.
+The input data can be bounded via parameters to the constructor.
+
+
+## Logger
+
+Polygraphy also includes a global logger which can control the verbosity not only of messages emitted by Polygraphy,
+but also of those emitted by underlying frameworks, like TensorRT. For example, the `EXTRA_VERBOSE` verbosity corresponds
+to TensorRT's `kVERBOSE` logging mode.
+
+To set the verbosity of the logger, use:
+```py
+G_LOGGER.module_severity = severity
+```
+
+For example:
+```py
+G_LOGGER.module_severity = G_LOGGER.EXTRA_VERBOSE
+```
+
+By default logs are emitted to the stdout/stderr stream. If you would like to have them stored in a file, set the `log_file` property to your log file:
+```py
+G_LOGGER.log_file = "your_log_file.log"
+```
+
+If you would like to have logs being emitted to python `logging` module, set the following flag:
+```py
+G_LOGGER.use_python_logging_system = True
+```
+
+After that, you can define your logging configuration using the `logging` module.
+
+## Putting It All Together
+
+Now that you know the basic components of Polygraphy, let's take a look at how they fit together.
+
+In this example, we will write a script that:
+1. Builds a TensorRT engine from an ONNX model
+2. Bounds input values in the range `[0, 2]`
+3. Runs inference using ONNX-Runtime and TensorRT
+4. Compares the results and checks that they match
+
+```python
+from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.backend.trt import TrtRunner, EngineFromNetwork, NetworkFromOnnxPath
+from polygraphy.comparator import Comparator, DataLoader
+
+model_path = "/path/to/model.onnx"
+
+build_onnxrt_session = SessionFromOnnx(model_path)
+build_engine = EngineFromNetwork(NetworkFromOnnxPath(model_path))
+
+runners = [
+    OnnxrtRunner(build_onnxrt_session),
+    TrtRunner(build_engine),
+]
+
+data_loader = DataLoader(val_range=(0, 2))
+
+run_results = Comparator.run(runners, data_loader=data_loader)
+assert bool(Comparator.compare_accuracy(run_results))
+```
+
+## Enabling PyTorch
+
+In order to enable PyTorch, you need to provide three things to the `PytRunner`:
+1. A model loader: In the simplest case, this can be a callable that returns a `torch.nn.Module`.
+
+2. `input_metadata`: A `TensorMetadata` describing the inputs of the model. This maps input names to their shapes and data types. As with other runners, `None` may be used to indicate dynamic dimensions.
+
+    **NOTE:** Other runners are able to automatically determine input metadata by inspecting the model definition, but because of the way PyTorch is implemented, it is difficult to write a generic function to determine model inputs from a `torch.nn.Module`.
+
+3. `output_names`: A list of output names. This is used by the `Comparator` to match `PytRunner` outputs to those of other runners.
+
+
+## Examples
+
+You can find complete code examples that use the Polygraphy Python API [here](../examples/api).
+
+
+## Python API Reference Documentation
+
+For more details, see the [Polygraphy Python API reference documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/polygraphy/index.html).
+
+### Building Python API Documentation Locally
+
+To build the API documentation, first install required packages:
+
+```bash
+python -m pip install -r docs/requirements.txt
+```
+
+and then use the `make` target to build docs:
+
+```bash
+make docs
+```
+
+The HTML documentation will be generated under `build/docs`
+To view the docs, open `build/docs/index.html` in a browser or HTML viewer.
+
+## Deprecation Policy
+
+When removing or modifying public APIs in breaking ways, Polygraphy follows the following procedure:
+
+1. The API is marked to be removed in some future version of Polygraphy (let's call it version `N`).
+
+2. The API continues to work as normal until *at least* version `N`, but will issue deprecation warnings when used.
+    These warnings typically suggest which new APIs should be used instead.
+
+3. The API may be removed in any Polygraphy version >= `N`.

tools/Polygraphy/polygraphy/__init__.py

@@ -0,0 +1,3 @@
+import polygraphy.config
+
+__version__ = "0.49.27"

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.")

tools/Polygraphy/polygraphy/common/__init__.py

@@ -0,0 +1,2 @@
+from polygraphy.exception import *
+from polygraphy.common.struct import *

tools/Polygraphy/polygraphy/common/interface.py

@@ -0,0 +1,192 @@
+#
+# 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 util
+from polygraphy.logger import G_LOGGER
+
+#
+# NOTE: These classes intentionally don't inherit from the built-in collections (dict, list, etc.)
+# because doing so prevents us from providing custom JSON serialization methods, since the default
+# encoder implementation can handle most of the built-in collections and therefore doesn't dispatch
+# to custom implementations.
+#
+
+
+def TypedDict(key_type_func, value_type_func):
+    """
+    Returns a class (not an instance) that will provide a dictionary-like
+    interface with runtime type checks.
+
+    Note: The types are provided lazily via a callable to avoid unnecessary dependencies
+    on types from external packages at import-time.
+
+    Args:
+        key_type_func (Callable() -> type):
+                A callable that returns the expected key type.
+        value_type_func (Callable() -> type):
+                A callable that returns the expected value type.
+    """
+
+    class Interface:
+        def __init__(self, dct=None):
+            self.dct = OrderedDict(util.default(dct, {}))
+            self.key_type = key_type_func()
+            self.value_type = value_type_func()
+
+        def _check_types(self, key, val):
+            if not isinstance(key, self.key_type):
+                G_LOGGER.critical(
+                    f"Unsupported key type in {self}. Key: {repr(key)} is type `{type(key).__name__}` but {type(self).__name__} expects type `{self.key_type.__name__}`"
+                )
+            if not isinstance(val, self.value_type):
+                G_LOGGER.critical(
+                    f"Unsupported value type in {self}. Value: {repr(val)} for key: {repr(key)} is type `{type(val).__name__}` but {type(self).__name__} expects type `{self.value_type.__name__}`"
+                )
+
+        def keys(self):
+            return self.dct.keys()
+
+        def values(self):
+            return self.dct.values()
+
+        def items(self):
+            return self.dct.items()
+
+        def update(self, other):
+            for key, val in other.items():
+                self._check_types(key, val)
+            return self.dct.update(other)
+
+        def __contains__(self, key):
+            return key in self.dct
+
+        def __getitem__(self, key):
+            return self.dct[key]
+
+        def __setitem__(self, key, val):
+            self._check_types(key, val)
+            self.dct[key] = val
+
+        def __str__(self):
+            return str(self.dct)
+
+        def __repr__(self):
+            return repr(self.dct)
+
+        def __len__(self):
+            return len(self.dct)
+
+        def __eq__(self, other):
+            return self.dct == other.dct
+
+        def __iter__(self):
+            return self.dct.__iter__()
+
+        def __copy__(self):
+            new_dict = type(self)()
+            new_dict.__dict__.update(self.__dict__)
+            new_dict.dct = copy.copy(self.dct)
+            return new_dict
+
+        def __deepcopy__(self, memo):
+            new_dict = type(self)()
+            new_dict.__dict__.update(self.__dict__)
+            new_dict.dct = copy.deepcopy(self.dct)
+            return new_dict
+
+    return Interface
+
+
+def TypedList(elem_type_func):
+    """
+    Returns a class (not an instance) that will provide a list-like
+    interface with runtime type checks.
+
+    Note: The types are provided lazily via a callable to avoid unnecessary dependencies
+    on types from external packages at import-time.
+
+    Args:
+        elem_type_func (Callable() -> type):
+                A callable that returns the expected list-element type.
+    """
+
+    class Interface:
+        def __init__(self, lst=None):
+            self.lst = util.default(lst, [])
+            self.elem_type = elem_type_func()
+
+        def _check_type(self, elem):
+            if not isinstance(elem, self.elem_type):
+                G_LOGGER.critical(
+                    f"Unsupported element type type in {type(self).__name__}. Element: {repr(elem)} is type: {type(elem).__name__} but type: {self.elem_type.__name__} was expected"
+                )
+
+        def __contains__(self, key):
+            return key in self.lst
+
+        def __getitem__(self, index):
+            return self.lst[index]
+
+        def __setitem__(self, index, elem):
+            self._check_type(elem)
+            self.lst[index] = elem
+
+        def __str__(self):
+            return str(self.lst)
+
+        def __repr__(self):
+            return repr(self.lst)
+
+        def append(self, elem):
+            self._check_type(elem)
+            return self.lst.append(elem)
+
+        def extend(self, elems):
+            for elem in elems:
+                self._check_type(elem)
+            return self.lst.extend(elems)
+
+        def __iadd__(self, elems):
+            for elem in elems:
+                self._check_type(elem)
+            self.lst += elems
+
+        def __len__(self):
+            return len(self.lst)
+
+        def __eq__(self, other):
+            return self.lst == other.lst
+
+        def __iter__(self):
+            return self.lst.__iter__()
+
+        def __copy__(self):
+            new_list = type(self)()
+            new_list.__dict__.update(self.__dict__)
+            new_list.lst = copy.copy(self.lst)
+            return new_list
+
+        def __deepcopy__(self, memo):
+            new_list = type(self)()
+            new_list.__dict__.update(self.__dict__)
+            new_list.lst = copy.deepcopy(self.lst)
+            return new_list
+
+    return Interface

tools/Polygraphy/polygraphy/common/struct.py

@@ -0,0 +1,203 @@
+#
+# 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.common.interface import TypedDict
+from polygraphy.datatype import DataType
+from polygraphy.json import Decoder, Encoder, add_json_methods
+
+
+class BoundedShape(list):
+    """
+    Represents a shape with min/max bounds.
+    """
+
+    def __init__(self, shape, min=None, max=None):
+        super().__init__(shape)
+        self.min = min
+        self.max = max
+
+    def __repr__(self):
+        return f"BoundedShape({list(self)}, min={self.min}, max={self.max})"
+
+
+class MetadataTuple:
+    def __init__(self, dtype, shape, docstring):
+        self.dtype = dtype
+        self.shape = shape
+        self.docstring = docstring
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @dtype.setter
+    def dtype(self, new):
+        self._dtype = DataType.from_dtype(new) if new is not None else None
+
+    def __iter__(self):
+        yield from [self.dtype, self.shape]
+
+    def __repr__(self):
+        return f"MetadataTuple({self.dtype}, {self.shape}, {self.docstring})"
+
+    def __str__(self):
+        ret = ""
+        meta_items = []
+        if self.dtype is not None:
+            meta_items.append(f"dtype={self.dtype}")
+        if self.shape is not None:
+            meta_items.append(f"shape={tuple(self.shape)}")
+        if self.docstring:
+            meta_items.append(self.docstring)
+        if meta_items:
+            ret += "[" + ", ".join(meta_items) + "]"
+        return ret
+
+    def __eq__(self, other):
+        return self.shape == other.shape and self.dtype == other.dtype
+
+
+@mod.export()
+class TensorMetadata(TypedDict(lambda: str, lambda: MetadataTuple)):
+    """
+    An OrderedDict[str, MetadataTuple] that maps input names to their data types and shapes.
+
+    Shapes may include negative values, ``None``, or strings to indicate dynamic dimensions.
+
+    Example:
+    ::
+
+        shape = tensor_meta["input0"].shape
+        dtype = tensor_meta["input0"].dtype
+    """
+
+    @staticmethod
+    def from_feed_dict(feed_dict):
+        """
+        Constructs a new TensorMetadata using information from the provided feed_dict.
+
+        Args:
+            feed_dict (OrderedDict[str, Union[numpy.ndarray, torch.tensor]]):
+                    A mapping of input tensor names to corresponding input arrays.
+
+        Returns:
+            TensorMetadata
+        """
+        meta = TensorMetadata()
+        for name, arr in feed_dict.items():
+            meta.add(name, util.array.dtype(arr), util.array.shape(arr))
+        return meta
+
+    def add(self, name, dtype, shape, min_shape=None, max_shape=None, docstring=None):
+        """
+        Convenience function for adding entries.
+
+        Args:
+            name (str): The name of the input.
+            dtype (Any):
+                    The data type of the input.
+                    This can be any type that can be converted to a Polygraphy DataType.
+            shape (Sequence[Union[int, str]]]):
+                    The shape of the input. Dynamic dimensions may
+                    be indicated by negative values, ``None``, or a string.
+
+            min_shape (Sequence[int]):
+                    The minimum valid shape for the input.
+                    If provided, this shape should not include any dynamic dimensions.
+            max_shape (Sequence[int]):
+                    The maximum valid shape for the input.
+                    If provided, this shape should not include any dynamic dimensions.
+            docstring (str):
+                    Any additional information associated with a tensor.
+
+        Returns:
+            The newly added entry.
+        """
+        self[name] = MetadataTuple(
+            dtype,
+            (
+                BoundedShape(shape, min=min_shape, max=max_shape)
+                if shape is not None
+                else None
+            ),
+            docstring,
+        )
+        return self
+
+    def __repr__(self):
+        ret = "TensorMetadata()"
+        for name, metadata_tuple in self.items():
+            (dtype, shape) = metadata_tuple
+            ret += util.make_repr(
+                ".add",
+                name,
+                dtype,
+                list(shape),
+                min_shape=shape.min,
+                max_shape=shape.max,
+                docstring=metadata_tuple.docstring,
+            )[0]
+        return ret
+
+    def __str__(self):
+        sep = ",\n "
+        elems = [f"{name} {meta_tuple}".strip() for name, meta_tuple in self.items()]
+        return "{" + sep.join(elems) + "}"
+
+
+@mod.export()
+@add_json_methods("formatted array")
+class FormattedArray:
+    """
+    [EXPERIMENTAL, UNTESTED] This API is experimental and untested and may be significantly
+    modified in future releases. Use with caution!
+
+    Representes an array whose semantic shape differs from its physical size in memory.
+
+    For example, consider an ``NCHW`` tensor of shape ``(1, 3, 28, 28)``. If we use a vectorized format
+    like ``N(C/4)HW4``, then the physical size of the array would be ``(1, 1, 28, 28 * 4)`` since
+    the channel dimension would be padded to a multiple of 4. However, we still need a way to keep
+    track of the semantic shape for things like shape inference.
+
+    This class provides a mechanism to specify the shape of an array independently of
+    the underlying array.
+    """
+
+    def __init__(self, array, shape):
+        """
+        Args:
+            array (Union[np.ndarray, polygraphy.cuda.DeviceView]):
+                    The array. In most cases, this will be a raw byte-array.
+            shape (Sequence[int]):
+                    The semantic shape of the data.
+        """
+        self.array = array
+        self.shape = shape
+
+
+@Encoder.register(FormattedArray)
+def encode(farray):
+    return {
+        "array": farray.array,
+        "shape": farray.shape,
+    }
+
+
+@Decoder.register(FormattedArray)
+def decode(dct):
+    return FormattedArray(dct["array"], dct["shape"])

tools/Polygraphy/polygraphy/comparator/__init__.py

@@ -0,0 +1,5 @@
+from polygraphy.comparator.comparator import *
+from polygraphy.comparator.compare import *
+from polygraphy.comparator.data_loader import *
+from polygraphy.comparator.postprocess import *
+from polygraphy.comparator.struct import *

tools/Polygraphy/polygraphy/comparator/comparator.py

@@ -0,0 +1,423 @@
+#
+# 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 queue
+from multiprocessing import Process, Queue
+
+from polygraphy import mod, util
+from polygraphy.common import TensorMetadata
+from polygraphy.comparator import util as comp_util
+from polygraphy.comparator.compare import CompareFunc
+from polygraphy.comparator.data_loader import DataLoader, DataLoaderCache
+from polygraphy.comparator.struct import AccuracyResult, IterationResult, RunResults
+from polygraphy.logger import G_LOGGER, LogMode
+
+
+@mod.export()
+class Comparator:
+    """
+    Compares inference outputs.
+    """
+
+    @staticmethod
+    def run(
+        runners,
+        data_loader=None,
+        warm_up=None,
+        use_subprocess=None,
+        subprocess_timeout=None,
+        subprocess_polling_interval=None,
+        save_inputs_path=None,
+    ):
+        """
+        Runs the supplied runners sequentially.
+
+        Args:
+            runners (List[BaseRunner]):
+                    A list of runners to run.
+            data_loader (Sequence[OrderedDict[str, numpy.ndarray]]):
+                    A generator or iterable that yields a dictionary that maps input names to input numpy buffers.
+                    In the simplest case, this can be a `List[Dict[str, numpy.ndarray]]` .
+
+                    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 an `TensorMetadata`
+                    instance by this function.
+                    Note that this does not work for generators or lists.
+
+                    The number of iterations run by this function is controlled by the number of items supplied
+                    by the data loader.
+
+                    Defaults to an instance of `DataLoader`.
+            warm_up (int):
+                    The number of warm up runs to perform for each runner before timing.
+                    Defaults to 0.
+            use_subprocess (bool):
+                    Whether each runner should be run in a subprocess. This allows each runner to have exclusive
+                    access to the GPU. When using a subprocess, runners and loaders will never be modified.
+            subprocess_timeout (int):
+                    The timeout before a subprocess is killed automatically. This is useful for handling processes
+                    that never terminate. A value of None disables the timeout. Defaults to None.
+            subprocess_polling_interval (int):
+                    The polling interval, in seconds, for checking whether a subprocess has completed or crashed.
+                    In rare cases, omitting this parameter when subprocesses are enabled may cause this function
+                    to hang indefinitely if the subprocess crashes.
+                    A value of 0 disables polling. Defaults to 30 seconds.
+            save_inputs_path (str):
+                    Path at which to save inputs used during inference. This will include all inputs generated by
+                    the provided data_loader, and will be saved as a JSON List[Dict[str, numpy.ndarray]].
+
+        Returns:
+            RunResults:
+                    A mapping of runner names to the results of their inference.
+                    The ordering of `runners` is preserved in this mapping.
+        """
+        warm_up = util.default(warm_up, 0)
+        data_loader = util.default(data_loader, DataLoader())
+        use_subprocess = util.default(use_subprocess, False)
+        subprocess_polling_interval = util.default(subprocess_polling_interval, 30)
+        loader_cache = DataLoaderCache(data_loader, save_inputs_path=save_inputs_path)
+
+        def execute_runner(runner, loader_cache):
+            with runner as active_runner:
+                # DataLoaderCache will ensure that the feed_dict does not contain any extra entries
+                # based on the provided input_metadata.
+                loader_cache.set_input_metadata(
+                    active_runner.get_input_metadata(use_numpy_dtypes=False)
+                )
+
+                if warm_up:
+                    G_LOGGER.start(
+                        f"{active_runner.name:35} | Running {warm_up} warm-up run(s)"
+                    )
+                    try:
+                        feed_dict = loader_cache[0]
+                    except IndexError:
+                        G_LOGGER.warning(
+                            f"{warm_up} warm-up run(s) were requested, but data loader did not supply any data. Skipping warm-up run(s)"
+                        )
+                    else:
+                        G_LOGGER.ultra_verbose(
+                            f"Warm-up Input Buffers:\n{util.indent_block(feed_dict)}"
+                        )
+                        # First do a few warm-up runs, and don't time them.
+                        for _ in range(warm_up):
+                            active_runner.infer(feed_dict=feed_dict)
+                    G_LOGGER.finish(
+                        f"{active_runner.name:35} | Finished {warm_up} warm-up run(s)"
+                    )
+
+                # Then, actual iterations.
+                index = 0
+                iteration_results = []
+                iterations_num = len(loader_cache)
+                total_runtime = 0
+                for index, feed_dict in enumerate(loader_cache):
+                    G_LOGGER.info(
+                        f"{active_runner.name:35}\n---- Inference Input(s) ----\n{TensorMetadata().from_feed_dict(feed_dict)}",
+                        mode=LogMode.ONCE,
+                    )
+
+                    G_LOGGER.extra_verbose(
+                        lambda: f"{active_runner.name:35} | Feeding inputs:\n{util.indent_block(dict(feed_dict))}"
+                    )
+                    outputs = active_runner.infer(feed_dict=feed_dict)
+
+                    runtime = active_runner.last_inference_time()
+                    total_runtime += runtime
+
+                    # Only make a deep copy if we have more than one iteration.
+                    # For single iteration case, we can use the outputs directly since they won't be reused.
+                    # This allows running with a large number of outputs (e.g. for accuracy debugging) without memory explosion.
+                    iteration_results.append(
+                        IterationResult(
+                            outputs=copy.deepcopy(outputs) if iterations_num > 1 else outputs,
+                            runtime=runtime,
+                            runner_name=active_runner.name,
+                        )
+                    )
+
+                    G_LOGGER.info(
+                        f"{active_runner.name:35}\n---- Inference Output(s) ----\n{TensorMetadata().from_feed_dict(outputs)}",
+                        mode=LogMode.ONCE,
+                    )
+                    G_LOGGER.extra_verbose(
+                        lambda: f"{active_runner.name:35} | Inference Time: {runtime * 1000.0:.3f} ms | Received outputs:\n{util.indent_block(dict(outputs))}"
+                    )
+
+                total_runtime_ms = total_runtime * 1000.0
+                G_LOGGER.finish(
+                    f"{active_runner.name:35} | Completed {index + 1} iteration(s) in {total_runtime_ms:.4g} ms | Average inference time: {total_runtime_ms / float(index + 1):.4g} ms."
+                )
+                return iteration_results
+
+        # Wraps execute_runner to use a queue.
+        def execute_runner_with_queue(runner_queue, runner, loader_cache):
+            iteration_results = None
+            try:
+                iteration_results = execute_runner(runner, loader_cache)
+            except:
+                # Cannot necessarily send the exception back over the queue.
+                G_LOGGER.backrace()
+            util.try_send_on_queue(runner_queue, iteration_results)
+            # After finishing, send the updated loader_cache back.
+            util.try_send_on_queue(runner_queue, loader_cache)
+
+        # Do all inferences in one loop, then comparisons at a later stage.
+        # We run each runner in a separate process so that we can provide exclusive GPU access for each runner.
+        run_results = RunResults()
+
+        if not runners:
+            G_LOGGER.warning(
+                "No runners were provided to Comparator.run(). Inference will not be run, and run results will be empty."
+            )
+
+        for runner in runners:
+            G_LOGGER.start(f"{runner.name:35} | Activating and starting inference")
+            if use_subprocess:
+                runner_queue = Queue()
+                process = Process(
+                    target=execute_runner_with_queue,
+                    args=(runner_queue, runner, loader_cache),
+                )
+                process.start()
+
+                # If a subprocess hangs in a certain way, then process.join could block forever. Hence,
+                # we need to keep polling the process to make sure it really is alive.
+                iteration_results = None
+                while process.is_alive() and iteration_results is None:
+                    try:
+                        iteration_results = util.try_receive_on_queue(
+                            runner_queue, timeout=subprocess_polling_interval / 2
+                        )
+                        # Receive updated loader cache, or fall back if it could not be sent.
+                        loader_cache = util.try_receive_on_queue(
+                            runner_queue, timeout=subprocess_polling_interval / 2
+                        )
+                    except queue.Empty:
+                        G_LOGGER.extra_verbose("Polled subprocess - still running")
+
+                try:
+                    assert iteration_results is not None
+                    run_results.append((runner.name, iteration_results))
+                    process.join(subprocess_timeout)
+                except:
+                    G_LOGGER.critical(
+                        f"{runner.name:35} | Terminated prematurely. Check the exception logged above. If there is no exception logged above, make sure not to use the --use-subprocess flag or set use_subprocess=False in Comparator.run()."
+                    )
+                finally:
+                    process.terminate()
+
+                if loader_cache is None:
+                    G_LOGGER.critical(
+                        "Could not send data loader cache to runner subprocess. Please try disabling subprocesses "
+                        "by removing the --use-subprocess flag, or setting use_subprocess=False in Comparator.run()"
+                    )
+            else:
+                run_results.append((runner.name, execute_runner(runner, loader_cache)))
+
+        G_LOGGER.verbose(f"Successfully ran: {[r.name for r in runners]}")
+        return run_results
+
+    @staticmethod
+    def postprocess(run_results, postprocess_func):
+        """
+        Applies post processing to all the outputs in the provided run results.
+        This is a convenience function to avoid the need for manual iteration over the run_results dictionary.
+
+        Args:
+            run_results (RunResults): The result of Comparator.run().
+            postprocess_func (Callable(IterationResult) -> IterationResult):
+                    The function to apply to each ``IterationResult``.
+
+        Returns:
+            RunResults: The updated run results.
+        """
+        G_LOGGER.start(
+            f"Applying post-processing to outputs: {postprocess_func.__name__}"
+        )
+        for _, iteration_results in run_results:
+            for index, iter_res in enumerate(iteration_results):
+                iteration_results[index] = postprocess_func(iter_res)
+        G_LOGGER.finish("Finished applying post-processing")
+        return run_results
+
+    @staticmethod
+    def default_comparisons(run_results):
+        # Sets up default comparisons - which is to compare each runner to the subsequent one.
+        return [(i, i + 1) for i in range(len(run_results) - 1)]
+
+    @staticmethod
+    def compare_accuracy(
+        run_results, fail_fast=False, comparisons=None, compare_func=None
+    ):
+        """
+        Args:
+            run_results (RunResults): The result of Comparator.run()
+
+
+            fail_fast (bool): Whether to exit after the first failure
+            comparisons (List[Tuple[int, int]]):
+                    Comparisons to perform, specified by runner indexes. For example, [(0, 1), (1, 2)]
+                    would compare the first runner with the second, and the second with the third.
+                    By default, this compares each result to the subsequent one.
+            compare_func (Callable(IterationResult, IterationResult) -> OrderedDict[str, bool]):
+                    A function that takes in two IterationResults, and returns a dictionary that maps output
+                    names to a boolean (or anything convertible to a boolean) indicating whether outputs matched.
+                    The order of arguments to this function is guaranteed to be the same as the ordering of the
+                    tuples contained in `comparisons`.
+
+        Returns:
+            AccuracyResult:
+                    A summary of the results of the comparisons. The order of the keys (i.e. runner pairs) is
+                    guaranteed to be the same as the order of `comparisons`. For more details, see the AccuracyResult
+                    docstring (e.g. help(AccuracyResult)).
+        """
+
+        def find_mismatched(match_dict):
+            return [name for name, matched in match_dict.items() if not bool(matched)]
+
+        compare_func = util.default(compare_func, CompareFunc.simple())
+        comparisons = util.default(
+            comparisons, Comparator.default_comparisons(run_results)
+        )
+
+        accuracy_result = AccuracyResult()
+        for runner0_index, runner1_index in comparisons:
+            (runner0_name, results0), (runner1_name, results1) = (
+                run_results[runner0_index],
+                run_results[runner1_index],
+            )
+
+            G_LOGGER.start(f"Accuracy Comparison | {runner0_name} vs. {runner1_name}")
+            with G_LOGGER.indent():
+                runner_pair = (runner0_name, runner1_name)
+                accuracy_result[runner_pair] = []
+
+                num_iters = min(len(results0), len(results1))
+                for iteration, (result0, result1) in enumerate(zip(results0, results1)):
+                    if num_iters > 1:
+                        G_LOGGER.info(f"Iteration: {iteration}")
+                    with contextlib.ExitStack() as stack:
+                        if num_iters > 1:
+                            stack.enter_context(G_LOGGER.indent())
+                        iteration_match_dict = compare_func(result0, result1)
+                        accuracy_result[runner_pair].append(iteration_match_dict)
+
+                        mismatched_outputs = find_mismatched(iteration_match_dict)
+                        if fail_fast and mismatched_outputs:
+                            return accuracy_result
+
+                G_LOGGER.extra_verbose(
+                    f"Finished comparing {runner0_name} with {runner1_name}"
+                )
+
+            passed, _, total = accuracy_result.stats(runner_pair)
+            pass_rate = accuracy_result.percentage(runner_pair) * 100.0
+
+            msg = f"Accuracy Summary | {runner0_name} vs. {runner1_name} | Passed: {passed}/{total} iterations | Pass Rate: {pass_rate}%"
+            if passed == total:
+                G_LOGGER.finish(msg)
+            else:
+                G_LOGGER.error(msg)
+
+        return accuracy_result
+
+    @staticmethod
+    def validate(run_results, check_inf=None, check_nan=None, fail_fast=None):
+        """
+        Checks output validity.
+
+        Args:
+            run_results (Dict[str, List[IterationResult]]): The result of Comparator.run().
+            check_inf (bool): Whether to fail on Infs. Defaults to False.
+            check_nan (bool): Whether to fail on NaNs. Defaults to True.
+            fail_fast (bool): Whether to fail after the first invalid value. Defaults to False.
+
+        Returns:
+            bool: True if all outputs were valid, False otherwise.
+        """
+        check_inf = util.default(check_inf, False)
+        check_nan = util.default(check_nan, True)
+        fail_fast = util.default(fail_fast, False)
+
+        def is_finite(output):
+            non_finite = util.array.logical_not(util.array.isfinite(output))
+            if util.array.any(non_finite):
+                G_LOGGER.error(
+                    "Inf Detected | One or more non-finite values were encountered in this output"
+                )
+                G_LOGGER.info(
+                    "Note: Use -vv or set logging verbosity to EXTRA_VERBOSE to display non-finite values",
+                    mode=LogMode.ONCE,
+                )
+                G_LOGGER.extra_verbose(f"Note: non-finite values at:\n{non_finite}")
+                G_LOGGER.extra_verbose(
+                    f"Note: non-finite values:\n{output[non_finite]}"
+                )
+                return False
+            return True
+
+        def is_not_nan(output):
+            nans = util.array.isnan(output)
+            if util.array.any(nans):
+                G_LOGGER.error(
+                    "NaN Detected | One or more NaNs were encountered in this output"
+                )
+                G_LOGGER.info(
+                    "Note: Use -vv or set logging verbosity to EXTRA_VERBOSE to display locations of NaNs",
+                    mode=LogMode.ONCE,
+                )
+                G_LOGGER.extra_verbose(f"Note: NaNs at:\n{nans}")
+                return False
+            return True
+
+        def validate_output(runner_name, output_name, output):
+            G_LOGGER.start(
+                f"{runner_name:35} | Validating output: {output_name} (check_inf={check_inf}, check_nan={check_nan})"
+            )
+            with G_LOGGER.indent():
+                comp_util.log_output_stats(output)
+
+                output_valid = True
+                if check_nan:
+                    output_valid &= is_not_nan(output)
+                if check_inf:
+                    output_valid &= is_finite(output)
+
+                if output_valid:
+                    G_LOGGER.finish(f"PASSED | Output: {output_name} is valid")
+                else:
+                    G_LOGGER.error(f"FAILED | Errors detected in output: {output_name}")
+                return output_valid
+
+        all_valid = True
+        G_LOGGER.start(f"Output Validation | Runners: {list(run_results.keys())}")
+        with G_LOGGER.indent():
+            for runner_name, results in run_results:
+                for result in results:
+                    for output_name, output in result.items():
+                        all_valid &= validate_output(runner_name, output_name, output)
+                        if fail_fast and not all_valid:
+                            return False
+
+            if all_valid:
+                G_LOGGER.finish("PASSED | Output Validation")
+            else:
+                G_LOGGER.error("FAILED | Output Validation")
+
+        return all_valid

tools/Polygraphy/polygraphy/comparator/data_loader.py

@@ -0,0 +1,569 @@
+#
+# 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
+from collections import OrderedDict
+
+from polygraphy import constants, func, mod, util
+from polygraphy.comparator.struct import RunResults
+from polygraphy.datatype import DataType
+from polygraphy.exception import DataTypeConversionException, PolygraphyException
+from polygraphy.json import save_json
+from polygraphy.logger import G_LOGGER, LogMode
+
+np = mod.lazy_import("numpy")
+torch = mod.lazy_import("torch")
+
+
+class ArraySampler:
+    def __init__(self, data_loader_backend_module, seed):
+        """
+        Args:
+            data_loader_backend_module (str):
+                    The module specifying the array type to use to generate arrays.
+                    Can be either "numpy" or "torch".
+            seed (int):
+                    The seed to use when generating random inputs.
+        """
+        self.rng = None
+        VALID_ARRAY_MODULES = ["numpy", "torch"]
+        if data_loader_backend_module not in VALID_ARRAY_MODULES:
+            G_LOGGER.critical(
+                f"Invalid `data_loader_backend_module`. Note: got: {data_loader_backend_module} but valid modules are: {VALID_ARRAY_MODULES}"
+            )
+
+        self.data_loader_backend_module = data_loader_backend_module
+
+        if self.data_loader_backend_module == "numpy":
+            self.rng = np.random.RandomState(seed)
+        elif self.data_loader_backend_module == "torch":
+            self.rng = torch.Generator()
+            self.rng.manual_seed(seed)
+
+    def sample_integer(self, shape, dtype, low, high):
+        """
+        Samples an array containing integral values in the range [low, high], inclusive
+        """
+        dtype = (
+            DataType.to_dtype(
+                DataType.from_dtype(dtype), self.data_loader_backend_module
+            )
+            if dtype is not None
+            else dtype
+        )
+        if self.data_loader_backend_module == "numpy":
+            return np.array(
+                self.rng.randint(low=low, high=high + 1, size=shape, dtype=dtype)
+            )
+        elif self.data_loader_backend_module == "torch":
+            return torch.randint(low, high + 1, shape, generator=self.rng, dtype=dtype)
+
+    def sample_float(self, shape, dtype, fmin, fmax):
+        """
+        Samples an array containing float values in the range [fmin, fmax], inclusive
+        """
+        # Special handling for infinite lower/upper bounds
+        # Without this, two infinities will collapse into a NaN, resulting in no infinities
+        # in the final output.
+        scale = fmax - fmin
+        shift = fmin
+        if util.is_inf(fmin):
+            scale = fmin
+            shift = 0
+        if util.is_inf(fmax):
+            scale = fmax
+
+        dtype = (
+            DataType.to_dtype(
+                DataType.from_dtype(dtype), self.data_loader_backend_module
+            )
+            if dtype is not None
+            else dtype
+        )
+        if self.data_loader_backend_module == "numpy":
+            return np.array(
+                (self.rng.random_sample(size=shape) * scale + shift).astype(dtype)
+            )
+        elif self.data_loader_backend_module == "torch":
+            return torch.rand(shape, generator=self.rng, dtype=dtype)
+
+    def constant_array(self, shape, dtype):
+        dtype = (
+            DataType.to_dtype(
+                DataType.from_dtype(dtype), self.data_loader_backend_module
+            )
+            if dtype is not None
+            else dtype
+        )
+        if self.data_loader_backend_module == "numpy":
+            return np.array(shape, dtype=dtype)
+        elif self.data_loader_backend_module == "torch":
+            return torch.tensor(shape, dtype=dtype)
+
+
+@mod.export()
+class DataLoader:
+    """
+    Generates synthetic input data.
+    """
+
+    def __init__(
+        self,
+        seed=None,
+        iterations=None,
+        input_metadata=None,
+        int_range=None,
+        float_range=None,
+        val_range=None,
+        data_loader_backend_module=None,
+    ):
+        """
+        Args:
+            seed (int):
+                    The seed to use when generating random inputs.
+                    Defaults to ``util.constants.DEFAULT_SEED``.
+            iterations (int):
+                    The number of iterations for which to supply data.
+                    Defaults to 1.
+            input_metadata (TensorMetadata):
+                    A mapping of input names to their corresponding shapes and data types.
+                    This will be used to determine what shapes to supply for inputs with dynamic shape, as
+                    well as to set the data type of the generated inputs.
+                    If either dtype or shape are None, then the value will be automatically determined.
+                    For input shape tensors, i.e. inputs whose *value* describes a shape in the model, the
+                    provided shape will be used to populate the values of the inputs, rather than to determine
+                    their shape.
+            val_range (Union[Tuple[number], Dict[str, Tuple[number]]]):
+                    A tuple containing exactly 2 numbers, indicating the minimum and maximum values (inclusive)
+                    the data loader should generate.
+                    If either value in the tuple is None, the default will be used for that value.
+                    If None is provided instead of a tuple, then the default values will be used for both the
+                    minimum and maximum.
+                    This can be specified on a per-input basis using a dictionary. In that case,
+                    use an empty string ("") as the key to specify default range for inputs not explicitly listed.
+                    Defaults to (0.0, 1.0).
+            data_loader_backend_module (str):
+                    A string denoting what module to use to construct the input data arrays. Currently supports
+                    "numpy" and "torch".
+                    Defaults to "numpy".
+
+            int_range (Tuple[int]):
+                    [DEPRECATED - Use val_range instead]
+                    A tuple containing exactly 2 integers, indicating the minimum and maximum integer values (inclusive)
+                    the data loader should generate. If either value in the tuple is None, the default will be used
+                    for that value.
+                    If None is provided instead of a tuple, then the default values will be used for both the
+                    minimum and maximum.
+            float_range (Tuple[float]):
+                    [DEPRECATED - Use val_range instead]
+                    A tuple containing exactly 2 floats, indicating the minimum and maximum float values (inclusive)
+                    the data loader should generate. If either value in the tuple is None, the default will be used
+                    for that value.
+                    If None is provided instead of a tuple, then the default values will be used for both the
+                    minimum and maximum.
+        """
+
+        def default_tuple(tup, default):
+            if tup is None or (
+                not isinstance(tup, tuple) and not isinstance(tup, list)
+            ):
+                return default
+            new_tup = []
+            for elem, default_elem in zip(tup, default):
+                new_tup.append(util.default(elem, default_elem))
+            return tuple(new_tup)
+
+        self.seed = util.default(seed, constants.DEFAULT_SEED)
+        self.iterations = util.default(iterations, 1)
+        self.user_input_metadata = util.default(input_metadata, {})
+        self.data_loader_backend_module = util.default(
+            data_loader_backend_module, "numpy"
+        )
+
+        self._int_range_set = int_range is not None
+        if self._int_range_set:
+            mod.warn_deprecated(
+                "The int_range parameter in DataLoader", "val_range", remove_in="0.50.0"
+            )
+        self._int_range = default_tuple(int_range, (1, 25))
+
+        self._float_range_set = float_range is not None
+        if self._float_range_set:
+            mod.warn_deprecated(
+                "The float_range parameter in DataLoader",
+                "val_range",
+                remove_in="0.50.0",
+            )
+        self._float_range = default_tuple(float_range, (-1.0, 1.0))
+
+        self.input_metadata = None
+        self.default_val_range = default_tuple(val_range, (0.0, 1.0))
+        self.val_range = util.default(val_range, self.default_val_range)
+
+        if self.user_input_metadata:
+            G_LOGGER.info(
+                f"Will generate inference input data according to provided TensorMetadata: {self.user_input_metadata}"
+            )
+
+    def __repr__(self):
+        return util.make_repr(
+            "DataLoader",
+            seed=self.seed,
+            iterations=self.iterations,
+            input_metadata=self.user_input_metadata or None,
+            int_range=self._int_range,
+            float_range=self._float_range,
+            val_range=self.val_range,
+            data_loader_backend_module=self.data_loader_backend_module,
+        )[0]
+
+    def _get_range(self, name, cast_type):
+        if cast_type == int and self._int_range_set:
+            return self._int_range
+        elif cast_type == float and self._float_range_set:
+            return self._float_range
+
+        tup = util.value_or_from_dict(self.val_range, name, self.default_val_range)
+        return tuple(cast_type(val) for val in tup)
+
+    def __getitem__(self, index):
+        """
+        Generates random input data.
+
+        May update the DataLoader's `input_metadata` attribute.
+
+        Args:
+            index (int):
+                    Since this class behaves like an iterable, it takes an index parameter.
+                    Generated data is guaranteed to be the same for the same index.
+
+        Returns:
+            OrderedDict[str, Union[numpy.ndarray, torch.Tensor]]: A mapping of input names to input numpy buffers.
+        """
+        if index >= self.iterations:
+            raise IndexError()
+
+        G_LOGGER.verbose(f"Generating data using numpy seed: {self.seed + index}")
+
+        array_sampler = ArraySampler(self.data_loader_backend_module, self.seed + index)
+
+        def get_static_shape(name, shape):
+            static_shape = shape
+            if util.is_shape_dynamic(shape):
+                if shape.min is not None:
+                    static_shape = shape.min
+                elif shape.max is not None:
+                    static_shape = shape.max
+                else:
+                    static_shape = util.override_dynamic_shape(shape)
+
+                if static_shape != shape:
+                    if not util.is_valid_shape_override(static_shape, shape):
+                        G_LOGGER.critical(
+                            f"Input tensor: {name} | Cannot override original shape: {shape} to {static_shape}"
+                        )
+                    G_LOGGER.warning(
+                        f"Input tensor: {name} [shape={shape}] | Will generate data of shape: {static_shape}.\n"
+                        f"If this is incorrect, please provide a custom data loader.",
+                        mode=LogMode.ONCE,
+                    )
+            return static_shape
+
+        # Whether the user provided the values for a shape tensor input,
+        # rather than the shape of the input.
+        # If the shape is 1D, and has a value equal to the rank of the provided default shape, it is
+        # likely to be a shape tensor, and so its value, not shape, should be overriden.
+        # Note that this is a hack needed for older versions of TensorRT. Ideally, we wouldn't care
+        # whether the input is a shape tensor or not.
+        def is_shape_tensor(name, dtype):
+            if name not in self.input_metadata or name not in self.user_input_metadata:
+                return False
+
+            _, shape = self.input_metadata[name]
+            if (
+                (dtype is not None and not DataType.from_dtype(dtype).is_integral)
+                or util.is_shape_dynamic(shape)
+                or len(shape) != 1
+            ):
+                return False
+
+            user_shape = self.user_input_metadata[name].shape
+            # Shape of shape cannot be dynamic.
+            return not util.is_shape_dynamic(user_shape) and len(user_shape) == shape[0]
+
+        def generate_buffer(name, dtype, shape):
+            if is_shape_tensor(name, dtype):
+                buffer = array_sampler.constant_array(shape, dtype)
+                G_LOGGER.info(
+                    f"Assuming {name} is a shape tensor. Setting input values to: {buffer}. "
+                    "If these values are not correct, please set it correctly in 'input_metadata' or by providing --input-shapes",
+                    mode=LogMode.ONCE,
+                )
+            elif dtype is not None and (
+                DataType.from_dtype(dtype).is_integral
+                or DataType.from_dtype(dtype) == DataType.BOOL
+            ):
+                imin, imax = self._get_range(
+                    name,
+                    cast_type=int if DataType.from_dtype(dtype).is_integral else bool,
+                )
+                G_LOGGER.verbose(
+                    f"Input tensor: {name} | Generating input data in range: [{imin}, {imax}]",
+                    mode=LogMode.ONCE,
+                )
+                # high is 1 greater than the max int drawn.
+                buffer = array_sampler.sample_integer(shape, dtype, imin, imax)
+            else:
+                fmin, fmax = self._get_range(name, cast_type=float)
+                G_LOGGER.verbose(
+                    f"Input tensor: {name} | Generating input data in range: [{fmin}, {fmax}]",
+                    mode=LogMode.ONCE,
+                )
+
+                buffer = array_sampler.sample_float(shape, dtype, fmin, fmax)
+
+            return buffer
+
+        if self.input_metadata is None and self.user_input_metadata is not None:
+            self.input_metadata = self.user_input_metadata
+
+        buffers = OrderedDict()
+        # Expand wildcards to inputs in user_input_metadata
+        inp_to_user_inp, unmatched_user_inps = util.match_keys(list(self.user_input_metadata), list(self.input_metadata))
+        # Warn about unused metadata
+        if unmatched_user_inps:
+            for name in unmatched_user_inps:
+                if util.contains_wildcard(name):
+                    msg = f"Input tensor wildcard: {name} | Metadata was provided, but none of the matched inputs exist in one or more runners."
+                    G_LOGGER.warning(msg)
+                else:
+                    msg = f"Input tensor: {name} | Metadata was provided, but the input does not exist in one or more runners."
+                    close_match = util.find_str_in_iterable(
+                        name, self.input_metadata.keys()
+                    )
+                    if close_match:
+                        msg += f"\nMaybe you meant to set: {close_match}?"
+                    G_LOGGER.warning(msg)
+
+        for name, (dtype, shape) in self.input_metadata.items():
+            try:
+                dtype = (
+                    DataType.to_dtype(
+                        DataType.from_dtype(dtype), self.data_loader_backend_module
+                    )
+                    if dtype is not None
+                    else None
+                )
+            except DataTypeConversionException:
+                G_LOGGER.critical(
+                    f"Could not convert data type: {dtype} to {self.data_loader_backend_module}, so the default data loader cannot generate a {self.data_loader_backend_module} array for input: {name}. "
+                    f"Please use a custom data loader to provide inputs. "
+                )
+            if name in inp_to_user_inp:
+                user_dtype, user_shape = self.user_input_metadata[inp_to_user_inp[name]]
+
+                dtype = util.default(user_dtype, dtype)
+                dtype = (
+                    DataType.to_dtype(
+                        DataType.from_dtype(dtype), self.data_loader_backend_module
+                    )
+                    if dtype is not None
+                    else None
+                )
+                is_valid_shape_override = (
+                    user_shape is not None
+                    and util.is_valid_shape_override(user_shape, shape)
+                )
+
+                if util.is_shape_dynamic(user_shape):
+                    G_LOGGER.warning(
+                        f"Input tensor: {name} [shape={shape}] | Provided input shape: {user_shape} is dynamic.\nDynamic shapes cannot be used to generate inference data. Will use default shape instead.\nTo avoid this, please provide a fixed shape to the data loader. "
+                    )
+                elif not is_valid_shape_override and not is_shape_tensor(name, dtype):
+                    G_LOGGER.warning(
+                        f"Input tensor: {name} [shape={shape}] | Cannot use provided custom shape: {user_shape} to override tensor shape. Will use default shape instead.",
+                        mode=LogMode.ONCE,
+                    )
+                else:
+                    shape = util.default(user_shape, shape)
+
+            static_shape = get_static_shape(name, shape)
+            buffers[name] = generate_buffer(name, dtype, shape=static_shape)
+
+        # Warn about unused val_range
+        if not isinstance(self.val_range, tuple):
+            util.check_sequence_contains(
+                self.val_range.keys(),
+                list(self.input_metadata.keys()) + [""],
+                name="val_range",
+                log_func=G_LOGGER.warning,
+                check_missing=False,
+            )
+
+        return buffers
+
+
+# Caches data loaded by a DataLoader for use across multiple runners.
+class DataLoaderCache:
+    def __init__(self, data_loader, save_inputs_path=None):
+        self.data_loader = data_loader
+        self.cache = []  # List[OrderedDict[str, numpy.ndarray]]
+        self.save_inputs_path = save_inputs_path
+
+    @func.constantmethod
+    def __len__(self):
+        return len(self.cache)
+
+    @func.constantmethod
+    def __getitem__(self, iteration):
+        """
+        Load the specified iteration from the cache if present, or load it from the data loader.
+
+        Args:
+            iteration (int): The iteration whose data to retrieve.
+        """
+        if iteration >= len(self.cache):
+            raise IndexError()
+
+        # Attempts to match existing input buffers to the requested input_metadata
+        def coerce_cached_input(index, name, dtype, shape):
+            cached_feed_dict = self.cache[iteration]
+            cached_name = util.find_str_in_iterable(
+                name, cached_feed_dict.keys(), index
+            )
+            if cached_name is None:
+                G_LOGGER.critical(
+                    f"Input tensor: {name} | Does not exist in the data loader cache."
+                )
+
+            if cached_name != name:
+                G_LOGGER.warning(
+                    f"Input tensor: {name} | Buffer name ({cached_name}) does not match expected input name ({name})."
+                )
+
+            buffer = cached_feed_dict[cached_name]
+
+            if dtype != util.array.dtype(buffer):
+                G_LOGGER.warning(
+                    f"Input tensor: {name} | Buffer dtype ({util.array.dtype(buffer)}) does not match expected input dtype ({dtype}), attempting to cast. "
+                )
+
+                try:
+                    np_type = DataType.to_dtype(dtype, "numpy")
+                except:
+                    pass
+                else:
+                    type_info = None
+                    if dtype.is_integral:
+                        type_info = np.iinfo(np_type)
+                    elif dtype.is_floating:
+                        type_info = np.finfo(np_type)
+
+                    if type_info is not None and util.array.any(
+                        (buffer < type_info.min) | (buffer > type_info.max)
+                    ):
+                        G_LOGGER.warning(
+                            f"Some values in this input are out of range of {dtype}. Unexpected behavior may ensue!"
+                        )
+                buffer = util.array.cast(buffer, dtype)
+
+            if not util.is_valid_shape_override(util.array.shape(buffer), shape):
+                G_LOGGER.warning(
+                    f"Input tensor: {name} | Buffer shape ({util.array.shape(buffer)}) does not match expected input shape ({shape}). "
+                    f"Attempting to transpose/reshape. "
+                )
+                buffer = util.try_match_shape(buffer, shape)
+
+            if util.array.dtype(buffer) != dtype or not util.is_valid_shape_override(
+                util.array.shape(buffer), shape
+            ):
+                G_LOGGER.critical(
+                    f"Input tensor: {name} | Cannot reuse input data due to mismatch in shape or data type.\n"
+                    f"Note: Cached input: [dtype={util.array.dtype(buffer)}, shape={util.array.shape(buffer)}], "
+                    f"Requested input: [dtype={dtype}, shape={shape}]"
+                )
+            return buffer
+
+        feed_dict = OrderedDict()
+
+        # Reload from data loader if needed
+        data_loader_feed_dict = None
+
+        for index, (name, (dtype, shape)) in enumerate(self.input_metadata.items()):
+            try:
+                buffer = coerce_cached_input(
+                    index, name, DataType.from_dtype(dtype), shape
+                )
+            except PolygraphyException:
+                G_LOGGER.warning(
+                    f"Could not use buffer previously cached from data loader for input: {name}. Attempting to reload inputs from the data loader.\nNote that this will only work if the data loader supports random access.\nPlease refer to warnings above for details on why the previously generated input buffer didn't work. "
+                )
+                try:
+                    if data_loader_feed_dict is None:
+                        data_loader_feed_dict = self.data_loader[iteration]
+                    buffer = data_loader_feed_dict[name]
+                except:
+                    G_LOGGER.critical(
+                        "Could not reload inputs from data loader. Are the runners running the same model? "
+                        "If not, please rewrite the data loader to support random access."
+                    )
+            feed_dict[name] = buffer
+
+        return feed_dict
+
+    def set_input_metadata(self, input_metadata):
+        """
+        Set the input metadata for the data loader.
+
+        Args:
+            input_metadata (TensorMetadata):
+                    Input Metadata, including shape and type information. The cache may attempt to transform inputs to
+                    match the specified input_metadata when data already in the cache does not exactly match.
+        """
+        self.input_metadata = input_metadata
+        with contextlib.suppress(AttributeError):
+            self.data_loader.input_metadata = input_metadata
+
+        if not self.cache:
+            G_LOGGER.verbose("Loading inputs from data loader")
+            self.cache = list(self.data_loader)
+
+            def _is_feed_dict(inp):
+                if isinstance(inp, RunResults):
+                    return False
+
+                try:
+                    for name, _ in inp.items():
+                        if not isinstance(name, str):
+                            return False
+                except:
+                    return False
+                else:
+                    return True
+
+            if not self.cache:
+                G_LOGGER.warning("Data loader did not yield any input data.")
+            elif not _is_feed_dict(self.cache[0]):
+                G_LOGGER.critical(
+                    f"Data loader returned an object that cannot be recognized as a feed_dict (Dict[str, Union[np.ndarray, torch.Tensor, DeviceView]]):"
+                    f"\nNote: The object was:\n{self.cache[0]}.\n"
+                    f"\nHint: If this is a `RunReults` object (e.g. generated with `--save-outputs`), try using the "
+                    f"`data to-input` tool to convert it to a feed_dict compatible format. "
+                )
+
+            # Only save inputs the first time the cache is generated
+            if self.save_inputs_path is not None:
+                save_json(self.cache, self.save_inputs_path, "inference input data")

tools/Polygraphy/polygraphy/comparator/postprocess.py

@@ -0,0 +1,64 @@
+#
+# 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
+
+np = mod.lazy_import("numpy")
+
+
+@mod.export()
+class PostprocessFunc:
+    """
+    Provides functions that can apply post-processing to `IterationResult` s.
+    """
+
+    @staticmethod
+    # This function returns a top_k function that can be used as a postprocess_func.
+    def top_k(k=None):
+        """
+        Creates a function that applies a Top-K operation to a IterationResult.
+        Top-K will return the indices of the k largest values in the array.
+
+        Args:
+            k (Union[int, Tuple[int, int], Dict[str, int], Dict[str, Tuple[int, int]]]):
+                    The number of indices to keep and optionally the axis on which to operate.
+                    For example, a value of ``(5, 0)`` would keep the top 5 indices along axis 0.
+
+                    If this exceeds the axis length, it will be clamped.
+                    This can be specified on a per-output basis by providing a dictionary. In that case,
+                    use an empty string ("") as the key to specify default top-k value for outputs not explicitly listed.
+                    If no default is present, unspecified outputs will not be modified.
+                    Defaults to 10.
+
+        Returns:
+            Callable(IterationResult) -> IterationResult: The top-k function.
+        """
+        k = util.default(k, 10)
+        axis = -1
+
+        # Top-K implementation.
+        def top_k_impl(iter_result):
+            for name, output in iter_result.items():
+                k_val = util.value_or_from_dict(k, name)
+                if k_val:
+                    nonlocal axis
+                    if util.is_sequence(k_val):
+                        k_val, axis = k_val
+
+                    iter_result[name] = util.array.topk(output, k_val, axis)[1]
+            return iter_result
+
+        return top_k_impl

tools/Polygraphy/polygraphy/comparator/struct.py

@@ -0,0 +1,409 @@
+#
+# 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 OrderedDict
+
+from polygraphy import config, mod, util
+from polygraphy.common.interface import TypedDict, TypedList
+from polygraphy.json import Decoder, Encoder, add_json_methods, load_json, save_json
+from polygraphy.logger import G_LOGGER
+
+
+class LazyArray:
+    """
+    Represents a lazily loaded NumPy array or PyTorch Tensor.
+    For example, large arrays may be serialized to temporary files on the disk
+    to save memory.
+    """
+
+    def __init__(self, arr):
+        """
+        Args:
+            arr (Union[np.ndarray, torch.Tensor]): The array.
+        """
+        self.arr = None
+        self.tmpfile = None
+        if config.ARRAY_SWAP_THRESHOLD_MB >= 0 and util.array.nbytes(arr) > (
+            config.ARRAY_SWAP_THRESHOLD_MB << 20
+        ):
+            self.tmpfile = util.NamedTemporaryFile(suffix=".json")
+            G_LOGGER.extra_verbose(
+                f"Evicting large array ({util.array.nbytes(arr) / 1024.0 ** 2:.3f} MiB) from memory and saving to {self.tmpfile.name}"
+            )
+            save_json(arr, self.tmpfile.name)
+        else:
+            self.arr = arr
+
+    def load(self):
+        """
+        Load the array, deserializing from the disk if it was stored earlier.
+
+        Returns:
+            Union[np.ndarray, torch.Tensor]: The array
+        """
+        if self.arr is not None:
+            return self.arr
+
+        if self.tmpfile is None:
+            G_LOGGER.internal_error(
+                f"self.arr is None but self.tmpfile is also None; this should be impossible."
+            )
+        return load_json(self.tmpfile.name)
+
+
+@Encoder.register(LazyArray, alias="LazyNumpyArray")
+def encode(lazy_arr):
+    return {
+        "values": lazy_arr.load(),
+    }
+
+
+@Decoder.register(LazyArray, alias="LazyNumpyArray")
+def decode(dct):
+    return LazyArray(dct["values"])
+
+
+@mod.export()
+class IterationResult(TypedDict(lambda: str, lambda: LazyArray)):
+    """
+    An ordered dictionary containing the result of a running a single iteration of a runner.
+
+    This maps output names to arrays, and preserves the output ordering from the runner.
+
+    NOTE: The ``POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB`` environment variable can be set to enable
+    the arrays to be swapped to the disk.
+
+    Also includes additional fields indicating the name of the runner which produced the
+    outputs, and the time required to do so.
+    """
+
+    @staticmethod
+    def _to_lazy(nparray):
+        if isinstance(nparray, LazyArray):
+            return nparray
+        return LazyArray(nparray)
+
+    @staticmethod
+    def _to_lazy_dict(nparray_dict):
+        if nparray_dict is None:
+            return None
+
+        # Converts a Dict[str, np.ndarray] to a Dict[str, LazyArray]
+        lazy = OrderedDict()
+        for name, out in nparray_dict.items():
+            lazy[name] = IterationResult._to_lazy(out)
+        return lazy
+
+    def __init__(self, outputs=None, runtime=None, runner_name=None):
+        """
+        Args:
+            outputs (Dict[str, Union[np.array, torch.Tensor]]): The outputs of this iteration, mapped to their names.
+
+            runtime (float):
+                    The time required for this iteration, in seconds.
+                    Only used for logging purposes.
+            runner_name (str):
+                    The name of the runner that produced this output.
+                    If this is omitted, a default name is generated.
+        """
+        if outputs and config.ARRAY_SWAP_THRESHOLD_MB < 0:
+            total_size_gb = sum(
+                util.array.nbytes(arr)
+                for arr in outputs.values()
+                if util.array.is_torch(arr) or util.array.is_numpy(arr)
+            ) / (1024.0**3)
+            if total_size_gb >= 1:
+                G_LOGGER.warning(
+                    f"It looks like the outputs of this network are very large ({total_size_gb:.3f} GiB).\n"
+                    "To reduce memory usage, you may want to allow Polygraphy to swap these arrays to the disk "
+                    "using the POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB environment variable."
+                )
+
+        super().__init__(IterationResult._to_lazy_dict(outputs))
+        self.runtime = runtime
+        self.runner_name = util.default(runner_name, "custom_runner")
+
+    # Convenience methods to preserve np.ndarray in the interface.
+    def update(self, other):
+        return super().update(IterationResult._to_lazy_dict(other))
+
+    def __setitem__(self, name, arr):
+        return super().__setitem__(name, IterationResult._to_lazy(arr))
+
+    def values(self):
+        for arr in super().values():
+            yield arr.load()
+
+    def items(self):
+        for name, arr in super().items():
+            yield name, arr.load()
+
+    def __getitem__(self, name):
+        return super().__getitem__(name).load()
+
+    def __eq__(self, other):
+        if self.runtime != other.runtime or self.runner_name != other.runner_name:
+            return False
+
+        for key, val in self.items():
+            if key not in other:
+                return False
+
+            if not util.array.equal(val, other[key]):
+                return False
+
+        return True
+
+
+@Encoder.register(IterationResult)
+def encode(iter_result):
+    return {
+        "outputs": iter_result.dct,
+        "runtime": iter_result.runtime,
+        "runner_name": iter_result.runner_name,
+    }
+
+
+@Decoder.register(IterationResult)
+def decode(dct):
+    return IterationResult(
+        outputs=dct["outputs"], runtime=dct["runtime"], runner_name=dct["runner_name"]
+    )
+
+
+@mod.export()
+@add_json_methods("inference results")
+class RunResults(TypedList(lambda: tuple)):
+    """
+    Maps runners to per-iteration outputs (in the form of a ``List[IterationResult]``).
+
+    For example, if ``results`` is an instance of ``RunResults()``, then
+    to access the outputs of the first iteration from a specified runner, do:
+    ::
+
+        iteration = 0
+        runner_name = "trt-runner"
+        outputs = results[runner_name][iteration]
+
+        # `outputs` is a `Dict[str, np.ndarray]`
+
+
+    Note: Technically, this is a ``List[Tuple[str, List[IterationResult]]]``, but includes
+    helpers that make it behave like an OrderedDict that can contain duplicates.
+    """
+
+    def items(self):
+        """
+        Creates a generator that yields ``Tuple[str, List[IterationResult]]`` - runner names
+        and corresponding outputs.
+        """
+        for name, iteration_results in self.lst:
+            yield name, iteration_results
+
+    def keys(self):
+        """
+        Creates a generator that yields runner names (str).
+        """
+        for name, _ in self.lst:
+            yield name
+
+    def values(self):
+        """
+        Creates a generator that yields runner outputs (List[IterationResult]).
+        """
+        for _, iteration_results in self.lst:
+            yield iteration_results
+
+    def update(self, other):
+        """
+        Updates the results stored in this instance.
+
+        Args:
+            other (Union[Dict[str, List[IterationResult]], RunResults]):
+                    A dictionary or RunResults instance from which to update this one.
+        """
+        for name, iteration_results in other.items():
+            self.lst[name] = iteration_results
+        return self
+
+    def add(self, out_list, runtime=None, runner_name=None):
+        """
+        A helper to create a ``List[IterationResult]`` and map it to the specified runner_name.
+
+        This method cannot be used to modify an existing entry.
+
+        Calling this method is equivalent to:
+        ::
+
+            results[runner_name] = []
+            for out in out_list:
+                results[runner_name].append(IterationResult(out, runtime, runner_name))
+
+        Args:
+            out_list (List[Dict[str, np.array]]):
+                One or more set of outputs where each output is a dictionary
+                of output names mapped to NumPy arrays.
+
+            runtime (float):
+                    The time required for this iteration, in seconds.
+                    Only used for logging purposes.
+            runner_name (str):
+                    The name of the runner that produced this output.
+                    If this is omitted, a default name is generated.
+        """
+        runner_name = util.default(runner_name, "custom_runner")
+        iter_results = [IterationResult(out, runtime, runner_name) for out in out_list]
+        self[runner_name] = iter_results
+
+    def __getitem__(self, key):
+        if isinstance(key, int):
+            return self.lst[key]
+
+        for name, iteration_results in self.lst:
+            if name == key:
+                return iteration_results
+
+        G_LOGGER.critical(
+            f"{key:35} does not exist in this RunResults instance. Note: Available runners: {list(self.keys())}"
+        )
+
+    def __setitem__(self, key, value):
+        if isinstance(key, int):
+            self.lst[key] = value
+            return
+
+        for index, name in enumerate(self.keys()):
+            if name == key:
+                self.lst[index] = (key, value)
+                break
+        else:
+            self.append((key, value))
+
+    def __contains__(self, val):
+        if isinstance(val, str) or isinstance(val, bytes):
+            return val in list(self.keys())
+        return val in self.lst
+
+    def __eq__(self, other):
+        for (r0, its0), (r1, its1) in zip(self.lst, other.lst):
+            if r0 != r1:
+                return False
+
+            if its0 != its1:
+                return False
+        return True
+
+
+@Encoder.register(RunResults)
+def encode(results):
+    return {"lst": results.lst}
+
+
+@Decoder.register(RunResults)
+def decode(dct):
+    return RunResults(list(map(tuple, dct["lst"])))
+
+
+@mod.export()
+class AccuracyResult(TypedDict(lambda: tuple, lambda: list)):
+    """
+    An ordered dictionary including details about the result of ``Comparator.compare_accuracy``.
+
+    More specifically, it is an ``OrderedDict[Tuple[str, str], List[OrderedDict[str, bool]]]`` which maps a runner
+    pair (a tuple containing both runner names) to a list of dictionaries of booleans (or anything that can be
+    converted into a boolean, such as an ``OutputCompareResult``), indicating whether there was a match in the outputs of
+    the corresponding iteration. The ``List[OrderedDict[str, bool]]`` is constructed from the dictionaries returned
+    by ``compare_func`` in ``compare_accuracy``.
+
+    For example, to see if there's a match between ``runner0`` and
+    ``runner1`` during the 1st iteration for an output called ``output0``:
+    ::
+
+        runner_pair = ("runner0", "runner1")
+        iteration = 0
+        output_name = "output0"
+        match = bool(accuracy_result[runner_pair][iteration][output_name])
+
+    If there's a mismatch, you can inspect the outputs from
+    the results of ``Comparator.run()``, assumed here to be called ``run_results``:
+    ::
+
+        runner0_output = run_results["runner0"][iteration][output_name]
+        runner1_output = run_results["runner1"][iteration][output_name]
+    """
+
+    def __bool__(self):
+        """
+        Whether all outputs matched for every iteration.
+        You can use this function to avoid manually checking each output. For example:
+        ::
+
+            if accuracy_result:
+                print("All matched!")
+
+        Returns:
+            bool
+        """
+        return all(
+            [
+                bool(match)
+                for outs in self.values()
+                for out in outs
+                for match in out.values()
+            ]
+        )
+
+    def _get_runner_pair(self, runner_pair):
+        return util.default(runner_pair, list(self.keys())[0])
+
+    def percentage(self, runner_pair=None):
+        """
+        Returns the percentage of iterations that matched for the given pair of runners,
+        expressed as a decimal between 0.0 and 1.0.
+
+        Always returns 1.0 when the number of iterations is 0, or when there are no runner comparisons.
+
+        Args:
+            runner_pair (Tuple[str, str]):
+                    A pair of runner names describing which runners to check.
+                    Defaults to the first pair in the dictionary.
+        """
+        if not list(self.keys()):
+            return 1.0  # No data in this result.
+
+        matched, _, total = self.stats(runner_pair)
+        if not total:
+            return 1.0  # No iterations
+        return float(matched) / float(total)
+
+    def stats(self, runner_pair=None):
+        """
+        Returns the number of iterations that matched, mismatched, and the total number of iterations.
+
+        Args:
+            runner_pair (Tuple[str, str]):
+                    A pair of runner names describing which runners to check.
+                    Defaults to the first pair in the dictionary.
+
+        Returns:
+            Tuple[int, int, int]: Number of iterations that matched, mismatched, and total respectively.
+        """
+        runner_pair = self._get_runner_pair(runner_pair)
+        outs = self[runner_pair]
+        matched = sum([all([match for match in out.values()]) for out in outs])
+        total = len(outs)
+        return matched, total - matched, total

tools/Polygraphy/polygraphy/comparator/util.py

@@ -0,0 +1,405 @@
+#
+# 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 functools
+import math
+import os
+
+from polygraphy import config, mod, util
+from polygraphy.logger import G_LOGGER
+from polygraphy.datatype import DataType
+
+import math
+import os
+
+np = mod.lazy_import("numpy")
+plt = mod.lazy_import("matplotlib.pyplot")
+matplotlib = mod.lazy_import("matplotlib")
+
+
+def cast_up(buffer):
+    dtype = util.array.dtype(buffer)
+
+    if dtype == DataType.FLOAT16:
+        buffer = util.array.cast(buffer, DataType.FLOAT32)
+    elif dtype in [DataType.INT8, DataType.UINT8, DataType.INT16, DataType.UINT16]:
+        buffer = util.array.cast(buffer, DataType.INT32)
+    elif dtype == DataType.UINT32:
+        buffer = util.array.cast(buffer, DataType.INT64)
+    return buffer
+
+
+def use_higher_precision(func):
+    """
+    Decorator that will cast the input numpy buffer(s) to a higher precision before computation.
+    """
+
+    @functools.wraps(func)
+    def wrapped(*buffers):
+        if any(util.is_empty_shape(util.array.shape(buffer)) for buffer in buffers):
+            return 0
+
+        new_buffers = [cast_up(buffer) for buffer in buffers]
+        return func(*new_buffers)
+
+    return wrapped
+
+
+@use_higher_precision
+def compute_max(buffer):
+    return util.array.max(buffer)
+
+
+# Returns index of max value
+@use_higher_precision
+def compute_argmax(buffer):
+    return util.array.unravel_index(util.array.argmax(buffer), util.array.shape(buffer))
+
+
+@use_higher_precision
+def compute_min(buffer):
+    return util.array.min(buffer)
+
+
+# Returns index of min value
+@use_higher_precision
+def compute_argmin(buffer):
+    return util.array.unravel_index(util.array.argmin(buffer), util.array.shape(buffer))
+
+
+def compute_mean(buffer):
+    return util.array.mean(buffer, dtype=DataType.FLOAT32)
+
+
+@use_higher_precision
+def compute_std(buffer):
+    return util.array.std(buffer)
+
+
+@use_higher_precision
+def compute_variance(buffer):
+    return util.array.var(buffer)
+
+
+@use_higher_precision
+def compute_median(buffer):
+    return util.array.median(buffer)
+
+
+def compute_quantile(buffer, q):
+    return util.array.quantile(buffer, q)
+
+
+def compute_average_magnitude(buffer):
+    return util.array.mean(util.array.abs(buffer), dtype=DataType.FLOAT32)
+
+
+def str_histogram(output, hist_range=None):
+    if util.array.dtype(output) == DataType.BOOL:
+        return ""
+
+    try:
+        try:
+            hist, bin_edges = util.array.histogram(output, range=hist_range)
+        except (ValueError, RuntimeError) as err:
+            G_LOGGER.verbose(f"Could not generate histogram. Note: Error was: {err}")
+            return ""
+
+        max_num_elems = compute_max(hist)
+        if not max_num_elems:  # Empty tensor
+            return
+
+        bin_edges = [f"{bin:.3g}" for bin in bin_edges]
+        max_start_bin_width = max(len(bin) for bin in bin_edges)
+        max_end_bin_width = max(len(bin) for bin in bin_edges[1:])
+
+        MAX_WIDTH = 40
+        ret = "---- Histogram ----\n"
+        ret += f"{'Bin Range':{max_start_bin_width + max_end_bin_width + 5}}|  Num Elems | Visualization\n"
+        for num, bin_start, bin_end in zip(hist, bin_edges, bin_edges[1:]):
+            bar = "#" * int(MAX_WIDTH * float(num) / float(max_num_elems))
+            ret += f"({bin_start:<{max_start_bin_width}}, {bin_end:<{max_end_bin_width}}) | {num:10} | {bar}\n"
+        return ret
+    except Exception as err:
+        G_LOGGER.verbose(f"Could not generate histogram.\nNote: Error was: {err}")
+        if config.INTERNAL_CORRECTNESS_CHECKS:
+            raise
+        return ""
+
+
+def str_output_stats(output, runner_name=None):
+    ret = ""
+    if runner_name:
+        ret += f"{runner_name} | Stats: "
+
+    try:
+        ret += f"mean={compute_mean(output):.5g}, std-dev={compute_std(output):.5g}, var={compute_variance(output):.5g}, median={compute_median(output):.5g}, min={compute_min(output):.5g} at {compute_argmin(output)}, max={compute_max(output):.5g} at {compute_argmax(output)}, avg-magnitude={compute_average_magnitude(output):.5g}"
+
+        # np.quantile doesn't work with boolean input, so we don't show quantile error if the output type is boolean
+        if output.dtype == bool:
+            ret += "\n"
+        else:
+            ret += f", p90={compute_quantile(output, 0.9):.5g}, p95={compute_quantile(output, 0.95):.5g}, p99={compute_quantile(output, 0.99):.5g}\n"
+    except Exception as err:
+        G_LOGGER.verbose(f"Could not generate statistics.\nNote: Error was: {err}")
+        ret += "<Error while computing statistics>"
+        if config.INTERNAL_CORRECTNESS_CHECKS:
+            raise
+    return ret
+
+
+def log_output_stats(output, info_hist=False, runner_name=None, hist_range=None):
+    ret = str_output_stats(output, runner_name)
+    G_LOGGER.info(ret)
+    with G_LOGGER.indent():
+        # For small outputs, show the entire output instead of just a histogram.
+        SMALL_OUTPUT_THRESHOLD = 100
+        if util.array.size(output) <= SMALL_OUTPUT_THRESHOLD:
+            G_LOGGER.log(
+                lambda: f"---- Values ----\n{util.indent_block(output)}",
+                severity=G_LOGGER.INFO if info_hist else G_LOGGER.VERBOSE,
+            )
+        G_LOGGER.log(
+            lambda: str_histogram(output, hist_range),
+            severity=G_LOGGER.INFO if info_hist else G_LOGGER.VERBOSE,
+        )
+
+
+def build_heatmaps(
+    arr, min_val, max_val, prefix, save_dir=None, show=None, use_lognorm=None
+):
+    """
+    Display an array as an image or set of images. The last two dimensions are interpreted as
+    the height and width and the leading dimensions are flattened and treated as the number
+    of images to display.
+
+    Args:
+        arr (Union[torch.Tensor, numpy.ndarray]): The input array or tensor.
+        min_val (float): The minimum value in the input array
+        max_val (float): The maximum value in the input array
+        prefix (str): The prefix to use when displaying titles for figures.
+        save_dir (Optional[str]): Path to a directory in which to save images of the heatmaps.
+        show (Optional[bool]): Whether to display the heatmap.
+        use_lognorm (bool): Whether to use LogNorm instead of Normalize when displaying values.
+    """
+    G_LOGGER.start(f"Building heatmaps for {prefix}. This may take a while...")
+    with G_LOGGER.indent():
+        MAX_HEIGHT = 1080
+        MAX_WIDTH = 1920
+        MAX_NUM_ROWS = 14
+        MAX_NUM_COLS = 7
+        FONT_SIZE = "xx-small"
+
+        shape = util.array.shape(arr)
+        if len(shape) < 3:
+            arr = util.array.view(
+                arr,
+                dtype=util.array.dtype(arr),
+                shape=([1] * (3 - len(shape))) + list(shape),
+            )
+
+        original_shape = util.array.shape(arr)
+        arr = util.array.view(
+            arr,
+            dtype=util.array.dtype(arr),
+            shape=(-1, original_shape[-2], original_shape[-1]),
+        )
+
+        shape = util.array.shape(arr)
+        num_images = shape[0]
+
+        def coord_str_from_img_idx(img_idx):
+            coord = []
+            for dim in reversed(original_shape[:-2]):
+                coord.insert(0, img_idx % dim)
+                img_idx //= dim
+            return f"({','.join(map(str, coord))},0:{shape[-2]},0:{shape[-1]})"
+
+        # We treat each 2D slice of the array as a separate image.
+        # Multiple images may be displayed on a single figure (in a grid) and we may have multiple figures.
+        num_rows = min(MAX_HEIGHT // shape[-2], MAX_NUM_ROWS)
+        num_cols = min(MAX_WIDTH // shape[-1], MAX_NUM_COLS)
+
+        # Remove any excess images per figure
+        if num_images < num_rows * num_cols:
+            num_cols = min(num_images, num_cols)
+            num_rows = math.ceil(num_images / num_cols)
+
+        num_images_per_figure = num_rows * num_cols
+        num_figures = math.ceil(num_images / num_images_per_figure)
+
+        # Populate each image in each figure.
+        for fig_idx in range(num_figures):
+            fig, axs = plt.subplots(
+                num_rows, num_cols, squeeze=False, dpi=200, constrained_layout=True
+            )
+            base_img_idx = fig_idx * num_images_per_figure
+
+            try:
+                # When the error is all the same, we can't use LogNorm.
+                if use_lognorm and min_val != max_val:
+                    norm = matplotlib.colors.LogNorm(vmin=min_val, vmax=max_val)
+                    prefix += " (Log Scale)"
+                else:
+                    norm = matplotlib.colors.Normalize(vmin=min_val, vmax=max_val)
+
+                fig_title = f"{prefix}: {coord_str_from_img_idx(base_img_idx)} to {coord_str_from_img_idx(min(base_img_idx + num_images_per_figure, num_images) - 1)}"
+                fig.suptitle(fig_title, fontsize=FONT_SIZE)
+
+                G_LOGGER.extra_verbose(f"Building heatmaps for {fig_title}")
+
+                images = []
+                for row in range(num_rows):
+                    for col in range(num_cols):
+                        img_idx = base_img_idx + (col + row * num_cols)
+
+                        ax = axs[row, col]
+                        ax.set_axis_off()
+
+                        if img_idx < shape[0]:
+                            img = arr[img_idx]
+                            title = f"{coord_str_from_img_idx(img_idx)}"
+                        else:
+                            img = np.zeros(shape=(shape[-2:]))
+                            title = "Out Of Bounds"
+                        ax.set_title(title, fontsize=FONT_SIZE)
+
+                        images.append(
+                            ax.imshow(
+                                img, cmap="plasma", filternorm=False, resample=False
+                            )
+                        )
+
+                for im in images:
+                    im.set_norm(norm)
+
+                fig.colorbar(images[0], ax=axs, shrink=0.7)
+
+                if save_dir is not None:
+                    path = os.path.join(
+                        save_dir, f"{util.sanitize_filename(fig_title)}.svg"
+                    )
+                    util.makedirs(path)
+                    G_LOGGER.info(f"Saving '{prefix}' heatmap to: '{path}'")
+                    fig.savefig(path)
+
+                if show:
+                    plt.show()
+            finally:
+                plt.close(fig)
+
+
+def scatter_plot_error_magnitude(
+    absdiff,
+    reldiff,
+    reference_output,
+    min_reldiff,
+    max_reldiff,
+    runner0_name,
+    runner1_name,
+    out0_name,
+    out1_name,
+    save_dir=None,
+    show=None,
+):
+    """
+    Display a plot of absolute/relative difference against the magnitude of the output.
+
+    Args:
+        absdiff (Union[torch.Tensor, numpy.ndarray]): The absolute difference.
+        reldiff (Union[torch.Tensor, numpy.ndarray]): The relative difference.
+        reference_output (Union[torch.Tensor, numpy.ndarray]): The output to consider as the reference output.
+        min_reldiff (float): The minimum relative difference
+        max_reldiff (float): The maximum relative difference
+        runner0_name (str): The name of the first runner.
+        runner1_name (str): The name of the second runner.
+        out0_name (str): The name of the output of the first runner.
+        out1_name (str): The name of the output of the second runner.
+        save_dir (Optional[str]): Path to a directory in which to save images of the plots.
+        show (Optional[bool]): Whether to display the error metrics plot.
+    """
+    G_LOGGER.start(
+        f"Building error metrics plot for {out0_name}. This may take a while..."
+    )
+    with G_LOGGER.indent():
+        title = f"Error metrics between output0 and output1\noutput0: {runner0_name:35} | {out0_name}\noutput1: {runner1_name:35} | {out1_name}"
+        fname = util.sanitize_filename(f"error_metrics_{out0_name}.png")
+        TICK_FONT_SIZE = 6
+        TITLE_FONT_SIZE = 7
+        NUM_X_TICKS = 20
+        NUM_Y_LINEAR_TICKS = 10
+
+        def set_ax_properties(ax):
+            ax.tick_params(axis="x", labelrotation=90)
+            ax.tick_params(axis="both", labelsize=TICK_FONT_SIZE)
+            ax.grid(linestyle="--")
+            ax.xaxis.label.set_fontsize(TITLE_FONT_SIZE)
+            ax.yaxis.label.set_fontsize(TITLE_FONT_SIZE)
+
+        def set_linear_ax(ax):
+            xticks = ax.get_xticks()
+            yticks = ax.get_yticks()
+            ax.set_xticks(np.linspace(0, xticks[-1], NUM_X_TICKS))
+            ax.set_yticks(np.linspace(0, yticks[-1], NUM_Y_LINEAR_TICKS))
+            set_ax_properties(ax)
+
+        def set_log_ax(ax, min_diff, max_diff):
+            ax.set_yscale("log")
+            xticks = ax.get_xticks()
+
+            # Add a very small epsilon to prevent division by 0:
+            eps = 1e-15
+            yrange = np.log10(np.array([min_diff + eps, max_diff + eps]))
+            yrange[0] = math.floor(yrange[0])
+            yrange[1] = math.ceil(yrange[1])
+
+            ax.set_xticks(np.linspace(0, xticks[-1], NUM_X_TICKS))
+            ax.set_yticks(np.power(10, np.arange(yrange[0], yrange[1], 1)))
+            set_ax_properties(ax)
+
+        magnitude = util.array.abs(reference_output)
+        fig, axs = plt.subplots(2, sharex=True, constrained_layout=True)
+
+        try:
+            fig.suptitle(title, fontsize=TITLE_FONT_SIZE)
+
+            axs[0].scatter(magnitude, absdiff, s=1)
+            axs[0].set(ylabel="Absolute error")
+            set_linear_ax(axs[0])
+
+            axs[1].scatter(magnitude, reldiff, s=1)
+            label_suffix = ""
+            # When the range of the data is 0, we can't use log scale.
+            if min_reldiff != max_reldiff:
+                set_log_ax(axs[1], min_reldiff, max_reldiff)
+                label_suffix = " (log scale)"
+            else:
+                set_linear_ax(axs[1])
+            axs[1].set(
+                xlabel="output1 magnitude", ylabel=f"Relative error{label_suffix}"
+            )
+
+            if save_dir is not None:
+                path = os.path.join(save_dir, fname)
+                util.makedirs(path)
+                G_LOGGER.info(f"Saving error metrics plot to: '{path}'")
+                fig.savefig(path, dpi=1200, bbox_inches="tight")
+
+            if show:
+                plt.show()
+
+        finally:
+            plt.close(fig)

tools/Polygraphy/polygraphy/config.py

@@ -0,0 +1,65 @@
+#
+# 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 os
+import sys
+
+INTERNAL_CORRECTNESS_CHECKS = bool(
+    os.environ.get("POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS", "0") != "0"
+)
+"""
+bool: Whether internal correctness checks are enabled.
+This can be configured by setting the 'POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS' environment variable.
+"""
+
+AUTOINSTALL_DEPS = bool(os.environ.get("POLYGRAPHY_AUTOINSTALL_DEPS", "0") != "0")
+"""
+bool: Whether Polygraphy will automatically install required Python packages at runtime.
+This can be configured by setting the 'POLYGRAPHY_AUTOINSTALL_DEPS' environment variable.
+"""
+
+ASK_BEFORE_INSTALL = bool(os.environ.get("POLYGRAPHY_ASK_BEFORE_INSTALL", "0" != "0"))
+"""
+bool: Whether Polygraphy should ask before automatically installing required Python packages.
+Has no effect if AUTOINSTALL_DEPS is not enabled.
+This can be configured by setting the 'POLYGRAPHY_ASK_BEFORE_INSTALL' environment variable.
+"""
+
+INSTALL_CMD = os.environ.get(
+    "POLYGRAPHY_INSTALL_CMD", f"{sys.executable} -m pip install"
+).split()
+"""
+List[str]: The command to use to automatically install dependencies. Only relevant when
+AUTOINSTALL_DEPS is enabled. Defaults to ``["python", "-m", "pip", "install"]``.
+This can be configured by setting the 'POLYGRAPHY_INSTALL_CMD' environment variable to a
+string containing the command; for example: ``python3 -m pip install``.
+"""
+
+ARRAY_SWAP_THRESHOLD_MB = int(
+    os.environ.get("POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB", "-1")
+)
+"""
+int: The threshold, in megabytes, above which Polygraphy will evict an array from memory and swap it to disk.
+A negative value disables swapping and a value of 0 causes all arrays to be saved to disk.
+Disabled by default.
+This can be configured by setting the 'POLYGRAPHY_ARRAY_SWAP_THRESHOLD_MB' environment variable.
+"""
+
+USE_TENSORRT_RTX = bool(os.environ.get("POLYGRAPHY_USE_TENSORRT_RTX", "0") != "0")
+"""
+bool: Whether to use TensorRT RTX as the TensorRT backend.
+This can be configured by setting the 'POLYGRAPHY_USE_TENSORRT_RTX' environment variable.
+"""

tools/Polygraphy/polygraphy/constants.py

@@ -0,0 +1,32 @@
+#
+# 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.
+#
+
+# For legacy purposes
+from polygraphy.config import AUTOINSTALL_DEPS, INTERNAL_CORRECTNESS_CHECKS
+
+DEFAULT_SHAPE_VALUE = 1
+DEFAULT_SEED = 1
+
+TAB = " " * 4  # The one true tab
+
+MARK_ALL = "mark-all"
+"""
+Special value for ModifyOutputs loaders indicating that all values should be marked as outputs
+"""
+
+LEGACY_TYPE_MARKER = "polygraphy_serialized_json_type"
+TYPE_MARKER = "polygraphy_class"

tools/Polygraphy/polygraphy/cuda/__init__.py

@@ -0,0 +1 @@
+from polygraphy.cuda.cuda import *

tools/Polygraphy/polygraphy/cuda/cuda.py

@@ -0,0 +1,521 @@
+#
+# 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 ctypes
+import os
+import sys
+
+from polygraphy import func, mod, util
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER
+
+np = mod.lazy_import("numpy")
+
+
+def void_ptr(val=None):
+    return ctypes.c_void_p(val)
+
+
+@mod.export()
+class MemcpyKind:
+    """
+    Enumerates different kinds of copy operations.
+    """
+
+    HostToHost = ctypes.c_int(0)
+    """Copies from host memory to host memory"""
+    HostToDevice = ctypes.c_int(1)
+    """Copies from host memory to device memory"""
+    DeviceToHost = ctypes.c_int(2)
+    """Copies from device memory to host memory"""
+    DeviceToDevice = ctypes.c_int(3)
+    """Copies from device memory to device memory"""
+    Default = ctypes.c_int(4)
+
+
+@mod.export()
+class Cuda:
+    """
+    NOTE: Do *not* construct this class manually.
+    Instead, use the ``wrapper()`` function to get the global wrapper.
+
+    Wrapper that exposes low-level CUDA functionality.
+    """
+
+    def __init__(self):
+        self.handle = None
+
+        fallback_lib = None
+        if sys.platform.startswith("win"):
+            cuda_paths = [os.environ.get("CUDA_PATH", "")]
+            cuda_paths += os.environ.get("PATH", "").split(os.path.pathsep)
+            lib_pat = "cudart64_*.dll"
+        else:
+            cuda_paths = [
+                *os.environ.get("LD_LIBRARY_PATH", "").split(os.path.pathsep),
+                os.path.join("/", "usr", "local", "cuda", "lib64"),
+                os.path.join("/", "usr", "lib"),
+                os.path.join("/", "lib"),
+            ]
+            lib_pat = "libcudart.so*"
+            fallback_lib = "libcudart.so"
+
+        cuda_paths = list(
+            filter(lambda x: x, cuda_paths)
+        )  # Filter out empty paths (i.e. "")
+
+        candidates = util.find_in_dirs(lib_pat, cuda_paths)
+        if not candidates:
+            log_func = G_LOGGER.critical if fallback_lib is None else G_LOGGER.warning
+            log_func(
+                f"Could not find the CUDA runtime library.\nNote: Paths searched were:\n{cuda_paths}"
+            )
+
+            lib = fallback_lib
+            G_LOGGER.warning(f"Attempting to load: '{lib}' using default loader paths")
+        else:
+            G_LOGGER.verbose(f"Found candidate CUDA libraries: {candidates}")
+            lib = candidates[0]
+
+        self.handle = ctypes.CDLL(lib)
+
+        if not self.handle:
+            G_LOGGER.critical(
+                "Could not load the CUDA runtime library. Is it on your loader path?"
+            )
+
+    @func.constantmethod
+    def check(self, status):
+        if status != 0:
+            G_LOGGER.critical(
+                f"CUDA Error: {status}. To figure out what this means, refer to https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html#group__CUDART__TYPES_1g3f51e3575c2178246db0a94a430e0038"
+            )
+
+    @func.constantmethod
+    def create_stream(self):
+        # Signature: () -> int
+        ptr = void_ptr()
+        self.check(self.handle.cudaStreamCreate(ctypes.byref(ptr)))
+        return ptr.value
+
+    @func.constantmethod
+    def stream_synchronize(self, ptr):
+        # Signature: int -> None
+        self.check(self.handle.cudaStreamSynchronize(void_ptr(ptr)))
+
+    @func.constantmethod
+    def destroy_stream(self, ptr):
+        # Signature: int -> None
+        self.check(self.handle.cudaStreamDestroy(void_ptr(ptr)))
+
+    @func.constantmethod
+    def malloc(self, nbytes):
+        """
+        Allocates memory on the GPU.
+
+        Args:
+            nbytes (int): The number of bytes to allocate.
+
+        Returns:
+            int: The memory address of the allocated region, i.e. a device pointer.
+
+        Raises:
+            PolygraphyException: If an error was encountered during the allocation.
+        """
+        ptr = void_ptr()
+        nbytes = ctypes.c_size_t(nbytes)  # Required to prevent overflow
+        self.check(self.handle.cudaMalloc(ctypes.byref(ptr), nbytes))
+        return ptr.value
+
+    @func.constantmethod
+    def free(self, ptr):
+        """
+        Frees memory allocated on the GPU.
+
+        Args:
+            ptr (int): The memory address, i.e. a device pointer.
+
+        Raises:
+            PolygraphyException: If an error was encountered during the free.
+        """
+        self.check(self.handle.cudaFree(void_ptr(ptr)))
+
+    @func.constantmethod
+    def memcpy(self, dst, src, nbytes, kind, stream_ptr=None):
+        """
+        Copies data between host and device memory.
+
+        Args:
+            dst (int):
+                    The memory address of the destination, i.e. a pointer.
+            src (int):
+                    The memory address of the source, i.e. a pointer.
+            nbytes (int):
+                    The number of bytes to copy.
+            kind (MemcpyKind):
+                    The kind of copy to perform.
+            stream_ptr (int):
+                    The memory address of a CUDA stream, i.e. a pointer.
+                    If this is not provided, a synchronous copy is performed.
+
+        Raises:
+            PolygraphyException: If an error was encountered during the copy.
+        """
+        nbytes = ctypes.c_size_t(nbytes)  # Required to prevent overflow
+        if stream_ptr is not None:
+            self.check(
+                self.handle.cudaMemcpyAsync(
+                    void_ptr(dst), void_ptr(src), nbytes, kind, void_ptr(stream_ptr)
+                )
+            )
+        else:
+            self.check(
+                self.handle.cudaMemcpy(void_ptr(dst), void_ptr(src), nbytes, kind)
+            )
+
+
+G_CUDA = None
+
+
+@mod.export()
+def wrapper():
+    """
+    Returns the global Polygraphy CUDA wrapper.
+
+    Returns:
+        Cuda: The global CUDA wrapper.
+    """
+    global G_CUDA
+    if G_CUDA is None:
+        G_CUDA = Cuda()
+    return G_CUDA
+
+
+@mod.export()
+class Stream:
+    """
+    High-level wrapper for a CUDA stream.
+    """
+
+    def __init__(self):
+        self.ptr = wrapper().create_stream()
+        """int: The memory address of the underlying CUDA stream"""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Frees the underlying CUDA stream.
+        """
+        self.free()
+
+    def free(self):
+        """
+        Frees the underlying CUDA stream.
+
+        You can also use a context manager to manage the stream lifetime.
+        For example:
+        ::
+
+            with Stream() as stream:
+                ...
+        """
+        wrapper().destroy_stream(self.ptr)
+        self.handle = ctypes.c_void_p(None)
+
+    def synchronize(self):
+        """
+        Synchronizes the stream.
+        """
+        wrapper().stream_synchronize(self.ptr)
+
+
+def try_get_stream_handle(stream):
+    if stream is None:
+        return None
+    return stream.ptr
+
+
+@mod.export()
+class DeviceView:
+    """
+    A read-only view of a GPU memory region.
+    """
+
+    def __init__(self, ptr, shape, dtype):
+        """
+        Args:
+            ptr (int): A pointer to the region of memory.
+
+            shape (Tuple[int]): The shape of the region.
+            dtype (DataType): The data type of the region.
+        """
+        self.ptr = int(ptr)
+        """int: The memory address of the underlying GPU memory"""
+        self.shape = shape
+        """Tuple[int]: The shape of the device buffer"""
+        self.itemsize = None
+        self.dtype = dtype
+        """DataType: The data type of the device buffer"""
+
+    def _check_host_buffer(self, host_buffer, copying_from):
+        if util.array.dtype(host_buffer) != self._dtype:
+            G_LOGGER.error(
+                f"Host buffer type: {util.array.dtype(host_buffer)} does not match the type of this device buffer: {self._dtype}. This may cause CUDA errors!"
+            )
+
+        if not util.array.is_contiguous(host_buffer):
+            G_LOGGER.critical(
+                "Provided host buffer is not contiguous in memory.\n"
+                "Hint: Use `util.make_contiguous()` or `np.ascontiguousarray()` to make the array contiguous in memory."
+            )
+
+        # If the host buffer is an input, the device buffer should be large enough to accomodate it.
+        # Otherwise, the host buffer needs to be large enough to accomodate the device buffer.
+        if copying_from:
+            if util.array.nbytes(host_buffer) > self.nbytes:
+                G_LOGGER.critical(
+                    f"Provided host buffer is larger than device buffer.\n"
+                    f"Note: host buffer is {util.array.nbytes(host_buffer)} bytes but device buffer is only {self.nbytes} bytes.\n"
+                    f"Hint: Use `resize()` to resize the device buffer to the correct shape."
+                )
+        else:
+            if util.array.nbytes(host_buffer) < self.nbytes:
+                G_LOGGER.critical(
+                    f"Provided host buffer is smaller than device buffer.\n"
+                    f"Note: host buffer is only {util.array.nbytes(host_buffer)} bytes but device buffer is {self.nbytes} bytes.\n"
+                    f"Hint: Use `util.array.resize_or_reallocate()` to resize the host buffer to the correct shape."
+                )
+
+    @property
+    def dtype(self):
+        try:
+            # For backwards compatibility
+            mod.warn_deprecated(
+                "Using NumPy data types in DeviceView/DeviceArray",
+                use_instead=None,
+                remove_in="0.50.0",
+            )
+            G_LOGGER.warning(
+                f"In the future, you will need to use `DataType.from_dtype(device_view.dtype).numpy()` to retrieve the NumPy data type"
+            )
+            return DataType.to_dtype(self._dtype, "numpy")
+        except:
+            return self._dtype
+
+    @dtype.setter
+    def dtype(self, new):
+        self._dtype = DataType.from_dtype(new)
+        self.itemsize = self._dtype.itemsize
+
+    @property
+    def nbytes(self):
+        """
+        The number of bytes in the memory region.
+        """
+        return util.volume(self.shape) * self.itemsize
+
+    @func.constantmethod
+    def copy_to(self, host_buffer, stream=None):
+        """
+        Copies from this device buffer to the provided host buffer.
+
+        Args:
+            host_buffer (Union[numpy.ndarray, torch.Tensor]):
+                    The host buffer to copy into. The buffer must be contiguous in
+                    memory (see np.ascontiguousarray or torch.Tensor.contiguous) and
+                    large enough to accomodate the device buffer.
+            stream (Stream):
+                    A Stream instance. Performs a synchronous copy if no stream is provided.
+
+        Returns:
+            np.ndarray: The host buffer
+        """
+        if not self.nbytes:
+            return host_buffer
+
+        self._check_host_buffer(host_buffer, copying_from=False)
+        wrapper().memcpy(
+            dst=util.array.data_ptr(host_buffer),
+            src=self.ptr,
+            nbytes=self.nbytes,
+            kind=MemcpyKind.DeviceToHost,
+            stream_ptr=try_get_stream_handle(stream),
+        )
+        return host_buffer
+
+    @func.constantmethod
+    def numpy(self):
+        """
+        Create a new NumPy array containing the contents of this device buffer.
+
+        Returns:
+            np.ndarray: The newly created NumPy array.
+        """
+        arr = np.empty(self.shape, dtype=DataType.to_dtype(self._dtype, "numpy"))
+        self.copy_to(arr)
+        return arr
+
+    def __str__(self):
+        return f"DeviceView[(dtype={self._dtype.name}, shape={self.shape}), ptr={hex(self.ptr)}]"
+
+    def __repr__(self):
+        return util.make_repr(
+            "DeviceView", ptr=self.ptr, shape=self.shape, dtype=self._dtype
+        )[0]
+
+
+@mod.export()
+class DeviceArray(DeviceView):
+    """
+    An array on the GPU.
+    """
+
+    def __init__(self, shape=None, dtype=None):
+        """
+        Args:
+            shape (Tuple[int]): The initial shape of the buffer.
+            dtype (DataType): The data type of the buffer.
+        """
+        super().__init__(
+            ptr=0,
+            shape=util.default(shape, tuple()),
+            dtype=util.default(dtype, DataType.FLOAT32),
+        )
+        self.allocated_nbytes = 0
+        self.resize(self.shape)
+
+    def __enter__(self):
+        return self
+
+    @staticmethod
+    def raw(shape=None):
+        """
+        Creates an untyped device array of the specified shape.
+
+        Args:
+            shape (Tuple[int]):
+                The initial shape of the buffer, in units of bytes.
+                For example, a shape of ``(4, 4)`` would allocate a 16 byte array.
+
+        Returns:
+            DeviceArray: The raw device array.
+        """
+        return DeviceArray(shape=shape, dtype=DataType.UINT8)
+
+    def resize(self, shape):
+        """
+        Resizes or reshapes the array to the specified shape.
+
+        If the allocated memory region is already large enough,
+        no reallocation is performed.
+
+        Args:
+            shape (Tuple[int]): The new shape.
+
+        Returns:
+            DeviceArray: self
+        """
+        nbytes = util.volume(shape) * self.itemsize
+        if nbytes > self.allocated_nbytes:
+            self.free()
+            self.ptr = wrapper().malloc(nbytes)
+            self.allocated_nbytes = nbytes
+        self.shape = shape
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Frees the underlying memory of this DeviceArray.
+        """
+        self.free()
+
+    def free(self):
+        """
+        Frees the GPU memory associated with this array.
+
+        You can also use a context manager to ensure that memory is freed. For example:
+        ::
+
+            with DeviceArray(...) as arr:
+                ...
+        """
+        wrapper().free(self.ptr)
+        self.shape = tuple()
+        self.allocated_nbytes = 0
+        self.ptr = 0
+
+    def copy_from(self, host_buffer, stream=None):
+        """
+        Copies from the provided host buffer into this device buffer.
+
+        Args:
+            host_buffer (Union[numpy.ndarray, torch.Tensor]):
+                    The host buffer to copy from. The buffer must be contiguous in
+                    memory (see np.ascontiguousarray or torch.Tensor.contiguous) and not
+                    larger than this device buffer.
+            stream (Stream):
+                    A Stream instance. Performs a synchronous copy if no stream is provided.
+
+        Returns:
+            DeviceArray: self
+        """
+        if not util.array.nbytes(host_buffer):
+            return self
+
+        self._check_host_buffer(host_buffer, copying_from=True)
+        wrapper().memcpy(
+            dst=self.ptr,
+            src=util.array.data_ptr(host_buffer),
+            nbytes=util.array.nbytes(host_buffer),
+            kind=MemcpyKind.HostToDevice,
+            stream_ptr=try_get_stream_handle(stream),
+        )
+        return self
+
+    def view(self, shape=None, dtype=None):
+        """
+        Creates a read-only DeviceView from this DeviceArray.
+
+        Args:
+            shape (Sequence[int]):
+                    The desired shape of the view.
+                    Defaults to the shape of this array or view.
+            dtype (DataType):
+                    The desired data type of the view.
+                    Defaults to the data type of this array or view.
+
+        Returns:
+            DeviceView: A view of this arrays data on the device.
+        """
+        shape = util.default(shape, self.shape)
+        dtype = util.default(dtype, self._dtype)
+        view = DeviceView(self.ptr, shape, dtype)
+
+        if view.nbytes > self.nbytes:
+            G_LOGGER.critical(
+                "A view cannot exceed the number of bytes of the original array.\n"
+                f"Note: Original array has shape: {self.shape} and dtype: {self._dtype}, which requires {self.nbytes} bytes, "
+                f"while the view has shape: {shape} and dtype: {dtype}, which requires {view.nbytes} bytes, "
+            )
+        return view
+
+    def __str__(self):
+        return f"DeviceArray[(dtype={self._dtype.name}, shape={self.shape}), ptr={hex(self.ptr)}]"
+
+    def __repr__(self):
+        return util.make_repr("DeviceArray", shape=self.shape, dtype=self._dtype)[0]

tools/Polygraphy/polygraphy/datatype/__init__.py

@@ -0,0 +1,6 @@
+from polygraphy.datatype.datatype import *
+from polygraphy.datatype.numpy import *
+from polygraphy.datatype.onnx import *
+from polygraphy.datatype.onnxrt import *
+from polygraphy.datatype.tensorrt import *
+from polygraphy.datatype.torch import *

tools/Polygraphy/polygraphy/datatype/datatype.py

@@ -0,0 +1,273 @@
+#
+# 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 functools
+from textwrap import dedent
+
+from polygraphy import mod
+from polygraphy.exception import DataTypeConversionException
+from polygraphy.logger import G_LOGGER, LogMode
+
+
+import enum
+
+
+class _SkipImporterException(Exception):
+    pass
+
+
+class _DataTypeKind(enum.Enum):
+    FLOATING_POINT = 0
+    INTEGRAL = 1
+    _OTHER = 3
+
+
+@mod.export()
+class DataTypeEntry:
+    """
+    Represents a data type.
+    Can be transformed to and from data type classes of external modules, like NumPy.
+
+    Do *not* construct objects of this type directly. Instead, use the predefined data types
+    provided in the ``DataType`` class.
+    """
+
+    def __init__(self, name, itemsize, type: _DataTypeKind):
+        self.name = name
+        """The human-readable name of the data type"""
+        self.itemsize = itemsize
+        """The size in bytes of a single element of this data type"""
+
+        # self._type describes the basic kind of the type we have.
+        # For example, this can
+        self._type = type
+
+    @property
+    def is_floating(self):
+        return self._type == _DataTypeKind.FLOATING_POINT
+
+    @property
+    def is_integral(self):
+        return self._type == _DataTypeKind.INTEGRAL
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return f"DataType.{self.name.upper()}"
+
+
+@mod.export()
+class DataType:
+    # Docstring will be populated by the loop below
+    """
+    Aggregates supported Polygraphy data types. Each data type is accessible
+    via this class as a class member of type ``DataTypeEntry``.
+
+    Members:
+    """
+    _IMPORTER_FUNCS = {}
+    _EXPORTER_FUNCS = {}
+
+    __members__ = {
+        "FLOAT64": DataTypeEntry("float64", 8, _DataTypeKind.FLOATING_POINT),
+        "FLOAT32": DataTypeEntry("float32", 4, _DataTypeKind.FLOATING_POINT),
+        "FLOAT16": DataTypeEntry("float16", 2, _DataTypeKind.FLOATING_POINT),
+        "FLOAT4": DataTypeEntry("float4", 0.5, _DataTypeKind.FLOATING_POINT),
+        "INT16": DataTypeEntry("int16", 2, _DataTypeKind.INTEGRAL),
+        "INT32": DataTypeEntry("int32", 4, _DataTypeKind.INTEGRAL),
+        "INT64": DataTypeEntry("int64", 8, _DataTypeKind.INTEGRAL),
+        "INT8": DataTypeEntry("int8", 1, _DataTypeKind.INTEGRAL),
+        "INT4": DataTypeEntry("int4", 0.5, _DataTypeKind.INTEGRAL),
+        "UINT16": DataTypeEntry("uint16", 2, _DataTypeKind.INTEGRAL),
+        "UINT32": DataTypeEntry("uint32", 4, _DataTypeKind.INTEGRAL),
+        "UINT64": DataTypeEntry("uint64", 8, _DataTypeKind.INTEGRAL),
+        "UINT8": DataTypeEntry("uint8", 1, _DataTypeKind.INTEGRAL),
+        "BOOL": DataTypeEntry("bool", 1, _DataTypeKind._OTHER),
+        "STRING": DataTypeEntry("string", 0, _DataTypeKind._OTHER),
+        "BFLOAT16": DataTypeEntry("bfloat16", 2, _DataTypeKind.FLOATING_POINT),
+        "FLOAT8E4M3FN": DataTypeEntry("float8e4m3fn", 1, _DataTypeKind.FLOATING_POINT),
+        "FLOAT8E4M3FNUZ": DataTypeEntry(
+            "float8e4m3fnuz", 1, _DataTypeKind.FLOATING_POINT
+        ),
+        "FLOAT8E5M2": DataTypeEntry("float8e5m2", 1, _DataTypeKind.FLOATING_POINT),
+        "FLOAT8E5M2FNUZ": DataTypeEntry(
+            "float8e5m2fnuz", 1, _DataTypeKind.FLOATING_POINT
+        ),
+    }
+
+    @staticmethod
+    def from_dtype(dtype, source_module=None):
+        """
+        Converts a data type from any known external libraries to a corresponding
+        Polygraphy data type.
+
+        Args:
+            dtype (Any): A data type from an external library.
+            source_module (str):
+                    The name of the module from where the provided `dtype` originates.
+                    If this is not provided, Polygraphy will attempt to guess the module
+                    in order to convert the data type.
+
+        Returns:
+            DataTypeEntry: The corresponding Polygraphy data type.
+
+        Raises:
+            PolygraphyException: If the data type could not be converted.
+        """
+        if dtype is None:
+            G_LOGGER.critical(f"Could not convert: {dtype} to a Polygraphy data type")
+
+        if isinstance(dtype, DataTypeEntry):
+            return dtype
+
+        if source_module is not None:
+            if source_module not in DataType._IMPORTER_FUNCS:
+                G_LOGGER.critical(
+                    f"Could not find source module: {source_module} in known importers. "
+                    f"Note: Importer functions have been registered for the following modules: {list(DataType._IMPORTER_FUNCS.keys())}"
+                )
+            try:
+                return DataType._IMPORTER_FUNCS[source_module](dtype)
+            except _SkipImporterException:
+                pass
+        else:
+            for func in DataType._IMPORTER_FUNCS.values():
+                try:
+                    ret = func(dtype)
+                except _SkipImporterException:
+                    pass
+                else:
+                    return ret
+
+        msg = f"Could not convert: {dtype} to a corresponding Polygraphy data type. Leaving this type in its source format."
+        G_LOGGER.warning(msg, mode=LogMode.ONCE)
+        G_LOGGER.internal_error(msg)
+        return dtype
+
+    @staticmethod
+    def to_dtype(dtype, target_module):
+        """
+        Converts a Polygraphy data type to one from any known external libraries.
+
+        Args:
+            dtype (DataType):
+                    A Polygraphy data type. If something other than a Polygraphy data type is provided,
+                    then this function will return it without modifying it.
+            target_module (str):
+                    The name of the module whose data type class to convert this data type to.
+
+        Returns:
+            Any: The corresponding data type from the target module.
+
+        Raises:
+            PolygraphyException: If the data type could not be converted.
+        """
+        if not isinstance(dtype, DataTypeEntry):
+            G_LOGGER.internal_error(
+                f"Received input of type other than DataType: {dtype}"
+            )
+            return dtype
+
+        if target_module not in DataType._EXPORTER_FUNCS:
+            G_LOGGER.critical(
+                f"Could not find target module: {target_module} in known exporters. "
+                f"Note: Exporter functions have been registered for the following modules: {list(DataType._EXPORTER_FUNCS.keys())}"
+            )
+        return DataType._EXPORTER_FUNCS[target_module](dtype)
+
+
+DataType.__doc__ = dedent(DataType.__doc__)
+for name, value in DataType.__members__.items():
+    setattr(DataType, name, value)
+    DataType.__doc__ += f"\t- {name}\n"
+
+
+def register_dtype_importer(source_module):
+    """
+    Registers an importer function with the DataType class.
+
+    IMPORTANT: You *must* ensure that the importer function does not attempt to automatically install
+    or import modules which are not already installed.
+    With a lazily imported module, `module.is_installed()/is_importable()` is an easy way to guard the code against this.
+    We do not want to automatically install heavy modules like PyTorch or TensorRT just for the sake of DataType.
+
+    For example:
+    ::
+
+        @register_dtype_importer("numpy")
+        def func(dtype):
+            ...
+
+    The importer function should return `None` if no corresponding data type could be found
+    or if the input type did not match what was expected.
+
+    The newly registered function is then usable via `from_dtype`:
+    ::
+
+        dtype = DataType.from_dtype(np.int64)
+    """
+
+    def register_importer_impl(func):
+        @functools.wraps(func)
+        def new_func(dtype):
+            val = func(dtype)
+            if val is None:
+                # We raise an exception to indicate that `from_dtype` should skip this importer and try a different one.
+                # We have to do it this way since we don't necessarily know which importer is the right one to use.
+                raise _SkipImporterException()
+            return val
+
+        DataType._IMPORTER_FUNCS[source_module] = new_func
+        return new_func
+
+    return register_importer_impl
+
+
+def register_dtype_exporter(target_module):
+    """
+    Registers an exporter function with the DataType class.
+
+    For example:
+    ::
+
+        @register_dtype_exporter("numpy")
+        def func(dtype):
+            ...
+
+    The newly registered function is then accessible with, for example:
+    ::
+
+        np_dtype = DataType.FLOAT32.numpy()
+    """
+
+    def register_exporter_impl(func):
+        @functools.wraps(func)
+        def new_func(dtype):
+            val = func(dtype)
+            if val is None:
+                G_LOGGER.critical(
+                    f"Could not convert Polygraphy data type: {dtype} to a corresponding {target_module} data type. ",
+                    ExceptionType=DataTypeConversionException,
+                )
+            return val
+
+        new_func.__name__ = target_module
+        setattr(DataTypeEntry, new_func.__name__, new_func)
+        DataType._EXPORTER_FUNCS[target_module] = new_func
+        return new_func
+
+    return register_exporter_impl

tools/Polygraphy/polygraphy/datatype/numpy.py

@@ -0,0 +1,77 @@
+#
+# 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.datatype.datatype import (
+    DataType,
+    register_dtype_importer,
+    register_dtype_exporter,
+)
+
+np = mod.lazy_import("numpy")
+
+
+def _get_mapping():
+    DATATYPE_FROM_NUMPY = {
+        np.double: DataType.FLOAT64,
+        np.float32: DataType.FLOAT32,
+        np.float16: DataType.FLOAT16,
+        np.int16: DataType.INT16,
+        np.int32: DataType.INT32,
+        np.int64: DataType.INT64,
+        np.int8: DataType.INT8,
+        np.uint16: DataType.UINT16,
+        np.uint32: DataType.UINT32,
+        np.uint64: DataType.UINT64,
+        np.uint8: DataType.UINT8,
+        np.bool_: DataType.BOOL,
+        np.str_: DataType.STRING,
+    }
+    return {np.dtype(key): val for key, val in DATATYPE_FROM_NUMPY.items()}
+
+
+@register_dtype_importer("numpy")
+def from_numpy(numpy_type):
+    """
+    Converts a NumPy data type to a Polygraphy DataType.
+
+    Args:
+        numpy_type (np.dtype): The NumPy data type.
+
+    Returns:
+        DataType: The Polygraphy data type.
+    """
+    if not np.is_installed() or not np.is_importable():
+        return None
+
+    try:
+        dtype = np.dtype(numpy_type)
+    except TypeError:
+        return None
+
+    return _get_mapping().get(dtype)
+
+
+@register_dtype_exporter("numpy")
+def from_datatype(self):
+    """
+    Converts this Polygraphy DataType to a NumPy data type.
+
+    Returns:
+        np.dtype: The NumPy data type.
+    """
+    return util.invert_dict(_get_mapping()).get(self)

tools/Polygraphy/polygraphy/datatype/onnx.py

@@ -0,0 +1,86 @@
+#
+# 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.datatype.datatype import (
+    DataType,
+    register_dtype_importer,
+    register_dtype_exporter,
+)
+
+onnx = mod.lazy_import("onnx")
+
+
+def _get_mapping():
+    DATATYPE_FROM_ONNX = {
+        "DOUBLE": DataType.FLOAT64,
+        "FLOAT": DataType.FLOAT32,
+        "FLOAT16": DataType.FLOAT16,
+        "INT16": DataType.INT16,
+        "INT32": DataType.INT32,
+        "INT64": DataType.INT64,
+        "INT8": DataType.INT8,
+        "INT4": DataType.INT4,
+        "UINT16": DataType.UINT16,
+        "UINT32": DataType.UINT32,
+        "UINT64": DataType.UINT64,
+        "UINT8": DataType.UINT8,
+        "BOOL": DataType.BOOL,
+        "STRING": DataType.STRING,
+        "BFLOAT16": DataType.BFLOAT16,
+        "FLOAT8E4M3FN": DataType.FLOAT8E4M3FN,
+        "FLOAT8E4M3FNUZ": DataType.FLOAT8E4M3FNUZ,
+        "FLOAT8E5M2": DataType.FLOAT8E5M2,
+        "FLOAT8E5M2FNUZ": DataType.FLOAT8E5M2FNUZ,
+    }
+    if None in DATATYPE_FROM_ONNX:
+        del DATATYPE_FROM_ONNX[None]
+
+    onnx_type_map = dict(onnx.TensorProto.DataType.items())
+    return {
+        onnx_type_map[key]: val
+        for key, val in DATATYPE_FROM_ONNX.items()
+        if key in onnx_type_map
+    }
+
+
+@register_dtype_importer("onnx")
+def from_onnx(onnx_type):
+    """
+    Converts an ONNX data type to a Polygraphy DataType.
+
+    Args:
+        onnx_type (onnx.TensorProto.DataType): The ONNX data type.
+
+    Returns:
+        DataType: The Polygraphy data type.
+    """
+    if not onnx.is_installed() or not onnx.is_importable():
+        return None
+
+    return _get_mapping().get(onnx_type)
+
+
+@register_dtype_exporter("onnx")
+def from_datatype(self):
+    """
+    Converts this Polygraphy DataType to an ONNX data type.
+
+    Returns:
+        onnx.TensorProto.DataType: The ONNX data type.
+    """
+    return util.invert_dict(_get_mapping()).get(self)

tools/Polygraphy/polygraphy/datatype/onnxrt.py

@@ -0,0 +1,72 @@
+#
+# 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.datatype.datatype import (
+    DataType,
+    register_dtype_importer,
+    register_dtype_exporter,
+)
+
+__DATATYPE_FROM_ONNXRT = {
+    "tensor(double)": DataType.FLOAT64,
+    "tensor(float)": DataType.FLOAT32,
+    "tensor(float16)": DataType.FLOAT16,
+    "tensor(int16)": DataType.INT16,
+    "tensor(int32)": DataType.INT32,
+    "tensor(int64)": DataType.INT64,
+    "tensor(int8)": DataType.INT8,
+    "tensor(int4)": DataType.INT4,
+    "tensor(uint16)": DataType.UINT16,
+    "tensor(uint32)": DataType.UINT32,
+    "tensor(uint64)": DataType.UINT64,
+    "tensor(uint8)": DataType.UINT8,
+    "tensor(bool)": DataType.BOOL,
+    "tensor(string)": DataType.STRING,
+    "tensor(bfloat16)": DataType.BFLOAT16,
+    "tensor(float8e4m3fn)": DataType.FLOAT8E4M3FN,
+    "tensor(float8e4m3fnuz)": DataType.FLOAT8E4M3FNUZ,
+    "tensor(float8e5m2)": DataType.FLOAT8E5M2,
+    "tensor(float8e5m2fnuz)": DataType.FLOAT8E5M2FNUZ,
+}
+
+__ONNXRT_FROM_DATATYPE = util.invert_dict(__DATATYPE_FROM_ONNXRT)
+
+
+@register_dtype_importer("onnxruntime")
+def from_onnxrt(onnxrt_type):
+    """
+    Converts an ONNX-Runtime data type to a Polygraphy DataType.
+
+    Args:
+        onnxrt_type (str): The ONNX-Runtime data type.
+
+    Returns:
+        DataType: The Polygraphy data type.
+    """
+    return __DATATYPE_FROM_ONNXRT.get(onnxrt_type)
+
+
+@register_dtype_exporter("onnxruntime")
+def from_datatype(self):
+    """
+    Converts this Polygraphy DataType to an ONNX-Runtime data type.
+
+    Returns:
+        str: The ONNX-Runtime data type.
+    """
+    return __ONNXRT_FROM_DATATYPE.get(self)

tools/Polygraphy/polygraphy/datatype/tensorrt.py

@@ -0,0 +1,74 @@
+#
+# 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.datatype.datatype import (
+    DataType,
+    register_dtype_importer,
+    register_dtype_exporter,
+)
+from polygraphy.mod.trt_importer import lazy_import_trt
+
+trt = lazy_import_trt()
+
+
+def _get_mapping():
+    DATATYPE_FROM_TENSORRT = {
+        trt.float32: DataType.FLOAT32,
+        trt.float16: DataType.FLOAT16,
+        trt.int32: DataType.INT32,
+        trt.int8: DataType.INT8,
+        util.try_getattr(trt, "int64"): DataType.INT64,
+        util.try_getattr(trt, "uint8"): DataType.UINT8,
+        util.try_getattr(trt, "bool"): DataType.BOOL,
+        util.try_getattr(trt, "bfloat16"): DataType.BFLOAT16,
+        util.try_getattr(trt, "fp8"): DataType.FLOAT8E4M3FN,
+        util.try_getattr(trt, "int4"): DataType.INT4,
+        util.try_getattr(trt, "fp4"): DataType.FLOAT4,
+    }
+    if None in DATATYPE_FROM_TENSORRT:
+        del DATATYPE_FROM_TENSORRT[None]
+
+    return DATATYPE_FROM_TENSORRT
+
+
+@register_dtype_importer("tensorrt")
+def from_tensorrt(tensorrt_type):
+    """
+    Converts a TensorRT data type to a Polygraphy DataType.
+
+    Args:
+        tensorrt_type (tensorrt.DataType): The TensorRT data type.
+
+    Returns:
+        DataType: The Polygraphy data type.
+    """
+    if not trt.is_installed() or not trt.is_importable():
+        return None
+
+    return _get_mapping().get(tensorrt_type)
+
+
+@register_dtype_exporter("tensorrt")
+def from_datatype(self):
+    """
+    Converts this Polygraphy DataType to a TensorRT data type.
+
+    Returns:
+        tensorrt.DataType: The TensorRT data type.
+    """
+    return util.invert_dict(_get_mapping()).get(self)

tools/Polygraphy/polygraphy/datatype/torch.py

@@ -0,0 +1,68 @@
+#
+# 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.datatype.datatype import (
+    DataType,
+    register_dtype_importer,
+    register_dtype_exporter,
+)
+
+torch = mod.lazy_import("torch>=1.13.0")
+
+
+def _get_mapping():
+    return {
+        torch.float64: DataType.FLOAT64,
+        torch.float32: DataType.FLOAT32,
+        torch.float16: DataType.FLOAT16,
+        torch.int16: DataType.INT16,
+        torch.int32: DataType.INT32,
+        torch.int64: DataType.INT64,
+        torch.int8: DataType.INT8,
+        torch.uint8: DataType.UINT8,
+        torch.bool: DataType.BOOL,
+        torch.bfloat16: DataType.BFLOAT16,
+    }
+
+
+@register_dtype_importer("torch")
+def from_torch(torch_type):
+    """
+    Converts a PyTorch data type to a Polygraphy DataType.
+
+    Args:
+        torch_type (torch.dtype): The PyTorch data type.
+
+    Returns:
+        DataType: The Polygraphy data type.
+    """
+    if not torch.is_installed() or not torch.is_importable():
+        return None
+
+    return _get_mapping().get(torch_type)
+
+
+@register_dtype_exporter("torch")
+def from_datatype(self):
+    """
+    Converts this Polygraphy DataType to a PyTorch data type.
+
+    Returns:
+        torch.dtype: The PyTorch data type.
+    """
+    return util.invert_dict(_get_mapping()).get(self)

tools/Polygraphy/polygraphy/exception/__init__.py

@@ -0,0 +1 @@
+from polygraphy.exception.exception import *

tools/Polygraphy/polygraphy/exception/exception.py

@@ -0,0 +1,52 @@
+#
+# 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
+
+
+# Do not raise this exception manually. Instead, use G_LOGGER.critical().
+@mod.export()
+class PolygraphyException(Exception):
+    """
+    An exception raised by Polygraphy.
+    """
+
+    pass
+
+
+# Do not raise this exception manually. Instead, use G_LOGGER.internal_error().
+@mod.export()
+class PolygraphyInternalException(Exception):
+    """
+    An exception raised when a Polygraphy internal check is violated.
+    Polygraphy internal checks can be enabled by setting the ``POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS``
+    environment variable to ``1``.
+    This is *not* a child class of PolygraphyException because it
+    indicates a bug in Polygraphy itself.
+    """
+
+    pass
+
+
+# Do not raise this exception manually. Instead, use G_LOGGER.critical(..., ExceptionType=DataTypeConversionException).
+@mod.export()
+class DataTypeConversionException(PolygraphyException):
+    """
+    An exception during conversion to or from a Polygraphy DataType.
+    """
+
+    pass

tools/Polygraphy/polygraphy/func/__init__.py

@@ -0,0 +1 @@
+from polygraphy.func.func import *

tools/Polygraphy/polygraphy/func/func.py

@@ -0,0 +1,199 @@
+#
+# 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 functools
+import inspect
+
+from polygraphy import config, mod
+from polygraphy.logger import G_LOGGER
+
+
+def make_iterable(obj):
+    return obj if type(obj) == tuple else (obj,)
+
+
+@mod.export()
+def extend(extend_func):
+    """
+    A decorator that uses the function it decorates to extend the function
+    provided as a parameter.
+
+    This is best illustrated with an example:
+    ::
+
+        def x(a0, a1, a2):
+            rv0 = [a0, a1, a2]
+            rv1 = None
+            return rv0, rv1
+
+        @extend(x)
+        def y(rv0, rv1):
+            rv0.append(-1)
+
+        # We can now call `y` as if it were `x`, and we will receive
+        # the return values from `x` after any modifications by `y`
+        rv0, rv1 = y(1, 2, 3)
+        assert rv0 == [1, 2, 3, -1]
+        assert rv1 is None
+
+    In this case, ``extend`` is essentially syntactic sugar for:
+    ::
+
+        def y(a0, a1, a2):
+            rv0, rv1 = x(a0, a1, a2)
+
+            # Body of `y` from previous section
+            rv0.append(-1)
+
+            return rv0, rv1
+
+    If ``y`` does not return anything, or returns ``None``, then ``extend`` will
+    ensure that the return value of ``x`` is forwarded to the caller.
+    This means that ``y`` will provide exactly the same interface as ``x``.
+
+    If `y` returns something other than ``None``, then its return value will be
+    provided to the caller, and the return value of ``x`` will be discarded.
+
+    In some cases, it may be necessary to access the parameters of ``x``.
+    If ``y`` takes all of ``x``'s parameters prior to its usual parameters
+    (which are the return values of ``x``), then all arguments to ``x`` will be
+    forwarded to ``y``.
+    Note that if ``x`` modifies its arguments in place, then ``y`` will see the modified
+    arguments, `not` the original ones.
+    For example:
+    ::
+
+        def x(x_arg):
+            return x_arg + 1
+
+        def y(x_arg, x_ret):
+            # `y` can now see both the input argument as well as the return value of `x`
+            assert x_ret == x_arg + 1
+
+        assert y(5) == 6
+
+
+    NOTE: This function will automatically unpack tuples returned by the function
+    being extended. Thus, the following implementation of ``x`` would behave just like
+    the one mentioned above:
+    ::
+
+        def x(a0, a1, a2):
+            ret = (rv0, rv1)
+            return ret # Tuple will be unpacked, and `y` still sees 2 parameters
+
+    NOTE: The decorated function must not use variadic parameters like ``*args`` or ``**kwargs``.
+
+    Args:
+        extend_func (Callable): A callable to extend.
+    """
+
+    def extend_decorator(func):
+        @functools.wraps(func)
+        def extended_func(*args, **kwargs):
+            extend_func_retval = extend_func(*args, **kwargs)
+            extend_func_ret_tuple = make_iterable(extend_func_retval)
+
+            func_params = inspect.signature(func).parameters
+            # Special case for when the extended function does not return anything
+            if (
+                len(func_params) == 0
+                and len(extend_func_ret_tuple) == 1
+                and extend_func_ret_tuple[0] is None
+            ):
+                func_retval = func()
+            elif len(extend_func_ret_tuple) == len(func_params):
+                func_retval = func(*extend_func_ret_tuple)
+            elif len(func_params) == len(extend_func_ret_tuple) + len(args) + len(
+                kwargs
+            ):
+                # We need to turn `extend_func_ret_tuple` into keyword arguments so that it can
+                # be ordered after `**kwargs`.
+                ret_arg_names = [
+                    param.name
+                    for param in list(func_params.values())[
+                        -len(extend_func_ret_tuple) :
+                    ]
+                ]
+                ret_kwargs = dict(zip(ret_arg_names, extend_func_ret_tuple))
+                func_retval = func(*args, **kwargs, **ret_kwargs)
+            else:
+
+                def try_get_name(fn):
+                    try:
+                        return fn.__name__
+                    except:
+                        return fn
+
+                G_LOGGER.critical(
+                    f"Function: {try_get_name(func)} accepts {len(func_params)} parameter(s), "
+                    f"but needs to accept {len(extend_func_ret_tuple)} parameter(s) from: {try_get_name(extend_func)} instead."
+                    f"\nNote: Parameters should be: {tuple(map(type, extend_func_ret_tuple))}"
+                )
+
+            if func_retval is not None:
+                return func_retval
+            return extend_func_retval
+
+        return extended_func
+
+    return extend_decorator
+
+
+@mod.export()
+def constantmethod(func):
+    """
+    A decorator that denotes constant methods.
+
+    NOTE: This decorator does nothing if the POLYGRAPHY_INTERNAL_CORRECTNESS_CHECKS environment variable is not set to `1`
+
+    Example:
+    ::
+
+        class Dummy:
+            def __init__(self):
+                self.x = 1
+
+            @func.constantmethod
+            def modify_x(self):
+                self.x = 2
+
+        d = Dummy()
+        d.modify_x() # This will fail!
+
+
+    This provides only minimal protection against accidental mutation of instance attributes.
+    For example, if a class includes references (e.g. a numpy array member), this function cannot
+    ensure that the contents of that member (e.g. the values in a numpy array) will remain unchanged.
+    """
+    if not config.INTERNAL_CORRECTNESS_CHECKS:
+        return func
+
+    @functools.wraps(func)
+    def wrapper(self, *args, **kwargs):
+        old_dict = copy.copy(vars(self))
+        ret = None
+        try:
+            ret = func(self, *args, **kwargs)
+        finally:
+            if vars(self) != old_dict:
+                G_LOGGER.internal_error(
+                    f"{self} was mutated in a constant method! Note:\nOld state: {old_dict}\nNew state: {vars(self)}"
+                )
+        return ret
+
+    return wrapper

tools/Polygraphy/polygraphy/json/__init__.py

@@ -0,0 +1 @@
+from polygraphy.json.serde import *

tools/Polygraphy/polygraphy/json/serde.py

@@ -0,0 +1,442 @@
+#
+# 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 base64
+import functools
+import io
+import json
+from collections import OrderedDict
+
+from polygraphy import constants, mod, util
+from polygraphy.logger import G_LOGGER
+
+np = mod.lazy_import("numpy")
+torch = mod.lazy_import("torch>=1.13.0")
+
+
+def legacy_str_from_type(typ):
+    return "__polygraphy_encoded_" + typ.__name__
+
+
+def str_from_type(typ):
+    return typ.__name__
+
+
+class BaseCustomImpl:
+    """
+    Base class for Polygraphy's JSON encoder/decoder.
+    """
+
+    @classmethod
+    def register(cls, typ, alias=None):
+        """
+        Decorator that registers JSON encoding/decoding functions for types.
+
+        Args:
+            typ (type): The type to register
+            alias (str):
+                    An alias under which to also register the decoder function.
+                    This can be used to retain backwards-compatibility when a class
+                    name changes.
+
+        For the documentation that follows, assume we have a class:
+        ::
+
+            class Dummy:
+                def __init__(self, x):
+                    self.x = x
+
+        ========
+        Encoders
+        ========
+
+        Encoder functions should accept instances of the specified type and
+        return dictionaries.
+
+        For example:
+        ::
+
+            @Encoder.register(Dummy)
+            def encode(dummy):
+                return {"x": dummy.x}
+
+
+        To use the custom encoder, use the `to_json` helper:
+        ::
+
+            d = Dummy(x=1)
+            d_json = to_json(d)
+
+
+        ========
+        Decoders
+        ========
+
+        Decoder functions should accept dictionaries, and return instances of the
+        type.
+
+        For example:
+        ::
+
+            @Decoder.register(Dummy)
+            def decode(dct):
+                return Dummy(x=dct["x"])
+
+
+        To use the custom decoder, use the `from_json` helper:
+        ::
+
+            from_json(d_json)
+
+
+        Args:
+            typ (type): The type of the class for which to register the function.
+        """
+
+        def register_impl(func):
+            def add(key, val):
+                if key in cls.polygraphy_registered:
+                    G_LOGGER.critical(
+                        f"Duplicate serialization function for type: {key}.\nNote: Existing function: {cls.polygraphy_registered[key]}, New function: {func}"
+                    )
+                cls.polygraphy_registered[key] = val
+
+            if cls == Encoder:
+
+                def wrapped(obj):
+                    dct = func(obj)
+                    dct[constants.TYPE_MARKER] = str_from_type(typ)
+                    return dct
+
+                add(typ, wrapped)
+                return wrapped
+            elif cls == Decoder:
+
+                def wrapped(dct):
+                    if constants.TYPE_MARKER in dct:
+                        del dct[constants.TYPE_MARKER]
+
+                    type_name = legacy_str_from_type(typ)
+                    if type_name in dct:
+                        del dct[type_name]
+
+                    return func(dct)
+
+                add(legacy_str_from_type(typ), wrapped)
+                add(str_from_type(typ), wrapped)
+                if alias is not None:
+                    add(alias, wrapped)
+            else:
+                G_LOGGER.critical("Cannot register for unrecognized class type: ")
+
+        return register_impl
+
+
+@mod.export()
+class Encoder(BaseCustomImpl, json.JSONEncoder):
+    """
+    Polygraphy's custom JSON Encoder implementation.
+    """
+
+    polygraphy_registered = {}
+
+    def default(self, o):
+        if type(o) in self.polygraphy_registered:
+            return self.polygraphy_registered[type(o)](o)
+        return super().default(o)
+
+
+@mod.export()
+class Decoder(BaseCustomImpl):
+    """
+    Polygraphy's custom JSON Decoder implementation.
+    """
+
+    polygraphy_registered = {}
+
+    def __call__(self, pairs):
+        # The encoder will insert special key-value pairs into dictionaries encoded from
+        # custom types. If we find one, then we know to decode using the corresponding custom
+        # type function.
+        dct = OrderedDict(pairs)
+
+        # Handle legacy naming first - these keys should not be present in JSON generated by more recent versions of Polygraphy.
+        for type_str, func in self.polygraphy_registered.items():
+            if (
+                type_str in dct and dct[type_str] == constants.LEGACY_TYPE_MARKER
+            ):  # Found a custom type!
+                return func(dct)
+
+        type_name = dct.get(constants.TYPE_MARKER)
+        if type_name is not None:
+            if type_name not in self.polygraphy_registered:
+                user_type_name = {
+                    "Tensor": "torch.Tensor",
+                    "ndarray": "np.ndarray",
+                }.get(type_name, type_name)
+                G_LOGGER.critical(
+                    f"Could not decode serialized type: {user_type_name}. This could be because a required module is missing. "
+                )
+            return self.polygraphy_registered[type_name](dct)
+
+        return dct
+
+
+NUMPY_REGISTRATION_SUCCESS = False
+TORCH_REGISTRATION_SUCCESS = False
+COMMON_REGISTRATION_SUCCESS = False
+
+
+def try_register_common_json(func):
+    """
+    Decorator that attempts to register common JSON encode/decode methods
+    if the methods have not already been registered.
+
+    This needs to be attempted multiple times because dependencies may become available in the
+    middle of execution - for example, if using dependency auto-installation.
+    """
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        global NUMPY_REGISTRATION_SUCCESS
+        if not NUMPY_REGISTRATION_SUCCESS and np.is_installed() and np.is_importable():
+            # We define this alongside load_json/save_json so that it is guaranteed to be
+            # imported before we need to encode/decode NumPy arrays.
+            @Encoder.register(np.ndarray)
+            def encode(array):
+                outfile = io.BytesIO()
+                np.save(outfile, array, allow_pickle=False)
+                outfile.seek(0)
+                data = base64.b64encode(outfile.read()).decode()
+                return {"array": data}
+
+            @Decoder.register(np.ndarray)
+            def decode(dct):
+                def load(mode="base64"):
+                    if mode == "base64":
+                        data = base64.b64decode(dct["array"].encode(), validate=True)
+                    elif mode == "latin-1":
+                        data = dct["array"].encode(mode)
+                    else:
+                        assert False, f"Unsupported mode: {mode}"
+                    infile = io.BytesIO(data)
+                    return np.load(infile, allow_pickle=False)
+
+                try:
+                    arr = load()
+                except:
+                    arr = load("latin-1")  # For backwards compatibility
+                if isinstance(arr, np.ndarray):
+                    return arr
+                return list(arr.values())[0]  # For backwards compatibility
+
+            NUMPY_REGISTRATION_SUCCESS = True
+
+        global TORCH_REGISTRATION_SUCCESS
+        if (
+            not TORCH_REGISTRATION_SUCCESS
+            and torch.is_installed()
+            and torch.is_importable()
+        ):
+
+            @Encoder.register(torch.Tensor)
+            def encode(tensor):
+                outfile = io.BytesIO()
+                torch.save(tensor, outfile)
+                outfile.seek(0)
+                data = base64.b64encode(outfile.read()).decode()
+                return {"tensor": data}
+
+            @Decoder.register(torch.Tensor)
+            def decode(dct):
+                data = base64.b64decode(dct["tensor"].encode(), validate=True)
+                infile = io.BytesIO(data)
+                return torch.load(infile)
+
+            TORCH_REGISTRATION_SUCCESS = True
+
+        global COMMON_REGISTRATION_SUCCESS
+        if not COMMON_REGISTRATION_SUCCESS:
+            # Pull in some common types so that we can get their associated serialization/deserialization
+            # functions. This allows the user to avoid importing these manually.
+            # Note: We can only do this here for submodules with no external dependencies.
+            # That means, for example, nothing from `backend/` can be imported here.
+            from polygraphy.common import FormattedArray
+            from polygraphy.comparator import RunResults
+
+            COMMON_REGISTRATION_SUCCESS = True
+        return func(*args, **kwargs)
+
+    return wrapped
+
+
+@mod.export()
+@try_register_common_json
+def to_json(obj):
+    """
+    Encode an object to JSON.
+
+    NOTE: For Polygraphy objects, you should use the ``to_json()`` method instead.
+
+    Returns:
+        str: A JSON representation of the object.
+    """
+    return json.dumps(obj, cls=Encoder, indent=constants.TAB)
+
+
+@mod.export()
+@try_register_common_json
+def from_json(src):
+    """
+    Decode a JSON string to an object.
+
+    NOTE: For Polygraphy objects, you should use the ``from_json()`` method instead.
+
+    Args:
+        src (str):
+                The JSON representation of the object
+
+    Returns:
+        object: The decoded instance
+    """
+    return json.loads(src, object_pairs_hook=Decoder())
+
+
+@mod.export()
+@try_register_common_json
+def save_json(obj, dest, description=None):
+    """
+    Encode an object as JSON and save it to a file.
+
+    NOTE: For Polygraphy objects, you should use the ``save()`` method instead.
+
+    Args:
+        obj : The object to save.
+        src (Union[str, file-like]): The path or file-like object to save to.
+    """
+    util.save_file(to_json(obj), dest, mode="w", description=description)
+
+
+@mod.export()
+@try_register_common_json
+def load_json(src, description=None):
+    """
+    Loads a file and decodes the JSON contents.
+
+    NOTE: For Polygraphy objects, you should use the ``load()`` method instead.
+
+    Args:
+        src (Union[str, file-like]): The path or file-like object to load from.
+
+    Returns:
+        object: The object, or `None` if nothing could be read.
+    """
+    return from_json(util.load_file(src, mode="r", description=description))
+
+
+@mod.export()
+def add_json_methods(description=None):
+    """
+    Decorator that adds 4 JSON helper methods to a class:
+
+    - to_json(): Convert to JSON string
+    - from_json(): Convert from JSON string
+    - save(): Convert to JSON and save to file
+    - load(): Load from file and convert from JSON
+
+    Args:
+        description (str):
+                A description of what is being saved or loaded.
+    """
+
+    def add_json_methods_impl(cls):
+        # JSON methods
+
+        def check_decoded(obj):
+            if not isinstance(obj, cls):
+                G_LOGGER.critical(
+                    f"Provided JSON cannot be decoded into a {cls.__name__}.\nNote: JSON was decoded into a {type(obj)}:\n{obj}"
+                )
+            return obj
+
+        def _to_json_method(self):
+            """
+            Encode this instance as a JSON object.
+
+            Returns:
+                str: A JSON representation of this instance.
+            """
+            return to_json(self)
+
+        def _from_json_method(src):
+            return check_decoded(from_json(src))
+
+        _from_json_method.__doc__ = f"""
+            Decode a JSON object and create an instance of this class.
+
+            Args:
+                src (str):
+                        The JSON representation of the object
+
+            Returns:
+                {cls.__name__}: The decoded instance
+
+            Raises:
+                PolygraphyException:
+                        If the JSON cannot be decoded to an instance of {cls.__name__}
+            """
+
+        cls.to_json = _to_json_method
+        cls.from_json = staticmethod(_from_json_method)
+
+        # Save/Load methods
+
+        def _save_method(self, dest):
+            """
+            Encode this instance as a JSON object and save it to the specified path
+            or file-like object.
+
+            Args:
+                dest (Union[str, file-like]):
+                      The path or file-like object to write to.
+
+            """
+            save_json(self, dest, description=description)
+
+        def _load_method(src):
+            return check_decoded(load_json(src, description=description))
+
+        _load_method.__doc__ = f"""
+            Loads an instance of this class from a JSON file.
+
+            Args:
+                src (Union[str, file-like]): The path or file-like object to read from.
+
+            Returns:
+                {cls.__name__}: The decoded instance
+
+            Raises:
+                PolygraphyException:
+                        If the JSON cannot be decoded to an instance of {cls.__name__}
+            """
+
+        cls.save = _save_method
+        cls.load = staticmethod(_load_method)
+
+        return cls
+
+    return add_json_methods_impl

tools/Polygraphy/polygraphy/logger/__init__.py

@@ -0,0 +1 @@
+from polygraphy.logger.logger import G_LOGGER, LogMode

tools/Polygraphy/polygraphy/logger/logger.py

@@ -0,0 +1,700 @@
+#
+# 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 enum
+import functools
+import logging
+import os
+import sys
+import time
+import traceback
+
+COLORED_MODULE_PRESENT = None
+
+
+def has_colors():
+    global COLORED_MODULE_PRESENT
+    if COLORED_MODULE_PRESENT is None:
+        try:
+            import colored
+
+            COLORED_MODULE_PRESENT = True
+        except:
+            COLORED_MODULE_PRESENT = False
+            print(
+                "[W] 'colored' module is not installed, will not use colors when logging. "
+                "To enable colors, please install the 'colored' module: python3 -m pip install colored"
+            )
+    return COLORED_MODULE_PRESENT
+
+
+# Context manager to apply indentation to messages
+class LoggerIndent:
+    def __init__(self, logger, indent):
+        self.logger = logger
+        self.old_indent = self.logger.logging_indent
+        self.indent = indent
+
+    def __enter__(self):
+        self.logger.logging_indent = self.indent
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.logging_indent = self.old_indent
+
+
+# Context manager to temporarily set verbosity
+class LoggerVerbosity:
+    def __init__(self, logger, severity):
+        self.logger = logger
+        self.old_severity = copy.copy(self.logger.module_severity)
+        self.module_severity = severity
+
+    def __enter__(self):
+        self.logger.module_severity = self.module_severity
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.module_severity = self.old_severity
+
+
+class LogMode(enum.IntEnum):
+    """
+    Specifies how messages should be logged.
+    """
+
+    EACH = 0
+    """Log the message each time"""
+    ONCE = 1
+    """Log the message only once. The same message will not be logged again."""
+
+
+class SeverityTrie:
+    """
+    A trie that represents per-path logging verbosities.
+    """
+
+    def _split_path(self, path):
+        # Leading or duplicate slashes can create empty elements in the path components. We ignore those.
+        return list(filter(lambda x: x, path.split(os.path.sep)))
+
+    def __init__(self, severity_dict):
+        assert "" in severity_dict, "severity_dict must include default severity!"
+
+        self.trie = {}
+        for path, severity in severity_dict.items():
+            cur_dict = self.trie
+            for path_component in self._split_path(path):
+                if path_component not in cur_dict:
+                    cur_dict[path_component] = {}
+                cur_dict = cur_dict[path_component]
+            cur_dict[""] = severity
+
+        # Skip path checking if we don't have any path entries.
+        self.has_non_default_entries = len(self.trie) > 1
+
+    def get(self, path=None):
+        """
+        Get the logging verbosity for the given path.
+
+        Args:
+            path (str): The path
+
+        Returns:
+            int: The logging verbosity.
+        """
+        default_severity = self.trie[""]
+        if path is None or not self.has_non_default_entries:
+            return default_severity
+
+        cur_dict = self.trie
+
+        def get_value(dct):
+            return dct.get("", default_severity)
+
+        for path_component in self._split_path(path):
+            if path_component not in cur_dict:
+                return get_value(cur_dict)
+            cur_dict = cur_dict[path_component]
+        return get_value(cur_dict)
+
+    def __str__(self):
+        return str(self.trie)
+
+
+class Logger:
+    """
+    Global logging interface. Do **not** construct a logger manually.
+    Instead, use ``G_LOGGER``, the global logger.
+    """
+
+    ULTRA_VERBOSE = -20  # Cast it into the flames!
+    """Enable unreasonably verbose messages and above"""
+    SUPER_VERBOSE = -10
+    """Enable extremely verbose messages and above"""
+    EXTRA_VERBOSE = 0
+    """Enable extra verbose messages and above"""
+    VERBOSE = 10
+    """Enable verbose messages and above"""
+    INFO = 20
+    """Enable informative messages and above"""
+    START = 22
+    """Enable messages indicating when a task is started and above"""
+    FINISH = 28
+    """Enable messages indicating when a task is finished and above"""
+    WARNING = 30
+    """Enable only warning messages and above"""
+    ERROR = 40
+    """Enable only error messages and above"""
+    CRITICAL = 50
+    """Enable only critical/fatal error messages and above"""
+
+    SEVERITY_LETTER_MAPPING = {
+        ULTRA_VERBOSE: "[U]",
+        SUPER_VERBOSE: "[S]",
+        EXTRA_VERBOSE: "[X]",
+        VERBOSE: "[V]",
+        INFO: "[I]",
+        START: "[I]",
+        FINISH: "[I]",
+        WARNING: "[W]",
+        ERROR: "[E]",
+        CRITICAL: "[!]",
+    }
+
+    SEVERITY_COLOR_MAPPING = {
+        ULTRA_VERBOSE: "dark_gray",
+        SUPER_VERBOSE: "medium_violet_red",
+        EXTRA_VERBOSE: "medium_purple",
+        VERBOSE: "light_magenta",
+        INFO: None,
+        START: "light_cyan",
+        FINISH: "light_green",
+        WARNING: "light_yellow",
+        ERROR: "light_red",
+        CRITICAL: "light_red",
+    }
+
+    def __init__(
+        self, severity=INFO, colors=True, letter=True, timestamp=False, line_info=False
+    ):
+        """
+        Args:
+            severity (Union[int, Dict[str, int]]):
+                    Severity below which messages will be ignored.
+                    This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+                    logging severities. In this case, use the ``""`` to indicate the default severity.
+                    Paths should be relative to the `polygraphy/` directory.
+                    For example, `polygraphy/backend` can be specified with just `backend/`.
+                    For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+                    This is converted to a ``SeverityTrie`` on assignment.
+                    Defaults to G_LOGGER.INFO.
+            colors (bool):
+                    Whether to use colored output.
+                    Defaults to True.
+            letter (bool):
+                    Whether to prepend each logging message with a letter indicating it's severity.
+                    Defaults to True.
+            timestamp (bool):
+                    Whether to include a timestamp in the logging output.
+                    Defaults to False.
+            line_info (bool):
+                    Whether to include file and line number information in the logging output.
+                    Defaults to False.
+        """
+        self.logging_indent = 0
+        self.once_logged = set()
+        self.colors = colors
+        self.letter = letter
+        self.timestamp = timestamp
+        self.line_info = line_info
+        self.logger_callbacks = []
+        self.log_file = None
+        """
+        Path to a log file to write logging output from Polygraphy.
+        This will not include logging messages from libraries used by Polygraphy, like
+        TensorRT or ONNX-Runtime.
+        """
+        self.module_severity = severity
+        """
+        Severity below which messages will be ignored.
+        This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+        logging severities. In this case, use the ``""`` to indicate the default severity.
+        Paths should be relative to the `polygraphy/` directory.
+        For example, `polygraphy/backend` can be specified with just `backend/`.
+        For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+        This is converted to a ``SeverityTrie`` on assignment.
+        Defaults to G_LOGGER.INFO.
+        """
+        self._use_python_logging_system = False
+        """
+        A flag indicating whether to use the Python `logging` module for log emission.
+        By default, logs are emitted to `stdout` or `stderr`. When this flag is set to `True`,
+        the logger uses the Python `logging` module instead of `stdout`/`stderr`.
+        This allows logs to be integrated into a unified logging system which can be helpful
+        in advanced cases like using multiprocessing or when the user wants one logging system
+        to manage Polygraphy and other libraries logs.
+        """
+
+    @property
+    def log_file(self):
+        return self._log_path
+
+    @log_file.setter
+    def log_file(self, value):
+        self._log_path = value
+        self._log_file = None
+        if self._log_path:
+            dir_path = os.path.dirname(self._log_path)
+            if dir_path:
+                dir_path = os.path.realpath(dir_path)
+                os.makedirs(dir_path, exist_ok=True)
+            self._log_file = open(self._log_path, "w")
+
+    @property
+    def module_severity(self):
+        return self._module_severity
+
+    @module_severity.setter
+    def module_severity(self, value):
+        if isinstance(value, SeverityTrie):
+            self._module_severity = value
+        else:
+            if not isinstance(value, dict):
+                value = {"": value}
+            if "" not in value:
+                value[""] = Logger.INFO
+            self._module_severity = SeverityTrie(value)
+
+        self._run_callbacks()
+
+    @property
+    def severity(self):
+        print(
+            "Warning: Accessing the `severity` property of G_LOGGER is deprecated and will be removed in v0.50.0. Use `module_severity` instead"
+        )
+        return self._module_severity.get()
+
+    @severity.setter
+    def severity(self, value):
+        print(
+            "Warning: Accessing the `severity` property of G_LOGGER is deprecated and will be removed in v0.50.0. Use `module_severity` instead"
+        )
+        self.module_severity = value
+
+    @property
+    def use_python_logging_system(self):
+        return self._use_python_logging_system
+
+    @use_python_logging_system.setter
+    def use_python_logging_system(self, value):
+        self._use_python_logging_system = value
+        if value:
+            self.python_logger = logging.getLogger("Polygraphy")
+            self.severity_level_mapping = {
+                Logger.ULTRA_VERBOSE: 2,
+                Logger.SUPER_VERBOSE: 4,
+                Logger.EXTRA_VERBOSE: 6,
+            }
+
+    def module_path(self, path):
+        """
+        Converts a given path to a path relative to the Polygraphy root module.
+        If the path is not part of the Polygraphy module, returns a path relative to the common prefix.
+
+        Args:
+            path (str): The path
+
+        Returns:
+            str: The path relative to the Polygraphy root module or common prefix.
+        """
+        import polygraphy
+
+        module_root_dir = polygraphy.__path__[0]
+        file_path = os.path.relpath(path, module_root_dir)
+        if os.pardir in file_path:
+            common_path_len = len(os.path.commonpath([module_root_dir, path]))
+            file_path = path[common_path_len:].lstrip(os.path.sep)
+        return file_path
+
+    def _run_callbacks(self):
+        for callback in self.logger_callbacks:
+            callback(self._module_severity)
+
+    def register_callback(self, callback):
+        """
+        Registers a callback with the logger, which will be invoked when the logging severity is modified.
+        The callback is guaranteed to be called at least once in the register_callback function.
+
+        Args:
+            callback (Callable(SeverityTrie)):
+                    A callback that accepts the current logger severity trie.
+        """
+        callback(self._module_severity)
+        self.logger_callbacks.append(callback)
+
+    def indent(self, level=1):
+        """
+        Returns a context manager that indents all strings logged by the specified amount.
+
+        Args:
+            level (int): The indentation level
+        """
+        return LoggerIndent(self, level + self.logging_indent)
+
+    def verbosity(self, severity=CRITICAL):
+        """
+        Returns a context manager that temporarily changes the severity of the logger for its duration.
+
+        Args:
+            severity (Union[int, Dict[str, int]]):
+                    Severity below which messages will be ignored.
+                    This can be specified on a per-submodule/file basis by providing a dictionary of paths to
+                    logging severities. In this case, use the ``""`` to indicate the default severity.
+                    For example: ``{"": G_LOGGER.INFO, "backend/trt": G_LOGGER.VERBOSE}``
+                    Defaults to Logger.CRITICAL, which will suppress all messages.
+        """
+        return LoggerVerbosity(self, severity)
+
+    def log(self, message, severity, mode=LogMode.EACH, stack_depth=2, error_ok=False):
+        """
+        Logs a message to stdout.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            severity (int):
+                    The severity with which to log this message. If the severity is less than
+                    the logger's current severity, the message is suppressed. Provided callables
+                    will not be called in that case.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+            stack_depth (int):
+                    The stack depth to use to determine file and line information.
+                    Defaults to 2.
+            error_ok (bool):
+                    Whether to suppress errors encountered while logging.
+                    When this is True, in the event of an error, the message will not be
+                    logged, but the logger will recover and resume execution.
+                    When False, the logger will re-raise the exception.
+        """
+        from polygraphy import config, constants
+
+        def get_rel_file_path_and_lineno():
+            file_path = sys._getframe(stack_depth).f_code.co_filename
+            line_no = sys._getframe(stack_depth).f_lineno
+            # If we can't get a valid path, keep walking the stack until we can.
+            new_stack_depth = stack_depth
+            while not os.path.exists(file_path) and new_stack_depth > 0:
+                new_stack_depth -= 1
+                file_path = sys._getframe(new_stack_depth).f_code.co_filename
+                line_no = sys._getframe(new_stack_depth).f_lineno
+
+            return self.module_path(file_path), line_no
+
+        def process_message(message, file_path, line_no):
+            def get_prefix():
+                prefix = ""
+                if self.letter:
+                    prefix += Logger.SEVERITY_LETTER_MAPPING[severity] + " "
+                if self.timestamp:
+                    prefix += f"({time.strftime('%X')}) "
+                if self.line_info:
+                    prefix += f"[{file_path}:{line_no}] "
+                return prefix
+
+            def apply_indentation(prefix, message):
+                message_lines = str(message).splitlines()
+                tab = constants.TAB * self.logging_indent
+                newline_tab = "\n" + tab + " " * len(prefix)
+                return tab + newline_tab.join([line for line in message_lines])
+
+            def apply_color(message):
+                if self.colors and has_colors():
+                    import colored
+
+                    color = Logger.SEVERITY_COLOR_MAPPING[severity]
+
+                    if color:
+                        try:
+                            color = colored.fore(color)
+                        except:
+                            color = [colored.fg(color)]
+
+                        return colored.stylize(message, color)
+                return message
+
+            prefix = get_prefix()
+            message = apply_indentation(prefix, message)
+            return apply_color(f"{prefix}{message}")
+
+        file_path, line_no = None, None
+        if self.line_info or self.module_severity.has_non_default_entries:
+            file_path, line_no = get_rel_file_path_and_lineno()
+
+        def should_log(message):
+            if severity < self.module_severity.get(file_path):
+                return False
+
+            if mode == LogMode.ONCE:
+                message_hash = hash(message)
+                if message_hash in self.once_logged:
+                    return False
+                self.once_logged.add(message_hash)
+            return True
+
+        if not should_log(message):
+            return
+
+        if callable(message):
+            try:
+                message = message()
+            except Exception as err:
+                if not error_ok or config.INTERNAL_CORRECTNESS_CHECKS:
+                    raise
+                message = f"<Error while logging this message: {str(err)}>"
+
+        message = str(message)
+
+        # Use the warnings module in correctness checking mode so all warnings are
+        # visible in the test result summary.
+        if config.INTERNAL_CORRECTNESS_CHECKS and severity == Logger.WARNING:
+            import warnings
+
+            warnings.warn(message)
+
+        message = process_message(message, file_path, line_no)
+
+        if self._log_file is not None:
+            self._log_file.write(message + "\n")
+            self._log_file.flush()
+
+        if self._use_python_logging_system:
+            # python logging system does not handle negative levels, map them to positive values
+            level = severity if severity > 0 else self.severity_level_mapping[severity]
+            self.python_logger.log(level=level, msg=message)
+        else:
+            print(
+                message, file=sys.stdout if severity < Logger.CRITICAL else sys.stderr
+            )
+
+    def backtrace(self, depth=0, limit=None, severity=ERROR):
+        limit = (
+            limit if limit is not None else (3 - self.module_severity.get() // 10) * 2
+        )  # Info provides 1 stack frame
+        limit = max(limit, 0)
+        frame = sys._getframe(depth + 2)
+        self.log(
+            " ".join(traceback.format_stack(f=frame, limit=limit)), severity=severity
+        )
+
+    def ultra_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with ULTRA_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.ULTRA_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def super_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with SUPER_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.SUPER_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def extra_verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with EXTRA_VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.EXTRA_VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def verbose(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with VERBOSE severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.VERBOSE, mode=mode, stack_depth=3, error_ok=True)
+
+    def info(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with INFO severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.INFO, mode=mode, stack_depth=3)
+
+    def start(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with START severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.START, mode=mode, stack_depth=3)
+
+    def finish(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with FINISH severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.FINISH, mode=mode, stack_depth=3)
+
+    def warning(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with WARNING severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.WARNING, mode=mode, stack_depth=3)
+
+    def error(self, message, mode=LogMode.EACH):
+        """
+        Logs a message to stdout with ERROR severity.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            mode (LogMode):
+                    Controls how the message is logged.
+                    See LogMode for details.
+        """
+        self.log(message, Logger.ERROR, mode=mode, stack_depth=3)
+
+    def critical(self, message, ExceptionType=None):
+        """
+        Logs a message to stdout with CRITICAL severity and raises an exception.
+
+        Args:
+            message (Union[str, Callable() -> str]):
+                    A string or callable which returns a string of the message to log.
+            ExceptionType (type):
+                    The type of exception to raise.
+                    Defaults to PolygraphyException.
+
+        Raises:
+            ExceptionType
+        """
+        self.log(message, Logger.CRITICAL, stack_depth=3)
+        from polygraphy.exception import PolygraphyException
+
+        ExceptionType = ExceptionType or PolygraphyException
+        raise ExceptionType(message) from None
+
+    def internal_error(self, message):
+        from polygraphy import config
+
+        if not config.INTERNAL_CORRECTNESS_CHECKS:
+            return
+
+        self.log(message, Logger.CRITICAL, stack_depth=3)
+        from polygraphy.exception import PolygraphyInternalException
+
+        raise PolygraphyInternalException(message)
+
+    def _str_from_module_info(self, module, name=None):
+        ret = ""
+
+        def try_append(func):
+            nonlocal ret
+            try:
+                ret += func()
+            except:
+                pass
+
+        try_append(lambda: name or f"Loaded Module: {module.__name__}")
+        try_append(lambda: f" | Version: {module.__version__}")
+        try_append(lambda: f" | Path: {list(map(os.path.realpath, module.__path__))}")
+        return ret
+
+    def module_info(self, module, name=None, severity=VERBOSE):
+        message = self._str_from_module_info(module, name)
+        self.log(message, severity=severity, stack_depth=3, mode=LogMode.ONCE)
+
+    def log_exception(self, func):
+        """
+        Decorator that causes exceptions in a function to be logged.
+        This is useful in cases where the exception is caught by a caller, but should
+        still be logged.
+        """
+
+        @functools.wraps(func)
+        def wrapped(*args, **kwargs):
+            from polygraphy.exception import PolygraphyException
+
+            try:
+                return func(*args, **kwargs)
+            except PolygraphyException:
+                # `PolygraphyException`s are always logged.
+                raise
+            except Exception as err:
+                G_LOGGER.error(err)
+                raise
+
+        return wrapped
+
+
+G_LOGGER = Logger()
+"""The global logger. Use this instead of constructing a logger"""
+
+# For backwards compatibility
+G_LOGGER.exit = G_LOGGER.critical

tools/Polygraphy/polygraphy/mod/__init__.py

@@ -0,0 +1,3 @@
+from polygraphy.mod.importer import *
+from polygraphy.mod.exporter import *
+from polygraphy.mod.util import version

tools/Polygraphy/polygraphy/mod/exporter.py

@@ -0,0 +1,362 @@
+#
+# 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 inspect
+import sys
+import warnings
+from textwrap import dedent
+
+import polygraphy
+from polygraphy import config
+from polygraphy.logger import G_LOGGER
+from polygraphy.mod.util import version
+
+
+def _add_to_all(symbol, module):
+    if hasattr(module, "__all__"):
+        module.__all__.append(symbol)
+    else:
+        module.__all__ = [symbol]
+
+
+def _define_in_module(name, symbol, module):
+    assert name not in vars(module), "This symbol is already defined!"
+    vars(module)[name] = symbol
+    _add_to_all(name, module)
+
+
+def export(funcify=False, func_name=None):
+    """
+    Decorator that exports a symbol into the ``__all__`` attribute of
+    the caller's module. This makes the symbol visible in a ``*`` import
+    (e.g. ``from module import *``) and hides other symbols unless they are
+    also present in ``__all__``.
+
+    Args:
+        funcify (bool):
+                Whether to create and export a function that will call a decorated Polygraphy loader.
+                The decorated type *must* be a subclass of ``BaseLoader`` if ``funcify=True``.
+
+                This is useful to provide convenient short-hands to immediately evaluate loaders.
+                For example:
+                ::
+
+                    @mod.export(funcify=True)
+                    class SuperCoolModelFromPath(BaseLoader):
+                        def __init__(self, init_params):
+                            ...
+
+                        def call_impl(self, call_params):
+                            ...
+
+                    # We can now magically access an immediately evaluated functional
+                    # variant of the loader:
+                    model = super_cool_model_from_path(init_params, call_params)
+
+                    # Which is equivalent to:
+                    load_model = SuperCoolModelFromPath(init_params)
+                    model = load_model(call_params)
+
+
+                The signature of the generated function is a combination of the signatures
+                of ``__init__`` and ``call_impl``. Specifically, parameters without defaults will
+                precede those with defaults, and ``__init__`` parameters will precede ``call_impl``
+                parameters. Special parameters like ``*args`` and ``**kwargs`` will always be the last
+                parameters in the generated signature if they are present in the loader method signatures.
+                The return value(s) will always come from ``call_impl``.
+
+                For example:
+                ::
+
+                    # With __init__ signature:
+                    def __init__(a, b=0) -> None:
+
+                    # And call_impl signature:
+                    def call_impl(c, d=0) -> z:
+
+                    # The generated function will have a signature:
+                    def generated(a, c, b=0, d=0) -> z:
+
+        func_name (str):
+                If funcify is True, this controls the name of the generated function.
+                By default, the exported function will use the same name as the loader, but
+                ``snake_case`` instead of ``PascalCase``.
+    """
+    module = inspect.getmodule(sys._getframe(1))
+
+    # Find a method by wallking the inheritance hierarchy of a type:
+    def find_method(symbol, method):
+        hierarchy = inspect.getmro(symbol)
+        for ancestor in hierarchy:
+            if method in vars(ancestor):
+                return vars(ancestor)[method]
+
+        assert (
+            False
+        ), f"Could not find method: {method} in the inheritance hierarcy of: {symbol}"
+
+    def export_impl(func_or_cls):
+        _add_to_all(func_or_cls.__name__, module)
+
+        if funcify:
+            # We only support funcify-ing BaseLoaders, and only if __init__ and call_impl
+            # have no overlapping parameters.
+            from polygraphy.backend.base import BaseLoader
+
+            assert inspect.isclass(
+                func_or_cls
+            ), "Decorated type must be a loader to use funcify=True"
+            assert BaseLoader in inspect.getmro(
+                func_or_cls
+            ), "Decorated type must derive from BaseLoader to use funcify=True"
+
+            def get_params(method):
+                return list(
+                    inspect.signature(
+                        find_method(func_or_cls, method)
+                    ).parameters.values()
+                )[1:]
+
+            def is_variadic(param):
+                return param.kind in [param.VAR_POSITIONAL, param.VAR_KEYWORD]
+
+            def has_default(param):
+                return param.default != param.empty
+
+            def get_param_name(p):
+                # For variadic arguments, p.name will drop the *, **
+                return str(p) if is_variadic(p) else p.name
+
+            def param_names(params):
+                return [get_param_name(p) for p in params]
+
+            loader = func_or_cls
+
+            init_params = get_params("__init__")
+            call_impl_params = get_params("call_impl")
+
+            assert (
+                set(param_names(call_impl_params)) - set(param_names(init_params))
+            ) == set(
+                param_names(call_impl_params)
+            ), "Cannot funcify a type where call_impl and __init__ have the same argument names!"
+
+            # Dynamically generate a function with the right signature.
+
+            # To generate the signature, we use the init and call_impl arguments,
+            # but move required arguments (i.e. without default values) to the front.
+
+            def build_arg_list(should_include):
+                def str_from_param(p):
+                    return get_param_name(p) + (
+                        f"={p.default}" if has_default(p) else ""
+                    )
+
+                arg_list = [str_from_param(p) for p in init_params if should_include(p)]
+                arg_list += [
+                    str_from_param(p) for p in call_impl_params if should_include(p)
+                ]
+                return arg_list
+
+            non_default_args = build_arg_list(
+                should_include=lambda p: not is_variadic(p) and not has_default(p)
+            )
+            default_args = build_arg_list(
+                should_include=lambda p: not is_variadic(p) and has_default(p)
+            )
+            special_args = build_arg_list(should_include=is_variadic)
+
+            signature = ", ".join(non_default_args + default_args + special_args)
+
+            init_args = ", ".join(param_names(init_params))
+            call_impl_args = ", ".join(param_names(call_impl_params))
+
+            def pascal_to_snake(name):
+                return "".join(
+                    f"_{c.lower()}" if c.isupper() else c for c in name
+                ).lstrip("_")
+
+            nonlocal func_name
+            func_name = func_name or pascal_to_snake(loader.__name__)
+
+            func_code = dedent(
+                f"""
+                def {func_name}({signature}):
+                    return loader_binding({init_args})({call_impl_args})
+
+                func_var = {func_name}
+                """
+            )
+
+            new_locals = {}
+            exec(
+                func_code,
+                # Need to bind the loader this way, or it won't be accesible from func_code.
+                {"loader_binding": loader},
+                new_locals,
+            )
+            func = new_locals["func_var"]
+
+            # Next we setup the docstring so that it is a combination of the __init__
+            # and call_impl docstrings.
+            func.__doc__ = f"Immediately evaluated functional variant of :class:`{loader.__name__}` .\n"
+
+            def try_add_method_doc(method):
+                call_impl = find_method(loader, method)
+                if call_impl.__doc__:
+                    func.__doc__ += dedent(call_impl.__doc__)
+
+            try_add_method_doc("__init__")
+            try_add_method_doc("call_impl")
+
+            # Now that the function has been defined, we just need to add it into the module's
+            # __dict__ so it is accessible like a normal symbol.
+            _define_in_module(func_name, func, module)
+
+        # We don't actually want to modify the decorated object.
+        return func_or_cls
+
+    return export_impl
+
+
+def warn_deprecated(
+    name, use_instead, remove_in, module_name=None, always_show_warning=False
+):
+
+    if version(polygraphy.__version__) >= version(remove_in):
+        G_LOGGER.internal_error(
+            f"{name} should have been removed in version: {remove_in}"
+        )
+
+    full_obj_name = f"{module_name}.{name}" if module_name else name
+    msg = (
+        f"{full_obj_name} is deprecated and will be removed in Polygraphy {remove_in}."
+    )
+    if use_instead is not None:
+        msg += f" Use {use_instead} instead."
+
+    warnings.warn(msg, DeprecationWarning, stacklevel=3)
+    if always_show_warning:
+        G_LOGGER.warning(msg)
+
+
+def deprecate(remove_in, use_instead, module_name=None, name=None):
+    """
+    Decorator that marks a function or class as deprecated.
+    When the function or class is used, a warning will be issued.
+
+    Args:
+        remove_in (str):
+                The version in which the decorated type will be removed.
+        use_instead (str):
+                The function or class to use instead.
+        module_name (str):
+                The name of the containing module. This will be used to
+                generate more informative warnings.
+                Defaults to None.
+        name (str):
+                The name of the object being deprecated.
+                If not provided, this is automatically determined based on the decorated type.
+                Defaults to None.
+    """
+
+    def deprecate_impl(obj):
+        if config.INTERNAL_CORRECTNESS_CHECKS and version(
+            polygraphy.__version__
+        ) >= version(remove_in):
+            G_LOGGER.internal_error(
+                f"{obj} should have been removed in version: {remove_in}"
+            )
+
+        nonlocal name
+        name = name or obj.__name__
+
+        if inspect.ismodule(obj):
+
+            class DeprecatedModule:
+                def __getattr__(self, attr_name):
+                    warn_deprecated(name, use_instead, remove_in, module_name)
+                    self = obj
+                    return getattr(self, attr_name)
+
+                def __setattr__(self, attr_name, value):
+                    warn_deprecated(name, use_instead, remove_in, module_name)
+                    self = obj
+                    return setattr(self, attr_name, value)
+
+            DeprecatedModule.__doc__ = f"Deprecated: Use {use_instead} instead"
+            return DeprecatedModule()
+        elif inspect.isclass(obj):
+
+            class Deprecated(obj):
+                def __init__(self, *args, **kwargs):
+                    warn_deprecated(name, use_instead, remove_in, module_name)
+                    super().__init__(*args, **kwargs)
+
+            Deprecated.__doc__ = f"Deprecated: Use {use_instead} instead"
+            return Deprecated
+        elif inspect.isfunction(obj):
+
+            def wrapped(*args, **kwargs):
+                warn_deprecated(name, use_instead, remove_in, module_name)
+                return obj(*args, **kwargs)
+
+            wrapped.__doc__ = f"Deprecated: Use {use_instead} instead"
+            return wrapped
+        else:
+            G_LOGGER.internal_error(f"deprecate is not implemented for: {obj}")
+
+    return deprecate_impl
+
+
+def export_deprecated_alias(name, remove_in, use_instead=None):
+    """
+    Decorator that creates and exports a deprecated alias for
+    the decorated class or function.
+
+    The alias will behave like the decorated type, except it will
+    issue a deprecation warning when used.
+
+    To create a deprecated alias for an entire module, invoke the
+    function manually within the module like so:
+    ::
+
+        mod.export_deprecated_alias("old_mod_name", remove_in="0.0.0")(sys.modules[__name__])
+
+    Args:
+        name (str):
+                The name of the deprecated alias.
+        remove_in (str):
+                The version, as a string, in which the deprecated alias will be removed.
+        use_instead (str):
+                The name of the function, class, or module to use instead.
+                If this is ``None``, the new name will be automatically determined.
+                Defaults to None.
+    """
+    module = inspect.getmodule(sys._getframe(1))
+
+    def export_deprecated_alias_impl(obj):
+        new_obj = deprecate(
+            remove_in,
+            use_instead=use_instead or obj.__name__,
+            module_name=module.__name__,
+            name=name,
+        )(obj)
+        _define_in_module(name, new_obj, module)
+        return obj
+
+    return export_deprecated_alias_impl

tools/Polygraphy/polygraphy/mod/importer.py

@@ -0,0 +1,377 @@
+#
+# 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 importlib
+import importlib.util
+import os
+import subprocess as sp
+import sys
+from typing import List
+
+try:
+    # Available in Python 3.8+
+    import importlib.metadata
+except ModuleNotFoundError:
+    pass
+
+from polygraphy import constants
+from polygraphy.mod import util as mod_util
+
+# Tracks all of Polygraphy's lazy imports, excluding internal ones.
+_all_external_lazy_imports = set()
+
+# Sometimes the Python package name differs from the module name.
+_PKG_NAME_FROM_MODULE = {}
+
+# Some packages need additional flags to install correctly.
+_EXTRA_FLAGS_FOR_MODULE = {
+    "onnx_graphsurgeon": ["--extra-index-url=https://pypi.ngc.nvidia.com"],
+}
+
+
+LATEST_VERSION = "==latest"
+"""Indicates that the latest version of the package is preferred in lazy_import"""
+
+
+def _version_ok(ver, preferred):
+    if preferred == LATEST_VERSION:
+        return False
+
+    pref_ver = preferred.lstrip("<=>").strip()
+    cond = preferred.rstrip(pref_ver).strip()
+    check = {
+        "==": lambda x, y: x == y,
+        ">=": lambda x, y: x >= y,
+        ">": lambda x, y: x > y,
+        "<=": lambda x, y: x <= y,
+        "<": lambda x, y: x < y,
+    }[cond]
+    return check(mod_util.version(ver), mod_util.version(pref_ver))
+
+
+def lazy_import(
+    name: str,
+    log: bool = None,
+    pkg_name: str = None,
+    install_flags: List[str] = None,
+    requires: List[str] = None,
+):
+    """
+    Lazily import a module.
+
+    If config.AUTOINSTALL_DEPS is set to 1,
+    missing modules are automatically installed, and existing modules may be
+    upgraded if newer versions are required.
+
+    Args:
+        name (str):
+                The name of the module and optionally the preferred version of the package,
+                formatted as a version string. For example, ``'example_module>=0.5.0'`` or ``'example_module==1.8.0'``.
+        log (bool):
+                Whether to log information about the module.
+                Defaults to True.
+        pkg_name (str):
+                The name of the package that provides this module, if it is different from the module name.
+                Used only if automatic installation of dependencies is enabled.
+        install_flags (List[str]):
+                Additional flags to provide to the installation command.
+                Used only if automatic installation of dependencies is enabled.
+        requires (List[str]):
+                Additional dependencies required by the module which are *not* specified as dependencies.
+                This parameter should only be required when a module does not correctly specify dependencies.
+                Defaults to [].
+
+    Returns:
+        LazyModule:
+                A lazily loaded module. When an attribute is first accessed,
+                the module will be imported.
+    """
+
+    def issue_wrong_version_error(installed_version, version):
+        from polygraphy.logger import G_LOGGER, LogMode
+
+        G_LOGGER.error(
+            f"Module: '{name}' version '{installed_version}' is installed, but version '{version}' is required.\n"
+            f"Please install the required version or set POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables "
+            f"to allow Polygraphy to do so automatically.\n"
+            f"Attempting to continue with the currently installed version of this module, but note that this may cause errors!",
+            mode=LogMode.ONCE,
+        )
+
+    VERSION_CHARS = ["=", ">", "<"]
+
+    log = True if log is None else log
+    requires = [] if requires is None else requires
+
+    def split_name_version(inp):
+        version_char_indices = [
+            inp.index(char) for char in VERSION_CHARS if char in inp
+        ]
+        if not version_char_indices:
+            return inp, None
+
+        min_index = min(version_char_indices)
+        return inp[:min_index], inp[min_index:]
+
+    name, version = split_name_version(name)
+    all_required_mods = list(map(split_name_version, requires)) + [(name, version)]
+
+    if "polygraphy" not in name:
+        _all_external_lazy_imports.add(name)
+
+    def import_mod():
+        from polygraphy import config
+        from polygraphy.logger import G_LOGGER, LogMode
+
+        def install_mod(install_name, install_version, raise_error=True):
+            modname = install_name.split(".")[0]
+            pkg = (
+                pkg_name
+                if pkg_name is not None
+                else _PKG_NAME_FROM_MODULE.get(modname, modname)
+            )
+            extra_flags = (
+                install_flags
+                if install_flags is not None
+                else _EXTRA_FLAGS_FOR_MODULE.get(modname, [])
+            )
+
+            def fail():
+                log_func = G_LOGGER.critical if raise_error else G_LOGGER.warning
+                log_func(
+                    f"Could not automatically install required module: {pkg}. Please install it manually."
+                )
+
+            if config.ASK_BEFORE_INSTALL:
+                res = None
+                while res not in ["y", "n"]:
+                    res = input(
+                        f"Automatically install '{pkg}' (version: {install_version or 'any'}) ([Y]/n)? "
+                    )
+                    res = res.strip()[:1].lower() or "y"
+
+                if res == "n":
+                    fail()
+
+            if install_version == LATEST_VERSION:
+                extra_flags.append("--upgrade")
+            elif install_version is not None:
+                pkg += install_version
+
+            cmd = config.INSTALL_CMD + [pkg] + extra_flags
+            G_LOGGER.info(f"Running installation command: {' '.join(cmd)}")
+            status = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
+            if status.returncode != 0:
+                G_LOGGER.error(
+                    f"Error during installation:\n{constants.TAB}{status.stderr.decode()}"
+                )
+                fail()
+
+            mod = importlib.import_module(install_name)
+            return mod
+
+        mod = None
+        try:
+            mod = importlib.import_module(name)
+        except ImportError as err:
+            if config.AUTOINSTALL_DEPS:
+                for install_name, install_version in all_required_mods:
+                    G_LOGGER.info(
+                        f"Module: '{install_name}' is required, but not installed. Attempting to install now."
+                    )
+                    mod = install_mod(install_name, install_version)
+            else:
+                G_LOGGER.critical(
+                    f"Module: '{name}' is required but could not be imported.\nNote: Error was: {err}\n"
+                    f"You can set POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables to allow "
+                    f"Polygraphy to automatically install missing modules.\n"
+                )
+
+        # Auto-upgrade if necessary
+        for install_name, install_version in all_required_mods:
+            installed_mod = importlib.import_module(install_name)
+            if (
+                install_version is not None
+                and hasattr(installed_mod, "__version__")
+                and not _version_ok(installed_mod.__version__, install_version)
+            ):
+                if config.AUTOINSTALL_DEPS:
+                    G_LOGGER.info(
+                        f"Note: Module: '{install_name}' version '{installed_mod.__version__}' is installed, but version '{install_version}' is required.\n"
+                        f"Attempting to upgrade now."
+                    )
+                    # We can try to use the other version if install fails, so this is non-fatal.
+                    installed_mod = install_mod(
+                        install_name, install_version, raise_error=False
+                    )
+                    if install_name == name:
+                        mod = installed_mod
+
+                elif install_version != LATEST_VERSION:
+                    issue_wrong_version_error(
+                        installed_mod.__version__, install_version
+                    )
+
+        if log:
+            G_LOGGER.module_info(mod)
+
+        return mod
+
+    MODULE_VAR_NAME = "module"
+
+    class LazyModule:
+        def __init__(self):
+            super().__setattr__(MODULE_VAR_NAME, None)
+
+        def __polygraphy_import_mod(self):
+            if self.module is None:
+                super().__setattr__(MODULE_VAR_NAME, import_mod())
+            return self.module
+
+        def __getattr__(self, name):
+            module = self.__polygraphy_import_mod()
+            return getattr(module, name)
+
+        def __setattr__(self, name, value):
+            module = self.__polygraphy_import_mod()
+            return setattr(module, name, value)
+
+        def is_installed(self):
+            """
+            Checks whether any version of this module is installed.
+            The module will not be imported by this method.
+
+            Returns:
+                bool: Whether the module is installed.
+            """
+            global importlib
+
+            try:
+                return name in sys.modules or (
+                    importlib.util.find_spec(name) is not None
+                )
+            except:
+                return False
+
+        def is_importable(self):
+            """
+            Checks whether this module is importable. Note that a module may be installed but not importable.
+
+            Returns:
+                bool: Whether the module is importable.
+            """
+            try:
+                importlib.import_module(name)
+                return True
+            except:
+                return False
+
+    return LazyModule()
+
+
+def has_mod(modname):
+    """
+    Checks whether a module is installed without importing the module.
+
+    Args:
+        modname (str): The name of the module to check.
+
+    Returns:
+        bool: Whether the module is installed.
+    """
+    import warnings
+
+    import polygraphy
+    from polygraphy.logger import G_LOGGER
+
+    remove_in = "0.50.0"
+    if mod_util.version(polygraphy.__version__) >= mod_util.version(remove_in):
+        G_LOGGER.internal_error(
+            f"has_mod should have been removed in version: {remove_in}"
+        )
+    warnings.warn(
+        f"has_mod is deprecated and will be removed in Polygraphy {remove_in}",
+        DeprecationWarning,
+        stacklevel=3,
+    )
+
+    try:
+        return modname in sys.modules or (importlib.util.find_spec(modname) is not None)
+    except ValueError:
+        return False
+
+
+def autoinstall(lazy_mod):
+    """
+    If the config.AUTOINSTALL_DEPS is set to 1, automatically install or upgrade a module.
+    Does nothing if autoinstallation is disabled.
+
+    Args:
+        lazy_mod (LazyModule):
+                A lazy module, like that returned by ``lazy_import``.
+    """
+    from polygraphy import config
+
+    if not config.AUTOINSTALL_DEPS:
+        return
+
+    try:
+        # It doesn't matter which attribute we try to get as any call to `__getattr__` will
+        # trigger the automatic installation.
+        getattr(lazy_mod, "__fake_polygraphy_autoinstall_attr")
+    except:
+        pass
+
+
+def import_from_script(path, name):
+    """
+    Imports a specified symbol from a Python script.
+
+    Args:
+        path (str): A path to the Python script. The path must include a '.py' extension.
+        name (str): The name of the symbol to import from the script.
+
+    Returns:
+        object: The loaded symbol.
+    """
+    from polygraphy.logger import G_LOGGER
+
+    dir = os.path.dirname(path)
+    modname = os.path.splitext(os.path.basename(path))[0]
+
+    sys.path.insert(0, dir)
+
+    with contextlib.ExitStack() as stack:
+
+        def reset_sys_path():
+            del sys.path[0]
+            del sys.modules[modname]
+
+        stack.callback(reset_sys_path)
+
+        try:
+            importlib.invalidate_caches()
+            mod = importlib.import_module(modname)
+            return getattr(mod, name)
+        except Exception as err:
+            ext = os.path.splitext(path)[1]
+            err_msg = f"Could not import symbol: {name} from script: {path}"
+            if ext != ".py":
+                err_msg += f"\nThis could be because the extension of the file is not '.py'. Note: The extension is: {ext}"
+            err_msg += f"\nNote: Error was: {err}"
+            G_LOGGER.critical(err_msg)

tools/Polygraphy/polygraphy/mod/trt_importer.py

@@ -0,0 +1,39 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 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, mod
+
+def tensorrt_module_and_version_string():
+    """
+    Returns the name of the TensorRT module to import. This selects between
+    TensorRT and TensorRT-RTX based on the value of config.USE_TENSORRT_RTX,
+    and ensures that a consistent version of the module is imported.
+    """
+    if config.USE_TENSORRT_RTX:
+        return "tensorrt_rtx>=1.0"
+    else:
+        return "tensorrt>=8.5"
+
+def lazy_import_trt():
+    """
+    Returns either tensorrt or tensorrt_rtx based on config.USE_TENSORRT_RTX.
+    
+    Prefer to use this function instead of mod.lazy_import("tensorrt>=8.5") to
+    import TensorRT, so that your code can use TensorRT-RTX.
+    """
+
+    return mod.lazy_import(tensorrt_module_and_version_string())

tools/Polygraphy/polygraphy/mod/util.py

@@ -0,0 +1,41 @@
+#
+# 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.
+#
+
+
+def version(version_str):
+    def process_version_part(num):
+        suffix = None
+        if "+" in num:
+            num, suffix = num.split("+")
+
+        try:
+            num = int(num)
+        except ValueError:
+            VERSION_SUFFIXES = ["a", "b", "rc", "post", "dev"]
+            # One version part can only contain one of the above suffixes
+            for version_suffix in VERSION_SUFFIXES:
+                if version_suffix in num:
+                    num = num.partition(version_suffix)
+                    break
+
+        return [num, suffix] if suffix is not None else [num]
+
+    ver_list = []
+    for num in version_str.split("."):
+        ver_list.extend(process_version_part(num))
+
+    return tuple(ver_list)

tools/Polygraphy/polygraphy/tools/README.md

@@ -0,0 +1,182 @@
+# Polygraphy Command-line Toolkit User Guide
+
+## Table of Contents
+
+- [Introduction](#introduction)
+- [Common Use-Cases](#common-use-cases)
+    - [Inspecting A Model](#inspecting-a-model)
+    - [Converting A Model To TensorRT](#converting-a-model-to-tensorrt)
+    - [Linting An ONNX Model](#linting-an-onnx-model)
+    - [Sanitizing An ONNX Model](#sanitizing-an-onnx-model)
+    - [Comparing A Model Between Frameworks](#comparing-a-model-between-frameworks)
+    - [Modifying Input Shapes In An ONNX Model](#modifying-input-shapes-in-an-onnx-model)
+- [Advanced Topics](#advanced-topics)
+    - [Deterministic Engine Builds](#deterministic-engine-builds)
+    - [Defining A Custom TensorRT Network Or Builder Configuration](#defining-a-custom-tensorrt-network-or-builder-configuration)
+    - [Extracting A Subgraph Of An ONNX Model](#extracting-a-subgraph-of-an-onnx-model)
+    - [Debugging Intermittent TensorRT Failures](#debugging-intermittent-tensorrt-failures)
+    - [Reducing Failing ONNX Models](#reducing-failing-onnx-models)
+- [Examples](#examples)
+- [Adding New Tools](#adding-new-tools)
+
+
+## Introduction
+
+The Polygraphy command-line toolkit includes several tools covering a wide variety of prototyping
+and debugging use-cases. This guide provides a broad overview of the capabilities included.
+
+For more information about a specific tool, see the README in the corresponding directory here.
+
+All the tools provided by Polygraphy can be invoked using the polygraphy binary: `polygraphy`.
+
+For usage information on a specific tool, you can refer to the help output with: `polygraphy <subtool> -h`.
+
+*NOTE: Some of the tools included are still experimental. Any tool labeled `[EXPERIMENTAL]`*
+*may be subject to backwards-incompatible changes, or even complete removal at any point in time.*
+
+## Common Use-Cases
+
+
+### Inspecting A Model
+
+It is often useful to inspect various aspects of a model, such as layer names, attributes,
+or overall structure. The `inspect model` subtool aims to provide this functionality in a way
+that is conducive to programmatic analysis; that is, it provides a human-readable text representation
+of the model (and it's `grep`-friendly!).
+
+For more details, refer to the examples, which, among other things, show how to inspect:
+- [TensorRT Networks](../../examples/cli/inspect/01_inspecting_a_tensorrt_network/)
+- [TensorRT Engines](../../examples/cli/inspect/02_inspecting_a_tensorrt_engine/)
+- [ONNX models](../../examples/cli/inspect/03_inspecting_an_onnx_model/)
+
+You can find the complete listing of `inspect` examples [here](../../examples/cli/inspect).
+
+
+### Converting A Model To TensorRT
+
+The `convert` tool can be used to convert various types of models to other formats;
+for example, converting ONNX models to TensorRT.
+
+For more information, refer to the examples, which, among other things, show how to:
+- [Convert models with dynamic shapes to TensorRT](../../examples/cli/convert/03_dynamic_shapes_in_tensorrt/)
+- [Run TensorRT INT8 calibration during conversion](../../examples/cli/convert/01_int8_calibration_in_tensorrt)
+
+You can find the complete listing of `convert` examples [here](../../examples/cli/convert/).
+
+### Linting An ONNX Model
+
+Validating an ONNX model can be more than checking if it passes ONNX specifications.
+For example, it could be data-dependant, or depend on the underlying runtime.
+Moreover, an ONNX model can be broken in multiple places that may only be uncovered iteratively.
+
+The subtool `check lint` validates the model for specific use-cases (dynamic shape, custom data, custom weights etc.).
+When model is broken, it also attempts to catch all independent exceptions/warnings in one go.
+
+For more details, refer to [`check` example 01](../../examples/cli/check/01_linting_an_onnx_model/).
+
+### Sanitizing An ONNX Model
+
+Sometimes, TensorRT may be unable to import an ONNX model. In these cases, it often helps to
+sanitize the ONNX model, removing excess nodes and folding constants. The `surgeon sanitize`
+subtool provides this capability using ONNX-GraphSurgeon and ONNX-Runtime.
+
+For more details, refer to [`surgeon` example 02](../../examples/cli/surgeon/02_folding_constants/).
+
+You can find the complete listing of `surgeon` examples [here](../../examples/cli/surgeon/).
+
+
+### Comparing A Model Between Frameworks
+
+The `run` tool can run inference and compare outputs generated by one or more frameworks.
+This can be used to check the accuracy of a framework for a particular model and debug
+when the accuracy is poor.
+
+Moreover, it can be used to generate Python scripts that do the same.
+
+For more information, refer to the examples, which, among other things, show how to:
+- [Compare outputs between frameworks](../../examples/cli/run/01_comparing_frameworks/)
+- [Save and load outputs to compare across runs](../../examples/cli/run/02_comparing_across_runs)
+- [Generate comparison scripts](../../examples/cli/run/03_generating_a_comparison_script/)
+
+
+You can find the complete listing of `run` examples [here](../../examples/cli/run/).
+
+
+### Modifying Input Shapes In An ONNX Model
+
+The best way to modify input shapes in an ONNX model is to re-export the model with
+the desired input shapes. When this is not possible, you can use the `surgeon sanitize`
+tool to do this.
+
+For more information, refer to [`surgeon` example 03](../../examples/cli/surgeon/03_modifying_input_shapes/).
+
+
+## Advanced Topics
+
+### Deterministic Engine Builds
+
+Because of the way TensorRT works, the engine building process is non-deterministic.
+Even when using the same model with the same builder configuration, engines can vary slightly.
+
+To make engine building reproducible, all Polygraphy CLI tools that involve building TensorRT engines
+accept `--save-tactics` and `--load-tactics` options, which allow you to save the tactics selected
+for an engine and reload them during subsequent builds.
+
+For more information, refer to [`convert` example 02](../../examples/cli/convert/02_deterministic_engine_builds_in_tensorrt/).
+
+
+### Defining A Custom TensorRT Network Or Builder Configuration
+
+Many of the command-line tools involve creating TensorRT networks. Most of the time, networks
+are created by parsing a model from a framework (generally in ONNX format). However, it
+is also possible to define the TensorRT network manually using a Python script.
+
+This is useful if you want to modify the network in some way using the TensorRT Python
+API; For example, setting layer precisions, or per-layer device preferences.
+
+Similarly, it is also possible to provide a custom builder configuration.
+
+This is useful in cases where Polygraphy may not yet support certain TensorRT builder configuration options.
+
+For more information, refer to [`run` example 04](../../examples/cli/run/04_defining_a_tensorrt_network_or_config_manually).
+
+
+### Extracting A Subgraph Of An ONNX Model
+
+When debugging ONNX models, it can be useful to extract subgraphs to look at them in isolation.
+The `surgeon extract` tool allows you to do so.
+
+For more information, refer to [`surgeon` example 01](../../examples/cli/surgeon/01_isolating_subgraphs/).
+
+
+### Debugging Intermittent TensorRT Failures
+
+Since the TensorRT optimizer is inherently non-deterministic, rebuilding an engine may
+result in different tactics being selected. If a particular tactic is faulty, this may manifest
+as an intermittent failure. The `debug build` tools helps find faulty tactics in such cases
+and reproduce failures reliably.
+
+For more information, refer to [`debug` example 01](../../examples/cli/debug/01_debugging_flaky_trt_tactics/).
+
+
+### Reducing Failing ONNX Models
+
+When investigating bugs involving ONNX models, it can be useful to reduce the model to a minimum
+failing subgraph. This helps us pinpoint the issue and makes further debugging much easier.
+The `debug reduce` tools helps us automate this process.
+
+For more information, refer to [`debug` example 02](../../examples/cli/debug/02_reducing_failing_onnx_models/).
+
+
+## Examples
+
+For examples, see the corresponding subdirectory in [examples/cli](../../examples/cli)
+
+
+## Adding New Tools
+
+You can add a new tool by adding a new file in this directory, creating a
+class that extends the [`Tool` base class](./base/tool.py), and adding
+the new tool to the [registry](./registry.py).
+
+For details on developing tools, see [this example](../../examples/dev/01_writing_cli_tools/).

tools/Polygraphy/polygraphy/tools/__init__.py

@@ -0,0 +1,3 @@
+from polygraphy.tools.base import *
+from polygraphy.tools.registry import TOOL_REGISTRY
+from polygraphy.tools._main import *

tools/Polygraphy/polygraphy/tools/_main.py

@@ -0,0 +1,83 @@
+#
+# 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 import mod
+
+
+@mod.export()
+def main(run_opts = None):
+    """
+    The Polygraphy CLI Toolkit
+
+    Includes various subtools that can assist with prototyping and
+    debugging inference with deep learning models. See the help output
+    for details.
+    """
+    import argparse
+    import sys
+    import time
+
+    import polygraphy
+    from polygraphy.exception import PolygraphyException
+    from polygraphy.logger import G_LOGGER
+    from polygraphy.tools.registry import TOOL_REGISTRY
+
+    start_time = time.time()
+    cmd = f"{' '.join(sys.argv)}"
+
+    parser = argparse.ArgumentParser(
+        description="Polygraphy: A Deep Learning Debugging Toolkit",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "-v", "--version", action="version", version=G_LOGGER._str_from_module_info(polygraphy, name="Polygraphy")
+    )
+
+    subparsers = parser.add_subparsers(title="Tools", dest="tools")
+    subparsers.required = True
+
+    for tool in TOOL_REGISTRY:
+        tool.setup_parser(subparsers)
+
+    if run_opts is not None:
+        args, unknown = parser.parse_known_args(run_opts)
+    else:
+        args, unknown = parser.parse_known_args()
+
+    if unknown:
+        G_LOGGER.error(f"Unrecognized Options: {unknown}")
+        return 1
+
+    selected_tool = args.subcommand
+    show_start_end_logging = False
+
+    try:
+        selected_tool.parse(args)
+        show_start_end_logging = selected_tool.show_start_end_logging(args)
+
+        if show_start_end_logging:
+            G_LOGGER.start(f"RUNNING | Command: {cmd}")
+
+        status = selected_tool.run(args)
+    except PolygraphyException:
+        # `PolygraphyException`s indicate user error, so we need not display the stack trace.
+        status = 1
+
+    end_time = time.time()
+    if show_start_end_logging:
+        log_func = G_LOGGER.finish if status == 0 else G_LOGGER.error
+        log_func(f"{'PASSED' if status == 0 else 'FAILED'} | Runtime: {end_time - start_time:.3f}s | Command: {cmd}")
+
+    return status

tools/Polygraphy/polygraphy/tools/args/README.md

@@ -0,0 +1,43 @@
+# Command-Line Argument Groups
+
+## Introduction
+
+The Polygraphy command-line toolkit includes several tools, many of which require the same or similar functionality as other tools.
+To facilitate efficient code reuse, commonly used functionality is bundled into reusable components, which we'll refer
+to as `Argument Group`s.
+
+## Argument Groups
+
+An `Argument Group` combines a set of command-line options, parsing logic, and functionality related to those arguments.
+The interface is defined in [BaseArgs](./base.py). The most commonly used methods are:
+
+- `add_parser_args(self, parser)`: Adds options to an `argparse` argument parser.
+
+- `parse(self, args)`: Parses command-line arguments from the `args` object created by `argparse` and
+    populates attributes in the argument group.
+    For simple arguments, this amounts to assigning an attribute of the argument group to a
+    corresponding attribute in the `args` object.
+    For more complex arguments, such as input shapes, this may involve more complex parsing.
+
+- `add_to_script(self, script)`: Adds code to a Python script that will provide functionality
+    related to the argument group. For example, `OnnxLoadArgs`, which is responsible for
+    loading ONNX models, may populate the script with a `OnnxFromPath` loader.
+
+    *NOTE: You may be wondering why we add code to a script instead of just providing a*
+    *method that will perform the relevant function, e.g. `load_onnx` in this case.*
+    *The reason is that adding to a script allows tools like `polygraphy run` to compose*
+    *together complex behavior and generate Python scripts which can be edited and used for more advanced needs.*
+
+Many of the argument groups *also* provide immediately evaluated helper methods, such as `load_onnx()` in the case of
+`OnnxLoadArgs`, which can be used by tools that do not need to generate scripts.
+These helpers typically use the `run_script` method from `polygraphy.tools.args.util` to reuse the logic from their
+`add_to_script` method.
+
+## Usage
+
+Tools can subscribe to argument groups by implementing the `get_subscriptions()` interface defined in [tool.py](../base/tool.py).
+This will add all the command-line options provided by the argument group to the tool, and these will be parsed
+automatically before the tool's `run` method is called.
+The tool can then access the argument groups via `self.arg_groups`.
+
+For details, see [the example](../../../examples/dev/01_writing_cli_tools).

tools/Polygraphy/polygraphy/tools/args/__init__.py

@@ -0,0 +1,19 @@
+#
+# 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.tools.args.backend import *
+from polygraphy.tools.args.comparator import *
+from polygraphy.tools.args.logger import *
+from polygraphy.tools.args.model import *

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

@@ -0,0 +1,21 @@
+#
+# 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.tools.args.backend.onnx import *
+from polygraphy.tools.args.backend.onnxrt import *
+from polygraphy.tools.args.backend.pluginref import *
+from polygraphy.tools.args.backend.tf import *
+from polygraphy.tools.args.backend.trt import *
+from polygraphy.tools.args.backend.runner_select import *

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

@@ -0,0 +1,16 @@
+#
+# 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.tools.args.backend.onnx.loader import *

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

@@ -0,0 +1,754 @@
+#
+# 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 os
+
+from polygraphy import constants, mod, util
+from polygraphy.common import TensorMetadata
+from polygraphy.comparator import IterationResult
+from polygraphy.comparator.data_loader import DataLoaderCache
+from polygraphy.logger import G_LOGGER, LogMode
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.args.comparator.data_loader import DataLoaderArgs
+from polygraphy.tools.args.model import ModelArgs
+from polygraphy.tools.script import (
+    Script,
+    make_invocable,
+    make_invocable_if_nondefault_kwargs,
+)
+
+onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
+onnxrt_backend = mod.lazy_import("polygraphy.backend.onnxrt")
+
+
+@mod.export()
+class OnnxInferShapesArgs(BaseArgs):
+    """
+    ONNX Shape Inference: ONNX shape inference.
+
+    Depends on:
+
+        - OnnxLoadArgs
+        - DataLoaderArgs: if allow_force_fallback == True
+    """
+
+    def __init__(self, default: bool = None, allow_force_fallback: bool = None):
+        """
+        Args:
+            default (bool):
+                    Whether shape inference should be enabled by default.
+                    Defaults to False.
+            allow_force_fallback (bool):
+                    Whether fallback shape inference using ONNX-Runtime should be allowed.
+                    Defaults to False.
+        """
+        super().__init__()
+        self._default = util.default(default, False)
+        self._allow_force_fallback = util.default(allow_force_fallback, False)
+
+    def add_parser_args_impl(self):
+        shape_infer_group = self.group.add_mutually_exclusive_group()
+        if self._default:
+            shape_infer_group.add_argument(
+                "--no-shape-inference",
+                help="Disable ONNX shape inference when loading the model",
+                dest="do_shape_inference",
+                action="store_false",
+                default=True,
+            )
+        else:
+            shape_infer_group.add_argument(
+                "--shape-inference",
+                "--do-shape-inference",
+                help="Enable ONNX shape inference when loading the model",
+                dest="do_shape_inference",
+                action="store_true",
+                default=False,
+            )
+
+        if self._allow_force_fallback:
+            shape_infer_group.add_argument(
+                "--force-fallback-shape-inference",
+                help="Force Polygraphy to use ONNX-Runtime to determine metadata for "
+                "tensors in the graph. This can be useful in cases where ONNX shape inference does not generate correct information. "
+                "Note that this will cause dynamic dimensions to become static. ",
+                action="store_true",
+                default=None,
+            )
+
+        self.group.add_argument(
+            "--no-onnxruntime-shape-inference",
+            help="Disable using ONNX-Runtime's shape inference utilities. "
+            "This will force Polygraphy to use `onnx.shape_inference` instead. "
+            "Note that ONNX-Runtime's shape inference utilities may be more performant and memory-efficient. ",
+            dest="allow_onnxruntime",
+            action="store_false",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            do_shape_inference (bool): Whether to do shape inference.
+            force_fallback (bool): Whether to force fallback shape inference.
+            allow_onnxruntime (bool): Whether to allow ONNX-Runtime's shape inference utilities to be used.
+        """
+        self.do_shape_inference = args_util.get(args, "do_shape_inference")
+        self.force_fallback = args_util.get(args, "force_fallback_shape_inference")
+        self.allow_onnxruntime = args_util.get(args, "allow_onnxruntime")
+
+        # No point is running ONNX shape inference if we're going to use fallback inference.
+        if self.force_fallback:
+            self.do_shape_inference = False
+
+    def add_to_script_impl(self, script, loader_name):
+        """
+        Note that this method does not take fallback shape inference into account.
+        To support fallback shape inference, the tool must call `fallback_inference()` manually.
+
+        Args:
+            loader_name (str):
+                    The name of the loader which should be consumed by the ``InferShapes`` loader.
+
+        Returns:
+            str: The name of the ``InferShapes`` loader added to the script.
+        """
+        if self.do_shape_inference:
+            script.add_import(imports=["InferShapes"], frm="polygraphy.backend.onnx")
+            loader_name = script.add_loader(
+                make_invocable(
+                    "InferShapes",
+                    loader_name,
+                    external_data_dir=self.arg_groups[OnnxLoadArgs].external_data_dir,
+                    allow_onnxruntime=self.allow_onnxruntime,
+                ),
+                "infer_shapes",
+            )
+        return loader_name
+
+    def infer_shapes(self, model, force=None):
+        """
+        Run shape inference on an ONNX model if `do_shape_inference` is True
+        according to arguments provided on the command-line.
+
+        Args:
+            model (onnx.ModelProto): The model in which to infer shapes.
+            force (bool):
+                    Force shape inference to run even if `do_shape_inference` is False.
+                    Defaults to False.
+
+        Returns:
+            onnx.ModelProto: The model with shapes inferred.
+        """
+        force = util.default(force, False)
+        with util.TempAttrChange(
+            self, {"do_shape_inference": True if force else self.do_shape_inference}
+        ):
+            loader = args_util.run_script(self.add_to_script, model)
+            return util.invoke_if_callable(loader)[0]
+
+    def fallback_inference(self, onnx_model, outputs=None):
+        """
+        Run inference with ONNX-Runtime.
+
+        This can be used to retrieve values/shapes/data types for all
+        tensors in the model when other shape inference approaches fail.
+
+        Args:
+            onnx_model (onnx.ModelProto):
+                    The ONNX model in which to infer shapes.
+
+
+            outputs (List[str]):
+                    The names of the outputs to retrieved.
+                    Defaults to constants.MARK_ALL
+
+        Returns:
+            (IterationResult, TensorMetadata):
+                    A tuple containing two elements:
+                    1. Mapping of values for all tensors in the model, including inputs.
+                    2. Metadata for every tensor in the model.
+        """
+        outputs = util.default(outputs, constants.MARK_ALL)
+        with G_LOGGER.verbosity(
+            G_LOGGER.module_severity.get(G_LOGGER.module_path(__file__)) + 10
+        ):
+            load_model = onnx_backend.ModifyOutputs(
+                onnx_model, outputs=outputs, copy=True
+            )
+            with onnxrt_backend.OnnxrtRunner(
+                onnxrt_backend.SessionFromOnnx(onnx_backend.BytesFromOnnx(load_model))
+            ) as runner:
+                data_loader = self.arg_groups[DataLoaderArgs].get_data_loader()
+                loader_cache = DataLoaderCache(data_loader)
+                loader_cache.set_input_metadata(
+                    runner.get_input_metadata(use_numpy_dtypes=False)
+                )
+
+                feed_dict = loader_cache[0]
+
+                with G_LOGGER.verbosity(
+                    G_LOGGER.module_severity.get(G_LOGGER.module_path(__file__)) - 10
+                ):
+                    G_LOGGER.info(
+                        f"Running fallback shape inference using input metadata:\n{TensorMetadata.from_feed_dict(feed_dict)}"
+                    )
+
+                outputs = runner.infer(feed_dict)
+                # We include the inputs here so that we have values for all tensors in the model.
+                outputs.update(feed_dict)
+                # Use IterationResult here since it can handle very large tensors by saving to disk.
+                # Layerwise outputs might otherwise take up too much memory.
+                return IterationResult(outputs), TensorMetadata.from_feed_dict(outputs)
+
+
+@mod.export()
+class OnnxSaveArgs(BaseArgs):
+    """
+    ONNX Model Saving: saving ONNX models.
+
+    Depends on:
+
+        - OnnxInferShapesArgs: if allow_shape_inference == True
+    """
+
+    def __init__(
+        self,
+        allow_shape_inference: bool = None,
+        output_opt: str = None,
+        output_short_opt: str = None,
+        output_opt_required: bool = None,
+        output_default_path: str = None,
+        allow_multiple_models: bool = None,
+    ):
+        """
+        Args:
+            allow_shape_inference (bool):
+                    Whether to allow shape inference when saving models.
+                    Defaults to False.
+            output_opt (str):
+                    The name of the output path option.
+                    Defaults to "output".
+                    Use a value of ``False`` to disable the option.
+            output_short_opt (str):
+                    The short option to use for the output path.
+                    Defaults to "-o".
+                    Use a value of ``False`` to disable the short option.
+            output_opt_required (bool):
+                    Whether the output path is a required argument.
+                    Defaults to False.
+            output_default_path (str):
+                    The default value to use for the output path option.
+                    Defaults to None.
+            allow_multiple_models (bool):
+                    Whether to enable support for saving more than one model.
+                    If this is True, the output path is expected to be a directory.
+                    Defaults to False.
+        """
+        super().__init__()
+
+        self._allow_shape_inference = util.default(allow_shape_inference, False)
+        self._output_opt = util.default(output_opt, "output")
+        self._output_short_opt = util.default(output_short_opt, "-o")
+        self._output_opt_required = util.default(output_opt_required, False)
+        self._output_default_path = output_default_path
+        self._allow_multiple_models = util.default(allow_multiple_models, False)
+
+        # add_to_script should never be called when `allow_multiple_models` is enabled.
+        # The one exception is that `save_onnx` should be able to call it, which is why we need this escape hatch.
+        self._disable_add_to_script_check = False
+
+    def add_parser_args_impl(self):
+        if self._output_opt:
+            params = ([self._output_short_opt] if self._output_short_opt else []) + [
+                f"--{self._output_opt}"
+            ]
+            help_msg = "Path to save the ONNX model"
+            if self._allow_multiple_models:
+                help_msg = "Path to a directory in which to save ONNX model(s)"
+
+            self.group.add_argument(
+                *params,
+                help=help_msg,
+                dest="save_onnx",
+                default=self._output_default_path,
+                required=self._output_opt_required,
+            )
+
+        self.group.add_argument(
+            "--save-external-data",
+            "--external-data-path",
+            help="Whether to save weight data in external file(s). "
+            + (
+                "You may optionally provide a value to this argument which will be used as a suffix for the external data files"
+                if self._allow_multiple_models
+                else "To use a non-default path, supply the desired path as an argument. This is always a relative path; "
+                "external data is always written to the same directory as the model. "
+            ),
+            default=None,
+            action="append",
+            nargs="?",
+            dest="external_data_path",
+        )
+        self.group.add_argument(
+            "--external-data-size-threshold",
+            help="The 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. "
+            "Optionally, use a `K`, `M`, or `G` suffix to indicate KiB, MiB, or GiB respectively. "
+            "For example, `--external-data-size-threshold=16M` is equivalent to `--external-data-size-threshold=16777216`. "
+            "Has no effect if `--save-external-data` is not set. Defaults to 1024 bytes.",
+            default=None,
+        )
+        self.group.add_argument(
+            "--no-save-all-tensors-to-one-file",
+            help="Do not save all tensors to a single file when saving external data. "
+            "Has no effect if `--save-external-data` is not set",
+            dest="all_tensors_to_one_file",
+            default=None,
+            action="store_false",
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            path (str): The path at which to save the ONNX model.
+            external_data_path (str): The path at which to save external data.
+            size_threshold (int): The size threshold above which external data is saved.
+            all_tensors_to_one_file (bool): Whether all external data should be written to a single file.
+        """
+        self.path = args_util.get(args, "save_onnx")
+
+        external_data_path = args_util.get(args, "external_data_path")
+        if external_data_path is not None:
+            external_data_path = external_data_path[0] or ""
+        self.external_data_path = external_data_path
+
+        self.size_threshold = args_util.parse_num_bytes(
+            args_util.get(args, "external_data_size_threshold")
+        )
+        self.all_tensors_to_one_file = args_util.get(args, "all_tensors_to_one_file")
+
+    def add_to_script_impl(self, script, loader_name):
+        """
+        Args:
+            loader_name (str):
+                    The name of the loader which should be consumed by the ``SaveOnnx`` loader.
+
+        Returns:
+            str: The name of the ``SaveOnnx`` loader added to the script.
+        """
+        if self._allow_multiple_models and not self._disable_add_to_script_check:
+            G_LOGGER.internal_error(
+                "OnnxSaveArgs.add_to_script() should never be called when `allow_multiple_models` is enabled"
+            )
+
+        if self.path is None:
+            return loader_name
+
+        # Need to run shape inference again after processing the graph since it may have changed.
+        if self._allow_shape_inference:
+            loader_name = self.arg_groups[OnnxInferShapesArgs].add_to_script(
+                script, loader_name
+            )
+
+        script.add_import(imports=["SaveOnnx"], frm="polygraphy.backend.onnx")
+        loader_name = script.add_loader(
+            make_invocable(
+                "SaveOnnx",
+                loader_name,
+                path=self.path,
+                external_data_path=self.external_data_path,
+                size_threshold=self.size_threshold,
+                all_tensors_to_one_file=self.all_tensors_to_one_file,
+            ),
+            "save_onnx",
+        )
+
+        return loader_name
+
+    def save_onnx(self, model, path: str = None):
+        """
+        Saves an ONNX model according to arguments provided on the command-line.
+
+        Args:
+            model (onnx.ModelProto): The ONNX model to save.
+
+            path (str):
+                    The path at which to save the model.
+                    If no path is provided, it is determined from command-line arguments.
+
+        Returns:
+            onnx.ModelProto: The model that was saved.
+        """
+        attrs = {"path": path, "_disable_add_to_script_check": True}
+        if self._allow_multiple_models:
+            if self.external_data_path is not None:
+                attrs["external_data_path"] = os.path.basename(
+                    os.path.splitext(path)[0]
+                ) + (self.external_data_path or "_ext_data")
+
+        with util.TempAttrChange(self, attrs):
+            loader = args_util.run_script(self.add_to_script, model)
+            return loader()
+
+
+@mod.export()
+class OnnxLoadArgs(BaseArgs):
+    """
+    ONNX Model Loading: loading ONNX models.
+
+    Depends on:
+
+        - ModelArgs
+        - OnnxInferShapesArgs: if allow_shape_inference == True
+        - OnnxSaveArgs: if allow_saving == True
+        - OnnxFromTfArgs: if allow_from_tf == True
+    """
+
+    def __init__(
+        self,
+        allow_saving: bool = None,
+        outputs_opt_prefix: str = None,
+        allow_shape_inference: bool = None,
+        allow_from_tf: bool = None,
+        allow_setting_upper_bounds: bool = None,
+    ):
+        """
+        Args:
+            allow_saving (bool):
+                    Whether to allow loaded models to be saved.
+                    Defaults to False.
+            outputs_opt_prefix (str):
+                    The prefix to use for the outputs option, which controls which tensors are marked as outputs.
+                    Defaults to "onnx-".
+                    Use a value of ``False`` to disable the option.
+            allow_shape_inference (bool):
+                    Whether to allow shape inference when loading models.
+                    Defaults to True.
+            allow_from_tf (bool):
+                    Whether to allow conversion of TensorFlow models to ONNX.
+                    Defaults to False.
+            allow_setting_upper_bounds (bool):
+                    Whether to allow setting upper bounds for unbounded DDS.
+                    Defaults to False.
+        """
+        super().__init__()
+        self._allow_saving = util.default(allow_saving, False)
+        self._allow_shape_inference = util.default(allow_shape_inference, True)
+        self._outputs_opt_prefix = util.default(outputs_opt_prefix, "onnx-")
+        self._allow_from_tf = util.default(allow_from_tf, False)
+        self._allow_setting_upper_bounds = util.default(
+            allow_setting_upper_bounds, False
+        )
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--external-data-dir",
+            "--load-external-data",
+            "--ext",
+            dest="external_data_dir",
+            help="Path to a directory containing external data for the model. "
+            "Generally, this is only required if the external data is not stored in the model directory.",
+        )
+        self.group.add_argument(
+            "--ignore-external-data",
+            help="Ignore 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.",
+            action="store_true",
+            default=None,
+        )
+
+        if (
+            self._outputs_opt_prefix is not False
+        ):  # Empty strings should not disable the option
+            self.group.add_argument(
+                f"--{self._outputs_opt_prefix}outputs",
+                help="Name(s) of ONNX tensor(s) to mark as output(s). "
+                "Using the special value 'mark all' indicates that all tensors should be used as outputs",
+                nargs="+",
+                default=None,
+                dest="onnx_outputs",
+            )
+            self.group.add_argument(
+                f"--{self._outputs_opt_prefix}exclude-outputs",
+                help="[EXPERIMENTAL] Name(s) of ONNX output(s) to unmark as outputs.",
+                nargs="+",
+                default=None,
+                dest="onnx_exclude_outputs",
+            )
+
+        self.group.add_argument(
+            "--fp-to-fp16",
+            help="Convert all floating point tensors in an ONNX model to 16-bit precision. "
+            "This is *not* needed in order to use TensorRT's fp16 precision, but may be useful for other backends. "
+            "Requires onnxmltools. ",
+            action="store_true",
+            default=None,
+        )
+
+        if self._allow_setting_upper_bounds:
+            self.group.add_argument(
+                "--set-unbounded-dds-upper-bound",
+                help="""
+                Set upper bounds for tensors with unbounded DDS(data-dependent shape).
+                Tensors with unbounded DDS can make it difficult for TensorRT to optimize inference performance
+                and memory usage. In the worst case, they can cause TensorRT engine build failures. To fix this,
+                Polygraphy supports setting upper bounds for tensors with unbounded DDS by inserting the ONNX
+                min operator. To specify per-tensor upper bounds, use the format:
+                --set-unbounded-dds-upper-bound [<tensor_name>:]<upper_bound>.
+                If no tensor name is provided, the upper bound is used for any tensors with unbounded DDS that
+                are not explicitly specified. For example:
+                --set-unbounded-dds-upper-bound 10000 tensor_a:5000 tensor_b:4000.
+
+                Note that setting upper bounds only works for models that have been constant folded and have shapes inferred.
+                """,
+                nargs="+",
+                default=None,
+                dest="upper_bounds",
+            )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            outputs (List[str]): Names of output tensors.
+            exclude_outputs (List[str]): Names of tensors which should be unmarked as outputs.
+            external_data_dir (str): Path to a directory from which to load external data.
+            ignore_external_data (bool): Whether to ignore loading external data.
+            convert_to_fp16 (bool): Whether to convert the model to FP16.
+            upper_bounds (Union[int, Dict[str, int]]): The upper bounds for tensors with unbounded DDS.
+        """
+        self.outputs = args_util.get_outputs(args, "onnx_outputs")
+        self.exclude_outputs = args_util.get(args, "onnx_exclude_outputs")
+        self.external_data_dir = args_util.get(args, "external_data_dir")
+        self.ignore_external_data = args_util.get(args, "ignore_external_data")
+        self.convert_to_fp16 = args_util.get(args, "fp_to_fp16")
+        self.upper_bounds = args_util.parse_arglist_to_dict(
+            args_util.get(args, "upper_bounds")
+        )
+
+    def _add_modify_onnx_outputs(
+        self, script, loader_name, disable_custom_outputs: bool = None
+    ):
+        if disable_custom_outputs:
+            outputs = None
+            exclude_outputs = None
+        else:
+            outputs = args_util.get_outputs_for_script(script, self.outputs)
+            exclude_outputs = self.exclude_outputs
+
+        modify_outputs_loader = make_invocable_if_nondefault_kwargs(
+            "ModifyOnnxOutputs",
+            loader_name,
+            outputs=outputs,
+            exclude_outputs=exclude_outputs,
+        )
+        if modify_outputs_loader is not None:
+            script.add_import(
+                imports="ModifyOutputs",
+                frm="polygraphy.backend.onnx",
+                imp_as="ModifyOnnxOutputs",
+            )
+            loader_name = script.add_loader(
+                modify_outputs_loader,
+                "modify_outputs",
+            )
+
+        return loader_name
+
+    def add_to_script_impl(
+        self, script, disable_custom_outputs: bool = None, serialize_model: bool = None
+    ):
+        """
+        Args:
+            disable_custom_outputs (bool):
+                    Whether to disallow modifying outputs according to the `outputs` and `exclude_outputs` attributes.
+                    Defaults to False.
+            serialize_model (bool):
+                    Whether to serialize the model.
+                    Defaults to False.
+
+        Returns:
+            str: The name of the ONNX loader added in the script.
+        """
+        model_type = self.arg_groups[ModelArgs].model_type
+        if model_type.is_onnx():
+            loader_name = self.arg_groups[ModelArgs].path
+            if self._allow_shape_inference:
+                loader_name = self.arg_groups[OnnxInferShapesArgs].add_to_script(
+                    script, loader_name
+                )
+
+            if (
+                loader_name == self.arg_groups[ModelArgs].path
+            ):  # Shape inference loader isn't being used, have to load.
+                script.add_import(
+                    imports=["OnnxFromPath"], frm="polygraphy.backend.onnx"
+                )
+                loader_str = make_invocable(
+                    "OnnxFromPath",
+                    self.arg_groups[ModelArgs].path,
+                    external_data_dir=self.external_data_dir,
+                    ignore_external_data=self.ignore_external_data,
+                )
+                loader_name = script.add_loader(loader_str, "load_onnx")
+        elif model_type.is_tf() and self._allow_from_tf:
+            from polygraphy.tools.args.backend.onnx.loader import OnnxFromTfArgs
+
+            loader_name = self.arg_groups[OnnxFromTfArgs].add_to_script(script)
+        else:
+            G_LOGGER.critical(
+                f"Model type: {model_type} could not be converted to an ONNX model."
+            )
+
+        loader_name = self._add_modify_onnx_outputs(
+            script, loader_name, disable_custom_outputs=disable_custom_outputs
+        )
+
+        if self.convert_to_fp16:
+            script.add_import(imports=["ConvertToFp16"], frm="polygraphy.backend.onnx")
+            loader_name = script.add_loader(
+                make_invocable("ConvertToFp16", loader_name), "convert_to_fp16"
+            )
+
+        if self._allow_saving:
+            loader_name = self.arg_groups[OnnxSaveArgs].add_to_script(
+                script, loader_name
+            )
+
+        if serialize_model:
+            script.add_import(imports=["BytesFromOnnx"], frm="polygraphy.backend.onnx")
+            loader_name = script.add_loader(
+                make_invocable("BytesFromOnnx", loader_name), "serialize_onnx"
+            )
+
+        if self._allow_setting_upper_bounds and self.upper_bounds is not None:
+            script.add_import(imports=["SetUpperBound"], frm="polygraphy.backend.onnx")
+            loader_name = script.add_loader(
+                make_invocable(
+                    "SetUpperBound", loader_name, upper_bounds=self.upper_bounds
+                ),
+                "set_upper_bound",
+            )
+
+        return loader_name
+
+    def must_use_onnx_loader(self, disable_custom_outputs: bool = None):
+        """
+        Whether this model needs to be loaded via a Polygraphy ONNX loader, e.g., in case it
+        needs modifications.
+
+        Args:
+            disable_custom_outputs (bool):
+                    Whether to disallow modifying outputs according to the `outputs` and `exclude_outputs` attributes.
+
+        Returns:
+            bool
+        """
+        tmp_script = Script()
+        inp_loader = "check_needs_modify"
+        needs_modify = (
+            self._add_modify_onnx_outputs(
+                tmp_script, inp_loader, disable_custom_outputs
+            )
+            != inp_loader
+        )
+        needs_shape_inference = (
+            self._allow_shape_inference
+            and self.arg_groups[OnnxInferShapesArgs].do_shape_inference
+        )
+        needs_save = (
+            self._allow_saving and self.arg_groups[OnnxSaveArgs].path is not None
+        )
+        needs_fp16_conversion = self.convert_to_fp16
+        needs_setting_upper_bounds = (
+            self._allow_setting_upper_bounds and self.upper_bounds is not None
+        )
+        # Currently, other loaders do not support external data, so we must fall back to the ONNX loader if it's present.
+        return (
+            not self.arg_groups[ModelArgs].model_type.is_onnx()
+            or needs_modify
+            or self.external_data_dir
+            or needs_shape_inference
+            or needs_save
+            or needs_fp16_conversion
+            or needs_setting_upper_bounds
+        )
+
+    def load_onnx(self):
+        """
+        Loads an ONNX model according to arguments provided on the command-line.
+
+        Returns:
+            onnx.ModelProto: The model that was loaded.
+        """
+        loader = args_util.run_script(self.add_to_script)
+        return loader()
+
+
+@mod.export()
+class OnnxFromTfArgs(BaseArgs):
+    """
+    TensorFlow-ONNX Model Conversion: converting TensorFlow models to ONNX.
+
+    Depends on:
+
+        - TfLoadArgs
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--opset",
+            help="Opset to use when converting to ONNX",
+            default=None,
+            type=int,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            opset (int): The ONNX opset version to use during conversion.
+        """
+        self.opset = args_util.get(args, "opset")
+
+    def add_to_script_impl(self, script):
+        from polygraphy.tools.args.backend.tf.loader import TfLoadArgs
+
+        G_LOGGER.verbose(
+            "Attempting to load as a TensorFlow model, using TF2ONNX to convert to ONNX. "
+            "If this is not correct, please specify --model-type",
+            mode=LogMode.ONCE,
+        )
+        script.add_import(imports=["OnnxFromTfGraph"], frm="polygraphy.backend.onnx")
+        loader_str = make_invocable(
+            "OnnxFromTfGraph",
+            self.arg_groups[TfLoadArgs].add_to_script(
+                script, disable_custom_outputs=True
+            ),
+            opset=self.opset,
+        )
+        loader_name = script.add_loader(loader_str, "export_onnx_from_tf")
+        return loader_name

tools/Polygraphy/polygraphy/tools/args/backend/onnxrt/__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.tools.args.backend.onnxrt.runner import *
+from polygraphy.tools.args.backend.onnxrt.loader import *

tools/Polygraphy/polygraphy/tools/args/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
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.args.model import ModelArgs
+from polygraphy.tools.args.backend.onnx.loader import OnnxLoadArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class OnnxrtSessionArgs(BaseArgs):
+    """
+    ONNX-Runtime Session Creation: creating an ONNX-Runtime Inference Session
+
+    Depends on:
+
+        - OnnxLoadArgs
+        - ModelArgs
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--providers",
+            "--execution-providers",
+            dest="providers",
+            help="A list of execution providers to use in order of priority. "
+            "Each provider 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'",
+            nargs="+",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            providers (List[str]): A list of execution providers.
+        """
+        self.providers = args_util.get(args, "providers")
+
+    def add_to_script_impl(self, script, onnx_name=None):
+        if onnx_name is None:  # default behavior according to self.arg_groups
+            if self.arg_groups[OnnxLoadArgs].must_use_onnx_loader():
+                onnx_name = self.arg_groups[OnnxLoadArgs].add_to_script(
+                    script, serialize_model=True
+                )
+            else:
+                onnx_name = self.arg_groups[ModelArgs].path
+
+        script.add_import(imports=["SessionFromOnnx"], frm="polygraphy.backend.onnxrt")
+        loader_name = script.add_loader(
+            make_invocable("SessionFromOnnx", onnx_name, providers=self.providers),
+            "build_onnxrt_session",
+        )
+        return loader_name
+
+    def load_onnxrt_session(self, model=None):
+        """
+        Loads an ONNX-Runtime Inference Session according to arguments provided on the command-line.
+
+        Args:
+            model (Union[bytes, str]): The model bytes or path to a model. Defaults to None, in which case, the model specified on the command-line is used.
+
+        Returns:
+            onnxruntime.InferenceSession
+        """
+        loader = args_util.run_script(self.add_to_script, model)
+        return loader()

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

@@ -0,0 +1,42 @@
+#
+# 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.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.args.backend.onnxrt.loader import OnnxrtSessionArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class OnnxrtRunnerArgs(BaseRunnerArgs):
+    """
+    ONNX-Runtime Inference: running inference with ONNX-Runtime.
+
+    Depends on:
+
+        - OnnxrtSessionArgs
+    """
+
+    def get_name_opt_impl(self):
+        return "ONNX-Runtime", "onnxrt"
+
+    def add_to_script_impl(self, script):
+        script.add_import(imports=["OnnxrtRunner"], frm="polygraphy.backend.onnxrt")
+        script.add_runner(
+            make_invocable(
+                "OnnxrtRunner", self.arg_groups[OnnxrtSessionArgs].add_to_script(script)
+            )
+        )

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

@@ -0,0 +1,16 @@
+#
+# 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.tools.args.backend.pluginref.runner import *

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

@@ -0,0 +1,47 @@
+#
+# 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.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.args.backend.onnx import OnnxLoadArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class PluginRefRunnerArgs(BaseRunnerArgs):
+    """
+    Plugin Reference Runner Inference: running inference with the plugin reference runner.
+
+    Depends on:
+
+        - OnnxLoadArgs
+    """
+
+    def get_name_opt_impl(self):
+        return "Plugin CPU Reference", "pluginref"
+
+    def add_to_script_impl(self, script):
+        script.add_import(imports=["GsFromOnnx"], frm="polygraphy.backend.onnx")
+        script.add_import(
+            imports=["PluginRefRunner"], frm="polygraphy.backend.pluginref"
+        )
+
+        onnx_name = self.arg_groups[OnnxLoadArgs].add_to_script(script)
+        loader_name = script.add_loader(
+            make_invocable("GsFromOnnx", onnx_name), "pluginref"
+        )
+        script.add_runner(make_invocable("PluginRefRunner", loader_name))

tools/Polygraphy/polygraphy/tools/args/backend/runner_select.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.
+#
+import argparse
+
+from polygraphy import mod
+from polygraphy.common.interface import TypedList
+from polygraphy.logger import G_LOGGER
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs, BaseRunnerArgs
+
+
+def make_action_cls(runner_opt):
+    class StoreRunnerOrdered(argparse.Action):
+        def __call__(self, parser, namespace, values, option_string=None):
+            if not hasattr(namespace, "runners"):
+                namespace.runners = []
+            namespace.runners.append(runner_opt)
+
+    return StoreRunnerOrdered
+
+
+class RunnerOptList(TypedList(lambda: tuple)):
+    def keys(self):
+        for opt, _ in self.lst:
+            yield opt
+
+    def values(self):
+        for _, name in self.lst:
+            yield name
+
+    def items(self):
+        yield from self.lst
+
+
+@mod.export()
+class RunnerSelectArgs(BaseArgs):
+    """
+    Runners: selecting runners to use for inference.
+    """
+
+    def add_parser_args_impl(self):
+        self._opt_to_group_map = {}
+        for arg_group in self.arg_groups.values():
+            if not isinstance(arg_group, BaseRunnerArgs):
+                continue
+
+            name, opt = arg_group.get_name_opt()
+            extra_help = arg_group.get_extra_help_text()
+            # Use opt as the key since it's guaranteed to be unique.
+            self._opt_to_group_map[opt] = arg_group
+            self.group.add_argument(
+                f"--{opt}",
+                help=f"Run inference using {name}. {extra_help}",
+                action=make_action_cls(opt),
+                dest="runners",
+                default=[],
+                nargs=0,
+            )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            runners (List[Tuple[str, str]]):
+                    A list of tuples mapping runner option strings to human readable names for all selected runners,
+                    in the order they were selected in.
+                    For example:
+                    ::
+
+                        [("trt", "TensorRT"), ("onnxrt", "ONNX-Runtime")]
+        """
+        runner_opts = args_util.get(args, "runners")
+
+        self.runners = RunnerOptList()
+        for opt in runner_opts:
+            self.runners.append((opt, self._opt_to_group_map[opt].get_name_opt()[0]))
+
+    def add_to_script_impl(self, script):
+        """
+        Adds all selected runners to the script.
+
+        Returns:
+            str: The name of the list of runners in the script.
+        """
+        if not self.runners:
+            G_LOGGER.warning(
+                "No runners have been selected. Inference will not be run!"
+            )
+
+        for opt in self.runners.keys():
+            self._opt_to_group_map[opt].add_to_script(script)
+        return script.get_runners()

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

@@ -0,0 +1,18 @@
+#
+# 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.tools.args.backend.tf.config import *
+from polygraphy.tools.args.backend.tf.loader import *
+from polygraphy.tools.args.backend.tf.runner import *

tools/Polygraphy/polygraphy/tools/args/backend/tf/config.py

@@ -0,0 +1,77 @@
+#
+# 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.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.script import make_invocable_if_nondefault
+
+
+@mod.export()
+class TfConfigArgs(BaseArgs):
+    """
+    TensorFlow Session Configuration: creating the TensorFlow SessionConfig.
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--gpu-memory-fraction",
+            help="Maximum percentage of GPU memory TensorFlow can allocate per process",
+            type=float,
+            default=None,
+        )
+        self.group.add_argument(
+            "--allow-growth",
+            help="Allow GPU memory allocated by TensorFlow to grow",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--xla",
+            help="[EXPERIMENTAL] Attempt to run graph with xla",
+            action="store_true",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            gpu_memory_fraction (float): The maximum percentage of GPU memory TensorFlow can allocate per session.
+            allow_growth (bool): Whether TensorFlow can dynamically allocate additional GPU memory.
+            xla (bool): Whether to enable XLA.
+        """
+        self.gpu_memory_fraction = args_util.get(args, "gpu_memory_fraction")
+        self.allow_growth = args_util.get(args, "allow_growth")
+        self.xla = args_util.get(args, "xla")
+
+    def add_to_script_impl(self, script):
+        config_loader_str = make_invocable_if_nondefault(
+            "CreateConfig",
+            gpu_memory_fraction=self.gpu_memory_fraction,
+            allow_growth=self.allow_growth,
+            use_xla=self.xla,
+        )
+        if config_loader_str is not None:
+            script.add_import(imports=["CreateConfig"], frm="polygraphy.backend.tf")
+            config_loader_name = script.add_loader(
+                config_loader_str, "create_tf_config"
+            )
+        else:
+            config_loader_name = None
+        return config_loader_name

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

@@ -0,0 +1,266 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from polygraphy import mod, util
+from polygraphy.logger import G_LOGGER, LogMode
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.args.model import ModelArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class TfTrtArgs(BaseArgs):
+    """
+    [UNTESTED] TensorFlow-TensorRT Integration: TensorFlow-TensorRT.
+
+    Depends on:
+
+        - TrtConfigArgs
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--tftrt",
+            "--use-tftrt",
+            help="Enable TF-TRT integration",
+            action="store_true",
+            default=None,
+            dest="use_tftrt",
+        )
+        self.group.add_argument(
+            "--minimum-segment-size",
+            help="Minimum length of a segment to convert to TensorRT",
+            type=int,
+            default=None,
+        )
+        self.group.add_argument(
+            "--dynamic-op",
+            help="Enable dynamic mode (defers engine build until runtime)",
+            action="store_true",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            use_tftrt (bool): Whether to use TF-TRT.
+            minimum_segment_size (int): The minimum size of segments offloaded to TRT.
+            dynamic_op (bool): Whether to enable dynamic mode, which defers engine building until runtime.
+        """
+        self.use_tftrt = args_util.get(args, "use_tftrt")
+        self.minimum_segment_size = args_util.get(args, "minimum_segment_size")
+        self.dynamic_op = args_util.get(args, "dynamic_op")
+
+    def add_to_script_impl(self, script, loader_name=None, suffix=None):
+        """
+        Args:
+            loader_name (str): The name of the loader which should be consumed by the ``UseTfTrt`` loader.
+        """
+        if self.use_tftrt:
+            from polygraphy.tools.args.backend.trt import TrtConfigArgs
+
+            script.add_import(imports=["UseTfTrt"], frm="polygraphy.backend.tf")
+            loader_str = make_invocable(
+                "UseTfTrt",
+                loader_name,
+                max_workspace_size=self.arg_groups[TrtConfigArgs]._workspace,
+                fp16=self.arg_groups[TrtConfigArgs].fp16,
+                int8=self.arg_groups[TrtConfigArgs].int8,
+                is_dynamic_op=self.dynamic_op,
+                minimum_segment_size=self.minimum_segment_size,
+            )
+            loader_name = script.add_loader(loader_str, "use_tftrt", suffix=suffix)
+        return loader_name
+
+
+@mod.export()
+class TfLoadArgs(BaseArgs):
+    """
+    TensorFlow Model Loading: loading TensorFlow models.
+
+    Depends on:
+
+        - ModelArgs
+        - TfTrtArgs: if allow_tftrt == True
+        - TrtSaveEngineBytesArgs: if allow_tftrt == True
+    """
+
+    def __init__(
+        self,
+        allow_artifacts: bool = None,
+        allow_custom_outputs: bool = None,
+        allow_tftrt: bool = None,
+    ):
+        """
+        Args:
+            allow_artifacts (bool):
+                    Whether to allow saving artifacts to the disk, like frozen models or TensorBoard visualizations.
+                    Defaults to True.
+            allow_custom_outputs (bool):
+                    Whether to allow marking custom output tensors.
+                    Defaults to True.
+            allow_tftrt (bool):
+                    Whether to allow applying TF-TRT.
+                    Defaults to False.
+
+        """
+        super().__init__()
+        self._allow_artifacts = util.default(allow_artifacts, True)
+        self._allow_custom_outputs = util.default(allow_custom_outputs, True)
+        self._allow_tftrt = util.default(allow_tftrt, False)
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--ckpt",
+            help="[EXPERIMENTAL] Name of the checkpoint to load. Required if the `checkpoint` file is missing. Should not include file extension "
+            "(e.g. to load `model.meta` use `--ckpt=model`)",
+            default=None,
+        )
+        if self._allow_custom_outputs:
+            self.group.add_argument(
+                "--tf-outputs",
+                help="Name(s) of TensorFlow output(s). "
+                "Using '--tf-outputs mark all' indicates that all tensors should be used as outputs",
+                nargs="+",
+                default=None,
+            )
+
+        if self._allow_artifacts:
+            self.group.add_argument(
+                "--save-pb",
+                help="Path to save the TensorFlow frozen graphdef",
+                default=None,
+                dest="save_frozen_graph_path",
+            )
+            self.group.add_argument(
+                "--save-tensorboard",
+                help="[EXPERIMENTAL] Path to save a TensorBoard visualization",
+                default=None,
+                dest="save_tensorboard_path",
+            )
+
+        self.group.add_argument(
+            "--freeze-graph",
+            help="[EXPERIMENTAL] Attempt to freeze the graph",
+            action="store_true",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            ckpt (str): Name of the checkpoint.
+            outputs (List[str]): Names of output tensors.
+            save_frozen_graph_path (str): The path at which the frozen graph will be saved.
+            save_tensorboard_path (str): The path at which the TensorBoard visualization will be saved.
+            freeze_graph (bool): Whether to attempt to freeze the graph.
+        """
+        self.ckpt = args_util.get(args, "ckpt")
+        self.outputs = args_util.get_outputs(args, "tf_outputs")
+        self.save_frozen_graph_path = args_util.get(args, "save_frozen_graph_path")
+        self.save_tensorboard_path = args_util.get(args, "save_tensorboard_path")
+        self.freeze_graph = args_util.get(args, "freeze_graph")
+
+    def add_to_script_impl(self, script, disable_custom_outputs=None):
+        """
+        Args:
+            disable_custom_outputs (bool):
+                    Whether to disallow modifying outputs according to the `outputs` attribute.
+                    Defaults to False.
+        """
+
+        model_file = self.arg_groups[ModelArgs].path
+        model_type = self.arg_groups[ModelArgs].model_type
+
+        if model_type == "ckpt":
+            G_LOGGER.verbose(
+                f"Loading a TensorFlow checkpoint from {model_file}. Please ensure you are not using the --use-subprocess flag",
+                mode=LogMode.ONCE,
+            )
+            script.add_import(imports=["GraphFromCkpt"], frm="polygraphy.backend.tf")
+            loader_id = "load_ckpt"
+            loader_str = make_invocable("GraphFromCkpt", model_file, self.ckpt)
+        elif model_type == "keras":
+            script.add_import(imports=["GraphFromKeras"], frm="polygraphy.backend.tf")
+            loader_id = "load_keras"
+            loader_str = make_invocable("GraphFromKeras", model_file)
+        elif model_type == "frozen":
+            script.add_import(imports=["GraphFromFrozen"], frm="polygraphy.backend.tf")
+            G_LOGGER.verbose(
+                "Attempting to load as a frozen graph. If this is not correct, please specify --model-type",
+                mode=LogMode.ONCE,
+            )
+            loader_id = "load_frozen"
+            loader_str = make_invocable("GraphFromFrozen", model_file)
+        else:
+            G_LOGGER.critical(
+                f"Model type: {model_type} cannot be imported with TensorFlow."
+            )
+
+        loader_name = script.add_loader(loader_str, loader_id)
+
+        if self.freeze_graph:
+            script.add_import(imports=["OptimizeGraph"], frm="polygraphy.backend.tf")
+            loader_name = script.add_loader(
+                make_invocable("OptimizeGraph", loader_name), "optimize_graph"
+            )
+
+        engine_dir = None
+        if self._allow_tftrt:
+            from polygraphy.tools.args.backend.trt import TrtSaveEngineBytesArgs
+
+            loader_name = self.arg_groups[TfTrtArgs].add_to_script(script, loader_name)
+            engine_dir = self.arg_groups[TrtSaveEngineBytesArgs].path
+
+        MODIFY_TF = "ModifyGraphOutputs"
+        outputs = (
+            None
+            if disable_custom_outputs
+            else args_util.get_outputs_for_script(script, self.outputs)
+        )
+        modify_tf_str = make_invocable(MODIFY_TF, loader_name, outputs=outputs)
+        if modify_tf_str != make_invocable(MODIFY_TF, loader_name):
+            script.add_import(imports=[MODIFY_TF], frm="polygraphy.backend.tf")
+            loader_name = script.add_loader(modify_tf_str, "modify_tf")
+
+        WRITE_TF = "SaveGraph"
+        write_tf_str = make_invocable(
+            WRITE_TF,
+            loader_name,
+            path=self.save_frozen_graph_path,
+            tensorboard_dir=self.save_tensorboard_path,
+            engine_dir=engine_dir,
+        )
+        if write_tf_str != make_invocable(WRITE_TF, loader_name):
+            script.add_import(imports=[WRITE_TF], frm="polygraphy.backend.tf")
+            loader_name = script.add_loader(write_tf_str, "save_tf")
+
+        return loader_name
+
+    def load_graph(self):
+        """
+        Loads a TensorFlow graph according to arguments provided on the command-line.
+
+        Returns:
+            tf.Graph
+        """
+        loader = args_util.run_script(self.add_to_script)
+        return loader()

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

@@ -0,0 +1,70 @@
+#
+# 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.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.args.backend.tf.config import TfConfigArgs
+from polygraphy.tools.args.backend.tf.loader import TfLoadArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class TfRunnerArgs(BaseRunnerArgs):
+    """
+    TensorFlow Inference: running inference with TensorFlow.
+
+    Depends on:
+
+        - TfConfigArgs
+        - TfLoadArgs
+    """
+
+    def get_name_opt_impl(self):
+        return "TensorFlow", "tf"
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--save-timeline",
+            help="[EXPERIMENTAL] Directory to save timeline JSON files for profiling inference (view at chrome://tracing)",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            timeline_path (str): Path at which to save timeline files for profiling.
+        """
+        self.timeline_path = args_util.get(args, "save_timeline")
+
+    def add_to_script_impl(self, script):
+        script.add_import(imports=["TfRunner"], frm="polygraphy.backend.tf")
+
+        graph_name = self.arg_groups[TfLoadArgs].add_to_script(script)
+        config_name = self.arg_groups[TfConfigArgs].add_to_script(script)
+
+        script.add_import(imports=["SessionFromGraph"], frm="polygraphy.backend.tf")
+        loader_name = script.add_loader(
+            make_invocable("SessionFromGraph", graph_name, config=config_name),
+            "build_tf_session",
+        )
+
+        script.add_runner(
+            make_invocable("TfRunner", loader_name, timeline_path=self.timeline_path)
+        )

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

@@ -0,0 +1,18 @@
+#
+# 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.tools.args.backend.trt.config import *
+from polygraphy.tools.args.backend.trt.loader import *
+from polygraphy.tools.args.backend.trt.runner import *

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

@@ -0,0 +1,995 @@
+#
+# 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
+
+from polygraphy import constants, mod, util
+from polygraphy.common import TensorMetadata
+from polygraphy import config as polygraphy_config
+from polygraphy.logger import G_LOGGER, LogMode
+from polygraphy.mod.trt_importer import tensorrt_module_and_version_string
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.backend.trt.helper import make_trt_enum_val
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.args.comparator.data_loader import DataLoaderArgs
+from polygraphy.tools.args.model import ModelArgs
+from polygraphy.tools.script import (
+    inline,
+    inline_identifier,
+    make_invocable,
+    make_invocable_if_nondefault,
+    safe,
+)
+
+
+def parse_profile_shapes(default_shapes, min_args, opt_args, max_args):
+    """
+    Parses TensorRT profile options from command-line arguments.
+
+    Args:
+        default_shapes (TensorMetadata): The inference input shapes.
+
+    Returns:
+        List[OrderedDict[str, Tuple[Shape]]]:
+            A list of profiles where each profile is a dictionary that maps
+            input names to a tuple of (min, opt, max) shapes.
+    """
+
+    def get_shapes(lst, idx):
+        # Overwrite a copy of default_shapes with the shapes for min, opt, or max (if applicable)
+        nonlocal default_shapes
+        default_shapes = copy.copy(default_shapes)
+        if idx < len(lst):
+            default_shapes.update(args_util.parse_meta(lst[idx], includes_dtype=False))
+
+        # Don't care about dtype, and need to override dynamic dimensions
+        shapes = {
+            name: util.override_dynamic_shape(shape)
+            for name, (_, shape) in default_shapes.items()
+        }
+
+        for name, shape in shapes.items():
+            if tuple(default_shapes[name].shape) != tuple(shape):
+                G_LOGGER.warning(
+                    f"Input tensor: {name} | For TensorRT profile, overriding dynamic shape: {default_shapes[name].shape} to: {shape}",
+                    mode=LogMode.ONCE,
+                )
+
+        return shapes
+
+    num_profiles = max(len(min_args), len(opt_args), len(max_args))
+
+    # For cases where input shapes are provided, we have to generate a profile
+    if not num_profiles and default_shapes:
+        num_profiles = 1
+
+    profiles = []
+    for idx in range(num_profiles):
+        min_shapes = get_shapes(min_args, idx)
+        opt_shapes = get_shapes(opt_args, idx)
+        max_shapes = get_shapes(max_args, idx)
+        if sorted(min_shapes.keys()) != sorted(opt_shapes.keys()):
+            G_LOGGER.critical(
+                f"Mismatch in input names between minimum shapes ({list(min_shapes.keys())}) and optimum shapes ({list(opt_shapes.keys())})"
+            )
+        elif sorted(opt_shapes.keys()) != sorted(max_shapes.keys()):
+            G_LOGGER.critical(
+                f"Mismatch in input names between optimum shapes ({list(opt_shapes.keys())}) and maximum shapes ({list(max_shapes.keys())})"
+            )
+
+        profile = {
+            name: (min_shapes[name], opt_shapes[name], max_shapes[name])
+            for name in min_shapes.keys()
+        }
+        profiles.append(profile)
+    return profiles
+
+
+@mod.export()
+class TrtConfigArgs(BaseArgs):
+    """
+    TensorRT Builder Configuration: creating the TensorRT BuilderConfig.
+
+    Depends on:
+
+        - DataLoaderArgs
+        - ModelArgs: if allow_custom_input_shapes == True
+    """
+
+    def __init__(
+        self,
+        precision_constraints_default: bool = None,
+        allow_random_data_calib_warning: bool = None,
+        allow_custom_input_shapes: bool = None,
+        allow_engine_capability: bool = None,
+        allow_tensor_formats: bool = None,
+        allow_compute_capabilities: bool = None,
+    ):
+        """
+        Args:
+            precision_constraints_default (str):
+                    The default value to use for the precision constraints option.
+                    Defaults to "none".
+            allow_random_data_calib_warning (bool):
+                    Whether to issue a warning when randomly generated data is being used for calibration.
+                    Defaults to True.
+            allow_custom_input_shapes (bool):
+                    Whether to allow custom input shapes when randomly generating data.
+                    Defaults to True.
+            allow_engine_capability (bool):
+                    Whether to allow engine capability to be specified.
+                    Defaults to False.
+            allow_tensor_formats (bool):
+                    Whether to allow tensor formats and related options to be set.
+                    Defaults to False.
+            allow_compute_capabilities (bool):
+                    Whether to allow compute capabilities options to be set.
+                    Defaults to False.
+        """
+        super().__init__()
+        self._precision_constraints_default = util.default(
+            precision_constraints_default, "none"
+        )
+        self._allow_random_data_calib_warning = util.default(
+            allow_random_data_calib_warning, True
+        )
+        self._allow_custom_input_shapes = util.default(allow_custom_input_shapes, True)
+        self._allow_engine_capability = util.default(allow_engine_capability, False)
+        self._allow_tensor_formats = util.default(allow_tensor_formats, False)
+        self._allow_compute_capabilities = util.default(allow_compute_capabilities, False)
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--trt-min-shapes",
+            action="append",
+            help="The minimum shapes the optimization profile(s) will support. "
+            "Specify this option once for each profile. If not provided, inference-time input shapes are used. "
+            "Format: --trt-min-shapes <input0>:[D0,D1,..,DN] .. <inputN>:[D0,D1,..,DN]",
+            nargs="+",
+            default=[],
+        )
+        self.group.add_argument(
+            "--trt-opt-shapes",
+            action="append",
+            help="The shapes for which the optimization profile(s) will be most performant. "
+            "Specify this option once for each profile. If not provided, inference-time input shapes are used. "
+            "Format: --trt-opt-shapes <input0>:[D0,D1,..,DN] .. <inputN>:[D0,D1,..,DN]",
+            nargs="+",
+            default=[],
+        )
+        self.group.add_argument(
+            "--trt-max-shapes",
+            action="append",
+            help="The maximum shapes the optimization profile(s) will support. "
+            "Specify this option once for each profile. If not provided, inference-time input shapes are used. "
+            "Format: --trt-max-shapes <input0>:[D0,D1,..,DN] .. <inputN>:[D0,D1,..,DN]",
+            nargs="+",
+            default=[],
+        )
+
+        self.group.add_argument(
+            "--tf32",
+            help="Enable tf32 precision in TensorRT",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--fp16",
+            help="Enable fp16 precision in TensorRT",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--bf16",
+            help="Enable bf16 precision in TensorRT",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--fp8",
+            help="Enable fp8 precision in TensorRT",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--int8",
+            help="Enable int8 precision in TensorRT. "
+            "If calibration is required but no calibration cache is provided, this option will cause TensorRT to run "
+            "int8 calibration using the Polygraphy data loader to provide calibration data. "
+            "If calibration is run and the model has dynamic shapes, the last optimization profile will be "
+            "used as the calibration profile. ",
+            action="store_true",
+            default=None,
+        )
+
+        precision_constraints_group = self.group.add_mutually_exclusive_group()
+        precision_constraints_group.add_argument(
+            "--precision-constraints",
+            help=f"If set to `prefer`, TensorRT will restrict available tactics to layer precisions specified in the network unless no implementation exists with the preferred layer constraints, in which case it will issue a warning and use the fastest available implementation. If set to `obey`, TensorRT will instead fail to build the network if no implementation exists with the preferred layer constraints. Defaults to `{self._precision_constraints_default}`",
+            choices=("prefer", "obey", "none"),
+            default=self._precision_constraints_default,
+        )
+        self.group.add_argument(
+            "--sparse-weights",
+            help="Enable optimizations for sparse weights in TensorRT",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--version-compatible",
+            help="Builds an engine designed to be forward TensorRT version compatible.",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--exclude-lean-runtime",
+            help="Exclude the lean runtime from the plan when version compatibility is enabled. ",
+            action="store_true",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--calibration-cache",
+            help="Path to load/save a calibration cache. "
+            "Used to store calibration scales to speed up the process of int8 calibration. "
+            "If the provided path does not yet exist, int8 calibration scales will be calculated and written to it during engine building. "
+            "If the provided path does exist, it will be read and int8 calibration will be skipped during engine building. ",
+            default=None,
+        )
+        self.group.add_argument(
+            "--calib-base-cls",
+            "--calibration-base-class",
+            dest="calibration_base_class",
+            help="The name of the calibration base class to use. For example, 'IInt8MinMaxCalibrator'. ",
+            default=None,
+        )
+        self.group.add_argument(
+            "--quantile",
+            type=float,
+            help="The quantile to use for IInt8LegacyCalibrator. Has no effect for other calibrator types.",
+            default=None,
+        )
+        self.group.add_argument(
+            "--regression-cutoff",
+            type=float,
+            help="The regression cutoff to use for IInt8LegacyCalibrator. Has no effect for other calibrator types.",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--load-timing-cache",
+            help="Path to load tactic timing cache. "
+            "Used to cache tactic timing information to speed up the engine building process. "
+            "If the file specified by --load-timing-cache does not exist, Polygraphy will emit a warning and fall back to "
+            "using an empty timing cache.",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--error-on-timing-cache-miss",
+            help="Emit error when a tactic being timed is not present in the timing cache.",
+            action="store_true",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--disable-compilation-cache",
+            help="Disable caching JIT-compiled code",
+            action="store_true",
+            default=None,
+        )
+
+        replay_group = self.group.add_mutually_exclusive_group()
+        replay_group.add_argument(
+            "--save-tactics",
+            "--save-tactic-replay",
+            help="Path to save a Polygraphy tactic replay file. "
+            "Details about tactics selected by TensorRT will be recorded and stored at this location as a JSON file. ",
+            dest="save_tactics",
+            default=None,
+        )
+        replay_group.add_argument(
+            "--load-tactics",
+            "--load-tactic-replay",
+            help="Path to load a Polygraphy tactic replay file, such as one created by --save-tactics. "
+            "The tactics specified in the file will be used to override TensorRT's default selections. ",
+            dest="load_tactics",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--tactic-sources",
+            help="Tactic sources to enable. This controls which libraries "
+            "(e.g. cudnn, cublas, etc.) TensorRT is allowed to load tactics from. "
+            "Values come from the names of the values in the trt.TacticSource enum and are case-insensitive. "
+            "If no arguments are provided, e.g. '--tactic-sources', then all tactic sources are disabled."
+            "Defaults to TensorRT's default tactic sources.",
+            nargs="*",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--trt-config-script",
+            help="Path to a Python script that defines a function that creates a "
+            "TensorRT IBuilderConfig. The function should take a builder and network as parameters and return a "
+            "TensorRT builder configuration. When this option is specified, all other config arguments are ignored. "
+            "By default, Polygraphy looks for a function called `load_config`. You can specify a custom function name "
+            "by separating it with a colon. For example: `my_custom_script.py:my_func`",
+            default=None,
+        )
+        self.group.add_argument(
+            "--trt-config-func-name",
+            help="[DEPRECATED - function name can be specified with --trt-config-script like so: `my_custom_script.py:my_func`]"
+            "When using a trt-config-script, this specifies the name of the function "
+            "that creates the config. Defaults to `load_config`. ",
+            default=None,
+        )
+        self.group.add_argument(
+            "--trt-config-postprocess-script",
+            "--trt-cpps",
+            help="[EXPERIMENTAL] Path to a Python script that defines a function that modifies a TensorRT IBuilderConfig. "
+            "This function will be called after Polygraphy has finished created the builder configuration and should take a builder, "
+            "network, and config as parameters and modify the config in place. "
+            "Unlike `--trt-config-script`, all other config arguments will be reflected in the config passed to the function."
+            "By default, Polygraphy looks for a function called `postprocess_config`. You can specify a custom function name "
+            "by separating it with a colon. For example: `my_custom_script.py:my_func`",
+            default=None,
+        )
+        self.group.add_argument(
+            "--trt-safety-restricted",
+            help="Enable safety scope checking in TensorRT",
+            action="store_true",
+            default=None,
+            dest="restricted",
+        )
+        self.group.add_argument(
+            "--refittable",
+            help="Enable the engine to be refitted with new weights after it is built.",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--strip-plan",
+            help="Builds the engine with the refittable weights stripped.",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--use-dla",
+            help="[EXPERIMENTAL] Use DLA as the default device type",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--allow-gpu-fallback",
+            help="[EXPERIMENTAL] Allow layers unsupported on the DLA to fall back to GPU. Has no effect if --use-dla is not set.",
+            action="store_true",
+            default=None,
+        )
+        self.group.add_argument(
+            "--pool-limit",
+            "--memory-pool-limit",
+            dest="memory_pool_limit",
+            help="Memory pool limits. Memory pool names come from the names of values in the `trt.MemoryPoolType` enum and are case-insensitive"
+            "Format: `--pool-limit <pool_name>:<pool_limit> ...`. For example, `--pool-limit dla_local_dram:1e9 workspace:16777216`. "
+            "Optionally, use a `K`, `M`, or `G` suffix to indicate KiB, MiB, or GiB respectively. "
+            "For example, `--pool-limit workspace:16M` is equivalent to `--pool-limit workspace:16777216`. ",
+            nargs="+",
+            default=None,
+        )
+        self.group.add_argument(
+            "--preview-features",
+            dest="preview_features",
+            help="Preview features to enable. Values come from the names of the values "
+            "in the trt.PreviewFeature enum, and are case-insensitive."
+            "If no arguments are provided, e.g. '--preview-features', then all preview features are disabled. "
+            "Defaults to TensorRT's default preview features.",
+            nargs="*",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--builder-optimization-level",
+            help="The builder optimization level. Setting 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. ",
+            type=int,
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--hardware-compatibility-level",
+            help="The hardware compatibility level to use for the engine. This allows engines built on one GPU architecture to work on GPUs "
+            "of other architectures. Values come from the names of values in the `trt.HardwareCompatibilityLevel` enum and are case-insensitive. "
+            "For example, `--hardware-compatibility-level ampere_plus` ",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--max-aux-streams",
+            help="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.execute_async_v3() call. "
+            "The default maximum number of auxiliary streams is determined by the heuristics in TensorRT on "
+            "whether enabling multi-stream would improve the performance. "
+            "Refer to the TensorRT API documentation for details.",
+            type=int,
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--quantization-flags",
+            dest="quantization_flags",
+            help="Int8 quantization flags to enable. Values come from the names of values "
+            "in the trt.QuantizationFlag enum, and are case-insensitive. "
+            "If no arguments are provided, e.g. '--quantization-flags', then all quantization flags are disabled. "
+            "Defaults to TensorRT's default quantization flags.",
+            nargs="*",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--profiling-verbosity",
+            help="The verbosity of NVTX annotations in the generated engine."
+            "Values come from the names of values in the `trt.ProfilingVerbosity` enum and are case-insensitive. "
+            "For example, `--profiling-verbosity detailed`. "
+            "Defaults to 'detailed'.",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--weight-streaming",
+            help="Build a weight streamable engine. Must be set with --strongly-typed. The weight streaming amount can be set with --weight-streaming-budget.",
+            action="store_true",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--runtime-platform",
+            help="The target 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. "
+            "Values come from the names of values in the `trt.RuntimePlatform` enum and are case-insensitive. "
+            "For example, `--runtime-platform same_as_build`, `--runtime-platform windows_amd64` ",
+            default=None,
+        )
+
+        if self._allow_engine_capability:
+            self.group.add_argument(
+                "--engine-capability",
+                help="The desired engine capability. "
+                "Possible values come from the names of the values in the trt.EngineCapability enum and are case-insensitive. ",
+                default=None,
+            )
+
+        if self._allow_tensor_formats:
+            self.group.add_argument(
+                "--direct-io",
+                help="Disallow reformatting layers at network input/output tensors which have user-specified formats. ",
+                action="store_true",
+                default=None,
+            )
+
+        self.group.add_argument(
+            "--tiling-optimization-level",
+            help="The tiling optimization level. Setting a higher optimization "
+            "level allows TensorRT to spend more building time for more tiling strategies. "
+            "Values come from the names of values in the `trt.TilingOptimizationLevel` enum and are case-insensitive.",
+            default=None,
+        )
+
+        if polygraphy_config.USE_TENSORRT_RTX and self._allow_compute_capabilities:
+            compute_capabilities_group = self.group.add_mutually_exclusive_group()
+
+            compute_capabilities_group.add_argument(
+                "--use-gpu",
+                help="Use the current GPU device as target for engine compilation. "
+                "Equivalent to setting ComputeCapability.CURRENT.",
+                action="store_true",
+                default=None,
+            )
+
+            compute_capabilities_group.add_argument(
+                "--compute-capabilities",
+                help="Specify target compute capabilities for engine compilation. "
+                "Values should be major.minor versions (e.g., '7.5 8.0 8.6').",
+                nargs="+",
+                default=None,
+            )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            profile_dicts (List[OrderedDict[str, Tuple[Shape]]]):
+                A list of profiles where each profile is a dictionary that maps
+                input names to a tuple of (min, opt, max) shapes.
+            tf32 (bool): Whether to enable TF32.
+            fp16 (bool): Whether to enable FP16.
+            bf16 (bool): Whether to enable BF16.
+            fp8  (bool): Whether to enable FP8.
+            int8 (bool): Whether to enable INT8.
+            precision_constraints (str): The precision constraints to apply.
+            restricted (bool): Whether to enable safety scope checking in the builder.
+            calibration_cache (str): Path to the calibration cache.
+            calibration_base_class (str): The name of the base class to use for the calibrator.
+            sparse_weights (bool): Whether to enable sparse weights.
+            load_timing_cache (str): Path from which to load a timing cache.
+            load_tactics (str): Path from which to load a tactic replay file.
+            save_tactics (str): Path at which to save a tactic replay file.
+            tactic_sources (List[str]): Strings representing enum values of the tactic sources to enable.
+            trt_config_script (str): Path to a custom TensorRT config script.
+            trt_config_func_name (str): Name of the function in the custom config script that creates the config.
+            trt_config_postprocess_script (str): Path to a TensorRT config postprocessing script.
+            trt_config_postprocess_func_name (str): Name of the function in the config postprocessing script that applies the post-processing.
+            use_dla (bool): Whether to enable DLA.
+            allow_gpu_fallback (bool): Whether to allow GPU fallback when DLA is enabled.
+            memory_pool_limits (Dict[str, int]): Mapping of strings representing memory pool enum values to memory limits in bytes.
+            engine_capability (str): The desired engine capability.
+            direct_io (bool): Whether to disallow reformatting layers at network input/output tensors which have user-specified formats.
+            preview_features (List[str]): Names of preview features to enable.
+            refittable (bool): Whether the engine should be refittable.
+            strip_plan (bool): Whether the engine should be built with the refittable weights stripped.
+            builder_optimization_level (int): The builder optimization level.
+            hardware_compatibility_level (str): A string representing a hardware compatibility level enum value.
+            profiling_verbosity (str): A string representing a profiling verbosity level enum value.
+            max_aux_streams (int): The maximum number of auxiliary streams that TensorRT is allowed to use.
+            version_compatible (bool): Whether or not to build a TensorRT forward-compatible.
+            exclude_lean_runtime (bool): Whether to exclude the lean runtime from a version compatible plan.
+            quantization_flags (List[str]): Names of quantization flags to enable.
+            error_on_timing_cache_miss (bool): Whether to emit error when a tactic being timed is not present in the timing cache.
+            disable_compilation_cache (bool): Whether to disable caching JIT-compiled code.
+            weight_streaming (bool): Whether to enable weight streaming for the TensorRT Engine.
+            runtime_platform (str): A string representing the target runtime platform enum value.
+            tiling_optimization_level (str): The tiling optimization level.
+            use_gpu (bool): Whether to use the current GPU device as target for engine compilation.
+            compute_capabilities (List[Tuple[int, int]]): List of (major, minor) compute capability tuples to target for engine compilation.
+        """
+
+        trt_min_shapes = args_util.get(args, "trt_min_shapes", default=[])
+        trt_max_shapes = args_util.get(args, "trt_max_shapes", default=[])
+        trt_opt_shapes = args_util.get(args, "trt_opt_shapes", default=[])
+
+        default_shapes = TensorMetadata()
+        if self._allow_custom_input_shapes:
+            if not hasattr(self.arg_groups[ModelArgs], "input_shapes"):
+                G_LOGGER.internal_error(
+                    "ModelArgs must be parsed before TrtConfigArgs!"
+                )
+            default_shapes = self.arg_groups[ModelArgs].input_shapes
+
+        self.profile_dicts = parse_profile_shapes(
+            default_shapes, trt_min_shapes, trt_opt_shapes, trt_max_shapes
+        )
+
+        self.tf32 = args_util.get(args, "tf32")
+        self.fp16 = args_util.get(args, "fp16")
+        self.bf16 = args_util.get(args, "bf16")
+        self.int8 = args_util.get(args, "int8")
+        self.fp8 = args_util.get(args, "fp8")
+        self.precision_constraints = args_util.get(args, "precision_constraints")
+
+        if self.precision_constraints == "none":
+            self.precision_constraints = None
+
+        self.restricted = args_util.get(args, "restricted")
+        self.refittable = args_util.get(args, "refittable")
+        self.strip_plan = args_util.get(args, "strip_plan")
+
+        self.calibration_cache = args_util.get(args, "calibration_cache")
+        calib_base = args_util.get(args, "calibration_base_class")
+        self.calibration_base_class = None
+        if calib_base is not None:
+            self.calibration_base_class = inline(
+                safe("trt.{:}", inline_identifier(calib_base))
+            )
+
+        self._quantile = args_util.get(args, "quantile")
+        self._regression_cutoff = args_util.get(args, "regression_cutoff")
+
+        self.sparse_weights = args_util.get(args, "sparse_weights")
+
+        self.load_timing_cache = args_util.get(args, "load_timing_cache")
+
+        self.load_tactics = args_util.get(args, "load_tactics")
+        self.save_tactics = args_util.get(args, "save_tactics")
+
+        tactic_sources = args_util.get(args, "tactic_sources")
+        self.tactic_sources = None
+        if tactic_sources is not None:
+            self.tactic_sources = [
+                make_trt_enum_val("TacticSource", source) for source in tactic_sources
+            ]
+
+        self.trt_config_script, self.trt_config_func_name = (
+            args_util.parse_script_and_func_name(
+                args_util.get(args, "trt_config_script"),
+                default_func_name="load_config",
+            )
+        )
+        (
+            self.trt_config_postprocess_script,
+            self.trt_config_postprocess_func_name,
+        ) = args_util.parse_script_and_func_name(
+            args_util.get(args, "trt_config_postprocess_script"),
+            default_func_name="postprocess_config",
+        )
+
+        func_name = args_util.get(args, "trt_config_func_name")
+        if func_name is not None:
+            mod.warn_deprecated(
+                "--trt-config-func-name",
+                "the config script argument",
+                "0.50.0",
+                always_show_warning=True,
+            )
+            self.trt_config_func_name = func_name
+
+        self.use_dla = args_util.get(args, "use_dla")
+        self.allow_gpu_fallback = args_util.get(args, "allow_gpu_fallback")
+
+        memory_pool_limits = args_util.parse_arglist_to_dict(
+            args_util.get(args, "memory_pool_limit"),
+            cast_to=args_util.parse_num_bytes,
+            allow_empty_key=False,
+        )
+        self.memory_pool_limits = None
+        if memory_pool_limits is not None:
+            self.memory_pool_limits = {
+                make_trt_enum_val("MemoryPoolType", pool_type): pool_size
+                for pool_type, pool_size in memory_pool_limits.items()
+            }
+
+        preview_features = args_util.get(args, "preview_features")
+        self.preview_features = None
+        if preview_features is not None:
+            self.preview_features = [
+                make_trt_enum_val("PreviewFeature", feature)
+                for feature in preview_features
+            ]
+
+        engine_capability = args_util.get(args, "engine_capability")
+        self.engine_capability = None
+        if engine_capability is not None:
+            self.engine_capability = make_trt_enum_val(
+                "EngineCapability", engine_capability
+            )
+
+        self.direct_io = args_util.get(args, "direct_io")
+        self.builder_optimization_level = args_util.get(
+            args, "builder_optimization_level"
+        )
+
+        self.hardware_compatibility_level = None
+        hardware_compatibility_level = args_util.get(
+            args, "hardware_compatibility_level"
+        )
+        if hardware_compatibility_level is not None:
+            self.hardware_compatibility_level = make_trt_enum_val(
+                "HardwareCompatibilityLevel", hardware_compatibility_level
+            )
+
+        self.runtime_platform = None
+        runtime_platform = args_util.get(
+            args, "runtime_platform"
+        )
+        if runtime_platform is not None:
+            self.runtime_platform = make_trt_enum_val(
+                "RuntimePlatform", runtime_platform
+            )
+
+        self.profiling_verbosity = None
+        profiling_verbosity = args_util.get(args, "profiling_verbosity")
+        if profiling_verbosity is not None:
+            self.profiling_verbosity = make_trt_enum_val(
+                "ProfilingVerbosity", profiling_verbosity
+            )
+
+        self.max_aux_streams = args_util.get(args, "max_aux_streams")
+        self.version_compatible = args_util.get(args, "version_compatible")
+        self.exclude_lean_runtime = args_util.get(args, "exclude_lean_runtime")
+
+        quantization_flags = args_util.get(args, "quantization_flags")
+        self.quantization_flags = None
+        if quantization_flags is not None:
+            self.quantization_flags = [
+                make_trt_enum_val("QuantizationFlag", flag)
+                for flag in quantization_flags
+            ]
+
+        if self.exclude_lean_runtime and not self.version_compatible:
+            G_LOGGER.critical(
+                f"`--exclude-lean-runtime` requires `--version-compatible` to be enabled."
+            )
+
+        self.error_on_timing_cache_miss = args_util.get(
+            args, "error_on_timing_cache_miss"
+        )
+
+        self.disable_compilation_cache = args_util.get(
+            args, "disable_compilation_cache"
+        )
+
+        self.weight_streaming = args_util.get(args, "weight_streaming")
+
+        self.tiling_optimization_level = None
+        tiling_optimization_level = args_util.get(
+            args, "tiling_optimization_level"
+        )
+        if tiling_optimization_level is not None:
+            self.tiling_optimization_level = make_trt_enum_val(
+                "TilingOptimizationLevel", tiling_optimization_level
+            )
+
+        # Parse compute capabilities arguments if enabled and TensorRT-RTX is available
+        self.use_gpu = False
+        self.compute_capabilities = None
+
+        if self._allow_compute_capabilities and polygraphy_config.USE_TENSORRT_RTX:
+            self.use_gpu = args_util.get(args, "use_gpu", default=False)
+            compute_capabilities_list = args_util.get(args, "compute_capabilities")
+
+            if compute_capabilities_list:
+                # Parse compute capabilities from list of strings
+                try:
+                    capabilities = []
+                    for cap_str in compute_capabilities_list:
+                        major, minor = map(int, cap_str.split('.'))
+                        capabilities.append((major, minor))
+                    self.compute_capabilities = capabilities
+                except ValueError:
+                    G_LOGGER.critical(f"Invalid compute capabilities format: {compute_capabilities_list}. "
+                                      "Expected format: space-separated 'major.minor' versions (e.g., '7.5 8.0').")
+
+    def add_to_script_impl(self, script):
+        profiles = []
+        for profile_dict in self.profile_dicts:
+            profile_str = "Profile()"
+            for name in profile_dict.keys():
+                profile_str += safe(
+                    ".add({:}, min={:}, opt={:}, max={:})", name, *profile_dict[name]
+                ).unwrap()
+            profiles.append(profile_str)
+        if profiles:
+            script.add_import(imports=["Profile"], frm="polygraphy.backend.trt")
+            profiles = safe(
+                "[\n{tab}{:}\n]",
+                inline(safe(f",\n{constants.TAB}".join(profiles))),
+                tab=inline(safe(constants.TAB)),
+            )
+            profile_name = script.add_loader(profiles, "profiles")
+        else:
+            profile_name = None
+
+        calibrator = None
+        if (
+            any(
+                arg is not None
+                for arg in [self.calibration_cache, self.calibration_base_class]
+            )
+            and not self.int8
+        ):
+            G_LOGGER.warning(
+                "Some int8 calibrator options were set, but int8 precision is not enabled. "
+                "Calibration options will be ignored. Please set --int8 to enable calibration. "
+            )
+
+        if self.int8:
+            script.add_import(imports=["Calibrator"], frm="polygraphy.backend.trt")
+            script.add_import(imports=["DataLoader"], frm="polygraphy.comparator")
+            data_loader_name = self.arg_groups[DataLoaderArgs].add_to_script(script)
+            if self.calibration_base_class:
+                script.add_import(imports=tensorrt_module_and_version_string(), imp_as="trt")
+
+            if (
+                self.arg_groups[DataLoaderArgs].is_using_random_data()
+                and (
+                    not self.calibration_cache
+                    or not os.path.exists(self.calibration_cache)
+                )
+                and self._allow_random_data_calib_warning
+            ):
+                G_LOGGER.warning(
+                    "Int8 Calibration is using randomly generated input data.\n"
+                    "This could negatively impact accuracy if the inference-time input data is dissimilar "
+                    "to the randomly generated calibration data.\n"
+                    "You may want to consider providing real data via the --data-loader-script option."
+                )
+
+            calibrator = make_invocable(
+                "Calibrator",
+                data_loader=(
+                    data_loader_name
+                    if data_loader_name
+                    else inline(safe("DataLoader()"))
+                ),
+                cache=self.calibration_cache,
+                BaseClass=self.calibration_base_class,
+                quantile=self._quantile,
+                regression_cutoff=self._regression_cutoff,
+            )
+
+        algo_selector = None
+        if self.load_tactics is not None:
+            script.add_import(imports=["TacticReplayer"], frm="polygraphy.backend.trt")
+            algo_selector = make_invocable("TacticReplayer", replay=self.load_tactics)
+        elif self.save_tactics is not None:
+            script.add_import(imports=["TacticRecorder"], frm="polygraphy.backend.trt")
+            algo_selector = make_invocable("TacticRecorder", record=self.save_tactics)
+
+        # Add a `tensorrt` import if any argument requires direct access to the module.
+        if any(
+            arg is not None
+            for arg in [
+                self.tactic_sources,
+                self.memory_pool_limits,
+                self.preview_features,
+                self.engine_capability,
+                self.profiling_verbosity,
+                self.hardware_compatibility_level,
+                self.runtime_platform,
+                self.quantization_flags,
+                self.tiling_optimization_level,
+                self.use_gpu,
+                self.compute_capabilities,
+            ]
+        ):
+            script.add_import(imports=tensorrt_module_and_version_string(), imp_as="trt")
+
+        if self.trt_config_script is not None:
+            script.add_import(
+                imports=["InvokeFromScript"], frm="polygraphy.backend.common"
+            )
+            config_loader_str = make_invocable(
+                "InvokeFromScript",
+                self.trt_config_script,
+                name=self.trt_config_func_name,
+            )
+        else:
+            # Use CreateConfigRTX if TensorRT-RTX is enabled, otherwise use CreateConfig
+            if polygraphy_config.USE_TENSORRT_RTX:
+                config_class = "CreateConfigRTX"
+                config_alias = "CreateTrtConfigRTX"
+                extra_args = {
+                    "use_gpu": self.use_gpu,
+                    "compute_capabilities": self.compute_capabilities,
+                }
+            else:
+                config_class = "CreateConfig"
+                config_alias = "CreateTrtConfig"
+                extra_args = {
+                    "tf32": self.tf32,
+                    "fp16": self.fp16,
+                    "bf16": self.bf16,
+                    "int8": self.int8,
+                    "fp8": self.fp8,
+                    "calibrator": calibrator,
+                    "use_dla": self.use_dla,
+                    "allow_gpu_fallback": self.allow_gpu_fallback,
+                }
+
+            config_loader_str = make_invocable_if_nondefault(
+                config_alias,
+                precision_constraints=self.precision_constraints,
+                restricted=self.restricted,
+                profiles=profile_name,
+                load_timing_cache=self.load_timing_cache,
+                algorithm_selector=algo_selector,
+                sparse_weights=self.sparse_weights,
+                tactic_sources=self.tactic_sources,
+                memory_pool_limits=self.memory_pool_limits,
+                refittable=self.refittable,
+                strip_plan=self.strip_plan,
+                preview_features=self.preview_features,
+                engine_capability=self.engine_capability,
+                direct_io=self.direct_io,
+                builder_optimization_level=self.builder_optimization_level,
+                hardware_compatibility_level=self.hardware_compatibility_level,
+                profiling_verbosity=self.profiling_verbosity,
+                max_aux_streams=self.max_aux_streams,
+                version_compatible=self.version_compatible,
+                exclude_lean_runtime=self.exclude_lean_runtime,
+                quantization_flags=self.quantization_flags,
+                error_on_timing_cache_miss=self.error_on_timing_cache_miss,
+                disable_compilation_cache=self.disable_compilation_cache,
+                weight_streaming=self.weight_streaming,
+                runtime_platform=self.runtime_platform,
+                tiling_optimization_level=self.tiling_optimization_level,
+                **extra_args
+            )
+            
+            if config_loader_str is None and polygraphy_config.USE_TENSORRT_RTX:
+                config_loader_str = make_invocable(config_alias)
+
+            if config_loader_str is not None:
+                if polygraphy_config.USE_TENSORRT_RTX:
+                    script.add_import(
+                        imports=config_class,
+                        frm="polygraphy.backend.tensorrt_rtx",
+                        imp_as=config_alias,
+                    )
+                else:
+                    script.add_import(
+                        imports=config_class,
+                        frm="polygraphy.backend.trt",
+                        imp_as=config_alias,
+                    )
+
+        if config_loader_str is not None:
+            config_loader_name = script.add_loader(
+                config_loader_str, "create_trt_config"
+            )
+        else:
+            config_loader_name = None
+
+        if self.trt_config_postprocess_script is not None:
+            # Need to set up a default config if there isn't one since `PostprocessConfig` will require a config.
+            if config_loader_name is None:
+                script.add_import(
+                    imports="CreateConfig",
+                    frm="polygraphy.backend.trt",
+                    imp_as="CreateTrtConfig",
+                )
+                config_loader_name = script.add_loader(
+                    make_invocable("CreateTrtConfig"), "create_trt_config"
+                )
+
+            script.add_import(
+                imports=["InvokeFromScript"], frm="polygraphy.backend.common"
+            )
+            script.add_import(
+                imports=["PostprocessConfig"],
+                frm="polygraphy.backend.trt",
+                imp_as="PostprocessTrtConfig",
+            )
+            func = make_invocable(
+                "InvokeFromScript",
+                self.trt_config_postprocess_script,
+                name=self.trt_config_postprocess_func_name,
+            )
+            config_loader_name = script.add_loader(
+                make_invocable("PostprocessTrtConfig", config_loader_name, func=func),
+                "postprocess_trt_config",
+            )
+
+        return config_loader_name
+
+    def create_config(self, builder, network):
+        """
+        Creates a TensorRT BuilderConfig according to arguments provided on the command-line.
+
+        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.
+        """
+        # Use CreateConfigRTX if TensorRT-RTX is enabled, otherwise use CreateConfig
+        if polygraphy_config.USE_TENSORRT_RTX:
+            from polygraphy.backend.tensorrt_rtx import CreateConfigRTX
+            default_loader = CreateConfigRTX()
+        else:
+            from polygraphy.backend.trt import CreateConfig
+            default_loader = CreateConfig()
+
+        loader = util.default(args_util.run_script(self.add_to_script), default_loader)
+        return loader(builder, network)

tools/Polygraphy/polygraphy/tools/args/backend/trt/helper.py

@@ -0,0 +1,28 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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.
+#
+
+
+# This file would have been called `util.py` but if we do that, then for some reason Python thinks
+# that this is the file we want when importing `polygraphy.tools.args.util`.
+from polygraphy.tools.script import inline, inline_identifier, safe
+
+
+def make_trt_enum_val(enum_name, value):
+    """
+    Helper function to create inline TRT enums for usage across various TRT classes.
+    """
+    return inline(safe("trt.{:}.{:}", inline_identifier(enum_name), inline_identifier(value.upper())))

tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py

@@ -0,0 +1,806 @@
+#
+# 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
+
+from polygraphy import mod, util
+from polygraphy.logger import G_LOGGER
+from polygraphy.mod.trt_importer import tensorrt_module_and_version_string
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.backend.onnx.loader import OnnxLoadArgs
+from polygraphy.tools.args.backend.trt.config import TrtConfigArgs
+from polygraphy.tools.args.backend.trt.helper import make_trt_enum_val
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.args.model import ModelArgs
+from polygraphy.tools.script import (
+    inline,
+    inline_identifier,
+    make_invocable,
+    make_invocable_if_nondefault_kwargs,
+    safe,
+)
+
+
+@mod.export()
+class TrtLoadPluginsArgs(BaseArgs):
+    """
+    TensorRT Plugin Loading: loading TensorRT plugins.
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--plugins",
+            help="Path(s) of plugin libraries to load",
+            nargs="+",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            plugins (List[str]): Path(s) to plugin libraries.
+        """
+        self.plugins = args_util.get(args, "plugins")
+
+    # If plugins are present, wrap the provided loader/object with LoadPlugins
+    def add_to_script_impl(self, script, loader_name: str):
+        """
+        Args:
+            loader_name (str):
+                    The name of the loader which should be consumed by the ``LoadPlugins`` loader.
+        """
+        if self.plugins:
+            script.add_import(imports=["LoadPlugins"], frm="polygraphy.backend.trt")
+            loader_str = make_invocable(
+                "LoadPlugins", plugins=self.plugins, obj=loader_name
+            )
+            loader_name = script.add_loader(loader_str, "load_plugins")
+        return loader_name
+
+
+@mod.export()
+class TrtOnnxFlagArgs(BaseArgs):
+    """
+    ONNX-TRT Parser Flags: setting flags for TensorRT's ONNX parser
+
+    Depends on:
+
+        - TrtConfigArgs: If NATIVE_INSTANCENORM should be automatically enabled in VC/HC mode
+    """
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--onnx-flags",
+            help="Flag(s) for adjusting the default parsing behavior of the ONNX parser."
+            "Flag values come from the `trt.OnnxParserFlag` enum and are case-insensitve."
+            "For example: --onnx-flags native_instancenorm ",
+            nargs="+",
+            default=None,
+        )
+        self.group.add_argument(
+            "--plugin-instancenorm",
+            help="Switch to clear the `trt.OnnxParserFlag.NATIVE_INSTANCENORM` flag and"
+            "force the usage of the plugin implementation of ONNX InstanceNorm."
+            "Note that `trt.OnnxParserFlag.NATIVE_INSTANCENORM` is ON by default since TensorRT 10.0.",
+            action="store_true",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            flags (List[str]): flags for onnxparser
+        """
+        self._flags = args_util.get(args, "onnx_flags", default=[])
+        self._plugin_instancenorm = args_util.get(
+            args, "plugin_instancenorm", default=None
+        )
+
+    def get_flags(self):
+        """
+        Updates and returns the ONNX parser flags as necessary.
+        This must be called only in `add_to_script_impl`.
+        Flags should not be accessed directly.
+        """
+        flags = copy.copy(self._flags) or []
+        if (
+            TrtConfigArgs in self.arg_groups
+            and (
+                self.arg_groups[TrtConfigArgs].hardware_compatibility_level is not None
+                or self.arg_groups[TrtConfigArgs].version_compatible
+            )
+            and "native_instancenorm" not in [f.lower() for f in flags]
+        ):
+            G_LOGGER.warning(
+                f"Version or hardware compatibility mode is enabled. Automatically enabling `NATIVE_INSTANCENORM` ONNX parser flag."
+            )
+            flags.append("native_instancenorm")
+
+        return (
+            [make_trt_enum_val("OnnxParserFlag", f) for f in flags] or None,
+            self._plugin_instancenorm,
+        )
+
+
+@mod.export()
+class TrtLoadNetworkArgs(BaseArgs):
+    """
+    TensorRT Network Loading: loading TensorRT networks.
+
+    Depends on:
+
+        - ModelArgs
+        - TrtLoadPluginsArgs
+        - OnnxLoadArgs: if allow_onnx_loading == True
+        - TrtOnnxFlagArgs
+    """
+
+    def __init__(
+        self,
+        allow_custom_outputs: bool = None,
+        allow_onnx_loading: bool = None,
+        allow_tensor_formats: bool = None,
+    ):
+        """
+        Args:
+            allow_custom_outputs (bool):
+                    Whether to allow marking custom output tensors.
+                    Defaults to True.
+            allow_onnx_loading (bool):
+                    Whether to allow parsing networks from an ONNX model.
+                    Defaults to True.
+            allow_tensor_formats (bool):
+                    Whether to allow tensor formats and related options to be set.
+                    Defaults to False.
+        """
+        super().__init__()
+        self._allow_custom_outputs = util.default(allow_custom_outputs, True)
+        self._allow_onnx_loading = util.default(allow_onnx_loading, True)
+        self._allow_tensor_formats = util.default(allow_tensor_formats, False)
+
+    def add_parser_args_impl(self):
+        if self._allow_custom_outputs:
+            self.group.add_argument(
+                "--trt-outputs",
+                help="Name(s) of TensorRT output(s). "
+                "Using '--trt-outputs mark all' indicates that all tensors should be used as outputs",
+                nargs="+",
+                default=None,
+            )
+            self.group.add_argument(
+                "--trt-exclude-outputs",
+                help="[EXPERIMENTAL] Name(s) of TensorRT output(s) to unmark as outputs.",
+                nargs="+",
+                default=None,
+            )
+
+        self.group.add_argument(
+            "--layer-precisions",
+            help="Compute precision to use for each layer. This should be specified on a per-layer basis, using the format: "
+            "--layer-precisions <layer_name>:<layer_precision>. Precision values come from the TensorRT data type aliases, like "
+            "float32, float16, int8, bool, etc. For example: --layer-precisions example_layer:float16 other_layer:int8. "
+            "When this option is provided, you should also set --precision-constraints to either 'prefer' or 'obey'. ",
+            nargs="+",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--tensor-dtypes",
+            "--tensor-datatypes",
+            help="Data type to use for each network I/O tensor. This should be specified on a per-tensor basis, using the format: "
+            "--tensor-datatypes <tensor_name>:<tensor_datatype>. Data type values come from the TensorRT data type aliases, like "
+            "float32, float16, int8, bool, etc. For example: --tensor-datatypes example_tensor:float16 other_tensor:int8. ",
+            nargs="+",
+            default=None,
+        )
+
+        if self._allow_tensor_formats:
+            self.group.add_argument(
+                "--tensor-formats",
+                help="Formats to allow for each network I/O tensor. This should be specified on a per-tensor basis, using the format: "
+                "--tensor-formats <tensor_name>:[<tensor_formats>,...]. Format values come from the `trt.TensorFormat` enum "
+                "and are case-insensitve. "
+                "For example: --tensor-formats example_tensor:[linear,chw4] other_tensor:[chw16]. ",
+                nargs="+",
+                default=None,
+            )
+
+        self.group.add_argument(
+            "--trt-network-func-name",
+            help="[DEPRECATED - function name can be specified alongside the script like so: `my_custom_script.py:my_func`] "
+            "When using a trt-network-script instead of other model types, this specifies the name "
+            "of the function that loads the network. Defaults to `load_network`.",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--trt-network-postprocess-script",
+            "--trt-npps",
+            help="[EXPERIMENTAL] Specify a post-processing script to run on the parsed TensorRT network. The script file may "
+            "optionally be suffixed with the name of the callable to be invoked.  For example: "
+            "`--trt-npps process.py:do_something`. If no callable is specified, then by default "
+            "Polygraphy uses the callable name `postprocess`. "
+            "The callable is expected to take a named argument `network` of type `trt.INetworkDefinition`. "
+            "Multiple scripts may be specified, in which case they are executed in the order given.",
+            nargs="+",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--strongly-typed",
+            help="Mark the network as being strongly typed.",
+            action="store_true",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--mark-debug",
+            help="Specify list of names of tensors to be marked as debug tensors."
+            "For example, `--mark-debug tensor1 tensor2 tensor3`. ",
+            nargs="+",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--mark-unfused-tensors-as-debug-tensors",
+            help="Mark unfused tensors as debug tensors.",
+            action="store_true",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            outputs (List[str]): Names of output tensors.
+            exclude_outputs (List[str]): Names of tensors which should be unmarked as outputs.
+            trt_network_func_name (str): The name of the function in a custom network script that creates the network.
+            layer_precisions (Dict[str, str]): Layer names mapped to their desired compute precision, in string form.
+            tensor_datatypes (Dict[str, str]): Tensor names mapped to their desired data types, in string form.
+            tensor_formats (Dict[str, List[str]]): Tensor names mapped to their desired formats, in string form.
+            postprocess_scripts (List[Tuple[str, str]]):
+                    A list of tuples specifying a path to a network postprocessing script and the name of the postprocessing function.
+            strongly_typed (bool): Whether to mark the network as being strongly typed.
+            mark_debug (List[str]): Names of tensors which should be marked as debug tensors.
+            mark_unfused_tensors_as_debug_tensors (bool): Whether to mark unfused tensors as debug tensors.
+        """
+        self.outputs = args_util.get_outputs(args, "trt_outputs")
+
+        self.exclude_outputs = args_util.get(args, "trt_exclude_outputs")
+
+        self.trt_network_func_name = args_util.get(args, "trt_network_func_name")
+
+        layer_precisions = args_util.parse_arglist_to_dict(
+            args_util.get(args, "layer_precisions"), allow_empty_key=False
+        )
+        self.layer_precisions = None
+        if layer_precisions is not None:
+            self.layer_precisions = {
+                name: inline(safe("trt.{}", inline_identifier(value)))
+                for name, value in layer_precisions.items()
+            }
+
+        tensor_datatypes = args_util.parse_arglist_to_dict(
+            args_util.get(args, "tensor_dtypes"), allow_empty_key=False
+        )
+        self.tensor_datatypes = None
+        if tensor_datatypes is not None:
+            self.tensor_datatypes = {
+                name: inline(safe("trt.{}", inline_identifier(value)))
+                for name, value in tensor_datatypes.items()
+            }
+
+        tensor_formats = args_util.parse_arglist_to_dict(
+            args_util.get(args, "tensor_formats"), allow_empty_key=False
+        )
+        self.tensor_formats = None
+        if tensor_formats is not None:
+            self.tensor_formats = {
+                name: [
+                    inline(
+                        safe("trt.TensorFormat.{}", inline_identifier(value.upper()))
+                    )
+                    for value in values
+                ]
+                for name, values in tensor_formats.items()
+            }
+
+        pps = args_util.parse_arglist_to_tuple_list(
+            args_util.get(args, "trt_network_postprocess_script"),
+            treat_missing_sep_as_val=False,
+        )
+        if pps is None:
+            pps = []
+
+        self.postprocess_scripts = []
+        for script_path, func in pps:
+            if not func:
+                func = "postprocess"
+            if not os.path.isfile(script_path):
+                G_LOGGER.warning(f"Could not find postprocessing script {script_path}")
+            self.postprocess_scripts.append((script_path, func))
+
+        self.strongly_typed = args_util.get(args, "strongly_typed")
+
+        self.mark_debug = args_util.get(args, "mark_debug")
+        self.mark_unfused_tensors_as_debug_tensors = args_util.get(
+            args, "mark_unfused_tensors_as_debug_tensors"
+        )
+
+    def add_to_script_impl(self, script):
+        network_func_name = self.arg_groups[ModelArgs].extra_model_info
+        if self.trt_network_func_name is not None:
+            mod.warn_deprecated(
+                "--trt-network-func-name",
+                "the model argument",
+                "0.50.0",
+                always_show_warning=True,
+            )
+            network_func_name = self.trt_network_func_name
+
+        model_file = self.arg_groups[ModelArgs].path
+        model_type = self.arg_groups[ModelArgs].model_type
+        outputs = args_util.get_outputs_for_script(script, self.outputs)
+        parser_flags, plugin_instancenorm = self.arg_groups[TrtOnnxFlagArgs].get_flags()
+
+        if any(
+            arg is not None
+            for arg in [
+                self.layer_precisions,
+                self.tensor_datatypes,
+                self.tensor_formats,
+                parser_flags,
+                plugin_instancenorm,
+            ]
+        ):
+            script.add_import(imports=tensorrt_module_and_version_string(), imp_as="trt")
+
+        if model_type == "trt-network-script":
+            script.add_import(
+                imports=["InvokeFromScript"], frm="polygraphy.backend.common"
+            )
+            loader_str = make_invocable(
+                "InvokeFromScript",
+                model_file,
+                name=network_func_name,
+            )
+            loader_name = script.add_loader(loader_str, "load_network")
+        elif self._allow_onnx_loading:
+            if self.arg_groups[OnnxLoadArgs].must_use_onnx_loader(
+                disable_custom_outputs=True
+            ):
+                # When loading from ONNX, we need to disable custom outputs since TRT requires dtypes on outputs,
+                # which our marking function doesn't guarantee.
+                script.add_import(
+                    imports=["NetworkFromOnnxBytes"], frm="polygraphy.backend.trt"
+                )
+                onnx_loader = self.arg_groups[OnnxLoadArgs].add_to_script(
+                    script, disable_custom_outputs=True, serialize_model=True
+                )
+                loader_str = make_invocable(
+                    "NetworkFromOnnxBytes",
+                    self.arg_groups[TrtLoadPluginsArgs].add_to_script(
+                        script, onnx_loader
+                    ),
+                    flags=parser_flags,
+                    plugin_instancenorm=plugin_instancenorm,
+                    strongly_typed=self.strongly_typed,
+                    mark_unfused_tensors_as_debug_tensors=self.mark_unfused_tensors_as_debug_tensors,
+                )
+                loader_name = script.add_loader(loader_str, "parse_network_from_onnx")
+            else:
+                script.add_import(
+                    imports=["NetworkFromOnnxPath"], frm="polygraphy.backend.trt"
+                )
+                loader_str = make_invocable(
+                    "NetworkFromOnnxPath",
+                    self.arg_groups[TrtLoadPluginsArgs].add_to_script(
+                        script, model_file
+                    ),
+                    flags=parser_flags,
+                    plugin_instancenorm=plugin_instancenorm,
+                    strongly_typed=self.strongly_typed,
+                    mark_unfused_tensors_as_debug_tensors=self.mark_unfused_tensors_as_debug_tensors,
+                )
+                loader_name = script.add_loader(loader_str, "parse_network_from_onnx")
+        else:
+            G_LOGGER.internal_error(
+                "Loading from ONNX is not enabled and a network script was not provided!"
+            )
+
+        def add_loader_if_nondefault(loader, result_var_name, **kwargs):
+            loader_str = make_invocable_if_nondefault_kwargs(
+                loader, loader_name, **kwargs
+            )
+            if loader_str is not None:
+                script.add_import(imports=[loader], frm="polygraphy.backend.trt")
+                return script.add_loader(loader_str, result_var_name)
+            return loader_name
+
+        for i, (script_path, func_name) in enumerate(self.postprocess_scripts):
+            script.add_import(
+                imports=["InvokeFromScript"], frm="polygraphy.backend.common"
+            )
+            pps = make_invocable("InvokeFromScript", script_path, name=func_name)
+            loader_name = add_loader_if_nondefault(
+                "PostprocessNetwork",
+                f"postprocess_step_{i}",
+                func=pps,
+                name=f"{script_path}:{func_name}",
+            )
+
+        loader_name = add_loader_if_nondefault(
+            "ModifyNetworkOutputs",
+            "set_network_outputs",
+            outputs=outputs,
+            exclude_outputs=self.exclude_outputs,
+        )
+        loader_name = add_loader_if_nondefault(
+            "SetLayerPrecisions",
+            "set_layer_precisions",
+            layer_precisions=self.layer_precisions,
+        )
+        loader_name = add_loader_if_nondefault(
+            "SetTensorDatatypes",
+            "set_tensor_datatypes",
+            tensor_datatypes=self.tensor_datatypes,
+        )
+        loader_name = add_loader_if_nondefault(
+            "SetTensorFormats", "set_tensor_formats", tensor_formats=self.tensor_formats
+        )
+        loader_name = add_loader_if_nondefault(
+            "MarkDebug", "mark_debug", mark_debug=self.mark_debug
+        )
+
+        return loader_name
+
+    def load_network(self):
+        """
+        Loads a TensorRT Network model according to arguments provided on the command-line.
+
+        Returns:
+            tensorrt.INetworkDefinition
+        """
+        loader = args_util.run_script(self.add_to_script)
+        return loader()
+
+
+@mod.export()
+class TrtSaveEngineBytesArgs(BaseArgs):
+    """
+    TensorRT Engine Saving: saving TensorRT engines.
+
+    Saves a serialized engine. This should be preferred over `TrtSaveEngineArgs()` since as of TensorRT 8.6,
+    version compatible engines cannot be re-serialized after they have been initially deserialized.
+    """
+
+    def __init__(self, output_opt: str = None, output_short_opt: str = None):
+        """
+        Args:
+            output_opt (str):
+                    The name of the output path option.
+                    Defaults to "output".
+                    Use a value of ``False`` to disable the option.
+            output_short_opt (str):
+                    The short option to use for the output path.
+                    Defaults to "-o".
+                    Use a value of ``False`` to disable the short option.
+        """
+        super().__init__()
+        self._output_opt = util.default(output_opt, "output")
+        self._output_short_opt = util.default(output_short_opt, "-o")
+
+    def add_parser_args_impl(self):
+        if self._output_opt:
+            params = ([self._output_short_opt] if self._output_short_opt else []) + [
+                f"--{self._output_opt}"
+            ]
+            self.group.add_argument(
+                *params,
+                help="Path to save the TensorRT Engine",
+                dest="save_engine",
+                default=None,
+            )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            path (str): The path at which to save the TensorRT engine.
+        """
+        self.path = args_util.get(args, "save_engine")
+
+    def add_to_script_impl(self, script, loader_name):
+        """
+        Args:
+            loader_name (str):
+                    The name of the loader which will generate the serialized engine.
+
+        Returns:
+            str: The name of the loader added to the script.
+        """
+        if self.path is None:
+            return loader_name
+
+        script.add_import(imports=["SaveBytes"], frm="polygraphy.backend.common")
+        return script.add_loader(
+            make_invocable("SaveBytes", loader_name, path=self.path),
+            "save_engine_bytes",
+        )
+
+    def save_engine_bytes(self, engine_bytes, path=None):
+        """
+        Saves a serialized TensorRT engine according to arguments provided on the command-line.
+
+        Args:
+            engine_bytes (bytes): The serialized TensorRT engine to save.
+
+            path (str):
+                    The path at which to save the engine.
+                    If no path is provided, it is determined from command-line arguments.
+
+        Returns:
+            bytes: The serialized engine that was saved.
+        """
+        with util.TempAttrChange(self, {"path": path}):
+            loader = args_util.run_script(self.add_to_script, engine_bytes)
+            return loader()
+
+
+@mod.deprecate(remove_in="0.55.0", use_instead="TrtSaveEngineBytesArgs")
+@mod.export()
+class TrtSaveEngineArgs(BaseArgs):
+    """
+    TensorRT Engine Saving: saving TensorRT engines.
+
+    Depends on:
+
+        - TrtSaveEngineBytesArgs
+    """
+
+    # For backwards-compatibility
+    @property
+    def path(self):
+        return self.arg_groups[TrtSaveEngineBytesArgs].path
+
+    def add_to_script_impl(self, script, loader_name):
+        """
+        Args:
+            loader_name (str):
+                    The name of the loader which will generate the engine.
+
+        Returns:
+            str: The name of the loader added to the script.
+        """
+        path = self.arg_groups[TrtSaveEngineBytesArgs].path
+
+        if path is None:
+            return loader_name
+
+        script.add_import(imports=["BytesFromEngine"], frm="polygraphy.backend.trt")
+        loader_name = script.add_loader(
+            make_invocable("BytesFromEngine", loader_name, path=path),
+            "bytes_from_engine",
+        )
+        return self.arg_groups[TrtSaveEngineArgs].add_to_script(script, loader_name)
+
+    def save_engine(self, engine, path=None):
+        """
+        Saves a TensorRT engine according to arguments provided on the command-line.
+
+        Args:
+            engine (trt.ICudaEngine): The TensorRT engine to save.
+
+            path (str):
+                    The path at which to save the engine.
+                    If no path is provided, it is determined from command-line arguments.
+
+        Returns:
+            tensorrt.ICudaEngine: The engine that was saved.
+        """
+        with util.TempAttrChange(self, {"path": path}):
+            loader = args_util.run_script(self.add_to_script, engine)
+            return loader()
+
+
+@mod.export()
+class TrtLoadEngineBytesArgs(BaseArgs):
+    """
+    TensorRT Engine: loading or building TensorRT engines.
+
+    Depends on:
+
+        - ModelArgs
+        - TrtLoadPluginsArgs
+        - TrtLoadNetworkArgs: if support for building engines is required
+        - TrtConfigArgs: if support for building engines is required
+        - TrtSaveEngineBytesArgs: if allow_saving == True
+    """
+
+    def __init__(self, allow_saving: bool = None):
+        """
+        Args:
+            allow_saving (bool):
+                    Whether to allow loaded models to be saved.
+                    Defaults to False.
+        """
+        super().__init__()
+        self._allow_saving = util.default(allow_saving, False)
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--save-timing-cache",
+            help="Path to save tactic timing cache if building an engine. "
+            "Existing caches will be appended to with any new timing information gathered. ",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            save_timing_cache (str): Path at which to save the tactic timing cache.
+        """
+        self.save_timing_cache = args_util.get(args, "save_timing_cache")
+
+    def add_to_script_impl(self, script, network_name=None):
+        """
+        Args:
+            network_name (str): The name of a variable in the script pointing to a network loader.
+        """
+        if self.arg_groups[ModelArgs].model_type == "engine":
+            script.add_import(
+                imports=["BytesFromPath"], frm="polygraphy.backend.common"
+            )
+
+            return script.add_loader(
+                make_invocable("BytesFromPath", self.arg_groups[ModelArgs].path),
+                "load_engine_bytes",
+            )
+
+        network_loader_name = network_name
+        if network_loader_name is None:
+            network_loader_name = self.arg_groups[TrtLoadNetworkArgs].add_to_script(
+                script
+            )
+
+        script.add_import(
+            imports=["EngineBytesFromNetwork"], frm="polygraphy.backend.trt"
+        )
+        config_loader_name = self.arg_groups[TrtConfigArgs].add_to_script(script)
+
+        script.add_import(
+            imports=["EngineBytesFromNetwork"], frm="polygraphy.backend.trt"
+        )
+        loader_str = make_invocable(
+            "EngineBytesFromNetwork",
+            self.arg_groups[TrtLoadPluginsArgs].add_to_script(
+                script, network_loader_name
+            ),
+            config=config_loader_name,
+            save_timing_cache=self.save_timing_cache,
+        )
+        loader_name = script.add_loader(loader_str, "build_engine")
+
+        if self._allow_saving:
+            loader_name = self.arg_groups[TrtSaveEngineBytesArgs].add_to_script(
+                script, loader_name
+            )
+        return loader_name
+
+    def load_engine_bytes(self, network=None):
+        """
+        Loads a TensorRT engine according to arguments provided on the command-line.
+
+        Args:
+            network (Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]):
+                    A tuple containing a TensorRT builder, network and optionally parser.
+
+        Returns:
+            tensorrt.ICudaEngine: The engine.
+        """
+        loader = args_util.run_script(self.add_to_script, network)
+        return loader()
+
+
+@mod.export()
+class TrtLoadEngineArgs(BaseArgs):
+    """
+    TensorRT Engine: loading TensorRT engines.
+
+    Depends on:
+
+        - TrtLoadEngineBytesArgs
+        - TrtLoadPluginsArgs
+    """
+
+    # For backwards-compatibility
+    @property
+    def save_timing_cache(self):
+        return self.arg_groups[TrtLoadEngineBytesArgs].save_timing_cache
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--load-runtime",
+            help="Path from which to load a runtime that can be used to load a version compatible "
+            "engine that excludes the lean runtime. ",
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            load_runtime (str):
+                    Path rom which to load a runtime that can be used to load a
+                    version compatible engine that excludes the lean runtime.
+        """
+        self.load_runtime = args_util.parse_path(
+            args_util.get(args, "load_runtime"), "Runtime"
+        )
+
+    def add_to_script_impl(self, script, network_name=None):
+        """
+        Args:
+            network_name (str): The name of a variable in the script pointing to a network loader.
+        """
+        load_serialized_engine = self.arg_groups[TrtLoadEngineBytesArgs].add_to_script(
+            script, network_name
+        )
+
+        script.add_import(imports=["EngineFromBytes"], frm="polygraphy.backend.trt")
+
+        runtime_loader = None
+        if self.load_runtime is not None:
+            script.add_import(imports=["LoadRuntime"], frm="polygraphy.backend.trt")
+            runtime_loader = script.add_loader(
+                make_invocable("LoadRuntime", self.load_runtime), "load_runtime"
+            )
+
+        return script.add_loader(
+            make_invocable(
+                "EngineFromBytes",
+                self.arg_groups[TrtLoadPluginsArgs].add_to_script(
+                    script, load_serialized_engine
+                ),
+                runtime=runtime_loader,
+            ),
+            "deserialize_engine",
+        )
+
+    def load_engine(self, network=None):
+        """
+        Loads a TensorRT engine according to arguments provided on the command-line.
+
+        Args:
+            network (Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]):
+                    A tuple containing a TensorRT builder, network and optionally parser.
+
+        Returns:
+            tensorrt.ICudaEngine: The engine.
+        """
+        loader = args_util.run_script(self.add_to_script, network)
+        return loader()

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

@@ -0,0 +1,108 @@
+#
+# 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.tools.args import util as args_util
+from polygraphy.tools.args.backend.trt.loader import TrtLoadEngineArgs
+from polygraphy.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.script import make_invocable
+
+
+@mod.export()
+class TrtRunnerArgs(BaseRunnerArgs):
+    """
+    TensorRT Inference: running inference with TensorRT.
+
+    Depends on:
+
+        - TrtLoadEngineArgs
+    """
+
+    def get_name_opt_impl(self):
+        return "TensorRT", "trt"
+
+    def add_parser_args_impl(self):
+        self.group.add_argument(
+            "--optimization-profile",
+            help="The index of optimization profile to use for inference",
+            type=int,
+            default=None,
+            dest="optimization_profile",
+        )
+        self.group.add_argument(
+            "--allocation-strategy",
+            help="The way activation memory is allocated. "
+            "static: Pre-allocate based on the max possible size across all profiles. "
+            "profile: Allocate what's needed for the profile to use."
+            "runtime: Allocate what's needed for the current input shapes.",
+            type=str,
+            default=None,
+            dest="allocation_strategy",
+            choices=["static", "profile", "runtime"],
+        )
+        self.group.add_argument(
+            "--weight-streaming-budget",
+            help="The amount of GPU memory in bytes that TensorRT can use for weights at runtime. The engine must be built with weight streaming enabled. It can take on the following values: "
+            "None or -2: Disables weight streaming at runtime. "
+            "-1: TensorRT will decide the streaming budget automatically. "
+            "0 to 100%%: The percentage of weights that TRT keeps on the GPU. 0%% will stream the maximum number of weights."
+            ">=0B: The exact amount of streamable weights that reside on the GPU (unit suffixes are supported).",
+            type=str,
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            optimization_profile (int): The index of the optimization profile to initialize the runner with.
+            allocation_strategy (str): The way activation memory is allocated.
+            weight_streaming_budget (int): The size of the weights on the GPU in bytes.
+            weight_streaming_percent (float): The percentage of weights on the GPU.
+        """
+        self.optimization_profile = args_util.get(args, "optimization_profile")
+        self.allocation_strategy = args_util.get(args, "allocation_strategy")
+        self.weight_streaming_budget = None
+        self.weight_streaming_percent = None
+
+        ws_arg = args_util.get(args, "weight_streaming_budget")
+        if ws_arg and ws_arg.endswith("%"):
+            percent = float(ws_arg[:-1])
+            assert (
+                0 <= percent <= 100
+            ), "Invalid percentage for --weight-streaming-budget!"
+            self.weight_streaming_percent = percent
+        elif ws_arg:
+            budget = args_util.parse_num_bytes(ws_arg)
+            assert (
+                budget == -2 or budget == -1 or budget >= 0
+            ), "Invalid amount for --weight-streaming-budget!"
+            self.weight_streaming_budget = budget
+
+    def add_to_script_impl(self, script):
+        script.add_import(imports=["TrtRunner"], frm="polygraphy.backend.trt")
+        loader_name = self.arg_groups[TrtLoadEngineArgs].add_to_script(script)
+        script.add_runner(
+            make_invocable(
+                "TrtRunner",
+                loader_name,
+                optimization_profile=self.optimization_profile,
+                allocation_strategy=self.allocation_strategy,
+                weight_streaming_budget=self.weight_streaming_budget,
+                weight_streaming_percent=self.weight_streaming_percent,
+            )
+        )

tools/Polygraphy/polygraphy/tools/args/base.py

@@ -0,0 +1,212 @@
+#
+# 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 typing import Tuple
+
+from polygraphy import util
+from polygraphy.common.interface import TypedDict
+from polygraphy.logger import G_LOGGER
+
+
+class ArgGroups(TypedDict(lambda: type, lambda: BaseArgs)):
+    """
+    Maps argument group types to argument groups.
+    """
+
+    pass
+
+
+class BaseArgs:
+    """
+    Adds a arguments to a command-line parser, and provides capabilities to create
+    Polygraphy objects based on the arguments.
+
+    Child classes that add options must define a docstring that includes a section header for the argument group,
+    a brief description which should complete the sentence: "Options related to ...", and finally, any dependencies:
+    ::
+
+        Section Header: Description
+
+        Depends on:
+
+            - OtherArgs0
+            - OtherArgs1: <additional info: condition under which it is needed, or reason for dependency>
+
+        <Optional Additional Documentation>
+
+    For example:
+    ::
+
+        TensorRT Engine: loading TensorRT engines.
+
+        Depends on:
+
+            - ModelArgs
+            - TrtLoadPluginsArgs
+            - TrtLoadNetworkArgs: if building engines
+            - TrtConfigArgs: if building engines
+            - TrtSaveEngineBytesArgs: if allow_saving == True
+
+    The section header and description will be used to popluate the tool's help output.
+    """
+
+    def __init__(self):
+        self.group = None
+        """The ``argparse`` argument group associated with this argument group"""
+
+        # This is populated by the tool base class via ``register()``.
+        self.arg_groups = None
+        """ArgGroups: Maps argument group types to argument groups"""
+
+    # Implementation for `allows_abbreviation`. Derived classes should override this instead of `allows_abbreviation`.
+    def allows_abbreviation_impl(self):
+        return True
+
+    def allows_abbreviation(self):
+        """
+        Whether to allow abbreviated options. When this is enabled, a prefix of an option can be used instead of
+        specifying the entire option. For example, an ``--iterations`` could be specified with just ``--iter``.
+        This breaks ``argparse.REMAINDER``, so any argument groups using that should disable this.
+        The default implementation returns True.
+
+        Returns:
+            bool
+        """
+        return self.allows_abbreviation_impl()
+
+    def register(self, arg_groups):
+        """
+        Registers a dictionary of all available argument groups with this argument group.
+
+        Args:
+            arg_groups (ArgGroups): Maps argument group types to argument groups.
+        """
+        self.arg_groups = arg_groups
+
+    # Implementation for `add_parser_args`. Derived classes should override this instead of `add_parser_args`.
+    # The `self.group` attribute will be populated with an argparse argument group before this is called.
+    def add_parser_args_impl(self):
+        pass
+
+    def add_parser_args(self, parser):
+        """
+        Add arguments to a command-line parser.
+
+        This method is guaranteed to only be called after `register`.
+
+        Args:
+            parser (argparse.ArgumentParser): The argument parser.
+        """
+        title, _, desc = self.__doc__.strip().splitlines()[0].rpartition(":")
+        if not title or not desc:
+            G_LOGGER.internal_error(
+                "Incorrect docstring format, expected 'Title: Description'.\n"
+                f"Note: Docstring was:\n{self.__doc__}.\n"
+                "See BaseArgs documentation for details."
+            )
+
+        num_prev_actions = len(parser._actions)
+
+        self.group = parser.add_argument_group(
+            title.strip(), f"Options related to {desc.strip()}"
+        )
+
+        self.add_parser_args_impl()
+        num_added_actions = len(parser._actions) - num_prev_actions
+
+        # Remove empty groups from the parser.
+        if self.group._action_groups:
+            G_LOGGER.internal_error("Argument groups should not create subgroups!")
+
+        # Remove empty groups from help text
+        if not num_added_actions:
+            parser._action_groups.remove(self.group)
+            self.group = None
+
+    # Implementation for `parse`. Derived classes should override this instead of `parse`.
+    def parse_impl(self, args):
+        pass
+
+    def parse(self, args):
+        """
+        Parses relevant arguments from command-line arguments and populates corresponding
+        attributes of this argument group.
+
+        This method is guaranteed to only be called after `add_parser_args`.
+
+        Args:
+            args: Arguments provided by argparse.
+        """
+        self.parse_impl(args)
+
+    # Implementation for `add_to_script`. Derived classes should override this instead of `add_to_script`.
+    def add_to_script_impl(self, script, *args, **kwargs):
+        raise NotImplementedError()
+
+    def add_to_script(self, script, *args, **kwargs) -> str:
+        """
+        Adds code to the given script that performs the functionality provided by this argument group.
+
+        For example, ``TrtConfigArgs`` would add a call to ``CreateConfig``.
+
+        This method is guaranteed to only be called after `parse`.
+
+        Args:
+            script (polygraphy.tools.script.Script):
+                    A script to which code should be added.
+
+        Returns:
+            str: The name of the variable that was modified or added in the script.
+        """
+        return self.add_to_script_impl(script, *args, **kwargs)
+
+
+class BaseRunnerArgs(BaseArgs):
+    """
+    Similar to BaseArgs, but meant specifically for argument groups dealing with runners.
+    """
+
+    def add_to_script(self, script) -> str:
+        """
+        Returns:
+            str: The name of the list of runners in the script.
+        """
+        self.add_to_script_impl(script)
+        return script.get_runners()
+
+    # Implementation for `get_name_opt`. Derived classes should override this instead of `get_name_opt`.
+    def get_name_opt_impl(self):
+        raise NotImplementedError()
+
+    def get_name_opt(self) -> Tuple[str, str]:
+        """
+        Returns a tuple containing a human readable name of the runner and the name of the command-line option (*without* leading dashes)
+        that should be used to select the runner controlled by this argument group.
+
+        For example: ``("TensorRT", "trt")``.
+        """
+        return self.get_name_opt_impl()
+
+    # Implementation for `get_extra_help_text`. Derived classes should override this instead of `get_extra_help_text`.
+    def get_extra_help_text_impl(self):
+        return ""
+
+    def get_extra_help_text(self) -> str:
+        """
+        Returns any extra help text to display in the tool help output for this runner.
+        """
+        return self.get_extra_help_text_impl()