chore: import upstream snapshot with attribution

tools/Polygraphy/tests/backend/base/test_loader.py

@@ -0,0 +1,32 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pytest
+from polygraphy.backend.base import BaseLoader
+from polygraphy.exception import PolygraphyInternalException
+
+
+def test_loader_checks_call_constant_method():
+    class BadLoader(BaseLoader):
+        def __init__(self):
+            self.x = 2
+
+        def call_impl(self):
+            self.x = 3
+
+    with pytest.raises(PolygraphyInternalException):
+        BadLoader()()

tools/Polygraphy/tests/backend/base/test_runner.py

@@ -0,0 +1,29 @@
+#
+# 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 numpy as np
+import pytest
+from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from tests.models.meta import ONNX_MODELS
+
+
+def test_infer_raises_if_runner_inactive():
+    runner = OnnxrtRunner(SessionFromOnnx(ONNX_MODELS["identity"].loader))
+    feed_dict = {"x": np.ones((1, 1, 2, 2), dtype=np.float32)}
+
+    with pytest.raises(PolygraphyException, match="Must be activated"):
+        runner.infer(feed_dict)

tools/Polygraphy/tests/backend/common/test_loader.py

@@ -0,0 +1,77 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from textwrap import dedent
+
+import pytest
+import tensorrt as trt
+from polygraphy import util
+from polygraphy.backend.common import InvokeFromScript, invoke_from_script
+from polygraphy.exception import PolygraphyException
+
+
+class TestImporter:
+    @pytest.mark.parametrize("loader", [InvokeFromScript, invoke_from_script])
+    def test_import_from_script(self, loader):
+        script = dedent(
+            """
+            from polygraphy.backend.trt import CreateNetwork
+            from polygraphy import func
+            import tensorrt as trt
+
+            @func.extend(CreateNetwork())
+            def load_network(builder, network):
+                inp = network.add_input("input", dtype=trt.float32, shape=(1, 1))
+                out = network.add_identity(inp).get_output(0)
+                network.mark_output(out)
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            if loader == InvokeFromScript:
+                load_network = loader(f.name, "load_network")
+                builder, network = load_network()
+            else:
+                builder, network = loader(f.name, "load_network")
+            with builder, network:
+                assert isinstance(builder, trt.Builder)
+                assert isinstance(network, trt.INetworkDefinition)
+                assert network.num_layers == 1
+                assert network.get_layer(0).type == trt.LayerType.IDENTITY
+
+    def test_import_non_existent(self):
+        script = dedent(
+            """
+            def example():
+                pass
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            with pytest.raises(
+                PolygraphyException, match="Could not import symbol: non_existent from"
+            ):
+                invoke_from_script(f.name, "non_existent")

tools/Polygraphy/tests/backend/onnx/test_loader.py

@@ -0,0 +1,380 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import os
+import tempfile
+
+import numpy as np
+import onnx
+import onnx_graphsurgeon as gs
+import pytest
+
+from polygraphy import constants
+from polygraphy.backend.onnx import (
+    ConvertToFp16,
+    FoldConstants,
+    ModifyOutputs,
+    OnnxFromBytes,
+    OnnxFromPath,
+    OnnxFromTfGraph,
+    SaveOnnx,
+    SetUpperBound,
+    extract_subgraph,
+    fold_constants,
+    gs_from_onnx,
+    infer_shapes,
+    onnx_from_path,
+)
+from polygraphy.common import TensorMetadata
+from polygraphy.logger import G_LOGGER
+from tests.helper import is_file_non_empty
+from tests.models.meta import ONNX_MODELS, TF_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestOnnxFromPath:
+    def test_basic(self):
+        loader = OnnxFromPath(ONNX_MODELS["identity"].path)
+        model = loader()
+        assert isinstance(model, onnx.ModelProto)
+        assert len(model.graph.node) == 1
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_loader_impl_methods):
+        check_warnings_on_loader_impl_methods(
+            OnnxFromPath(ONNX_MODELS["identity"].path)
+        )
+
+    def test_external_data(self):
+        model = ONNX_MODELS["ext_weights"]
+        loader = OnnxFromPath(model.path, model.ext_data)
+        assert isinstance(loader(), onnx.ModelProto)
+
+    def test_ignore_external_data(self):
+        model = ONNX_MODELS["ext_weights"]
+        loader = OnnxFromPath(model.path, ignore_external_data=True)
+        onnx_model = loader()
+
+        assert isinstance(onnx_model, onnx.ModelProto)
+        assert all(init.data_location == 1 for init in onnx_model.graph.initializer)
+
+
+class TestOnnxFromBytes:
+    def test_basic(self):
+        loader = OnnxFromBytes(ONNX_MODELS["identity"].loader)
+        model = loader()
+        assert isinstance(model, onnx.ModelProto)
+        assert len(model.graph.node) == 1
+
+
+class TestGsFromOnnx:
+    def test_basic(self):
+        graph = gs_from_onnx(OnnxFromPath(ONNX_MODELS["identity"].path))
+        assert isinstance(graph, gs.Graph)
+
+
+class TestExportOnnxFromTf:
+    def test_no_optimize(self):
+        pytest.importorskip("tensorflow")
+        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, optimize=False)
+        model = loader()
+
+    def test_opset(self):
+        pytest.importorskip("tensorflow")
+        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, opset=9)
+        model = loader()
+        assert model.opset_import[0].version == 9
+
+
+class TestModifyOnnx:
+    @pytest.mark.parametrize("copy", [True, False])
+    def test_layerwise(self, copy):
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        loader = ModifyOutputs(original_model, outputs=constants.MARK_ALL, copy=copy)
+        model = loader()
+        assert len(original_model.graph.output) == 1 or not copy
+        assert len(model.graph.output) == 2
+
+    @pytest.mark.parametrize("output", ["identity_out_0", "identity_out_2"])
+    def test_custom_outputs(self, output):
+        loader = ModifyOutputs(
+            OnnxFromPath(ONNX_MODELS["identity_identity"].path), outputs=[output]
+        )
+        model = loader()
+        assert len(model.graph.output) == 1
+        assert model.graph.output[0].name == output
+
+    def test_exclude_outputs_with_layerwise(self):
+        loader = ModifyOutputs(
+            OnnxFromPath(ONNX_MODELS["identity_identity"].path),
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["identity_out_2"],
+        )
+        model = loader()
+        assert len(model.graph.output) == 1
+        assert model.graph.output[0].name == "identity_out_0"
+
+
+@pytest.mark.parametrize("allow_onnxruntime", [True, False])
+class TestInferShapes:
+    def check_model(self, model):
+        # Find all intermediate tensors to check if they have shapes.
+        tensors = set()
+        for node in model.graph.node:
+            tensors.update(node.output)
+        tensors -= {out.name for out in model.graph.output}
+
+        assert len(model.graph.value_info) >= len(tensors)
+        for val in model.graph.value_info:
+            assert val.type.tensor_type.HasField("shape")
+
+    def test_model(self, allow_onnxruntime):
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        model = infer_shapes(original_model, allow_onnxruntime=allow_onnxruntime)
+        self.check_model(model)
+
+    def test_path(self, allow_onnxruntime):
+        model = infer_shapes(
+            ONNX_MODELS["identity_identity"].path, allow_onnxruntime=allow_onnxruntime
+        )
+        self.check_model(model)
+
+    @pytest.mark.parametrize("set_data_dir", [True, False])
+    def test_external_data(self, set_data_dir, allow_onnxruntime):
+        model = ONNX_MODELS["ext_weights_same_dir"]
+        model = infer_shapes(
+            model.path,
+            external_data_dir=model.ext_data if set_data_dir else None,
+            allow_onnxruntime=allow_onnxruntime,
+        )
+        self.check_model(model)
+
+    def test_save_to_disk_on_size_threshold(self, allow_onnxruntime):
+        model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
+        model = infer_shapes(
+            model, save_to_disk_threshold_bytes=0, allow_onnxruntime=allow_onnxruntime
+        )
+        self.check_model(model)
+
+
+class TestConvertToFp16:
+    @pytest.mark.parametrize("copy", [True, False])
+    def test_basic(self, copy):
+        # Precondition.
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        assert original_model.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT or not copy
+
+        # Under test.
+        loader = ConvertToFp16(original_model, copy=copy)
+        model = loader()
+
+        # Postcondition.
+        graph = gs_from_onnx(model)
+        graph.toposort()
+
+        assert graph.inputs[0].dtype == "float32"
+        assert graph.nodes[0].op == "Cast"
+        assert graph.nodes[1].op == "Identity"
+        assert graph.nodes[2].op == "Identity"
+        assert graph.nodes[3].op == "Cast"
+        assert graph.outputs[0].dtype == "float32"
+
+
+class TestFoldConstants:
+    @pytest.mark.parametrize("fold_shapes", [True, False])
+    @pytest.mark.parametrize("partitioning", [None, "basic", "recursive"])
+    @pytest.mark.parametrize("copy", [True, False])
+    @pytest.mark.parametrize("allow_onnxruntime_shape_inference", [True, False])
+    def test_basic(
+        self, partitioning, fold_shapes, copy, allow_onnxruntime_shape_inference
+    ):
+        original_model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
+        loader = FoldConstants(
+            original_model,
+            partitioning=partitioning,
+            fold_shapes=fold_shapes,
+            copy=copy,
+            error_ok=False,
+            allow_onnxruntime_shape_inference=allow_onnxruntime_shape_inference,
+        )
+        model = loader()
+        assert len(original_model.graph.node) != 1 or not copy
+        assert len(model.graph.node) == 1
+
+    @pytest.mark.parametrize(
+        "size_threshold, expect_folding",
+        [
+            (None, True),
+            (0, False),
+            (10 << 20, True),
+            (10 << 20 - 1, False),
+        ],
+    )
+    def test_size_threshold(self, size_threshold, expect_folding):
+        model = onnx_from_path(ONNX_MODELS["constant_fold_bloater"].path)
+        model = fold_constants(model, size_threshold=size_threshold)
+
+        if expect_folding:
+            assert len(model.graph.node) == 0
+        else:
+            assert len(model.graph.node) == 1
+            assert model.graph.node[0].op_type == "Tile"
+
+
+class TestSetUpperBound:
+    @pytest.mark.parametrize("global_upper_bound", [False, True])
+    @pytest.mark.parametrize("specified_upper_bound", [False, True])
+    def test_set_upper_bound(
+        self,
+        global_upper_bound,
+        specified_upper_bound,
+    ):
+        original_model = onnx_from_path(ONNX_MODELS["unbounded_dds"].path)
+        upper_bound_dict = {}
+        if not global_upper_bound and not specified_upper_bound:
+            upper_bound_dict[""] = 1000
+            upper_bound = 1000
+        if global_upper_bound:
+            upper_bound_dict[""] = 2000
+            upper_bound = 2000
+        if specified_upper_bound:
+            upper_bound_dict["cast_out_6"] = 4000
+            upper_bound = 4000
+
+        loader = SetUpperBound(
+            original_model,
+            upper_bounds=upper_bound_dict,
+        )
+
+        model = loader()
+        graph = gs_from_onnx(model)
+
+        # Check if there is a Min operator in the modified model
+        find_min = False
+        for node in graph.nodes:
+            if node.op == "Min":
+                find_min = True
+                # Check if the Min operator's second input is a constant tensor
+                assert isinstance(node.inputs[1], gs.Constant)
+
+                val = node.inputs[1].values
+                # Check if the constant value equals the target upper bound
+                assert val == upper_bound
+        assert find_min
+
+
+class TestSaveOnnx:
+    def test_save_onnx(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            outpath = os.path.join(outdir, "test", "nested")
+            loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path), path=outpath)
+            loader()
+            assert is_file_non_empty(outpath)
+
+    def test_external_data(self):
+        with tempfile.NamedTemporaryFile(dir=".") as path, tempfile.NamedTemporaryFile(dir=".") as data:
+            rpath_name = os.path.basename(data.name)
+            model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
+            loader = SaveOnnx(
+                model, path.name, external_data_path=rpath_name, size_threshold=0
+            )
+            loader()
+            assert is_file_non_empty(path.name)
+            assert is_file_non_empty(data.name)
+
+
+@pytest.fixture()
+def extract_model():
+    input_metadata = TensorMetadata().add("X", dtype=np.float32, shape=(64, 64))
+    output_metadata = TensorMetadata().add(
+        "identity_out_0", dtype=np.float32, shape=None
+    )
+    return (
+        onnx_from_path(ONNX_MODELS["identity_identity"].path),
+        input_metadata,
+        output_metadata,
+    )
+
+
+class TestExtractSubgraph:
+    def check_model(self, model):
+        graph = gs_from_onnx(model)
+        assert len(graph.nodes) == 1
+
+        assert len(graph.inputs) == 1
+        assert graph.inputs[0].name == "X"
+        assert graph.inputs[0].shape is not None
+        assert graph.inputs[0].dtype is not None
+
+        assert len(graph.outputs) == 1
+        assert graph.outputs[0].name == "identity_out_0"
+        assert graph.outputs[0].dtype is not None
+
+    def test_extract_onnx_model(self, extract_model):
+        original_model, input_meta, output_meta = extract_model
+        model = extract_subgraph(original_model, input_meta, output_meta)
+
+        assert original_model.graph.output[0].name == "identity_out_2"
+        self.check_model(model)
+
+    def test_extract_onnx_model_no_input_meta(self, extract_model):
+        model, _, output_meta = extract_model
+        model = extract_subgraph(model, output_metadata=output_meta)
+        self.check_model(model)
+
+    def test_extract_onnx_model_no_output_meta(self, extract_model):
+        model, input_meta, _ = extract_model
+        model = extract_subgraph(model, input_metadata=input_meta)
+        assert model.graph.output[0].name == "identity_out_2"
+
+    def test_extract_onnx_gs_graph(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        graph = gs_from_onnx(model)
+        subgraph = extract_subgraph(graph, input_meta, output_meta)
+        # Make sure original graph isn't modified.
+        assert len(graph.nodes) == 2
+
+        assert isinstance(subgraph, gs.Graph)
+        assert len(subgraph.nodes) == 1
+
+        assert len(subgraph.inputs) == 1
+        assert subgraph.inputs[0].name == "X"
+
+        assert len(subgraph.outputs) == 1
+        assert subgraph.outputs[0].name == "identity_out_0"
+
+    def test_extract_passes_no_input_shape(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        input_meta["X"].shape = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)
+
+    def test_extract_passes_no_input_dtype(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        input_meta["X"].dtype = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)
+
+    def test_extract_passes_no_output_shape(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        output_meta["identity_out_0"].shape = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)

tools/Polygraphy/tests/backend/onnx/test_util.py

@@ -0,0 +1,33 @@
+#
+# 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.backend.onnx import onnx_from_path, gs_from_onnx
+from polygraphy.backend.onnx import util as onnx_util
+from tests.models.meta import ONNX_MODELS
+
+
+def test_get_num_nodes():
+    model = onnx_from_path(ONNX_MODELS["scan"].path)
+    assert onnx_util.get_num_nodes(model) == 3  # Should count subgraph nodes.
+
+
+def test_get_unbounded_dds_tensors():
+    model = onnx_from_path(ONNX_MODELS["unbounded_dds"].path)
+    graph = gs_from_onnx(model)
+    tensors = onnx_util.get_unbounded_dds_tensors(graph)
+    assert len(tensors) == 1
+    assert tensors[0].name == "cast_out_6"

tools/Polygraphy/tests/backend/onnxrt/test_loader.py

@@ -0,0 +1,88 @@
+#
+# 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.backend.onnxrt import SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from tests.models.meta import ONNX_MODELS
+
+import onnxruntime as onnxrt
+import pytest
+
+
+class TestSessionFromOnnx:
+    def test_defaults(self):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == ["CPUExecutionProvider"]
+
+    @pytest.mark.parametrize(
+        "providers,expected",
+        [
+            (["cpu"], ["CPUExecutionProvider"]),
+            (["CPU"], ["CPUExecutionProvider"]),
+        ],
+    )
+    def test_provider_matching(self, providers, expected):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=providers)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == expected
+
+    @pytest.mark.skipif(
+        "TensorrtExecutionProvider" not in onnxrt.get_available_providers(),
+        reason="Skip test if TensorrtExecutionProvider is not available",
+    )
+    @pytest.mark.parametrize(
+        "providers,expected_dict",
+        [
+            # Searches for 'tensorrt' as the execution provider's name
+            (["tensorrt", "cpu"], {"TensorrtExecutionProvider": {}, "CPUExecutionProvider": {}}),
+            # Searches for the execution provider's name if the item is a tuple in the format (EP name, EP options)
+            (
+                    [("TensorrtExecutionProvider", {"trt_op_types_to_exclude": "Add"}), "CPUExecutionProvider"],
+                    {"TensorrtExecutionProvider": {"trt_op_types_to_exclude": "Add"}, "CPUExecutionProvider": {}}
+            ),
+        ],
+    )
+    def test_provider_with_options(self, providers, expected_dict):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=providers)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == list(expected_dict.keys())
+
+        provider_options = sess.get_provider_options()
+        for k, v in provider_options.items():
+            if expected_dict.get(k, None):
+                assert set(expected_dict[k].items()).issubset(v.items())
+
+    def test_invalid_providers_raise_errors(self):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=["cpu", "not_a_real_provider"])
+        with pytest.raises(
+            PolygraphyException,
+            match="Could not find specified ONNX-Runtime execution provider",
+        ):
+            loader()

