chore: import upstream snapshot with attribution

tools/polygraphy-extension-trtexec/.gitignore

@@ -0,0 +1,72 @@
+# Byte-compiled / optimized / DLL files
+*/__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Sphinx documentation
+docs/_build/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Crash dumps
+core

tools/polygraphy-extension-trtexec/CHANGELOG.md

@@ -0,0 +1,53 @@
+# Polygraphy Trtexec Extension Change Log
+
+Dates are in YYYY-MM-DD format.
+## vNext
+
+## v0.0.9 (2023-08-21)
+### Changed
+- Update API usage to accomodate `Polygraphy==0.49.0`.
+- Replace `workspace` with `memPoolSize` in `TrtexecRunner`.
+- Remove usage of `polygraphy.backend.trt.util` in `polygraphy_trtexec/backend/runner.py`.
+
+## v0.0.8 (2022-09-08)
+### Added
+- `trtexec-export-times`, `trtexec-export-output`, `trtexec-export-profile` and `trtexec-export-layer-info` flags
+- Report e2e Host Latency instead of GPU Compute time
+
+### Changed
+- Use latency as the inference metric instead of GPU Compute Time
+
+## v0.0.7 (2022-08-10)
+### Fixed
+- Fixed README rendering in wheel build
+
+## v0.0.6 (2020-07-21)
+### Changed
+- Bug fix in reporting inference time
+
+## v0.0.5 (2020-07-21)
+### Changed
+- Bug fix in handling `refit` flag for `polygraphy<'0.39.0`
+
+## v0.0.4 (2020-07-18)
+### Added
+- `trtexec-path` flag applies to allow specifying path to a custom `trtexec` binary
+
+## v0.0.3 (2020-07-15)
+### Changed
+- `refittable` flag applies to `trtexec` backend's `refit` only if `polygraphy>=0.39.0`
+- Renamed `no-builder-cache` to `trtexec-no-builder-cache`
+- Renamed `profiling-verbosity` to `trtexec-profiling-verbosity`
+
+## v0.0.2 (2020-07-05)
+### Added
+- `--trtexec-iterations` flag. Avoids unwanted behavior of running N times more iterations than expected when using the --iterations flag
+
+### Changed
+- Renamed `trtexec-warm-up` to `trtexec-warmup`
+
+### Fixed
+- Bug fix for failing pytests
+
+## v0.0.1 (2022-06-23)
+- Initial integration

tools/polygraphy-extension-trtexec/LICENSE.txt

@@ -0,0 +1,191 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2020 NVIDIA Corporation
+
+   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.

tools/polygraphy-extension-trtexec/Makefile

