chore: import upstream snapshot with attribution

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