tools/Polygraphy/tests/backend/onnxrt/test_runner.py

@@ -0,0 +1,105 @@
+#
+# 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 numpy as np
+import pytest
+import torch
+
+from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestOnnxrtRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = OnnxrtRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    def test_torch_tensors(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            arr = torch.ones((1, 1, 2, 2), dtype=torch.float32)
+            outputs = runner.infer({"x": arr})
+
+            assert isinstance(outputs["y"], torch.Tensor)
+            assert torch.equal(outputs["y"], arr)
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = ONNX_MODELS["identity"]
+        runner = OnnxrtRunner(SessionFromOnnx(model.loader))
+        check_warnings_on_runner_impl_methods(runner)
+
+    def test_shape_output(self):
+        model = ONNX_MODELS["reshape"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            model.check_runner(runner)
+
+    def test_dim_param_preserved(self):
+        model = ONNX_MODELS["dim_param"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            input_meta = runner.get_input_metadata(use_numpy_dtypes=False)
+            # In Polygraphy, we only use None to indicate a dynamic input dimension - not strings.
+            assert len(input_meta) == 1
+            for _, (_, shape) in input_meta.items():
+                assert shape == ["dim0", 16, 128]
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": np.ones(shape=(1, 1, 2, 2), dtype=np.int32)})
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": np.ones(shape=(1, 1, 3, 2), dtype=np.float32)})

tools/Polygraphy/tests/backend/pluginref/test_runner.py

@@ -0,0 +1,100 @@
+#
+# 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 numpy as np
+import pytest
+from polygraphy.backend.onnx import GsFromOnnx, OnnxFromPath
+from polygraphy.backend.pluginref import PluginRefRunner
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestPluginRefRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = PluginRefRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = ONNX_MODELS["identity"]
+        runner = PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path)))
+        check_warnings_on_runner_impl_methods(runner)
+
+    def test_works_on_multiple_nodes(self):
+        model = ONNX_MODELS["identity_identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            model.check_runner(runner)
+
+    def test_fail_on_unsupported_node(self):
+        model = ONNX_MODELS["and"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(
+                PolygraphyException,
+                match="does not have a reference implementation registered!",
+            ):
+                runner.infer(
+                    {
+                        "x": np.ones(shape=(3, 4), dtype=bool),
+                        "y": np.ones(shape=(3, 4), dtype=bool),
+                    }
+                )
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": np.ones(shape=(1, 1, 2, 2), dtype=np.int32)})
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": np.ones(shape=(1, 1, 3, 2), dtype=np.float32)})

tools/Polygraphy/tests/backend/tf/test_loader.py

@@ -0,0 +1,81 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import os
+import tempfile
+
+import pytest
+from polygraphy import constants, util
+from polygraphy.backend.tf import (
+    GraphFromFrozen,
+    ModifyGraphOutputs,
+    SaveGraph,
+    graph_from_frozen,
+)
+from polygraphy.logger import G_LOGGER
+from tests.helper import is_file_non_empty
+from tests.models.meta import TF_MODELS
+
+tf = pytest.importorskip("tensorflow")
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestFrozenGraphLoader:
+    def test_load_graph(self):
+        with tf.compat.v1.Graph().as_default() as graph:
+            inp = tf.placeholder(shape=(1, 1, 1, 1), dtype=tf.float32)
+            out = tf.identity(inp)
+
+        graph, outputs = graph_from_frozen(graph)
+        assert graph
+        assert outputs
+
+    def test_load_pb(self):
+        tf_loader = GraphFromFrozen(TF_MODELS["identity"].path)
+        tf_loader()
+
+
+class TestModifyGraph:
+    def test_layerwise(self):
+        load_frozen = GraphFromFrozen(TF_MODELS["identity"].path)
+        modify_tf = ModifyGraphOutputs(load_frozen, outputs=constants.MARK_ALL)
+
+        graph, outputs = modify_tf()
+        assert graph
+        assert outputs
+
+
+class TestSaveGraph:
+    def test_save_pb(self):
+        with util.NamedTemporaryFile() as outpath:
+            tf_loader = SaveGraph(
+                GraphFromFrozen(TF_MODELS["identity"].path), path=outpath.name
+            )
+            tf_loader()
+            assert is_file_non_empty(outpath.name)
+
+    def test_save_tensorboard(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            tf_loader = SaveGraph(
+                GraphFromFrozen(TF_MODELS["identity"].path), tensorboard_dir=outdir
+            )
+            tf_loader()
+            assert os.path.exists(tf_loader.tensorboard_dir)

tools/Polygraphy/tests/backend/tf/test_runner.py

@@ -0,0 +1,93 @@
+#
+# 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 numpy as np
+import pytest
+from polygraphy import util
+from polygraphy.backend.tf import SessionFromGraph, TfRunner
+from polygraphy.exception import PolygraphyException
+from tests.helper import is_file_non_empty
+from tests.models.meta import TF_MODELS
+
+pytest.importorskip("tensorflow")
+
+
+class TestTfRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = TfRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = TF_MODELS["identity"]
+        runner = TfRunner(SessionFromGraph(model.loader))
+        check_warnings_on_runner_impl_methods(runner)
+
+    @pytest.mark.skip(reason="Non-trivial to set up - requires CUPTI")
+    def test_save_timeline(self):
+        model = TF_MODELS["identity"]
+        with util.NamedTemporaryFile() as outpath:
+            with TfRunner(
+                SessionFromGraph(model.loader),
+                allow_growth=True,
+                save_timeline=outpath.name,
+            ) as runner:
+                model.check_runner(runner)
+                assert is_file_non_empty(outpath.name)
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "Input:0"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 15, 25, 30), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer(
+                    {"Input:0": np.ones(shape=(1, 15, 25, 30), dtype=np.int32)}
+                )
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer(
+                    {"Input:0": np.ones(shape=(1, 1, 25, 30), dtype=np.float32)}
+                )

tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py

@@ -0,0 +1,340 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from collections import namedtuple
+
+import pytest
+import tensorrt as trt
+from polygraphy import config, util
+from polygraphy.backend.trt import (
+    Algorithm,
+    TacticRecorder,
+    TacticReplayData,
+    TacticReplayer,
+    TensorInfo,
+)
+from polygraphy.exception import PolygraphyException
+
+# Skip all tests in this file if TensorRT-RTX is enabled
+if config.USE_TENSORRT_RTX:
+    pytest.skip("Algorithm selector tests are not compatible with TensorRT-RTX", allow_module_level=True)
+
+
+FakeAlgorithmContext = namedtuple(
+    "FakeAlgorithmContext", ["name", "num_inputs", "num_outputs"]
+)
+FakeAlgorithm = namedtuple("FakeAlgorithm", ["algorithm_variant", "io_info"])
+FakeAlgorithm.get_algorithm_io_info = lambda this, index: this.io_info[index]
+
+FakeAlgorithmVariant = namedtuple("FakeAlgorithmVariant", ["implementation", "tactic"])
+
+
+def fake_context(name):
+    return FakeAlgorithmContext(name=name, num_inputs=1, num_outputs=1)
+
+
+def make_tensor_info(
+    dtype=trt.float32,
+    strides=(1, 2, 3),
+    vectorized_dim=-1,
+    components_per_element=1,
+):
+    return TensorInfo(dtype, strides, vectorized_dim, components_per_element)
+
+
+def fake_algo(implementation=6, tactic=0, io=None):
+    io_info = [make_tensor_info()] * 2
+    if io:
+        io_info = []
+        for dtype, strides in io:
+            io_info.append(
+                TensorInfo(
+                    dtype=dtype,
+                    strides=strides,
+                    vectorized_dim=-1,
+                    components_per_element=1,
+                )
+            )
+
+    trt_algo = FakeAlgorithm(
+        algorithm_variant=FakeAlgorithmVariant(implementation, tactic), io_info=io_info
+    )
+    return trt_algo
+
+
+class TestTensorInfo:
+    @pytest.mark.parametrize(
+        "left, right, expected",
+        [
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                True,
+            ),
+            # Different data type
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float16, (1, 2, 3), -1, 1),
+                False,
+            ),
+            # Different vectotrization
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float32, (1, 2, 3), 0, 2),
+                False,
+            ),
+        ],
+    )
+    def test_equality(self, left, right, expected):
+        assert (left == right) == expected
+
+
+class TestAlgorithm:
+    @pytest.mark.parametrize(
+        "left, right, expected",
+        [
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                True,
+            ),  # Same
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    7,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different implementation
+            (
+                Algorithm(
+                    6,
+                    2,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different tactic
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.int8)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different input data type
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.int8)],
+                ),
+                False,
+            ),  # Different output data type
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)] * 2,
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different number of inputs
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)] * 2,
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different number of outputs
+        ],
+    )
+    def test_equality(self, left, right, expected):
+        assert (left == right) == expected
+
+
+@pytest.fixture(params=[True, False], ids=["path", "object"])
+def replay(request):
+    """
+    Returns:
+        Tuple[FakeAlgorithmContext, Algorithm, FakeAlgorithm,
+              Union[str, TacticReplayData], Union[str, TacticReplayData]]:
+                This fixture returns 5 things:
+                1. A fake TensorRT algorithm context
+                2. A Polygraphy Algorithm instance
+                3. A fake TensorRT algorithm (with the same information as (2))
+                4. An input tactic replay data, populated with the Polygraphy Algorithm (2), either
+                    as a ``TacticReplayData`` instance, or a path.
+                5. An output tactic replay data, empty, either as a ``TacticReplayData`` instance, or
+                    a path.
+    """
+    jsonify = request.param
+
+    name = "node_of_y"
+    context = fake_context(name)
+
+    trt_algo = fake_algo()
+    poly_algo = Algorithm.from_trt(context, trt_algo)
+
+    in_replay_data = TacticReplayData().add(name, poly_algo)
+    out_replay_data = TacticReplayData()
+    if jsonify:
+        inpath = util.NamedTemporaryFile("w")
+        in_replay_data.save(inpath.name)
+        in_replay_data = inpath.name
+
+        outpath = util.NamedTemporaryFile("r")
+        out_replay_data = outpath.name
+
+    yield context, poly_algo, trt_algo, in_replay_data, out_replay_data
+
+
+class TestReplayer:
+    def test_basic(self, replay):
+        context, _, algo, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        selected = replayer.select_algorithms(
+            context, [fake_algo(implementation=2), algo, fake_algo(tactic=1)]
+        )
+        assert selected == [1]
+
+    def test_new_layer_falls_back(self, replay):
+        _, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        selected = replayer.select_algorithms(
+            fake_context(name="new_layer"),
+            [fake_algo(2, 1), fake_algo(3, 4), fake_algo(5, 6)],
+        )
+        assert selected == [0, 1, 2]
+
+    def test_missing_algo_fails(self, replay):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        with pytest.raises(
+            PolygraphyException, match="was not provided by TensorRT as a choice"
+        ):
+            assert replayer.select_algorithms(context, [fake_algo(2, 1)]) == [0]
+
+    @pytest.mark.parametrize(
+        "algo",
+        [
+            fake_algo(2),
+            fake_algo(tactic=2),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.int8, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.int32, (1, 2)),
+                ]
+            ),
+        ],
+    )
+    def test_different_algo_fails(self, replay, algo):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        with pytest.raises(
+            PolygraphyException, match="was not provided by TensorRT as a choice"
+        ):
+            assert replayer.select_algorithms(context, [algo]) == [0]
+
+    def test_fails_if_wrong_selected(self, replay):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        # We should be able to check tactics even if we're not recording them.
+        with pytest.raises(
+            PolygraphyException, match="TensorRT selected a tactic different"
+        ):
+            replayer.report_algorithms([context], [fake_algo(implementation=9)])
+
+
+class TestRecorder:
+    def test_basic(self, replay):
+        context, poly_algo, algo, _, replay_data = replay
+        assert isinstance(replay_data, str) or not replay_data
+        replayer = TacticRecorder(replay_data)
+
+        replayer.report_algorithms([context], [algo])
+        assert len(replayer.data) == 1
+        assert replayer.data[context.name] == poly_algo