@@ -0,0 +1,35 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+.ONESHELL: # Run all rules in the same shell to make it easier to run tests with environment variables
+
+BUILD_DIR := build
+
+clean:
+	- rm -rf dist/ $(BUILD_DIR)/ polygraphy_trtexec.egg-info/ polygraphy_trtexec/__pycache__/ 
+
+build: clean
+	- python3 setup.py bdist_wheel
+
+install_deps: build
+	- python3 -m pip install colored wheel
+
+install: install_deps
+	- python3 -m pip install --force-reinstall --no-deps $(CURDIR)/dist/*.whl
+
+test: install
+	python3 -m pytest

tools/polygraphy-extension-trtexec/README.md

@@ -0,0 +1,30 @@
+# Extending `polygraphy run` to support `trtexec`
+
+## Introduction
+
+`polygraphy run` allows you to run inference with multiple backends, including TensorRT and ONNX-Runtime, and compare outputs.
+This extension adds support to run inference with `trtexec`.
+
+## Installation
+
+Follow the steps below to install the extension module. After installation, you should see the `trtexec` options in the help
+output of `polygraphy run`:
+
+1. Build using `setup.py`:
+
+    ```bash
+    python3 setup.py bdist_wheel
+    ```
+
+2. Install the wheel: The wheel is installed in the `dist` directory. Install the wheel by running the following command    
+    ```
+    python3 -m pip install dist/polygraphy_trtexec-*.whl \
+        --extra-index-url https://pypi.ngc.nvidia.com
+    ```
+    *NOTE: You may have to update the above command to install the appropriate version of the wheel*
+
+3. After the installation, you can run it on the trtexec backend by using the `--trtexec` flag as follows:
+
+    ```bash
+    polygraphy run sample.onnx --trtexec
+    ```

tools/polygraphy-extension-trtexec/polygraphy_trtexec/__init__.py

@@ -0,0 +1,18 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+__version__ = "0.0.9"

tools/polygraphy-extension-trtexec/polygraphy_trtexec/args/__init__.py

@@ -0,0 +1,18 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from polygraphy_trtexec.args.runner import *

tools/polygraphy-extension-trtexec/polygraphy_trtexec/args/runner.py

@@ -0,0 +1,366 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+This file defines the `TrtexecRunnerArgs` argument group, which manages
+command-line options that control the `TrtexecRunner` runner.
+
+The argument group implements the standard `BaseRunnerArgs` interface, which inherits from `BaseArgs`.
+"""
+
+import polygraphy
+from polygraphy import mod
+from polygraphy.tools.args import ModelArgs, TrtConfigArgs, TrtLoadPluginsArgs, TrtLoadNetworkArgs, TrtSaveEngineBytesArgs, util as args_util
+from polygraphy.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.script import make_invocable
+
+@mod.export()
+class TrtexecRunnerArgs(BaseRunnerArgs):
+    """
+    Trtexec Runner Inference: running inference with the trtexec backend.
+
+    Depends on:
+        ModelArgs
+        TrtConfigArgs
+        TrtLoadPluginsArgs
+        TrtSaveEngineBytesArgs
+    """
+
+    def get_name_opt_impl(self):
+        return "Trtexec Runner", "trtexec"
+
+    def add_parser_args_impl(self):
+        """
+        Add command-line arguments that trtexec supports
+        """
+
+        self.group.add_argument(
+            "--trtexec-path",
+            help="Path to find trtexec binary. By default, it expects to find it in PATH",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--use-cuda-graph",
+            help="Use CUDA graph to capture engine execution and then launch inference (default = disabled). This flag may be ignored if the graph capture fails",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--avg-runs",
+            help="Report performance measurements averaged over N consecutive iterations (default = 10)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--best",
+            help="Enable all precisions to achieve the best performance (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--duration",
+            help="Run performance measurements for at least N seconds wallclock time (default = 3)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--device",
+            help="Select cuda device N (default = 0)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--streams",
+            help="Instantiate N engines to use concurrently (default = 1)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--min-timing",
+            help="Set the minimum number of iterations used in kernel selection (default = 1)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--avg-timing",
+            help="Set the number of times averaged in each iteration for kernel selection (default = 8)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--expose-dma",
+            help="Serialize DMA transfers to and from device (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--no-data-transfers",
+            help="Disable DMA transfers to and from device (default = enabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--trtexec-warmup",
+            help="Run for N milliseconds on trtexec to warmup before measuring performance (default = 200)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--trtexec-iterations",
+            help="Run at least N inference iterations on trtexec (default = 10)",
+            default=None,
+            type=int,
+        )
+
+        self.group.add_argument(
+            "--trtexec-export-times",
+            help="Write the timing results in a json file",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--trtexec-export-output",
+            help="Write the output tensors to a json file",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--trtexec-export-profile",
+            help="Write the profile information per layer in a json file",
+            default=None,
+        )
+
+        self.group.add_argument(
+            "--trtexec-export-layer-info",
+            help="Write the layer information of the engine in a json file",
+            default=None,
+        )
+
+        # Optional
+
+        self.group.add_argument(
+            "--use-spin-wait",
+            help="Actively synchronize on GPU events. This option may decrease synchronization time but increase CPU usage and power (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--threads",
+            help="Enable multithreading to drive engines with independent threads or speed up refitting (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--use-managed-memory",
+            help="Use managed memory instead of separate host and device allocations (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--dump-refit",
+            help="Print the refittable layers and weights from a refittable engine",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--dump-output",
+            help="Print the output tensor(s) of the last inference iteration (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--dump-profile",
+            help="Print profile information per layer (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--dump-layer-info",
+            help="Print layer information of the engine to console (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--separate-profile-run",
+            help="Do not attach the profiler in the benchmark run; if profiling is enabled, a second profile run will be executed (default = disabled)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--trtexec-no-builder-cache",
+            help="Disable timing cache in builder (default is to enable timing cache)",
+            action="store_true",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--trtexec-profiling-verbosity",
+            help="Specify profiling verbosity. mode ::= layer_names_only|detailed|none (default = layer_names_only)",
+            default=False,
+        )
+
+        self.group.add_argument(
+            "--layer-output-types",
+            help="""Control per-layer output type constraints. Effective only when precisionConstraints is set to
+                              "obey" or "prefer". (default = none)"
+                              The specs are read left-to-right, and later ones override earlier ones. "*" can be used as a
+                              layerName to specify the default precision for all the unspecified layers. If a layer has more than
+                              one output, then multiple types separated by "+" can be provided for this layer.
+                              Per-layer output type spec ::= layerOutputTypes[","spec]
+                                                    layerOutputTypes ::= layerName":"type
+                                                    type ::= "fp32"|"fp16"|"int32"|"int8"["+"type]",
+                """,
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+        """
+
+        # Required options
+        self.trtexec_path = args_util.get(args, "trtexec_path")
+        self.use_cuda_graph = args_util.get(args, "use_cuda_graph")
+        self.avg_runs = args_util.get(args, "avg_runs")
+        self.best = args_util.get(args, "best")
+        self.duration = args_util.get(args, "duration")
+        self.device = args_util.get(args, "device")
+        self.streams = args_util.get(args, "streams")
+        self.min_timing = args_util.get(args, "min_timing")
+        self.avg_timing = args_util.get(args, "avg_timing")
+        self.expose_dma = args_util.get(args, "expose_dma")
+        self.no_data_transfers = args_util.get(args, "no_data_transfers")
+        self.trtexec_warmup = args_util.get(args, "trtexec_warmup")
+        self.trtexec_iterations = args_util.get(args, "trtexec_iterations")
+        self.trtexec_export_times = args_util.get(args, "trtexec_export_times")
+        self.trtexec_export_output = args_util.get(args, "trtexec_export_output")
+        self.trtexec_export_profile = args_util.get(args, "trtexec_export_profile")
+        self.trtexec_export_layer_info = args_util.get(args, "trtexec_export_layer_info")
+
+        # Optional options
+        self.use_spin_wait = args_util.get(args, "use_spin_wait")
+        self.threads = args_util.get(args, "threads")
+        self.use_managed_memory = args_util.get(args, "use_managed_memory")
+        self.dump_refit = args_util.get(args, "dump_refit")
+        self.dump_output = args_util.get(args, "dump_output")
+        self.dump_profile = args_util.get(args, "dump_profile")
+        self.dump_layer_info = args_util.get(args, "dump_layer_info")
+        self.separate_profile_run = args_util.get(args, "separate_profile_run")
+        self.trtexec_no_builder_cache = args_util.get(args, "trtexec_no_builder_cache")
+        self.trtexec_profiling_verbosity = args_util.get(args, "trtexec_profiling_verbosity")
+        self.layer_output_types = args_util.get(args, "layer_output_types")
+
+    def add_to_script_impl(self, script):
+        model_path = self.arg_groups[ModelArgs].path
+        model_type = self.arg_groups[ModelArgs].model_type
+        input_shapes = self.arg_groups[ModelArgs].input_shapes or None
+
+        profile_dicts = self.arg_groups[TrtConfigArgs].profile_dicts
+        tf32 = self.arg_groups[TrtConfigArgs].tf32
+        fp16 = self.arg_groups[TrtConfigArgs].fp16
+        int8 = self.arg_groups[TrtConfigArgs].int8
+        allow_gpu_fallback = self.arg_groups[TrtConfigArgs].allow_gpu_fallback
+        precision_constraints = self.arg_groups[TrtConfigArgs].precision_constraints
+        mem_pool_size = self.arg_groups[TrtConfigArgs].memory_pool_limits
+        use_dla = self.arg_groups[TrtConfigArgs].use_dla
+        if mod.version(polygraphy.__version__) >= mod.version('0.39.0'):
+            refit = self.arg_groups[TrtConfigArgs].refittable
+        else:
+            refit = None
+
+        plugins = self.arg_groups[TrtLoadPluginsArgs].plugins
+        layer_precisions = self.arg_groups[TrtLoadNetworkArgs].layer_precisions
+        if layer_precisions:
+            layer_precisions = {layer:str(precision) for (layer, precision) in layer_precisions.items()}
+
+        save_engine = self.arg_groups[TrtSaveEngineBytesArgs].path
+
+        # Add an import for TensorRT
+        script.add_import(imports=["tensorrt"], imp_as="trt")
+
+        # Add an import for the Trtexec runner.
+        script.add_import(imports=["TrtexecRunner"], frm="polygraphy_trtexec.backend")
+        # Add the Trtexec runner using the `Script.add_runner()` API.
+        script.add_runner(make_invocable(
+            "TrtexecRunner",
+            model_path=model_path,
+            model_type=model_type,
+
+            trtexec_path = self.trtexec_path,
+            use_cuda_graph=self.use_cuda_graph,
+            avg_runs=self.avg_runs,
+            best=self.best,
+            duration=self.duration,
+            device=self.device,
+            streams=self.streams,
+            min_timing=self.min_timing,
+            avg_timing=self.avg_timing,
+            expose_dma=self.expose_dma,
+            no_data_transfers=self.no_data_transfers,
+            trtexec_warmup=self.trtexec_warmup,
+            trtexec_iterations=self.trtexec_iterations,
+            trtexec_export_times=self.trtexec_export_times,
+            trtexec_export_output=self.trtexec_export_output,
+            trtexec_export_profile=self.trtexec_export_profile,
+            trtexec_export_layer_info=self.trtexec_export_layer_info,
+
+            # Optional
+            use_spin_wait=self.use_spin_wait,
+            threads=self.threads,
+            use_managed_memory=self.use_managed_memory,
+            dump_refit=self.dump_refit,
+            dump_output=self.dump_output,
+            dump_profile=self.dump_profile,
+            dump_layer_info=self.dump_layer_info,
+            refit=refit,
+            separate_profile_run=self.separate_profile_run,
+            trtexec_no_builder_cache=self.trtexec_no_builder_cache,
+            trtexec_profiling_verbosity=self.trtexec_profiling_verbosity,
+            layer_output_types=self.layer_output_types,
+
+            input_shapes=input_shapes,
+            profile_dicts=profile_dicts,
+            tf32=tf32,
+            fp16=fp16,
+            int8=int8,
+            allow_gpu_fallback=allow_gpu_fallback,
+            precision_constraints=precision_constraints,
+            mem_pool_size=mem_pool_size,
+            use_dla=use_dla,
+            layer_precisions=layer_precisions,
+            plugins=plugins,
+            save_engine=save_engine,
+            ))

tools/polygraphy-extension-trtexec/polygraphy_trtexec/backend/__init__.py

@@ -0,0 +1,20 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Since we've used `mod.export()`, we can simply `import *`.
+# Only objects which have been decorated by `mod.export()` will be visible.
+from polygraphy_trtexec.backend.runner import *

tools/polygraphy-extension-trtexec/polygraphy_trtexec/backend/runner.py

@@ -0,0 +1,395 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This file defines the `TrtexecRunner` runner, which takes an ONNX model and
+runs inference on the trtexec backend.
+
+The runner implements the standard `BaseRunner` interface.
+"""
+
+import subprocess
+import json
+import tempfile
+from collections import OrderedDict
+import shutil
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.common import TensorMetadata
+from polygraphy.datatype import DataType
+from polygraphy.logger import G_LOGGER
+from polygraphy.backend.onnx import onnx_from_path
+from polygraphy.backend.onnx.util import get_input_metadata
+from polygraphy.backend.trt import engine_from_bytes
+from polygraphy.backend.common import bytes_from_path
+
+trt = mod.lazy_import("tensorrt>=8.5")
+np = mod.lazy_import("numpy")
+
+TRTEXEC_DEFAULT_PATH = "trtexec"
+MiB = 1024 ** 2
+
+def which(path):
+    """
+    Check whether `path` is an executable available on PATH
+    """
+    return shutil.which(path) is not None
+
+def convert_shape(input_name, trt_spec):
+    """
+    Converts a trt shape spec to a trtexec shape spec
+    """
+    trtexec_spec = '{}:{}'.format(input_name, 'x'.join(map(str, trt_spec)))
+    return trtexec_spec
+
+def parse_input_shapes(input_shapes):
+    """
+    Generate necessary specs to pass to --minShapes, --optShapes, --maxShapes
+    for trtexec backend
+    """
+    if not input_shapes:
+        return None
+
+    trtexec_input_shapes = []
+    for input_name, shape_spec in input_shapes.items():
+        trtexec_input_shapes.append(convert_shape(input_name, shape_spec.shape))
+    return ','.join(trtexec_input_shapes)
+
+def parse_profile_dicts(profile_dicts):
+    """
+    Generate necessary specs to pass to --minShapes, --optShapes, --maxShapes
+    for trtexec backend
+    """
+    if not profile_dicts:
+        return None, None, None
+
+    trtexec_min_shapes, trtexec_opt_shapes, trtexec_max_shapes = [], [], []
+
+    for input_name, (min_shapes, opt_shapes, max_shapes) in profile_dicts[0].items():
+        trtexec_min_shapes.append(convert_shape(input_name, min_shapes))
+        trtexec_opt_shapes.append(convert_shape(input_name, opt_shapes))
+        trtexec_max_shapes.append(convert_shape(input_name, max_shapes))
+
+    return ','.join(trtexec_min_shapes), ','.join(trtexec_opt_shapes), ','.join(trtexec_max_shapes)
+
+def parse_layer_precisions(layer_precisions):
+    """
+    Generate necessary specs to pass to --layerPrecisions for trtexec backend
+    """
+    if not layer_precisions:
+        return None
+
+    trtexec_layer_precisions = []
+
+    for layer, precision in layer_precisions.items():
+        trtexec_precision = ""
+        if precision == "trt.float32":
+            trtexec_precision = "fp32"
+        elif precision == "trt.float16":
+            trtexec_precision = "fp16"
+        elif precision == "trt.int32":
+            trtexec_precision = "int32"
+        elif precision == "trt.int8":
+            trtexec_precision = "int8"
+        else:
+            G_LOGGER.critical(f"Unsupported precision type: {precision}")
+
+        trtexec_layer_precisions.append(f"{layer}:{trtexec_precision}")
+
+    return ','.join(trtexec_layer_precisions)
+
+def get_inference_time(perf_output):
+    """
+    Reads the output from the performance summary generated by the trtexec
+    binary to extract the required performance statistics
+    """
+    inference_time_stats = {}
+    for line in perf_output.split('\n'):
+        index = line.find('Latency:')
+        if index >= 0:
+            stats = line[index + len('Latency:'):].split(',')
+            for stat in stats:
+                metric, value = stat.split('=')
+                value = value.strip().split(' ')[0]
+                inference_time_stats[metric.strip()] = float(value)
+            return inference_time_stats
+    G_LOGGER.critical(f"Could not read inference time for trtexec backend. This "
+    "might cause polygraphy to misbehave")
+
+@mod.export()
+class TrtexecRunner(BaseRunner):
+    """
+    Runs inference using custom trtexec. It accepts all ONNX models.
+    """
+
+    def __init__(self, model_path, model_type=None,
+    trtexec_path=None, use_cuda_graph=None, avg_runs=None, best=None, duration=None, device=None, streams=None, min_timing=None, avg_timing=None, expose_dma=None, no_data_transfers=None, trtexec_warmup=None, trtexec_iterations=None, trtexec_export_times=None, trtexec_export_output=None, trtexec_export_profile=None, trtexec_export_layer_info=None,
+    use_spin_wait=None, threads=None, use_managed_memory=None, dump_refit=None, dump_output=None, dump_profile=None, dump_layer_info=None, refit=None, separate_profile_run=None, trtexec_no_builder_cache=None, trtexec_profiling_verbosity=None, layer_output_types=None, use_dla_core=None,
+    input_shapes=None, profile_dicts=None, tf32=None, fp16=None, int8=None, allow_gpu_fallback=None, precision_constraints=None, mem_pool_size=None, use_dla=None, layer_precisions=None, plugins=None, save_engine=None):
+        super().__init__(prefix="trtexec-runner")
+        self.model_path = model_path
+        self.model_type = model_type
+        self.trtexec_path = util.default(trtexec_path, TRTEXEC_DEFAULT_PATH)
+        if not which(self.trtexec_path):
+            G_LOGGER.critical(f"trtexec not found in given path: {self.trtexec_path}")
+
+        self.use_cuda_graph = use_cuda_graph
+        self.avg_runs = avg_runs
+        self.best = best
+        self.duration = duration
+        self.device = device
+        self.streams = streams
+        self.min_timing = min_timing
+        self.avg_timing = avg_timing
+        self.expose_dma = expose_dma
+        self.no_data_transfers = no_data_transfers
+        self.trtexec_warmup = trtexec_warmup
+        self.trtexec_iterations = trtexec_iterations
+        self.trtexec_export_times = trtexec_export_times
+        self.trtexec_export_output = trtexec_export_output
+        self.trtexec_export_profile = trtexec_export_profile
+        self.trtexec_export_layer_info = trtexec_export_layer_info
+
+        self.use_spin_wait = use_spin_wait
+        self.threads = threads
+        self.use_managed_memory = use_managed_memory
+        self.dump_refit = dump_refit
+        self.dump_output = dump_output
+        self.dump_profile = dump_profile
+        self.dump_layer_info = dump_layer_info
+        self.refit = refit
+        self.separate_profile_run = separate_profile_run
+        self.trtexec_no_builder_cache = trtexec_no_builder_cache
+        self.trtexec_profiling_verbosity = trtexec_profiling_verbosity
+        self.layer_output_types = layer_output_types
+        self.use_dla_core = use_dla_core
+        self.input_shapes = parse_input_shapes(input_shapes)
+        self.min_shapes, self.opt_shapes, self.max_shapes = parse_profile_dicts(profile_dicts)
+        self.no_tf32 = not tf32
+        self.fp16 = fp16
+        self.int8 = int8
+        self.allow_gpu_fallback = allow_gpu_fallback
+        self.precision_constraints = precision_constraints
+        self.use_dla = 0 if use_dla else None
+        self.plugins = plugins
+        self.layer_precisions = parse_layer_precisions(layer_precisions)
+        self.save_engine = save_engine
+        if mem_pool_size is None:
+            self.mem_pool_size = None
+        else:
+            self.mem_pool_size = ""
+            for k, v in mem_pool_size.items():
+                v = v / MiB # Convert bytes into MiB
+                if int(k) == 0:
+                    self.mem_pool_size += f"workspace:{v},"
+                elif int(k) == 1:
+                    self.mem_pool_size += f"dlaSRAM:{v},"
+                elif int(k) == 2:
+                    self.mem_pool_size += f"dlaLocalDRAM:{v},"
+                elif int(k) == 3:
+                    self.mem_pool_size += f"dlaGlobalDRAM:{v},"
+                else:
+                    pass
+            self.mem_pool_size = self.mem_pool_size.rstrip(',')
+
+
+    def activate_impl(self):
+        """
+        Initializes the construction of the command that needs to be run using
+        the `trtexec` backend.
+        """
+
+        self.cmd_args = [self.trtexec_path]
+        self.input_files = []
+
+        if self.trtexec_export_output:
+            self.export_output_file_handle = open(self.trtexec_export_output, 'w+')
+        else:
+            self.export_output_file_handle = tempfile.NamedTemporaryFile(delete=False)
+        self.export_output_file_name = self.trtexec_export_output or self.export_output_file_handle.name
+
+        model_type_mapping = self.get_model_type_mapping()
+
+        # Mapping the args of polygraphy run to that of trtexec
+        init_args_mapping = {
+            **model_type_mapping,
+            'exportOutput': self.export_output_file_name,
+
+            'useCudaGraph': self.use_cuda_graph,
+            'avgRuns': self.avg_runs,
+            'best': self.best,
+            'duration': self.duration,
+            'device': self.device,
+            'streams': self.streams,
+            'minTiming': self.min_timing,
+            'avgTiming': self.avg_timing,
+            'exposeDMA': self.expose_dma,
+            'noDataTransfers': self.no_data_transfers,
+            'warmUp': self.trtexec_warmup,
+            'iterations': self.trtexec_iterations,
+            'exportTimes': self.trtexec_export_times,
+            'exportProfile': self.trtexec_export_profile,
+            'exportLayerInfo': self.trtexec_export_layer_info,
+
+            'useSpinWait': self.use_spin_wait,
+            'threads': self.threads,
+            'useManagedMemory': self.use_managed_memory,
+            'dumpRefit': self.dump_refit,
+            'dumpOutput': self.dump_output,
+            'dumpProfile': self.dump_profile,
+            'dumpLayerInfo': self.dump_layer_info,
+            'refit': self.refit,
+            'separateProfileRun': self.separate_profile_run,
+            'noBuilderCache': self.trtexec_no_builder_cache,
+            'profilingVerbosity': self.trtexec_profiling_verbosity,
+            'layerPrecisions': self.layer_precisions,
+            'layerOutputTypes': self.layer_output_types,
+            'useDLACore': self.use_dla,
+
+            'shapes': self.input_shapes,
+            'minShapes': self.min_shapes,
+            'optShapes': self.opt_shapes,
+            'maxShapes': self.max_shapes,
+            'noTF32': self.no_tf32,
+            'fp16': self.fp16,
+            'int8': self.int8,
+            'allowGPUFallback': self.allow_gpu_fallback,
+            'precisionConstraints': self.precision_constraints,
+            'memPoolSize': self.mem_pool_size,
+            'plugins': self.plugins,
+            'saveEngine': self.save_engine,
+
+            'verbose': G_LOGGER.severity <= G_LOGGER.EXTRA_VERBOSE,
+        }
+
+        for arg, value in init_args_mapping.items():
+            self.add_cmd_args(arg, value)
+
+    def add_cmd_args(self, name, value=None):
+        """
+        Add the args to `self.cmd_args`. The function handles both
+        args and kwargs
+        """
+        if value is None:
+            return
+
+        if isinstance(value, bool):
+            # For a bool, add the arg only if the corresponding value is `True`
+            if value:
+                self.cmd_args.append('--{}'.format(name))
+        else:
+            self.cmd_args.append('--{}={}'.format(name, value))
+
+    def get_model_type_mapping(self):
+        """
+        Add the required args based on the model type
+        """
+        if self.model_type == 'onnx':
+            return {
+                'onnx': self.model_path
+            }
+
+        if self.model_type == 'engine':
+            return {
+                'loadEngine':self.model_path
+            }
+
+        G_LOGGER.critical(f"Unsupported model type: {self.model_type}. `trtexec` only supports TensorRT engines and ONNX models")
+
+    def generate_load_inputs_spec(self, feed_dict):
+        """
+        Reads the feed_dict metadata input dictionary and generates files to
+        pass as command line input to trtexec binary
+        """
+        load_inputs_spec = []
+        for input, values in feed_dict.items():
+            input_file = tempfile.NamedTemporaryFile(delete=False)
+            values.tofile(input_file.name)
+            load_inputs_spec.append('{}:{}'.format(input, input_file.name))
+            self.input_files.append(input_file)
+        self.load_inputs_spec = ','.join(load_inputs_spec)
+
+    def read_output_file(self):
+        """
+        Reads the output from the output file generated by the trtexec binary
+        """
+        outputs = OrderedDict()
+        content = json.load(self.export_output_file_handle)
+        for entry in content:
+            name, dimensions, values = entry['name'], entry['dimensions'], entry['values']
+            dimensions = [int(d) for d in dimensions.split('x')]
+            outputs[name] = np.array(values).reshape(*dimensions)
+        return outputs
+
+    def get_input_metadata_impl(self):
+        # Input metadata is used by Polygraphy's default data loader to
+        # determine the required shapes and datatypes of the input buffers.
+        if self.model_type == 'onnx':
+            model = onnx_from_path(self.model_path)
+            return get_input_metadata(model.graph)
+
+        if self.model_type =='engine':
+            engine = engine_from_bytes(bytes_from_path(self.model_path))
+            meta = TensorMetadata()
+            for idx in range(engine.num_io_tensors):
+                name = engine.get_tensor_name(idx)
+                if engine.get_tensor_mode(name) != trt.TensorIOMode.INPUT:
+                    continue
+                meta.add(name=name, dtype=DataType.from_dtype(engine.get_tensor_dtype(name), "tensorrt"), shape=engine.get_tensor_shape(name))
+            return meta
+
+
+    def infer_impl(self, feed_dict):
+        outputs = OrderedDict()
+
+        # Adds other args that need to generated during inference. For example,
+        # `feed_dict` is used to generate the args for `loadInputs`
+        self.construct_final_cmd(feed_dict)
+        G_LOGGER.info(f"The trtexec command being run: {self.cmd_args}")
+        perf_output = subprocess.run(self.cmd_args, stdout=subprocess.PIPE, text=True).\
+                            stdout
+        self.inference_time_stats = get_inference_time(perf_output)
+        G_LOGGER.verbose(f"Inference time statistics: {self.inference_time_stats}")
+
+        outputs = self.read_output_file()
+        # inference_time_stats records time in 'ms'. However, polygraphy
+        # expects time in seconds.
+        self.inference_time = self.inference_time_stats['median'] / 1000
+        return outputs
+
+    def last_inference_time_stats(self):
+        """
+        Provides the inference time statistics
+        """
+        return self.inference_time_stats
+
+    def construct_final_cmd(self, feed_dict):
+        """
+        Constructs the complete command to run inference on trtexec backend.
+        Adds any other args that need to generated during inference.
+        """
+
+        self.generate_load_inputs_spec(feed_dict)
+        self.add_cmd_args('loadInputs', self.load_inputs_spec)
+
+    def deactivate_impl(self):
+        # Close the temporary files that are created. Python automatically
+        # deletes temporary files after they are closed
+        self.export_output_file_handle.close()
+        for input_file in self.input_files:
+            input_file.close()

tools/polygraphy-extension-trtexec/polygraphy_trtexec/export.py

@@ -0,0 +1,28 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This file defines the entry point that will be exported by our extension module.
+`polygraphy run` will use this to add our custom argument groups.
+"""
+
+from polygraphy_trtexec.args import TrtexecRunnerArgs
+
+def export_argument_groups():
+    return [
+        TrtexecRunnerArgs(),
+    ]

tools/polygraphy-extension-trtexec/sample.onnx

@@ -0,0 +1,15 @@
+backend-test:[
+
+
x
y"Identity
+test_identityZ
+
x
+

+

+

+
+b
+
y
+

+

+

+

tools/polygraphy-extension-trtexec/setup.py

@@ -0,0 +1,56 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from setuptools import find_packages, setup
+import polygraphy_trtexec
+import os
+
+def main():
+    # We change to the project root directory so that `setup.py` is usable from any directory.
+    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+    os.chdir(ROOT_DIR)
+
+    setup(
+        name="polygraphy_trtexec",
+        version=polygraphy_trtexec.__version__,
+        description="Polygraphy Trtexec: Extension to run on trtexec backend",
+        long_description=open("README.md", "r", encoding="utf-8").read(),
+        long_description_content_type='text/markdown',
+        url="https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy",
+        author="NVIDIA",
+        author_email="svc_tensorrt@nvidia.com",
+        classifiers=[
+            "Intended Audience :: Developers",
+            "Programming Language :: Python :: 3",
+        ],
+        license="Apache 2.0",
+        install_requires=[
+            "polygraphy>=0.49.0",
+        ],
+        packages=find_packages(exclude=("tests", "tests.*")),
+        entry_points={
+            "polygraphy.run.plugins": [
+                "trtexec-runner=polygraphy_trtexec.export:export_argument_groups",
+            ]
+        },
+        zip_safe=True,
+        python_requires=">=3.6",
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/polygraphy-extension-trtexec/tests/__init__.py

@@ -0,0 +1,16 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

tools/polygraphy-extension-trtexec/tests/helper.py

@@ -0,0 +1,85 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import subprocess as sp
+from polygraphy import util
+import argparse
+import time
+
+TEST_DIR = os.path.abspath(os.path.dirname(os.path.abspath(__file__)))
+
+# Paramaters to run test on when testing trtexec args mapping
+TRTEXEC_PATH_PARAMS = ["trtexec"]
+TRTEXEC_EXPORT_TIMES_PARAMS = ["times.json"]
+TRTEXEC_EXPORT_OUTPUT_PARAMS = ["output.json"]
+TRTEXEC_EXPORT_PROFILE_PARAMS = ["profile.json"]
+TRTEXEC_EXPORT_LAYER_INFO_PARAMS = ["layer_info.json"]
+
+def poly_run(args):
+    """
+    Helper function to run the `polygraphy run` command with nececssary args
+    """
+    cmd_args = ["polygraphy", "run"] + args
+    output = sp.run(cmd_args, stdout=sp.PIPE).stdout.decode('utf-8')
+    return 'PASSED' in output
+
+def is_file_non_empty(path):
+    return os.stat(path).st_size != 0
+
+class ArgGroupTestHelper:
+    def __init__(self, arg_group, deps=None):
+        self.deps = util.default(deps, [])
+
+        self.arg_group = arg_group
+        self.parser = argparse.ArgumentParser()
+
+        arg_groups = {type(self.arg_group): self.arg_group}
+        arg_groups.update({type(dep): dep for dep in self.deps})
+
+        for dep in self.deps:
+            dep.register(arg_groups)
+        self.arg_group.register(arg_groups)
+
+        for dep in self.deps:
+            dep.add_parser_args(self.parser)
+        self.arg_group.add_parser_args(self.parser)
+
+    def parse_args(self, cli_args):
+        args = self.parser.parse_args(cli_args)
+        for dep in self.deps:
+            dep.parse(args)
+        self.arg_group.parse(args)
+        return args
+
+    def __getattr__(self, name):
+        if name in ["arg_group", "parser"]:
+            return super().__getattr__(name)
+        return getattr(self.arg_group, name)
+
+def time_func(func, warm_up=10, iters=50):
+    for _ in range(warm_up):
+        func()
+
+    total = 0
+    for _ in range(iters):
+        start = time.time()
+        func()
+        end = time.time()
+        total += end - start
+    return total / float(iters)
+

tools/polygraphy-extension-trtexec/tests/models/dynamic_identity.onnx

@@ -0,0 +1,15 @@
+backend_test:y
+
+
X
Y"Identityonnx_dynamic_identityZ&
+
X!
+

+

+
+height
+widthb&
+
Y!
+

+

+
+height
+widthB

tools/polygraphy-extension-trtexec/tests/models/identity.onnx

@@ -0,0 +1,15 @@
+backend-test:[
+
+
x
y"Identity
+test_identityZ
+
x
+

+

+

+
+b
+
y
+

+

+

+

tools/polygraphy-extension-trtexec/tests/models/meta.py

@@ -0,0 +1,26 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+
+META_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Path to models used for testing
+ONNX_MODELS_PATH = {
+    'identity': os.path.join(META_DIR, "identity.onnx"),
+    'dynamic_identity': os.path.join(META_DIR, "dynamic_identity.onnx"),
+}

tools/polygraphy-extension-trtexec/tests/requirements.txt

@@ -0,0 +1,7 @@
+colored
+polygraphy==0.49.0
+onnx==1.14.0
+onnxruntime==1.15.0
+protobuf==3.20.2
+pytest
+wheel

tools/polygraphy-extension-trtexec/tests/test_runner_args.py

@@ -0,0 +1,195 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+import pytest
+from polygraphy.tools.args import (
+    ModelArgs,
+    TrtConfigArgs,
+    TrtLoadPluginsArgs,
+    TrtLoadNetworkArgs,
+    TrtSaveEngineBytesArgs,
+)
+from polygraphy.tools.script import Script
+from tests.models.meta import ONNX_MODELS_PATH
+from tests.helper import *
+from polygraphy.tools.args import util as args_util
+from polygraphy_trtexec.args import TrtexecRunnerArgs
+
+@pytest.fixture()
+def trtexec_runner_args():
+    return ArgGroupTestHelper(
+        TrtexecRunnerArgs(),
+        deps=[
+            ModelArgs(),
+            TrtConfigArgs(),
+            TrtLoadPluginsArgs(),
+            TrtLoadNetworkArgs(),
+            TrtSaveEngineBytesArgs(),
+        ],
+    )
+
+
+class TestTrtexecRunnerArgs:
+    @pytest.mark.parametrize("trtexec_path_params", TRTEXEC_PATH_PARAMS)
+    def test_trtexec_path(self, trtexec_runner_args, trtexec_path_params):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-path={trtexec_path_params}"])
+        assert trtexec_runner_args.trtexec_path == trtexec_path_params
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_path == trtexec_path_params
+
+    def test_use_cuda_graph(self, trtexec_runner_args):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--use-cuda-graph"])
+        assert trtexec_runner_args.use_cuda_graph
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].use_cuda_graph
+
+    @pytest.mark.parametrize("num_avg_runs", range(1, 16, 5))
+    def test_avg_runs(self, trtexec_runner_args, num_avg_runs):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--avg-runs={num_avg_runs}"])
+        assert trtexec_runner_args.avg_runs == num_avg_runs
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].avg_runs == num_avg_runs
+
+    def test_best(self, trtexec_runner_args):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--best"])
+        assert trtexec_runner_args.best
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].best
+
+    @pytest.mark.parametrize("num_duration", range(1, 16, 5))
+    def test_duration(self, trtexec_runner_args, num_duration):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--duration={num_duration}"])
+        assert trtexec_runner_args.duration == num_duration
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].duration == num_duration
+
+    @pytest.mark.parametrize("num_device", range(0, 3))
+    def test_device(self, trtexec_runner_args, num_device):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--device={num_device}"])
+        assert trtexec_runner_args.device == num_device
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].device == num_device
+
+    @pytest.mark.parametrize("num_streams", range(1, 4))
+    def test_streams(self, trtexec_runner_args, num_streams):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--streams={num_streams}"])
+        assert trtexec_runner_args.streams == num_streams
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].streams == num_streams
+
+    @pytest.mark.parametrize("num_min_timing", range(1, 4))
+    def test_min_timing(self, trtexec_runner_args, num_min_timing):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--min-timing={num_min_timing}"])
+        assert trtexec_runner_args.min_timing == num_min_timing
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].min_timing == num_min_timing
+
+    @pytest.mark.parametrize("num_avg_timing", range(1, 4))
+    def test_avg_timing(self, trtexec_runner_args, num_avg_timing):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--avg-timing={num_avg_timing}"])
+        assert trtexec_runner_args.avg_timing == num_avg_timing
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].avg_timing == num_avg_timing
+
+    def test_expose_dma(self, trtexec_runner_args):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--expose-dma"])
+        assert trtexec_runner_args.expose_dma
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].expose_dma
+
+    def test_no_data_transfers(self, trtexec_runner_args):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--no-data-transfers"])
+        assert trtexec_runner_args.no_data_transfers
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].no_data_transfers
+
+    @pytest.mark.parametrize("num_warmup", range(100, 400, 100))
+    def test_trtexec_warmup(self, trtexec_runner_args, num_warmup):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-warmup={num_warmup}"])
+        assert trtexec_runner_args.trtexec_warmup == num_warmup
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_warmup == num_warmup
+
+    @pytest.mark.parametrize("num_iterations", range(100, 400, 100))
+    def test_trtexec_iterations(self, trtexec_runner_args, num_iterations):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-iterations={num_iterations}"])
+        assert trtexec_runner_args.trtexec_iterations == num_iterations
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_iterations == num_iterations
+
+    @pytest.mark.parametrize("trtexec_export_times_params", TRTEXEC_EXPORT_TIMES_PARAMS)
+    def test_trtexec_export_times(self, trtexec_runner_args, trtexec_export_times_params):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-export-times={trtexec_export_times_params}"])
+        assert trtexec_runner_args.trtexec_export_times == trtexec_export_times_params
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_export_times == trtexec_export_times_params
+
+    @pytest.mark.parametrize("trtexec_export_output_params", TRTEXEC_EXPORT_OUTPUT_PARAMS)
+    def test_trtexec_export_output_params(self, trtexec_runner_args, trtexec_export_output_params):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-export-output={trtexec_export_output_params}"])
+        assert trtexec_runner_args.trtexec_export_output == trtexec_export_output_params
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_export_output == trtexec_export_output_params
+
+    @pytest.mark.parametrize("trtexec_export_profile_params", TRTEXEC_EXPORT_PROFILE_PARAMS)
+    def test_trtexec_export_profile_params(self, trtexec_runner_args, trtexec_export_profile_params):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-export-profile={trtexec_export_profile_params}"])
+        assert trtexec_runner_args.trtexec_export_profile == trtexec_export_profile_params
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_export_profile == trtexec_export_profile_params
+
+    @pytest.mark.parametrize("trtexec_export_layer_info_params", TRTEXEC_EXPORT_LAYER_INFO_PARAMS)
+    def test_trtexec_export_layer_info_params(self, trtexec_runner_args, trtexec_export_layer_info_params):
+        trtexec_runner_args.parse_args([ONNX_MODELS_PATH["identity"], f"--trtexec-export-layer-info={trtexec_export_layer_info_params}"])
+        assert trtexec_runner_args.trtexec_export_layer_info == trtexec_export_layer_info_params
+
+        script = Script()
+        runners = args_util.run_script(trtexec_runner_args.add_to_script)
+        assert runners[0].trtexec_export_layer_info == trtexec_export_layer_info_params

tools/polygraphy-extension-trtexec/tests/test_runner_backend.py

@@ -0,0 +1,43 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from polygraphy_trtexec.backend.runner import TrtexecRunner, TRTEXEC_DEFAULT_PATH
+from tests.models.meta import ONNX_MODELS_PATH
+
+class TestTrtRunner:
+
+    def test_basic(self):
+        with TrtexecRunner(ONNX_MODELS_PATH["identity"], model_type="onnx") as runner:
+            assert runner.trtexec_path is TRTEXEC_DEFAULT_PATH
+            assert runner.use_cuda_graph is None
+            assert runner.avg_runs is None
+            assert runner.best is None
+            assert runner.duration is None
+            assert runner.device is None
+            assert runner.streams is None
+            assert runner.min_timing is None
+            assert runner.avg_timing is None
+            assert runner.expose_dma is None
+            assert runner.no_data_transfers is None
+            assert runner.trtexec_warmup is None
+            assert runner.trtexec_iterations is None
+            assert runner.trtexec_export_times is None
+            assert runner.trtexec_export_output is None
+            assert runner.trtexec_export_profile is None
+            assert runner.trtexec_export_layer_info is None
+            assert runner.is_active
+        assert not runner.is_active

tools/polygraphy-extension-trtexec/tests/test_trtexec.py

@@ -0,0 +1,85 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import os
+import onnx
+import pytest
+from polygraphy import util, mod
+from tests.helper import TEST_DIR, poly_run, is_file_non_empty
+from tests.models.meta import ONNX_MODELS_PATH
+
+class TestModelEngine:
+    def test_load_onnx_model(self):
+        # Load an onnx model and run inference
+        assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec"])
+
+    def test_save_load_engine(self):
+        # Save an engine and check if file exists. Run inference with saved engine
+        with util.NamedTemporaryFile() as outpath:
+            assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--save-engine", outpath.name])
+            assert is_file_non_empty(outpath.name)
+            assert poly_run(["--trtexec", outpath.name, "--model-type=engine"])
+
+class TestProfileShapes:
+    def test_input_shape(self):
+        assert poly_run([ONNX_MODELS_PATH['dynamic_identity'],
+            "--trtexec",
+            "--input-shapes", "X:[1,2,3,3]",
+        ])
+
+    def test_explicit_profile(self):
+        assert poly_run([ ONNX_MODELS_PATH['dynamic_identity'],
+            "--trtexec",
+            "--input-shapes", "X:[1,2,1,1]",
+            "--trt-min-shapes", "X:[1,2,1,1]",
+            "--trt-opt-shapes", "X:[1,2,1,1]",
+            "--trt-max-shapes", "X:[1,2,1,1]",
+        ])
+
+    def test_explicit_profile_opt_runtime_shapes_differ(self):
+        assert poly_run([ONNX_MODELS_PATH['dynamic_identity'],
+            "--trtexec",
+            "--input-shapes", "X:[1,2,2,2]",
+            "--trt-min-shapes", "X:[1,2,1,1]",
+            "--trt-opt-shapes", "X:[1,2,3,3]",
+            "--trt-max-shapes", "X:[1,2,4,4]",
+        ])
+
+class TestParameterMapping:
+    def test_args(self):
+        assert poly_run([ONNX_MODELS_PATH['identity'],
+            "--trtexec",
+            "--fp16", "--int8", "--best", "--tf32",
+            "--use-cuda-graph", "--no-data-transfers",
+        ])
+
+    def test_kwargs(self):
+        assert poly_run([ONNX_MODELS_PATH['identity'],
+            "--trtexec",
+            "--memory-pool-limit=workspace:1M", "--streams=2", "--device=0",
+            "--avg-runs=1", "--min-timing=1", "--avg-timing=3",
+            "--trtexec-iterations=1", "--trtexec-warmup=1", "--duration=1",
+        ])
+
+
+class TestAccuracyComparision:
+    def test_trtexec_trt(self):
+        # Compare inference results between trtexec and trt
+        assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--trt"])
+
+    def test_trtexec_onnxrt(self):
+        # Compare inference results between trtexec and onnxrt
+        assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--onnxrt"])