tools/Polygraphy/tests/backend/trt/test_calibrator.py

@@ -0,0 +1,330 @@
+#
+# 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 numpy as np
+import pytest
+import tensorrt as trt
+import torch
+
+from polygraphy import config, cuda, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    CreateConfig,
+    Profile,
+    engine_from_network,
+    get_trt_logger,
+    network_from_onnx_bytes,
+)
+from polygraphy.common import TensorMetadata
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from tests.helper import get_file_size, is_file_non_empty
+from tests.models.meta import ONNX_MODELS
+
+# Skip all tests in this file if TensorRT-RTX is enabled
+if config.USE_TENSORRT_RTX:
+    pytest.skip("Calibrator tests are not compatible with TensorRT-RTX", allow_module_level=True)
+
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def dynamic_identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["dynamic_identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def multi_input_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["reducable"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+def generate_data(num_batches):
+    for item in [np.ones((1, 1, 2, 2), dtype=np.float32)] * num_batches:
+        yield {"x": item}
+
+
+class TestCalibrator:
+    def check_calibrator_cleanup(self, calibrator):
+        # Calibrator buffers should be freed after the build
+        assert all([buf.allocated_nbytes == 0 for buf in calibrator.device_buffers.values()])
+
+    @pytest.mark.parametrize(
+        "BaseClass",
+        [
+            trt.IInt8Calibrator,
+            trt.IInt8LegacyCalibrator,
+            trt.IInt8EntropyCalibrator,
+            trt.IInt8EntropyCalibrator2,
+            trt.IInt8MinMaxCalibrator,
+        ],
+    )
+    def test_calibrator_basic(self, identity_builder_network, BaseClass):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}] * NUM_BATCHES
+        calibrator = Calibrator(data, BaseClass=BaseClass)
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_host_data_copied_to_device(self):
+        with Calibrator(generate_data(1)) as calibrator:
+            [ptr] = calibrator.get_batch(names=["x"])
+            v = cuda.DeviceView(ptr, shape=(1, 1, 2, 2), dtype=np.float32)
+            arr = v.numpy()
+            assert arr.shape == (1, 1, 2, 2)
+            assert np.all(arr == 1)
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_data_and_ordering_correct(self):
+        def generate_multidata(num_batches):
+            for _ in range(num_batches):
+                shape = (4, 5)
+                yield {
+                    "x0": np.zeros(shape, dtype=np.float32),
+                    "x1": cuda.DeviceArray(shape=shape, dtype=np.float32).copy_from(np.ones(shape, dtype=np.float32)),
+                    "x2": cuda.DeviceArray(shape=shape, dtype=np.float32)
+                    .copy_from(np.ones(shape, dtype=np.float32) * 2)
+                    .ptr,
+                }
+
+        NUM_BATCHES = 2
+        with Calibrator(generate_multidata(NUM_BATCHES)) as calibrator:
+            for _ in range(NUM_BATCHES):
+                ptrs = calibrator.get_batch(names=["x0", "x1", "x2"])
+                for index, ptr in enumerate(ptrs):
+                    v = cuda.DeviceView(ptr, shape=(4, 5), dtype=np.float32)
+                    assert np.all(v.numpy() == index)
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_generator_data(self, identity_builder_network):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        calibrator = Calibrator(generate_data(NUM_BATCHES))
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    # We should be able to mix DeviceView with NumPy arrays and PyTorch tensors.
+    @pytest.mark.parametrize("mode", ["array", "view", "pointer", "torch"])
+    def test_calibrator_device_buffers_multiinput(self, multi_input_builder_network, mode):
+        def generate_dev_data(num_batches):
+            with cuda.DeviceArray(shape=(1,), dtype=np.float32) as x:
+                for _ in range(num_batches):
+                    x.copy_from(np.ones((1,), dtype=np.float32))
+                    xdata = {
+                        "array": x,
+                        "view": cuda.DeviceView(x.ptr, x.shape, x.dtype),
+                        "pointer": x.ptr,
+                        "torch": torch.ones((1,), dtype=torch.float32),
+                    }[mode]
+                    yield {"X0": xdata, "Y0": np.zeros((1,), dtype=np.float32)}
+
+        builder, network = multi_input_builder_network
+        NUM_BATCHES = 2
+
+        calibrator = Calibrator(generate_dev_data(NUM_BATCHES))
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    # We want the calibrator to inter-op with TRT APIs seamlessly
+    def test_calibrator_outside_polygraphy(self, identity_builder_network):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        config = builder.create_builder_config()
+        config.set_flag(trt.BuilderFlag.INT8)
+        calibrator = Calibrator(generate_data(NUM_BATCHES))
+        config.int8_calibrator = calibrator
+        runtime = trt.Runtime(get_trt_logger())
+        engine = runtime.deserialize_cuda_engine(builder.build_serialized_network(network, config))
+
+        assert engine
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_with_path_name_cache(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile() as cache:
+            calibrator = Calibrator(data, cache=cache.name)
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                assert is_file_non_empty(cache.name)
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize("mode", ["wb+", "rb", "wb"])
+    def test_calibrator_with_file_object_cache(self, identity_builder_network, mode):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile(mode=mode) as cache:
+            calibrator = Calibrator(data, cache=cache)
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                if mode != "rb":
+                    assert is_file_non_empty(cache.name)
+        self.check_calibrator_cleanup(calibrator)
+
+    # read_calibration_cache should work even if an explicit cache is not provided
+    # This way, it is possible to calibrate quickly when calibrating multiple times.
+    def test_calibrator_caches_without_explicit_cache(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        calibrator = Calibrator(data)
+        # First, populate the cache
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            pass
+
+        # Check that the internal cache is populated
+        assert calibrator.read_calibration_cache()
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_rechecks_cache_on_reset(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile(mode="wb+") as cache:
+            calibrator = Calibrator(data, cache=cache.name)
+            # First, populate the cache
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                pass
+
+            # Ensure that now the calibrator will read from the cache when reset
+            calibrator.reset()
+            assert calibrator.cache_contents is None
+            assert len(calibrator.read_calibration_cache()) == get_file_size(cache.name)
+
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize(
+        "names",
+        [
+            (["fake-input", "x"]),
+            (["fake-input"]),
+        ],
+    )
+    def test_calibrator_invalid_input_fails(self, identity_builder_network, names):
+        builder, network = identity_builder_network
+
+        data = [{name: np.ones((1, 1, 2, 2), dtype=np.float32) for name in names}]
+        calibrator = Calibrator(data)
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+
+        with pytest.raises(PolygraphyException):
+            with engine_from_network((builder, network), create_config):
+                pass
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize(
+        "expected_meta,meta,should_pass",
+        [
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                True,
+            ),
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(-1, None, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                True,
+            ),
+            # Wrong data type
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float64, shape=(1, 3, 28, 28)),
+                False,
+            ),
+            # Wrong shape
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 2, 28, 28)),
+                False,
+            ),
+        ],
+    )
+    def test_calibrator_checks_input_metadata(self, expected_meta, meta, should_pass):
+        data = [
+            {
+                name: np.ones(shape=shape, dtype=DataType.to_dtype(dtype, "numpy"))
+                for name, (dtype, shape) in meta.items()
+            }
+        ]
+        calibrator = Calibrator(data)
+        calibrator.set_input_metadata(expected_meta)
+
+        with calibrator:
+            assert (calibrator.get_batch(list(expected_meta.keys())) is not None) == should_pass
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_forces_float32_data(self):
+        data_loader = DataLoader()
+
+        calibrator = Calibrator(data_loader)
+
+        meta = TensorMetadata().add("input", dtype=DataType.FLOAT16, shape=(1, 2, 3))
+        calibrator.set_input_metadata(meta)
+
+        data = data_loader[0]["input"]
+        # TRT requires all calibration inputs to be provided in FP32 regardless of the data type
+        # in the original model.
+        assert util.array.dtype(data) == DataType.FLOAT32
+
+    # TensorRT does not support changing input shapes during calibration
+    @pytest.mark.xfail
+    def test_calibrator_dynamic_shapes(self, dynamic_identity_builder_network):
+        builder, network = dynamic_identity_builder_network
+
+        SHAPES = [(1, 2, 1, 1), (1, 2, 3, 3)]
+
+        def generate_dynamic_shaped_data():
+            for shape in SHAPES:
+                yield {"X": np.ones(shape=shape, dtype=np.float32)}
+
+        calibrator = Calibrator(generate_dynamic_shaped_data())
+
+        create_config = CreateConfig(
+            int8=True,
+            calibrator=calibrator,
+            profiles=[Profile().add(name="X", min=(1, 2, 1, 1), opt=(1, 2, 2, 2), max=(1, 2, 4, 4))],
+        )
+        with engine_from_network((builder, network), create_config) as engine:
+            assert calibrator.num_batches == 2
+            assert engine
+        self.check_calibrator_cleanup(calibrator)

tools/Polygraphy/tests/backend/trt/test_config.py

@@ -0,0 +1,542 @@
+import contextlib
+import os
+import tempfile
+
+import pytest
+
+from polygraphy import mod, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    Profile,
+    network_from_onnx_bytes,
+    postprocess_config,
+)
+from polygraphy.common.struct import BoundedShape
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy import config as polygraphy_config
+from tests.helper import has_dla
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if polygraphy_config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+class TestCreateConfig:
+    def test_defaults(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig()
+        assert loader.timing_cache_path is None
+
+        with loader(builder, network) as config:
+            assert not config.get_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE)
+            with contextlib.suppress(AttributeError):
+                if polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.get_flag(trt.BuilderFlag.TF32)
+                else:
+                    assert not config.get_flag(trt.BuilderFlag.TF32)
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
+            assert not config.get_flag(trt.BuilderFlag.FP16)
+            assert not config.get_flag(trt.BuilderFlag.INT8)
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.BF16)
+            if mod.version(trt.__version__) >= mod.version("8.6"):
+                assert not config.get_flag(trt.BuilderFlag.FP8)
+                assert not config.get_flag(trt.BuilderFlag.VERSION_COMPATIBLE)
+                assert not config.get_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME)
+                if not polygraphy_config.USE_TENSORRT_RTX:
+                    assert (
+                        config.hardware_compatibility_level
+                        == trt.HardwareCompatibilityLevel.NONE
+                    )
+            if mod.version(trt.__version__) >= mod.version("10.2") and not polygraphy_config.USE_TENSORRT_RTX:
+                assert (
+                    config.runtime_platform
+                    == trt.RuntimePlatform.SAME_AS_BUILD
+                )
+            assert config.num_optimization_profiles == 1
+            if not polygraphy_config.USE_TENSORRT_RTX:
+                assert config.int8_calibrator is None
+            with contextlib.suppress(AttributeError):
+                if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.get_tactic_sources() == 24
+                elif mod.version(trt.__version__) >= mod.version("8.7"):
+                    assert config.get_tactic_sources() == 29
+                elif mod.version(trt.__version__) >= mod.version("8.5"):
+                    assert config.get_tactic_sources() == 31
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS)
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.DISABLE_COMPILATION_CACHE)
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
+            with contextlib.suppress(AttributeError):
+                if polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.engine_capability == trt.EngineCapability.STANDARD
+                else:
+                    assert config.engine_capability == trt.EngineCapability.DEFAULT
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.DIRECT_IO)
+
+    @pytest.mark.parametrize(
+        "engine_capability",
+        [
+            trt.EngineCapability.STANDARD,
+            trt.EngineCapability.SAFETY,
+            trt.EngineCapability.DLA_STANDALONE,
+        ],
+    )
+    def test_engine_capability(self, identity_builder_network, engine_capability):
+        builder, network = identity_builder_network
+        loader = CreateConfig(engine_capability=engine_capability)
+        with loader(builder, network) as config:
+            assert config.engine_capability == engine_capability
+
+    @pytest.mark.parametrize("flag", ["obey", "prefer", None])
+    def test_precision_constraints(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(precision_constraints=flag)
+        with loader(builder, network) as config:
+            obey_set = config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
+            prefer_set = config.get_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
+            if flag == "obey":
+                assert obey_set and not prefer_set
+            elif flag == "prefer":
+                assert not obey_set and prefer_set
+            else:
+                assert not obey_set and not prefer_set
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported before TRT 8.6",
+    )
+    @pytest.mark.parametrize(
+        "kwargs, expected_flag",
+        [
+            ({"version_compatible": True}, "VERSION_COMPATIBLE"),
+            (
+                {"version_compatible": True, "exclude_lean_runtime": True},
+                "EXCLUDE_LEAN_RUNTIME",
+            ),
+        ],
+    )
+    def test_version_compatibility_flags(
+        self, identity_builder_network, kwargs, expected_flag
+    ):
+        builder, network = identity_builder_network
+        loader = CreateConfig(**kwargs)
+        with loader(builder, network) as config:
+            assert config.get_flag(getattr(trt.BuilderFlag, expected_flag))
+
+    def test_direct_io(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig(direct_io=True)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.DIRECT_IO)
+
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_restricted(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(restricted=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.SAFETY_SCOPE) == flag
+
+    @pytest.mark.parametrize(
+        "arg_name, flag_type",
+        [
+            ("refittable", trt.BuilderFlag.REFIT),
+        ]
+        + (
+            [
+                (
+                    "disable_compilation_cache",
+                    trt.BuilderFlag.DISABLE_COMPILATION_CACHE,
+                ),
+            ]
+            if mod.version(trt.__version__) >= mod.version("9.0")
+            else []
+        )
+        + (
+            [
+                ("strip_plan", trt.BuilderFlag.STRIP_PLAN),
+            ]
+            if mod.version(trt.__version__) >= mod.version("10.0")
+            else []
+        )
+        + (
+            [
+                ("fp16", trt.BuilderFlag.FP16),
+                ("int8", trt.BuilderFlag.INT8),
+                ("allow_gpu_fallback", trt.BuilderFlag.GPU_FALLBACK),
+                ("tf32", trt.BuilderFlag.TF32),
+            ]
+            + (
+                [
+                    ("bf16", trt.BuilderFlag.BF16),
+                ]
+                if mod.version(trt.__version__) >= mod.version("8.7")
+                else []
+            )
+            + (
+                [
+                    ("fp8", trt.BuilderFlag.FP8),
+                ]
+                if mod.version(trt.__version__) >= mod.version("8.6")
+                else []
+            )
+            if not polygraphy_config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    @pytest.mark.parametrize("value", [True, False])
+    def test_flags(self, identity_builder_network, arg_name, flag_type, value):
+        builder, network = identity_builder_network
+        loader = CreateConfig(**{arg_name: value})
+        with loader(builder, network) as config:
+            assert config.get_flag(flag_type) == value
+
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_sparse_weights(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(sparse_weights=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS) == flag
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support DLA"
+    )
+    def test_use_dla(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig(use_dla=True)
+        with loader(builder, network) as config:
+            assert config.default_device_type == trt.DeviceType.DLA
+            if has_dla():
+                assert config.DLA_core == 0
+
+    with contextlib.suppress(AttributeError):
+        TACTIC_SOURCES_CASES = [
+            (None, 31),  # By default, all sources are enabled.
+            ([], 0),
+            ([trt.TacticSource.CUBLAS], 1),
+            ([trt.TacticSource.CUBLAS_LT], 2),
+            ([trt.TacticSource.CUDNN], 4),
+            ([trt.TacticSource.CUBLAS, trt.TacticSource.CUBLAS_LT], 3),
+            ([trt.TacticSource.CUBLAS, trt.TacticSource.CUDNN], 5),
+            ([trt.TacticSource.CUBLAS_LT, trt.TacticSource.CUDNN], 6),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                ],
+                7,
+            ),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                    trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
+                ],
+                15,
+            ),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                    trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
+                    trt.TacticSource.JIT_CONVOLUTIONS,
+                ],
+                31,
+            ),
+        ]
+
+        if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
+            TACTIC_SOURCES_CASES[0] = (None, 24)
+        elif mod.version(trt.__version__) >= mod.version("8.7"):
+            TACTIC_SOURCES_CASES[0] = (None, 29)
+
+    @pytest.mark.parametrize("sources, expected", TACTIC_SOURCES_CASES)
+    def test_tactic_sources(self, identity_builder_network, sources, expected):
+        builder, network = identity_builder_network
+        loader = CreateConfig(tactic_sources=sources)
+        with loader(builder, network) as config:
+            assert config.get_tactic_sources() == expected
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.7") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="API was added in TRT 8.7",
+    )
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_error_on_timing_cache_miss(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(error_on_timing_cache_miss=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS) == flag
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support calibrators"
+    )
+    def test_calibrator_metadata_set(self, identity_builder_network):
+        builder, network = identity_builder_network
+        calibrator = Calibrator(DataLoader())
+        loader = CreateConfig(int8=True, calibrator=calibrator)
+        with loader(builder, network) as config:
+            assert config.int8_calibrator
+            assert "x" in calibrator.data_loader.input_metadata
+            meta = calibrator.data_loader.input_metadata["x"]
+            assert meta.shape == BoundedShape((1, 1, 2, 2))
+            assert meta.dtype == DataType.FLOAT32
+
+    def test_multiple_profiles(self, identity_builder_network):
+        builder, network = identity_builder_network
+        profiles = [
+            Profile().add("x", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
+            Profile().add("x", (1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)),
+        ]
+        loader = CreateConfig(profiles=profiles)
+        with loader(builder, network) as config:
+            assert config.num_optimization_profiles == 2
+
+    @pytest.mark.parametrize("path_mode", [True, False], ids=["path", "file-like"])
+    def test_timing_cache(self, identity_builder_network, path_mode):
+        builder, network = identity_builder_network
+        with util.NamedTemporaryFile() as cache:
+            loader = CreateConfig(load_timing_cache=cache.name if path_mode else cache)
+            with loader(builder, network) as config:
+                assert config.get_timing_cache()
+
+    def test_fall_back_to_empty_timing_cache(self, identity_builder_network):
+        """Tests that passing in a nonexistent timing cache path is non-fatal"""
+        builder, network = identity_builder_network
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_name = os.path.join(tmpdir, "casper")
+            loader = CreateConfig(load_timing_cache=cache_name)
+            with loader(builder, network) as config:
+                assert config.get_timing_cache()
+
+    def test_empty_timing_cache_when_default(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig()
+        with loader(builder, network) as config:
+            cache = config.get_timing_cache()
+            with cache.serialize() as buffer:
+                cache_size = len(bytes(buffer))
+
+            cache.reset()
+            with cache.serialize() as buffer:
+                new_cache_size = len(bytes(buffer))
+            assert cache_size == new_cache_size
+
+    def test_profiling_verbosity(self, identity_builder_network):
+        builder, network = identity_builder_network
+        expected = trt.ProfilingVerbosity.NONE
+        loader = CreateConfig(profiling_verbosity=expected)
+        with loader(builder, network) as config:
+            assert config.profiling_verbosity == expected
+
+    with contextlib.suppress(AttributeError):
+        POOL_LIMITS = [
+            {trt.MemoryPoolType.WORKSPACE: 25},
+            {trt.MemoryPoolType.DLA_MANAGED_SRAM: 25},
+            {trt.MemoryPoolType.DLA_LOCAL_DRAM: 25},
+            {trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25},
+            # Multiple limits
+            {
+                trt.MemoryPoolType.DLA_LOCAL_DRAM: 20,
+                trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25,
+                trt.MemoryPoolType.WORKSPACE: 39,
+            },
+        ]
+
+        # @pytest.mark.skipif(
+        #     config.USE_TENSORRT_RTX,
+        #     reason="TensorRT-RTX does not support DLA memory pools"
+        # )
+        @pytest.mark.parametrize("pool_limits", POOL_LIMITS)
+        def test_memory_pool_limits(self, pool_limits, identity_builder_network):
+            if any("dla" in key.name.lower() for key in pool_limits) and not has_dla():
+                pytest.skip("DLA is not available on this system")
+
+            builder, network = identity_builder_network
+            loader = CreateConfig(memory_pool_limits=pool_limits)
+            with loader(builder, network) as config:
+                for pool_type, pool_size in pool_limits.items():
+                    assert config.get_memory_pool_limit(pool_type) == pool_size
+
+    @pytest.mark.parametrize(
+        "preview_features",
+        [
+            [trt.PreviewFeature.PROFILE_SHARING_0806]
+            if mod.version(trt.__version__) >= mod.version("10.0")
+            else (
+                [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
+                if polygraphy_config.USE_TENSORRT_RTX
+                else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
+            ),
+        ],
+    )
+    def test_preview_features(self, identity_builder_network, preview_features):
+        builder, network = identity_builder_network
+        loader = CreateConfig(preview_features=preview_features)
+        with loader(builder, network) as config:
+            # Check that only the enabled preview features are on.
+            for pf in trt.PreviewFeature.__members__.values():
+                expected = pf in preview_features
+                # TensorRT-RTX enables PROFILE_SHARING_0806 by default and can't be disabled
+                if polygraphy_config.USE_TENSORRT_RTX and pf == trt.PreviewFeature.PROFILE_SHARING_0806:
+                    expected = True
+                assert config.get_preview_feature(pf) == expected
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support quantization_flag API"
+    )
+    @pytest.mark.parametrize(
+        "quantization_flags",
+        [
+            [],
+            [trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION],
+        ],
+    )
+    def test_quantization_flags(self, identity_builder_network, quantization_flags):
+        builder, network = identity_builder_network
+        loader = CreateConfig(quantization_flags=quantization_flags)
+        with loader(builder, network) as config:
+            # Check that only the enabled quantization flags are on.
+            for qf in trt.QuantizationFlag.__members__.values():
+                assert config.get_quantization_flag(qf) == (qf in quantization_flags)
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported for TRT versions prior to 8.6",
+    )
+    @pytest.mark.parametrize("level", range(6))
+    def test_builder_optimization_level(self, identity_builder_network, level):
+        builder, network = identity_builder_network
+        loader = CreateConfig(builder_optimization_level=level)
+        with loader(builder, network) as config:
+            assert config.builder_optimization_level == level
+
+    if mod.version(trt.__version__) >= mod.version("8.6"):
+
+        @pytest.mark.parametrize(
+            "level",
+            [
+                trt.HardwareCompatibilityLevel.NONE,
+                trt.HardwareCompatibilityLevel.AMPERE_PLUS,
+            ],
+        )
+        def test_hardware_compatibility_level(self, identity_builder_network, level):
+            builder, network = identity_builder_network
+            loader = CreateConfig(hardware_compatibility_level=level)
+            with loader(builder, network) as config:
+                assert config.hardware_compatibility_level == level
+
+    if mod.version(trt.__version__) >= mod.version("10.2"):
+
+        @pytest.mark.parametrize(
+            "platform",
+            [
+                trt.RuntimePlatform.SAME_AS_BUILD,
+                trt.RuntimePlatform.WINDOWS_AMD64,
+            ],
+        )
+        def test_runtime_platform(self, identity_builder_network, platform):
+            builder, network = identity_builder_network
+            loader = CreateConfig(runtime_platform=platform)
+            with loader(builder, network) as config:
+                assert config.runtime_platform == platform
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported for TRT versions prior to 8.6",
+    )
+    @pytest.mark.parametrize("num_streams", range(3))
+    def test_max_aux_streams(self, identity_builder_network, num_streams):
+        builder, network = identity_builder_network
+        loader = CreateConfig(max_aux_streams=num_streams)
+        with loader(builder, network) as config:
+            assert config.max_aux_streams == num_streams
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("9.0") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="API was added in TRT 9.0",
+    )
+    def test_progress_monitor(self, identity_builder_network):
+        class DummyProgressMonitor(trt.IProgressMonitor):
+            def __init__(self):
+                trt.IProgressMonitor.__init__(self)
+
+            def phase_start(self, phase_name, parent_phase, num_steps):
+                pass
+
+            def phase_finish(self, phase_name):
+                pass
+
+            def step_complete(self, phase_name, step):
+                return True
+
+        builder, network = identity_builder_network
+        progress_monitor = DummyProgressMonitor()
+        loader = CreateConfig(progress_monitor=progress_monitor)
+        with loader(builder, network) as config:
+            assert config.progress_monitor == progress_monitor
+
+    if mod.version(trt.__version__) >= mod.version("10.8") and not polygraphy_config.USE_TENSORRT_RTX:
+        @pytest.mark.parametrize(
+            "level",
+            [
+                trt.TilingOptimizationLevel.NONE,
+                trt.TilingOptimizationLevel.FAST,
+                trt.TilingOptimizationLevel.MODERATE,
+                trt.TilingOptimizationLevel.FULL,
+            ],
+        )
+        def test_tiling_optimization_level(self, identity_builder_network, level):
+            builder, network = identity_builder_network
+            loader = CreateConfig(tiling_optimization_level=level)
+            with loader(builder, network) as config:
+                assert config.tiling_optimization_level == level
+
+
+class TestPostprocessConfig:
+    def test_with_config(self, identity_builder_network):
+        builder, network = identity_builder_network
+        config = CreateConfig()(builder, network)
+        assert not config.get_flag(trt.BuilderFlag.INT8)
+
+        config = postprocess_config(
+            config,
+            func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
+            builder=builder,
+            network=network,
+        )
+        assert config.get_flag(trt.BuilderFlag.INT8)
+
+    def test_with_config_callable(self, identity_builder_network):
+        builder, network = identity_builder_network
+        config = CreateConfig()
+
+        config = postprocess_config(
+            config,
+            func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
+            builder=builder,
+            network=network,
+        )
+        assert config.get_flag(trt.BuilderFlag.INT8)

tools/Polygraphy/tests/backend/trt/test_loader.py

@@ -0,0 +1,664 @@
+#
+# 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 __future__ import annotations
+
+import sys
+import pytest
+
+from polygraphy import config, constants, mod, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    EngineBytesFromNetwork,
+    EngineFromBytes,
+    EngineFromNetwork,
+    EngineFromPath,
+    LoadPlugins,
+    LoadRuntime,
+    ModifyNetworkOutputs,
+    NetworkFromOnnxBytes,
+    Profile,
+    SaveEngine,
+    buffer_from_engine,
+    bytes_from_engine,
+    create_config,
+    create_network,
+    engine_from_network,
+    get_trt_logger,
+    modify_network_outputs,
+    network_from_onnx_bytes,
+    network_from_onnx_path,
+    onnx_like_from_network,
+    postprocess_network,
+    set_layer_precisions,
+    set_tensor_datatypes,
+    set_tensor_formats,
+)
+from polygraphy.common.struct import BoundedShape
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from tests.helper import get_file_size, is_file_non_empty
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+##
+## Fixtures
+##
+
+
+@pytest.fixture(scope="session")
+def identity_engine():
+    network_loader = NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader)
+    engine_loader = EngineFromNetwork(network_loader)
+    with engine_loader() as engine:
+        yield engine
+
+
+@pytest.fixture(scope="session")
+def identity_vc_engine_bytes():
+    flags = [trt.OnnxParserFlag.NATIVE_INSTANCENORM]
+    config = CreateConfig(version_compatible=True)
+    network_loader = NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader, flags=flags)
+    engine_loader = EngineBytesFromNetwork(network_loader, config=config)
+    with engine_loader() as engine_bytes:
+        yield engine_bytes
+
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    yield builder, network
+
+
+@pytest.fixture(scope="session")
+def identity_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def identity_identity_network():
+    builder, network, parser = network_from_onnx_bytes(
+        ONNX_MODELS["identity_identity"].loader
+    )
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def reshape_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["reshape"].loader)
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def modifiable_network():
+    # Must return a loader since the network will be modified each time it's loaded.
+    return NetworkFromOnnxBytes(ONNX_MODELS["identity_identity"].loader)
+
+
+@pytest.fixture(scope="session")
+def modifiable_reshape_network():
+    # Must return a loader since the network will be modified each time it's loaded.
+    return NetworkFromOnnxBytes(ONNX_MODELS["reshape"].loader)
+
+
+##
+## Tests
+##
+
+
+class TestLoadPlugins:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="Plugin tests are not compatible with TensorRT-RTX"
+    )
+    def test_can_load_libnvinfer_plugins(self):
+        def get_plugin_names():
+            return [pc.name for pc in trt.get_plugin_registry().plugin_creator_list]
+
+        loader = LoadPlugins(
+            plugins=[
+                (
+                    "nvinfer_plugin.dll"
+                    if sys.platform.startswith("win")
+                    else "libnvinfer_plugin.so"
+                )
+            ]
+        )
+        loader()
+        assert get_plugin_names()
+
+
+class TestSerializedEngineLoader:
+    def test_serialized_engine_loader_from_lambda(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+
+            loader = EngineFromBytes(lambda: open(outpath.name, "rb").read())
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+        
+    def test_serialized_engine_loader_from_buffer(self, identity_engine):
+        with identity_engine.serialize() as buffer:
+            loader = EngineFromBytes(buffer)
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_custom_runtime(self, identity_engine):
+        with identity_engine.serialize() as buffer:
+            loader = EngineFromBytes(buffer, runtime=trt.Runtime(get_trt_logger()))
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 10.0"
+)
+class TestSerializedEngineLoaderFromDisk:
+    def test_serialized_engine_loader_from_lambda(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+
+            loader = EngineFromPath(lambda: outpath.name)
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_custom_runtime(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+            
+            loader = EngineFromPath(lambda: outpath.name, runtime=trt.Runtime(get_trt_logger()))
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("8.6"), reason="API was added in TRT 8.6"
+)
+@pytest.mark.skipif(
+    config.USE_TENSORRT_RTX,
+    reason="TensorRT-RTX does not have lean runtime shared objects"
+)
+class TestLoadRuntime:
+    def test_load_lean_runtime(self, nvinfer_lean_path):
+        loader = LoadRuntime(nvinfer_lean_path)
+        with loader() as runtime:
+            assert isinstance(runtime, trt.Runtime)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("8.6") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 8.6"
+)
+@pytest.mark.skipif(
+    config.USE_TENSORRT_RTX,
+    reason="TensorRT-RTX does not have libnvinfer_lean.so.1"
+)
+class TestSerializedVCEngineLoader:
+    def test_serialized_vc_engine_loader_from_lambda(self, identity_vc_engine_bytes):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f:
+                f.write(identity_vc_engine_bytes)
+
+            loader = EngineFromBytes(lambda: open(outpath.name, "rb").read())
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_from_buffer(self, identity_vc_engine_bytes):
+        loader = EngineFromBytes(identity_vc_engine_bytes)
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+
+class TestNetworkFromOnnxBytes:
+    def test_loader(self):
+        builder, network, parser = network_from_onnx_bytes(
+            ONNX_MODELS["identity"].loader
+        )
+        if not config.USE_TENSORRT_RTX:
+            assert not network.has_implicit_batch_dimension
+
+    @pytest.mark.parametrize(
+        "kwargs, flag",
+        (
+            [
+                (
+                    {"strongly_typed": True},
+                    trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED,
+                )
+            ]
+            if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    def test_network_flags(self, kwargs, flag):
+        builder, network, parser = network_from_onnx_bytes(
+            ONNX_MODELS["identity"].loader, **kwargs
+        )
+        assert network.get_flag(flag)
+
+
+class TestNetworkFromOnnxPath:
+    def test_loader(self):
+        builder, network, parser = network_from_onnx_path(ONNX_MODELS["identity"].path)
+        if not config.USE_TENSORRT_RTX:
+            assert not network.has_implicit_batch_dimension
+
+    @pytest.mark.parametrize(
+        "kwargs, flag",
+        (
+            [
+                (
+                    {"strongly_typed": True},
+                    trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED,
+                )
+            ]
+            if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    def test_network_flags(self, kwargs, flag):
+        builder, network, parser = network_from_onnx_path(
+            ONNX_MODELS["identity"].path, **kwargs
+        )
+        assert network.get_flag(flag)
+
+
+class TestModifyNetwork:
+    def test_mark_layerwise(self, modifiable_network):
+        load_network = ModifyNetworkOutputs(
+            modifiable_network, outputs=constants.MARK_ALL
+        )
+        builder, network, parser = load_network()
+
+        for layer in network:
+            for index in range(layer.num_outputs):
+                assert layer.get_output(index).is_network_output
+
+    def test_mark_custom_outputs(self, modifiable_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_network, outputs=["identity_out_0"]
+        )
+
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "identity_out_0"
+
+    def test_exclude_outputs_with_mark_layerwise(self, modifiable_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_network,
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["identity_out_2"],
+        )
+
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "identity_out_0"
+
+    def test_mark_shape_outputs(self, modifiable_reshape_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_reshape_network, outputs=["output", "reduce_prod_out_gs_2"]
+        )
+
+        assert network.num_outputs == 2
+        assert network.get_output(1).name == "reduce_prod_out_gs_2"
+
+    def test_unmark_shape_outputs(self, modifiable_reshape_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_reshape_network,
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["shape_out_gs_0", "reduce_prod_out_gs_2"],
+        )
+
+        assert network.num_outputs == 1
+
+    def test_mark_outputs_layer_with_optional_inputs(self):
+        builder, network = create_network()
+        inp = network.add_input("input", shape=(1, 3, 224, 224), dtype=trt.float32)
+        slice_layer = network.add_slice(
+            inp, (0, 0, 0, 0), (1, 3, 224, 224), (1, 1, 1, 1)
+        )
+
+        # Set a tensor for `stride` to increment `num_inputs` so we have some inputs
+        # which are `None` in between.
+        slice_layer.set_input(3, inp)
+        assert slice_layer.num_inputs == 4
+
+        slice = slice_layer.get_output(0)
+        slice.name = "Slice"
+
+        builder, network = modify_network_outputs((builder, network), outputs=["Slice"])
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "Slice"
+        assert network.get_output(0) == slice
+
+
+class TestPostprocessNetwork:
+    def test_basic(self, modifiable_network):
+        """Tests that the callback is actually invoked by Polygraphy."""
+        func_called = False
+
+        def func(network):
+            nonlocal func_called
+            func_called = True
+            assert isinstance(network, trt.INetworkDefinition)
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+        assert func_called
+
+    def test_kwargs(self, modifiable_network):
+        """Tests that callbacks that use **kwargs work as expected."""
+        func_called = False
+
+        def func(**kwargs):
+            nonlocal func_called
+            func_called = True
+            assert isinstance(kwargs["network"], trt.INetworkDefinition)
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+        assert func_called
+
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set"
+    )
+    def test_modify_network(self, modifiable_network):
+        """Tests that the network passed in is properly modified by the callback."""
+
+        # Performs the equivalent of set_layer_precisions
+        def func(network):
+            for layer in network:
+                if layer.name == "onnx_graphsurgeon_node_1":
+                    layer.precision = trt.float16
+                if layer.name == "onnx_graphsurgeon_node_3":
+                    layer.precision = trt.int8
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+
+        assert network[0].precision == trt.float16
+        assert network[1].precision == trt.int8
+
+    def test_negative_non_callable(self, modifiable_network):
+        """Tests that PostprocessNetwork properly rejects `func` objects that
+        are not callable."""
+
+        with pytest.raises(PolygraphyException, match=r"Object .* is not a callable"):
+            builder, network, parser = postprocess_network(modifiable_network, None)
+
+
+class TestSetLayerPrecisions:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set"
+    )
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_layer_precisions(
+            modifiable_network,
+            layer_precisions={
+                "onnx_graphsurgeon_node_1": trt.float16,
+                "onnx_graphsurgeon_node_3": trt.int8,
+            },
+        )
+
+        assert network[0].precision == trt.float16
+        assert network[1].precision == trt.int8
+
+
+class TestSetTensorDatatypes:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where tensor datatypes cannot be set"
+    )
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_tensor_datatypes(
+            modifiable_network,
+            tensor_datatypes={
+                "X": trt.float16,
+                "identity_out_2": trt.float16,
+            },
+        )
+
+        assert network.get_input(0).dtype == trt.float16
+        assert network.get_output(0).dtype == trt.float16
+
+
+class TestSetTensorFormats:
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_tensor_formats(
+            modifiable_network,
+            tensor_formats={
+                "X": [trt.TensorFormat.LINEAR, trt.TensorFormat.CHW4],
+                "identity_out_2": [trt.TensorFormat.HWC8],
+            },
+        )
+
+        assert network.get_input(0).allowed_formats == (
+            1 << int(trt.TensorFormat.LINEAR) | 1 << int(trt.TensorFormat.CHW4)
+        )
+        assert network.get_output(0).allowed_formats == 1 << int(trt.TensorFormat.HWC8)
+
+
+class TestEngineBytesFromNetwork:
+    def test_can_build(self, identity_network):
+        loader = EngineBytesFromNetwork(identity_network)
+        with loader() as serialized_engine:
+            assert isinstance(serialized_engine, trt.IHostMemory)
+
+
+class TestEngineFromNetwork:
+    def test_defaults(self, identity_network):
+        loader = EngineFromNetwork(identity_network)
+        assert loader.timing_cache_path is None
+
+    def test_can_build_with_parser_owning(self, identity_network):
+        loader = EngineFromNetwork(identity_network)
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    def test_can_build_without_parser_non_owning(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = EngineFromNetwork((builder, network))
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    def test_custom_runtime(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = EngineFromNetwork(
+            (builder, network), runtime=trt.Runtime(get_trt_logger())
+        )
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support calibrators"
+    )
+    @pytest.mark.parametrize(
+        "use_config_loader, set_calib_profile",
+        [(True, None), (False, False), (False, True)],
+    )
+    def test_can_build_with_calibrator(
+        self, identity_builder_network, use_config_loader, set_calib_profile
+    ):
+        builder, network = identity_builder_network
+        calibrator = Calibrator(DataLoader())
+
+        def check_calibrator():
+            # CreateConfig and EngineFromNetwork should set the input metadata for the calibrator,
+            # which in turn should be passed to the data loader.
+            assert calibrator.input_metadata is not None
+            assert "x" in calibrator.data_loader.input_metadata
+            meta = calibrator.data_loader.input_metadata["x"]
+            assert meta.shape == BoundedShape((1, 1, 2, 2))
+            assert meta.dtype == DataType.FLOAT32
+
+        if use_config_loader:
+            config = create_config(builder, network, int8=True, calibrator=calibrator)
+            check_calibrator()
+        else:
+            config = builder.create_builder_config()
+            config.set_flag(trt.BuilderFlag.INT8)
+            config.int8_calibrator = calibrator
+            # Since this network has static shapes, we shouldn't need to set a calibration profile.
+            if set_calib_profile:
+                calib_profile = (
+                    Profile().fill_defaults(network).to_trt(builder, network)
+                )
+                config.add_optimization_profile(calib_profile)
+                config.set_calibration_profile(calib_profile)
+
+        loader = EngineFromNetwork((builder, network), config)
+        with loader():
+            pass
+        check_calibrator()
+
+        # Calibrator buffers should be freed after the build
+        assert all(
+            [buf.allocated_nbytes == 0 for buf in calibrator.device_buffers.values()]
+        )
+
+    @pytest.mark.parametrize("path_mode", [True, False], ids=["path", "file-like"])
+    def test_timing_cache_generate_and_append(self, path_mode):
+        with util.NamedTemporaryFile() as total_cache, util.NamedTemporaryFile() as identity_cache:
+
+            def build_engine(model, cache):
+                if not path_mode:
+                    cache.seek(0)
+                network_loader = NetworkFromOnnxBytes(ONNX_MODELS[model].loader)
+                # In non-path_mode, use the file-like object directly.
+                # Must load the cache with CreateConfig so that new data is appended
+                # instead of overwriting the previous cache.
+                loader = EngineFromNetwork(
+                    network_loader,
+                    CreateConfig(load_timing_cache=cache.name),
+                    save_timing_cache=cache.name if path_mode else cache,
+                )
+                with loader():
+                    pass
+                if not path_mode:
+                    cache.seek(0)
+
+            assert not total_cache.read()
+
+            build_engine("const_foldable", total_cache)
+            const_foldable_cache_size = get_file_size(total_cache.name)
+
+            # Build this network twice. Once with a fresh cache so we can determine its size.
+            assert get_file_size(identity_cache.name) == 0
+            build_engine("identity", identity_cache)
+            identity_cache_size = get_file_size(identity_cache.name)
+
+            build_engine("identity", total_cache)
+            total_cache_size = get_file_size(total_cache.name)
+
+            # The total cache should be larger than either of the individual caches.
+            assert (
+                total_cache_size >= const_foldable_cache_size
+                and total_cache_size >= identity_cache_size
+            )
+            # The total cache should also be smaller than or equal to the sum of the individual caches since
+            # header information should not be duplicated.
+            assert total_cache_size <= (const_foldable_cache_size + identity_cache_size)
+
+
+class TestBytesFromEngine:
+    def test_serialize_engine(self, identity_network):
+        with engine_from_network(identity_network) as engine:
+            serialized_engine = bytes_from_engine(engine)
+            assert isinstance(serialized_engine, bytes)
+
+
+class TestBufferFromEngine:
+
+    def test_should_return_IHostMemory(self, identity_engine: trt.ICudaEngine) -> None:
+        # Precondition.
+        engine = identity_engine
+
+        # Under test.
+        buffer = buffer_from_engine(engine)
+
+        # Postcondition.
+        assert isinstance(buffer, trt.IHostMemory)
+
+    def test_should_content_match_engine(self, identity_engine: trt.ICudaEngine) -> None:
+        """Test that `BufferFromEngine` returns a buffer with the same content as the engine."""
+        # Precondition.
+        engine = identity_engine
+
+        # Under test.
+        buffer = buffer_from_engine(engine)
+
+        # Postcondition.
+        assert bytes(buffer) == bytes(engine.serialize())
+
+
+class TestSaveEngine:
+
+    def test_should_write_serialized_engine_to_file(self, identity_network: trt.ICudaEngine) -> None:
+        # Precondition.
+        with util.NamedTemporaryFile(mode="wb+") as out_file:
+            name = out_file.name
+            engine = engine_from_network(identity_network)
+
+            # Under test.
+            save_engine = SaveEngine(engine, path=out_file)
+            save_engine()
+            out_file.flush()
+
+            # Postcondition.
+            assert is_file_non_empty(out_file.name)
+            out_file.seek(0)
+            assert bytes(engine.serialize()) == bytes(out_file.read())
+
+
+class TestOnnxLikeFromNetwork:
+    @pytest.mark.parametrize(
+        "model_name",
+        [
+            "identity",
+            "empty_tensor_expand",
+            "const_foldable",
+            "and",
+            "scan",
+            "dim_param",
+            "tensor_attr",
+        ],
+    )
+    def test_onnx_like_from_network(self, model_name):
+        assert onnx_like_from_network(
+            NetworkFromOnnxBytes(ONNX_MODELS[model_name].loader)
+        )
+
+
+class TestDefaultPlugins:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="Plugin tests are not compatible with TensorRT-RTX"
+    )
+    def test_default_plugins(self):
+        network_loader = NetworkFromOnnxBytes(ONNX_MODELS["roialign"].loader)
+        engine_loader = EngineFromNetwork(network_loader)
+        engine = engine_loader()
+        assert engine is not None

tools/Polygraphy/tests/backend/trt/test_profile.py

@@ -0,0 +1,119 @@
+#
+# 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 numpy as np
+import pytest
+
+from polygraphy import config
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+else:
+    import tensorrt as trt
+from polygraphy.backend.trt import Profile, create_network, network_from_onnx_bytes
+from tests.models.meta import ONNX_MODELS
+
+
+@pytest.fixture(scope="session")
+def dynamic_identity_network():
+    builder, network, parser = network_from_onnx_bytes(
+        ONNX_MODELS["dynamic_identity"].loader
+    )
+    with builder, network, parser:
+        yield builder, network, parser
+
+
+class TestProfile:
+    def test_can_add(self):
+        profile = Profile()
+        min, opt, max = (1, 1), (2, 2), (4, 4)
+        assert profile.add("input", min=min, opt=opt, max=max) is profile
+        shape_tuple = profile["input"]
+        assert shape_tuple.min == min
+        assert shape_tuple.opt == opt
+        assert shape_tuple.max == max
+
+    def test_fill_defaults_does_not_overwrite(self, dynamic_identity_network):
+        _, network, _ = dynamic_identity_network
+        profile = Profile().add("X", (1, 1, 1, 1), (1, 1, 2, 2), (1, 1, 3, 3))
+
+        assert profile.fill_defaults(network) is profile
+        assert profile["X"].min == (1, 1, 1, 1)
+        assert profile["X"].opt == (1, 1, 2, 2)
+        assert profile["X"].max == (1, 1, 3, 3)
+
+    def test_fill_defaults_scalar_shape_tensor(self):
+        _, network = create_network()
+        fill_shape = network.add_input("fill_shape", shape=tuple(), dtype=trt.int32)
+
+        # Need to add some other operations so TensorRT treats `fill_shape` as a shape tensor.
+        if config.USE_TENSORRT_RTX:
+            fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE, trt.int32)
+        else:
+            fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE)
+        fill.set_input(0, fill_shape)
+        fill.set_input(
+            1,
+            network.add_constant(
+                shape=tuple(), weights=np.array(0).astype(np.int32)
+            ).get_output(0),
+        )
+        fill.set_input(
+            2,
+            network.add_constant(
+                shape=tuple(), weights=np.array(1).astype(np.int32)
+            ).get_output(0),
+        )
+
+        network.mark_output(fill.get_output(0))
+
+        assert fill_shape.is_shape_tensor
+
+        profile = Profile()
+        profile.fill_defaults(network)
+
+        assert profile[fill_shape.name].min == (1,)
+        assert profile[fill_shape.name].opt == (1,)
+        assert profile[fill_shape.name].max == (1,)
+
+    def test_to_trt(self, dynamic_identity_network):
+        builder, network, _ = dynamic_identity_network
+        profile = Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4))
+
+        trt_profile = profile.to_trt(builder, network)
+        trt_profile.get_shape("X") == ((1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4))
+
+    @pytest.mark.parametrize("name, should_match", [
+        ("inp_*",      [True for _ in range(12)]),
+        ("inp_?",      [False, False, False, *[True for _ in range(9)]]),
+        ("inp_[abc]",  [*[False for _ in range(6)], True, True, True, False, False, False]),
+        ("inp_[!abc]", [False, False, False, True, True, True, False, False, False, True, True, True]),
+    ])
+    def test_input_name_with_wildcards(self, name, should_match):
+        match_case = [
+            "inp_foo", "inp_bar", "inp_123", "inp_1", "inp_s", "inp_k",
+            "inp_a", "inp_b", "inp_c", "inp_d", "inp_e", "inp_f"
+        ]
+        builder, network = create_network()
+        for input in match_case:
+            network.add_input(input, shape=(-1, 2, 3), dtype=trt.float32)
+
+        profile = Profile()
+        profile.add(name, min=(2, 2, 3), opt=(2, 2, 3), max=(2, 2, 3))
+        profile.fill_defaults(network)
+        trt_prof = profile.to_trt(builder, network)
+        res = [trt_prof.get_shape(case) == [(2, 2, 3), (2, 2, 3), (2, 2, 3)] for case in match_case]
+        assert res == should_match 

tools/Polygraphy/tests/backend/trt/test_runner.py

@@ -0,0 +1,408 @@
+#
+# 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 threading
+
+import numpy as np
+import pytest
+import torch
+
+from polygraphy import config, cuda, mod
+from polygraphy.backend.trt import (
+    EngineFromNetwork,
+    NetworkFromOnnxBytes,
+    Profile,
+    TrtRunner,
+    engine_from_network,
+    network_from_onnx_bytes,
+)
+from polygraphy.backend.trt.runner import _get_array_on_cpu
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+@pytest.fixture(scope="class")
+def nonzero_engine():
+    model = ONNX_MODELS["nonzero"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+@pytest.fixture()
+def identity_engine():
+    model = ONNX_MODELS["identity"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+@pytest.fixture()
+def reducable_engine():
+    model = ONNX_MODELS["reducable"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+class TestTrtRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = TrtRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            assert runner.optimization_profile is None
+            assert runner.is_active
+            ONNX_MODELS["identity"].check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX has different warning output behavior")
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods, identity_engine):
+        runner = TrtRunner(identity_engine)
+        check_warnings_on_runner_impl_methods(runner)
+
+    @pytest.mark.parametrize(
+        "inp, expected",
+        [
+            ([1, 0, 1, 1], [[0, 2, 3]]),
+            ([1, 0, 0, 1], [[0, 3]]),
+            ([0, 0, 0, 1], [[3]]),
+        ],
+    )
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support data dependent shapes")
+    def test_data_dependent_shapes(self, nonzero_engine, inp, expected):
+        with TrtRunner(nonzero_engine) as runner:
+            outputs = runner.infer(
+                {
+                    "input": np.array(
+                        inp,
+                        dtype=(np.int32 if mod.version(trt.__version__) < mod.version("9.0") else np.int64),
+                    )
+                }
+            )
+            assert np.array_equal(outputs["nonzero_out_0"], np.array(expected, dtype=np.int32))
+
+    @pytest.mark.parametrize("copy_outputs_to_host", [True, False])
+    @pytest.mark.parametrize("device", ["cpu", "cuda"])
+    def test_torch_tensors(self, copy_outputs_to_host, identity_engine, device):
+        with TrtRunner(identity_engine) as runner:
+            arr = torch.ones([1, 1, 2, 2], dtype=torch.float32, device=device)
+            outputs = runner.infer({"x": arr}, copy_outputs_to_host=copy_outputs_to_host)
+            assert all(isinstance(t, torch.Tensor) for t in outputs.values())
+
+            assert torch.equal(outputs["y"].to("cpu"), arr.to("cpu"))
+
+            assert outputs["y"].device.type == ("cpu" if copy_outputs_to_host else "cuda")
+
+    def test_context(self, identity_engine):
+        with TrtRunner(identity_engine.create_execution_context) as runner:
+            ONNX_MODELS["identity"].check_runner(runner)
+
+    def test_device_buffer_order_matches_bindings(self, reducable_engine):
+        with TrtRunner(reducable_engine) as runner:
+            dev_buf_order = list(runner.device_input_buffers.keys())
+            for binding, dev_buf_name in zip(reducable_engine, dev_buf_order):
+                assert binding == dev_buf_name
+
+    def test_shape_output(self):
+        model = ONNX_MODELS["reshape"]
+        engine = engine_from_network(NetworkFromOnnxBytes(model.loader))
+        with engine, TrtRunner(engine.create_execution_context) as runner:
+            model.check_runner(runner)
+
+    def test_multithreaded_runners_from_engine(self, identity_engine):
+        with TrtRunner(identity_engine) as runner0, TrtRunner(identity_engine) as runner1:
+            t1 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner0,))
+            t2 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner1,))
+            t1.start()
+            t2.start()
+            t1.join()
+            t2.join()
+
+    @pytest.mark.parametrize("use_optimization_profile", [True, False])
+    @pytest.mark.skipif(
+        not config.USE_TENSORRT_RTX
+        and mod.version(trt.__version__) >= mod.version("8.6")
+        and mod.version(trt.__version__) < mod.version("8.7"),
+        reason="Bug in TRT 8.6",
+    )
+    def test_multiple_profiles(self, use_optimization_profile):
+        model = ONNX_MODELS["dynamic_identity"]
+        profile0_shapes = [
+            (1, 2, 1, 1),
+            (1, 2, 1, 1),
+            (1, 2, 1, 1),
+        ]  # Use min==opt==max to fix shapes in the engine.
+        profile1_shapes = [(1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)]
+        profile2_shapes = [(1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [
+            Profile().add("X", *profile0_shapes),
+            Profile().add("X", *profile1_shapes),
+            Profile().add("X", *profile2_shapes),
+        ]
+        config_loader = CreateConfig(profiles=profiles)
+        engine = engine_from_network(network_loader, config_loader)
+        context = engine.create_execution_context()
+
+        for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
+            with TrtRunner(
+                context,
+                optimization_profile=index if use_optimization_profile else None,
+            ) as runner:
+                if not use_optimization_profile:
+                    runner.set_profile(index)
+
+                assert runner.context.active_optimization_profile == index
+                for shape in shapes:
+                    model.check_runner(runner, {"X": shape})
+
+    @pytest.mark.skipif(
+        not config.USE_TENSORRT_RTX and mod.version(trt.__version__) < mod.version("10.0"),
+        reason="Feature not present before 10.0",
+    )
+    @pytest.mark.parametrize("allocation_strategy", [None, "static", "profile", "runtime"])
+    def test_allocation_strategies(self, allocation_strategy):
+        if config.USE_TENSORRT_RTX and allocation_strategy == "runtime":
+            pytest.skip("TensorRT-RTX issues with runtime allocation strategy")
+
+        model = ONNX_MODELS["residual_block"]
+        profile0_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (1, 3, 224, 224)]
+        profile1_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224)]
+        profile2_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (4, 3, 224, 224)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [
+            Profile().add("gpu_0/data_0", *profile0_shapes),
+            Profile().add("gpu_0/data_0", *profile1_shapes),
+            Profile().add("gpu_0/data_0", *profile2_shapes),
+        ]
+        config_loader = CreateConfig(profiles=profiles)
+        engine = engine_from_network(network_loader, config_loader)
+
+        for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
+            with TrtRunner(
+                engine,
+                optimization_profile=index,
+                allocation_strategy=allocation_strategy,
+            ) as runner:
+                for shape in shapes:
+                    model.check_runner(runner, {"gpu_0/data_0": shape})
+
+    def test_empty_tensor_with_dynamic_input_shape_tensor(self):
+        model = ONNX_MODELS["empty_tensor_expand"]
+        shapes = [(1, 2, 0, 3, 0), (2, 2, 0, 3, 0), (4, 2, 0, 3, 0)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [Profile().add("new_shape", *shapes)]
+        config_loader = CreateConfig(profiles=profiles)
+
+        with TrtRunner(EngineFromNetwork(network_loader, config_loader)) as runner:
+            for shape in shapes:
+                model.check_runner(runner, {"new_shape": shape})
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_name_feed_dict(self, names, err, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer({name: module.ones((1, 1, 2, 2), dtype=module.float32) for name in names})
+
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_dtype_feed_dict(self, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": module.ones((1, 1, 2, 2), dtype=module.int32)})
+
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_shape_feed_dict(self, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": module.ones((1, 1, 3, 2), dtype=module.float32)})
+
+    @pytest.mark.parametrize("use_view", [True, False])  # We should be able to use DeviceArray in place of DeviceView
+    def test_device_views(self, use_view, reducable_engine):
+        with TrtRunner(reducable_engine) as runner, cuda.DeviceArray((1,), dtype=np.float32) as x:
+            x.copy_from(np.ones((1,), dtype=np.float32))
+            outputs = runner.infer(
+                {
+                    "X0": x.view() if use_view else x,
+                    "Y0": np.ones((1,), dtype=np.float32),
+                }
+            )
+            assert outputs["identity_out_6"][0] == 2
+            assert outputs["identity_out_8"][0] == 2
+
+    def test_no_output_copy(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+            outputs = runner.infer({"x": inp}, copy_outputs_to_host=False)
+            assert isinstance(outputs["y"], cuda.DeviceView)
+            assert np.array_equal(outputs["y"].numpy(), inp)
+
+    def test_subsequent_infers_with_different_input_types(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+
+            def check(outputs):
+                assert np.all(outputs["y"] == inp)
+
+            check(runner.infer({"x": inp}))
+            check(runner.infer({"x": cuda.DeviceArray(shape=inp.shape, dtype=inp.dtype).copy_from(inp)}))
+
+            torch_outputs = runner.infer({"x": torch.from_numpy(inp)})
+            check({name: out.numpy() for name, out in torch_outputs.items()})
+            check(runner.infer({"x": inp}))
+
+    @pytest.mark.parametrize("use_view", [True, False])  # We should be able to use DeviceArray in place of DeviceView
+    def test_device_view_dynamic_shapes(self, use_view):
+        model = ONNX_MODELS["dynamic_identity"]
+        profiles = [
+            Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
+        ]
+        runner = TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader), CreateConfig(profiles=profiles)))
+        with runner, cuda.DeviceArray(shape=(1, 2, 3, 3), dtype=np.float32) as arr:
+            inp = np.random.random_sample(size=(1, 2, 3, 3)).astype(np.float32)
+            arr.copy_from(inp)
+            outputs = runner.infer({"X": (cuda.DeviceView(arr.ptr, arr.shape, arr.dtype) if use_view else arr)})
+            assert np.all(outputs["Y"] == inp)
+            assert outputs["Y"].shape == (1, 2, 3, 3)
+
+    def test_cannot_use_device_view_shape_tensor(self):
+        model = ONNX_MODELS["empty_tensor_expand"]
+        with TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader))) as runner, cuda.DeviceArray(
+            shape=(5,),
+            dtype=(
+                np.int32
+                if mod.version(trt.__version__) < mod.version("9.0") and not config.USE_TENSORRT_RTX
+                else np.int64
+            ),
+        ) as arr:
+            with pytest.raises(PolygraphyException, match="it must reside in host memory"):
+                runner.infer({"data": np.ones((2, 0, 3, 0), dtype=np.float32), "new_shape": arr})
+
+    @pytest.mark.parametrize("hwc_input", [True, False], ids=["hwc_input", "chw_input"])
+    @pytest.mark.parametrize("hwc_output", [True, False], ids=["hwc_output", "chw_output"])
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support custom I/O format networks")
+    def test_infer_chw_format(self, hwc_input, hwc_output):
+        model = ONNX_MODELS["identity_multi_ch"]
+        inp_shape = model.input_metadata["x"].shape
+        builder, network, parser = network_from_onnx_bytes(model.loader)
+
+        formats = 1 << int(trt.TensorFormat.HWC)
+        if hwc_input:
+            network.get_input(0).allowed_formats = formats
+        if hwc_output:
+            network.get_output(0).allowed_formats = formats
+
+        engine = engine_from_network((builder, network))
+
+        with TrtRunner(engine) as runner:
+            inp = np.random.normal(size=(inp_shape)).astype(np.float32)
+            if hwc_input:
+                inp = inp.transpose(0, 2, 3, 1)
+
+            outputs = runner.infer({"x": inp})
+            if hwc_input == hwc_output:  # output in CHW/HWC format and similarly shaped
+                assert np.allclose(outputs["y"], inp)
+            elif not hwc_input and hwc_output:  # output in HWC format and shaped (N, H, W, C)
+                assert np.allclose(outputs["y"].transpose(0, 3, 1, 2), inp)
+            else:  # hwc_input and not hwc_output: output in CHW format and shaped (N, C, H, W)
+                assert np.allclose(outputs["y"].transpose(0, 2, 3, 1), inp)
+
+    @pytest.mark.parametrize("use_torch", [True, False])
+    def test_get_array_on_cpu(self, use_torch):
+        shape = (4,)
+        with cuda.DeviceArray.raw(shape) as arr:
+            host_buffers = {}
+            stream = cuda.Stream()
+            host_arr = _get_array_on_cpu(arr, "test", host_buffers, stream, arr.nbytes, use_torch)
+
+            if use_torch:
+                assert isinstance(host_arr, torch.Tensor)
+            else:
+                assert isinstance(host_arr, np.ndarray)
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX,
+        reason="Feature not present before 10.0",
+    )
+    @pytest.mark.parametrize("budget", [None, -2, -1, 0, 0.5, 0.99, 1.0, 1000, np.inf])
+    def test_weight_streaming(self, budget):
+        model = ONNX_MODELS["matmul_2layer"]
+        network_loader = NetworkFromOnnxBytes(model.loader, strongly_typed=True)
+        config_loader = CreateConfig(weight_streaming=True)
+        engine = engine_from_network(network_loader, config_loader)
+
+        if budget == np.inf:
+            # set to max size - 1
+            budget = engine.streamable_weights_size - 1
+
+        kwargs = {"weight_streaming_budget": None, "weight_streaming_percent": None}
+        if budget is not None:
+            if 0 < budget <= 1:
+                kwargs["weight_streaming_percent"] = budget * 100
+            else:
+                kwargs["weight_streaming_budget"] = int(budget)
+
+        with TrtRunner(engine, optimization_profile=0, **kwargs) as runner:
+            model.check_runner(runner)
+
+    @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
+    def test_compute_capabilities_engine_building(self):
+        """Test compute capabilities integration with engine building"""
+        model = ONNX_MODELS["identity"]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+
+        # Test --use-gpu flag
+        config_loader = CreateConfig(use_gpu=True)
+        engine = engine_from_network(network_loader, config_loader)
+        with TrtRunner(engine) as runner:
+            model.check_runner(runner)
+
+        # Test --compute-capabilities flag
+        config_loader = CreateConfig(compute_capabilities=[(7, 5), (8, 0), (8, 6)])
+        engine = engine_from_network(network_loader, config_loader)
+        with TrtRunner(engine) as runner:
+            model.check_runner(runner)
+
+    @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
+    def test_compute_capabilities_mutual_exclusion(self):
+        """Test that use_gpu and compute_capabilities are mutually exclusive"""
+        # Test mutual exclusion - should raise an exception
+        with pytest.raises(PolygraphyException, match="use_gpu and compute_capabilities are mutually exclusive"):
+            CreateConfig(use_gpu=True, compute_capabilities=[(7, 5)])

tools/Polygraphy/tests/backend/trt/test_util.py

@@ -0,0 +1,322 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import contextlib
+from textwrap import dedent
+
+import pytest
+
+from polygraphy import config, mod
+from polygraphy.backend.trt import Profile, create_network
+from polygraphy.backend.trt import util as trt_util
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+@pytest.fixture(scope="session")
+def dummy_network():
+    builder, network = create_network()
+    network.add_input("X", dtype=trt.float32, shape=[-1])
+    with builder, network:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def layer_class_mapping():
+    return trt_util.get_layer_class_mapping()
+
+
+@pytest.mark.parametrize("layer_type", trt.LayerType.__members__.values())
+def test_all_layer_types_mapped(layer_class_mapping, layer_type):
+    waived_layers = [trt.LayerType.PLUGIN]
+    with contextlib.suppress(AttributeError):
+        waived_layers.append(trt.LayerType.PLUGIN_V3)
+    if layer_type in waived_layers:
+        pytest.skip("PLUGIN has no corresponding ILayer")
+    assert layer_type in layer_class_mapping
+
+
+# Can't use pytest.skip because we can't construct the test unless trt.MemoryPoolType exists.
+
+
+def adjust_memory_pool_limits_after_8_6(limits):
+    # Adjust tactic DRAM so we can match the output text reliably in update_expected_output.
+    if mod.version(trt.__version__) >= mod.version("8.6") or config.USE_TENSORRT_RTX:
+        limits[trt.MemoryPoolType.TACTIC_DRAM] = 1 << 30
+    return limits
+
+
+def update_expected_output(expected):
+    is_trt_10_plus = (
+        mod.version(trt.__version__) >= mod.version("10.0") or 
+        config.USE_TENSORRT_RTX
+    )
+    is_trt_8_6_plus = (
+        mod.version(trt.__version__) >= mod.version("8.6") or 
+        config.USE_TENSORRT_RTX
+    )
+    is_trt_8_7_plus = (
+        mod.version(trt.__version__) >= mod.version("8.7") or 
+        config.USE_TENSORRT_RTX
+    )
+    
+    if is_trt_8_6_plus:
+        if is_trt_10_plus:
+            expected = expected.replace(
+                "MiB]",
+                "MiB, TACTIC_DRAM: 1024.00 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]",
+            )
+        else:
+            expected = expected.replace("MiB]", "MiB, TACTIC_DRAM: 1024.00 MiB]")
+
+        if "Preview Features" not in expected:
+            if not is_trt_10_plus:
+                expected = (
+                    dedent(expected).strip()
+                    + "\nPreview Features       | [FASTER_DYNAMIC_SHAPES_0805, DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]"
+                )
+            else:
+                preview_features = "[PROFILE_SHARING_0806"
+                if config.USE_TENSORRT_RTX:
+                    preview_features += ", RUNTIME_ACTIVATION_RESIZE_10_10"
+                preview_features += "]"
+                expected = (
+                    dedent(expected).strip()
+                    + f"\nPreview Features       | {preview_features}"
+                )
+
+    if is_trt_8_7_plus:
+        # CUBLAS_LT is not longer enabled by default
+        expected = expected.replace("CUBLAS_LT, ", "")
+
+    if is_trt_10_plus:
+        expected = expected.replace(
+            "EngineCapability.DEFAULT", "EngineCapability.STANDARD"
+        )
+        expected = expected.replace("CUBLAS, ", "")
+        expected = expected.replace("CUDNN, ", "")
+
+    return expected
+
+
+@pytest.mark.parametrize(
+    "create_config, expected",
+    # NOTE: We set workspace sizes here so we can have predictable output
+    [
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                )
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                tactic_sources=[],
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | []
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 4 << 20}
+                )
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 4.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                **({} if config.USE_TENSORRT_RTX else {
+                    "fp16": True,
+                    "int8": True,
+                    "tf32": True,
+                }),
+                refittable=True,
+                precision_constraints="obey",
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}REFIT, TF32, OBEY_PRECISION_CONSTRAINTS]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format(
+                    "" if config.USE_TENSORRT_RTX else "FP16, INT8, ",
+                )
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                profiles=[
+                    Profile().add("X", [1], [1], [1]),
+                    Profile().add("X", [2], [2], [2]),
+                ],
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Optimization Profiles  | 2 profile(s)
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+    ] + ([] if config.USE_TENSORRT_RTX else [
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                use_dla=True,
+            ),
+            update_expected_output(
+                """
+                Flags                  | []
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB, DLA_MANAGED_SRAM: 0.00 MiB, DLA_LOCAL_DRAM: 1024.00 MiB, DLA_GLOBAL_DRAM: 512.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                DLA                    | Default Device Type: DeviceType.DLA, Core: -1
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """
+            ),
+        ),
+    ]) + [
+        (
+            (
+                CreateConfig(
+                    memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                        {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                    ),
+                    preview_features=[trt.PreviewFeature.PROFILE_SHARING_0806],
+                ),
+                update_expected_output(
+                    """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Preview Features       | [PROFILE_SHARING_0806]
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+                ),
+            )
+            if mod.version(trt.__version__) >= mod.version("10.0") or config.USE_TENSORRT_RTX
+            else (
+                CreateConfig(
+                    memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                        {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                    ),
+                    preview_features=(
+                        [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03] 
+                        if config.USE_TENSORRT_RTX 
+                        else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
+                    ),
+                ),
+                update_expected_output(
+                    """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Preview Features       | [{}]
+                """.format(
+                    "TF32" if config.USE_TENSORRT_RTX else "",
+                    "ALIASED_PLUGIN_IO_10_03" if config.USE_TENSORRT_RTX else "FASTER_DYNAMIC_SHAPES_0805"
+                )
+                ),
+            )
+        ),
+    ],
+    ids=[
+        "default",
+        "tactic-sources",
+        "memory-pool-limits",
+        "builder-flags" + ("-rtx" if config.USE_TENSORRT_RTX else ""),
+        "profiles",
+    ] + ([] if config.USE_TENSORRT_RTX else ["dla"]) + [
+        "preview-features",
+    ],
+)
+def test_str_from_config(create_config, expected, dummy_network):
+    config = create_config(*dummy_network)
+    actual = trt_util.str_from_config(config, dummy_network)
+    expected = dedent(expected).strip()
+    assert actual == expected
+
+
+def test_get_all_tensors_layer_with_null_inputs():
+    builder, network = create_network()
+    with builder, network:
+        inp = network.add_input("input", shape=(1, 3, 224, 224), dtype=trt.float32)
+        slice_layer = network.add_slice(
+            inp, (0, 0, 0, 0), (1, 3, 224, 224), (1, 1, 1, 1)
+        )
+
+        # Set a tensor for `stride` to increment `num_inputs` so we have some inputs
+        # which are `None` in between.
+        slice_layer.set_input(3, inp)
+        assert slice_layer.num_inputs == 4
+
+        slice = slice_layer.get_output(0)
+        slice.name = "Slice"
+        network.mark_output(slice)
+
+        assert trt_util.get_all_tensors(network) == {"input": inp, "Slice": slice}