chore: import upstream snapshot with attribution

tools/Polygraphy/examples/dev/01_writing_cli_tools/README.md

@@ -0,0 +1,35 @@
+# Writing Custom Command-Line Tools
+
+
+## Introduction
+
+Polygraphy includes various helper utilities to make it easier to write
+new command-line tools from scratch.
+
+In this example, we'll write a brand new tool called `gen-data` that will generate random data
+using Polygraphy's default data loader, and then write it to an output file. The user will
+be able to specify the number of values to generate as well as the output path.
+
+To do this, we'll create a child class of `Tool` and use the `DataLoaderArgs` argument
+group provided by Polygraphy.
+
+
+## Running The Example
+
+1. You can run the example tool from this directory. For example:
+
+    ```bash
+    ./gen-data -o data.json --num-values 25
+    ```
+
+2. We can even inspect the generated data with `inspect data`:
+
+    ```bash
+    polygraphy inspect data data.json -s
+    ```
+
+To see the other command-line options available in the example tool,
+run:
+```bash
+./gen-data -h
+```

tools/Polygraphy/examples/dev/01_writing_cli_tools/gen-data

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+Generates data and writes it to a file.
+"""
+
+from polygraphy import mod
+from polygraphy.common import TensorMetadata
+from polygraphy.json import save_json
+from polygraphy.tools import Tool
+from polygraphy.tools.args import DataLoaderArgs
+
+# Your tool should lazily import any external dependencies. By doing so,
+# we avoid creating hard dependencies on other packages.
+# Additionally, this allows Polygraphy to automatically install required packages
+# as they are needed, instead of requiring the user to do so up front.
+np = mod.lazy_import("numpy")
+
+
+class GenData(Tool):
+    # Polygraphy will use the docstring of the tool child class to generate
+    # the summary for the command-line help output.
+    """
+    Generate random data and write it to a file.
+    """
+
+    # First, we'll implement `get_subscriptions_impl()` to subscribe to argument groups
+    # that we're intrested in.
+    # All the argument groups we subscribe to will be stored in a member called
+    # arg_groups, which maps types to instances.
+    def get_subscriptions_impl(self):
+        return [DataLoaderArgs()]
+
+    # Next, we'll add custom arguments, beyond those provided by our subscribed
+    # argument groups, by defining `add_parser_args_impl`.
+    def add_parser_args_impl(self, parser):
+        parser.add_argument("-o", "--output", help="Path at which to write generated data.", required=True)
+        parser.add_argument("--num-values", help="The number of random values to generate.", default=1, type=int)
+
+    # Lastly, we implement `run`, which will implement the functionality of our tool.
+    def run(self, args):
+        # The DataLoaderArgs argument group provides a helper called `get_data_loader`, which
+        # will create a new data loader based on the command-line arguments provided by the user.
+        # See `polygraphy/tools/args/data_loader.py` for implementation details.
+        #
+        # To get data of the shape we want, we'll set the `input_metadata` parameter based on --num-values.
+        meta = TensorMetadata().add(name="data", dtype=np.float32, shape=(args.num_values,))
+        data_loader = self.arg_groups[DataLoaderArgs].get_data_loader(meta)
+
+        # data_loader behaves like a generator/iterable, so we can cast it to a `list` to
+        # generate all the data at once.
+        save_json(list(data_loader), dest=args.output, description="randomly generated numbers")
+
+
+# NOTE: To integrate a tool into Polygraphy, you will need to add it to the registry in
+# `polygraphy/tools/registry.py`.
+#
+# Alternatively, we can create a standalone tool by invoking the `main()` method, which will allow
+# our script to be used on the command-line.
+GenData().main()

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/README.md

@@ -0,0 +1,139 @@
+# Extending `polygraphy run`
+
+## Introduction
+
+`polygraphy run` allows you to run inference with multiple backends, including TensorRT and ONNX-Runtime, and compare outputs.
+While it does provide mechanisms to load and compare against custom outputs from unsupported backends,
+adding support for the backend via an extension module allows it to be integrated more seamlessly,
+providing a better user experience.
+
+In this example, we'll create an extension module for `polygraphy run` called `polygraphy_reshape_destroyer`,
+which will include the following:
+
+- A special loader that will replace no-op `Reshape` nodes in an ONNX model with `Identity` nodes.
+
+- A custom runner that supports ONNX models containing only `Identity` nodes.
+
+- Command-line options to:
+    - Enable or disable renaming nodes when a transformation is applied by the loader.
+    - Run the model in `slow`, `medium`, or `fast` mode.
+        In `slow` and `medium` modes, we'll inject a `time.sleep()` during inference
+        (this will result in massive performance gains in `fast` mode!).
+
+## Background
+
+Although this example is self-contained and concepts will be explained as you encounter them, it is still
+recommended that you first familiarize yourself with
+[Polygraphy's `Loader` and `Runner` APIs](../../../polygraphy/README.md),
+the [`Argument Group` Interface](../../../polygraphy/tools/args/README.md),
+as well as the [`Script` interface](../../../polygraphy/tools/script.py).
+
+After that, creating an extension module for `polygraphy run` is a simple matter of defining your
+custom `Loader`s/`Runner`s and `Argument Group`s and making them visible to Polygraphy via
+`setuptools`'s `entry_points` API.
+
+*NOTE: Defining a custom `Loader` is not strictly required, but will be covered in this example for the sake of completeness.*
+
+As a matter of convention, Polygraphy extension module names are prefixed with `polygraphy_`.
+
+## Reading The Example Code
+
+We've structured our example extension module such that it somewhat mirrors the structure of the Polygraphy repository.
+This should make it easier to see the parallels between functionality in the extension module and that provided by Polygraphy natively.
+The structure is:
+<!-- Polygraphy Test: Ignore Start -->
+```bash
+- extension_module/
+    - polygraphy_reshape_destroyer/
+        - backend/
+            - __init__.py   # Controls submodule-level exports
+            - loader.py     # Defines our custom loader.
+            - runner.py     # Defines our custom runner.
+        - args/
+            - __init__.py   # Controls submodule-level exports
+            - loader.py     # Defines command-line argument group for our custom loader.
+            - runner.py     # Defines command-line argument group for our custom runner.
+        - __init__.py       # Controls module-level exports
+        - export.py         # Defines the entry-point for `polygraphy run`.
+    - setup.py              # Builds our module
+```
+<!-- Polygraphy Test: Ignore End -->
+
+It is recommended that you read these files in the following order:
+
+1. [backend/loader.py](./extension_module/polygraphy_reshape_destroyer/backend/loader.py)
+2. [backend/runner.py](./extension_module/polygraphy_reshape_destroyer/backend/runner.py)
+3. [backend/\_\_init\_\_.py](./extension_module/polygraphy_reshape_destroyer/backend/__init__.py)
+4. [args/loader.py](./extension_module/polygraphy_reshape_destroyer/args/loader.py)
+5. [args/runner.py](./extension_module/polygraphy_reshape_destroyer/args/runner.py)
+6. [args/\_\_init\_\_.py](./extension_module/polygraphy_reshape_destroyer/args/__init__.py)
+7. [\_\_init\_\_.py](./extension_module/polygraphy_reshape_destroyer/__init__.py)
+8. [export.py](./extension_module/polygraphy_reshape_destroyer/export.py)
+9. [setup.py](./extension_module/setup.py)
+
+
+## Running The Example
+
+1. Build and install the extension module:
+
+    Build using `setup.py`:
+
+    ```bash
+    python3 extension_module/setup.py bdist_wheel
+    ```
+
+    Install the wheel:
+
+    ```bash
+    python3 -m pip install extension_module/dist/polygraphy_reshape_destroyer-0.0.1-py3-none-any.whl \
+        --extra-index-url https://pypi.ngc.nvidia.com
+    ```
+
+    *TIP: If you make changes to the example extension module, you can update your installed version by*
+    *rebuilding (by following step 1) and then running:*
+
+    ```bash
+    python3 -m pip install extension_module/dist/polygraphy_reshape_destroyer-0.0.1-py3-none-any.whl \
+        --force-reinstall --no-deps
+    ```
+
+2. Once the extension module is installed, you should see the options you added appear in the help output
+    of `polygraphy run`:
+
+    ```bash
+    polygraphy run -h
+    ```
+
+3. Next, we can try out our custom runner with an ONNX model containing a no-op Reshape:
+
+    ```bash
+    polygraphy run no_op_reshape.onnx --res-des
+    ```
+
+4. We can also try some of the other command-line options we added:
+
+    - Renaming replaced nodes:
+
+        ```bash
+        polygraphy run no_op_reshape.onnx --res-des --res-des-rename-nodes
+        ```
+
+    - Different inference speeds:
+
+        ```bash
+        polygraphy run no_op_reshape.onnx --res-des --res-des-speed=slow
+        ```
+
+        ```bash
+        polygraphy run no_op_reshape.onnx --res-des --res-des-speed=medium
+        ```
+
+        ```bash
+        polygraphy run no_op_reshape.onnx --res-des --res-des-speed=fast
+        ```
+
+5. Lastly, let's compare our implementation against ONNX-Runtime to make sure it is functionally correct:
+
+    ```bash
+    polygraphy run no_op_reshape.onnx --res-des --onnxrt
+    ```

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/README.md

@@ -0,0 +1,4 @@
+# Polygraphy Reshape Destroyer
+
+An example `polygraphy run` extension module that can replace no-op `Reshape`s in an ONNX
+model with `Identity` nodes and run inference.

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/__init__.py

@@ -0,0 +1,18 @@
+#
+# 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.
+#
+
+__version__ = "0.0.1"

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/args/__init__.py

@@ -0,0 +1,19 @@
+#
+# 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_reshape_destroyer.args.loader import *
+from polygraphy_reshape_destroyer.args.runner import *

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/args/loader.py

@@ -0,0 +1,115 @@
+#
+# 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.
+#
+
+"""
+This file defines the `ReplaceReshapeArgs` argument group, which manages
+command-line options that control the `ReplaceReshape` loader.
+
+The argument group implements the standard `BaseArgs` interface.
+"""
+
+
+from polygraphy import mod
+from polygraphy.tools.args import OnnxLoadArgs
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseArgs
+from polygraphy.tools.script import make_invocable
+
+
+# NOTE: Our argument groups can depend on any argument groups that `polygraphy run` subscribes to.
+#       See `polygraphy/tools/run/run.py` for the complete list.
+#       In this case, we'll take advantage of OnnxLoadArgs to load the ONNX model for us.
+@mod.export()
+class ReplaceReshapeArgs(BaseArgs):
+    # Argument groups employ a standardized format for their docstrings:
+    #
+    #  - The first line must include a title and description separated by a colon (':').
+    #   The description should answer the question: "What is this argument group responsible for?".
+    #
+    # - If our argument group depends on other argument groups, we must also add a `Depends on:` section
+    #   listing our dependencies.
+    #
+    # See the `BaseArgs` docstring for more details on the expected format.
+    #
+    """
+    ONNX Reshape Replacement: replacing no-op Reshape nodes with Identity in ONNX models
+
+    Depends on:
+
+        - OnnxLoadArgs
+    """
+
+    # Add any command-line options we want for our loader.
+    def add_parser_args_impl(self):
+        # The `BaseArgs` constructor will automatically set `self.group` to an `argparse`
+        # argument group to which we can add our command-line options.
+        #
+        # NOTE: In order to prevent collisions with other Polygraphy options, we'll prefix all the options
+        #       we add with `--res-des`, short for `REShape DEStroyer`.
+        self.group.add_argument(
+            "--res-des-rename-nodes",
+            help="Whether to rename nodes if they are replaced",
+            action="store_true",
+            default=None,
+        )
+
+    # Next, we'll implement parsing code for the arguments we added.
+    # This will allow our argument group to be used by other argument groups.
+    def parse_impl(self, args):
+        # The docstring for `parse_impl` must document which attributes it populates.
+        # These attributes are considered part of the public interface of the argument group
+        # and may be used by other argument groups and/or command-line tools.
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            rename_nodes (bool): Whether to rename nodes if they are replaced.
+        """
+        # We'll use `args_util.get` to retrieve attributes from `args`, which will return `None` if the attribute is not found.
+        # This will ensure that our argument group will continue to work even if a command-line option is disabled in the code.
+        self.rename_nodes = args_util.get(args, "res_des_rename_nodes")
+
+    # Finally, we can implement the logic which will add code to the script.
+    def add_to_script_impl(self, script):
+        # First, we'll use `OnnxLoadArgs` to add code to load the ONNX model.
+        # This will ensure that any options related to ONNX model loading are respected by our argument group.
+        # `OnnxLoadArgs`'s `add_to_script` method will return the name of a loader that loads an ONNX model.
+        loader_name = self.arg_groups[OnnxLoadArgs].add_to_script(script)
+
+        # Next, we'll add Polygraphy's `GsFromOnnx` loader so that we can convert the ONNX model to an
+        # ONNX-GraphSurgeon graph that can be fed to our custom loader.
+        #
+        # First, import the loader from Polygraphy:
+        script.add_import(imports=["GsFromOnnx"], frm="polygraphy.backend.onnx")
+        # Next, invoke the loader with arguments (in this case, the ONNX model loader name), and add it to the script.
+        loader_name = script.add_loader(
+            make_invocable("GsFromOnnx", loader_name), loader_id="gs_from_onnx"
+        )
+
+        # Finally, add the ReplaceReshapeArgs loader.
+        # Unlike the Polygraphy loaders, we'll need to import our loader from the extension module.
+        script.add_import(
+            imports=["ReplaceReshapes"], frm="polygraphy_reshape_destroyer.backend"
+        )
+        # Add the loader and return the ID so that it can be used by subsequent loaders or runners.
+        # NOTE: We can provide additional positional and keyword arguments to `make_invocable` to pass them on to the loader.
+        return script.add_loader(
+            make_invocable(
+                "ReplaceReshapes", loader_name, rename_nodes=self.rename_nodes
+            ),
+            loader_id="replace_reshapes",
+        )

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/args/runner.py

@@ -0,0 +1,84 @@
+#
+# 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.
+#
+"""
+This file defines the `IdentityOnlyRunnerArgs` argument group, which manages
+command-line options that control the `IdentityOnlyRunner` runner.
+
+The argument group implements the standard `BaseRunnerArgs` interface, which inherits from `BaseArgs`.
+"""
+
+from polygraphy import mod
+from polygraphy.tools.args import util as args_util
+from polygraphy.tools.args.base import BaseRunnerArgs
+from polygraphy.tools.script import make_invocable
+from polygraphy_reshape_destroyer.args.loader import ReplaceReshapeArgs
+
+
+# NOTE: Much like loader argument groups, runner argument groups may depend on other argument groups.
+@mod.export()
+class IdentityOnlyRunnerArgs(BaseRunnerArgs):
+    """
+    Identity-Only Runner Inference: running inference with the identity-only runner.
+
+    Depends on:
+
+        - ReplaceReshapeArgs
+    """
+
+    def get_name_opt_impl(self):
+        # Unlike regular `BaseArgs` argument groups, runner argument groups are also expected
+        # to provide a human readable name for the runner as well as a name for
+        # the option that will toggle the runner, not including leading dashes.
+        #
+        # We'll use "res-des" for the option, which will allow us to use the runner by setting `--res-des`.
+        return "Identity-Only Runner", "res-des"
+
+    def add_parser_args_impl(self):
+        # Once again, to prevent collisions with other Polygraphy options, we prefix our option with `res-des`.
+        self.group.add_argument(
+            "--res-des-speed",
+            help="Speed to run inference",
+            choices=["slow", "medium", "fast"],
+            # Since our runner uses `util.default`, we can use `None` as a universal default.
+            default=None,
+        )
+
+    def parse_impl(self, args):
+        """
+        Parses command-line arguments and populates the following attributes:
+
+        Attributes:
+            speed (str): Speed with which to run inference.
+        """
+        self.speed = args_util.get(args, "res_des_speed")
+
+    def add_to_script_impl(self, script):
+        # We'll rely on our ReplaceReshapeArgs argument group to create the ONNX-GraphSurgeon graph for us:
+        loader_name = self.arg_groups[ReplaceReshapeArgs].add_to_script(script)
+
+        # Next, we'll add an import for our runner.
+        script.add_import(
+            imports=["IdentityOnlyRunner"], frm="polygraphy_reshape_destroyer.backend"
+        )
+        # Lastly, we can add our runner using the `Script.add_runner()` API.
+        # Like in the loader implementation, additional arguments can be provided directly to `make_invocable`.
+        script.add_runner(
+            make_invocable("IdentityOnlyRunner", loader_name, speed=self.speed)
+        )
+
+        # NOTE: Unlike the `add_to_script_impl` method of regular `BaseArgs`, that of `BaseRunnerArgs`
+        #       is not expected to return anything.

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/backend/__init__.py

@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+# Since we've used `mod.export()`, we can simply `import *`.
+# Only objects which have been decorated by `mod.export()` will be visible.
+from polygraphy_reshape_destroyer.backend.loader import *
+from polygraphy_reshape_destroyer.backend.runner import *

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/backend/loader.py

@@ -0,0 +1,129 @@
+#
+# 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.
+#
+"""
+This file defines the `ReplaceReshapes` loader, which takes an ONNX-GraphSurgeon graph
+and replaces any no-op Reshapes with Identity nodes.
+
+The loader implements the standard `BaseLoader` interface.
+"""
+
+from typing import Callable, Union
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseLoader
+from polygraphy.logger import G_LOGGER
+
+# For external dependencies or any Polygraphy backends
+# (besides `polygraphy.backend.base`), you should use `mod.lazy_import`.
+#
+# This will enable Polygraphy to automatically install dependencies at runtime if required, and
+# will avoid creating a hard dependency on external packages.
+#
+# NOTE: As the name implies, `lazy_import` does *not* import the module until the first time it is
+#       accessed. Thus, you should be careful to avoid an antipattern like:
+#
+#   my_module = mod.lazy_import("my_module")
+#   submodule = my_module.submodule
+#
+# The second line will trigger an immediate import of `my_module`.
+# Instead, use something like:
+#
+#   submodule = mod.lazy_import("my_module.submodule")
+#
+gs = mod.lazy_import("onnx_graphsurgeon")
+
+
+# `mod.export()` adds the decorated class or function to this module's __all__ attribute.
+# When we do an `import *` from the `__init__.py` file in this submodule, this will ensure
+# that only the decorated objects are exported.
+#
+# NOTE: We use `funcify=True` so that an immediately evaluated functional loader (called `replace_reshapes`)
+#       will be automatically generated for us. This won't be used by the command-line toolkit, but could
+#       be useful if this module is ever used via the Python API.
+#
+@mod.export(funcify=True)
+class ReplaceReshapes(BaseLoader):
+    """
+    Functor that replaces no-op Reshape nodes in an ONNX-GraphSurgeon graph with Identity.
+    """
+
+    def __init__(
+        self, graph: Union[gs.Graph, Callable[[], gs.Graph]], rename_nodes: bool = None
+    ):
+        """
+        Replaces no-op Reshape nodes in an ONNX-GraphSurgeon graph with Identity.
+
+        Args:
+            graph (Union[gs.Graph, Callable() -> gs.Graph]):
+                    An ONNX-GraphSurgeon graph or a callable that returns one.
+            rename_nodes (bool):
+                    Whether to rename Reshape nodes when we convert them to Identity.
+                    Defaults to False.
+        """
+        # In addition to accepting a `gs.Graph` directly, we will also support callables, e.g. Polygraphy loaders.
+        # This will allow our loader to be composed together with other Polygraphy loaders.
+        #
+        # Since the `graph` parameter may be a callable, we'll assign it to a "private" member, i.e. prefixed with '_',
+        # to avoid conflating it with an actual `gs.Graph`.
+        #
+        self._graph = graph
+
+        # See the comment in `util.default` for details on why we use this approach rather than standard Python default parameters.
+        self.rename_nodes = util.default(rename_nodes, False)
+
+    # The `call_impl` method is responsible for doing the actual work of the loader.
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            gs.Graph: The graph with no-op Reshape nodes replaced by Identity.
+        """
+        # As mentioned before, `self._graph` could be a callable, so we invoke it if needed here.
+        #
+        # TIP: The second value returned by `invoke_if_callable` (unused here) is a boolean indicating
+        #      whether the argument was indeed a callable.
+        #
+        graph, _ = util.invoke_if_callable(self._graph)
+
+        for node in graph.nodes:
+            if node.op != "Reshape":
+                continue
+
+            # We can't determine that a Reshape is a no-op unless the new shape is known
+            # prior to inference-time, i.e. a constant.
+            if not isinstance(node.inputs[1], gs.Constant):
+                continue
+
+            # Reshape is only a no-op when the new shape is the same as the old shape.
+            new_shape = node.inputs[1].values
+            if list(node.inputs[0].shape) != list(new_shape):
+                continue
+
+            # Replace no-op reshape with an Identity. We can simply edit the operator name,
+            # clear any attributes, and then delete the second input.
+            G_LOGGER.info(f"Replacing no-op reshape: {node.name} with an Identity node")
+            if self.rename_nodes:
+                node.name += "_destroyed"
+                G_LOGGER.info(f"Renamed Identity node to: {node.name}")
+
+            node.op = "Identity"
+            node.attrs.clear()
+            del node.inputs[1]
+
+        # Finally, clean up the graph to remove any dangling tensors and return it.
+        graph.cleanup()
+        return graph

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/backend/runner.py

@@ -0,0 +1,125 @@
+#
+# 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.
+#
+
+"""
+This file defines the `IdentityOnlyRunner` runner, which takes an ONNX-GraphSurgeon graph
+containing only Identity nodes and runs inference.
+
+The runner implements the standard `BaseRunner` interface.
+"""
+
+import copy
+import time
+from collections import OrderedDict
+
+from polygraphy import mod, util
+from polygraphy.backend.base import BaseRunner
+from polygraphy.common import TensorMetadata
+from polygraphy.logger import G_LOGGER
+
+
+@mod.export()
+class IdentityOnlyRunner(BaseRunner):
+    """
+    Runs inference using custom Python code.
+    Only supports models containing only Identity nodes.
+    """
+
+    def __init__(self, graph, name=None, speed: str = None):
+        """
+        Args:
+            graph (Union[onnx_graphsurgeon.Graph, Callable() -> onnx_graphsurgeon.Graph]):
+                    An ONNX-GraphSurgeon graph or a callable that returns one.
+            name (str):
+                    The human-readable name prefix to use for this runner.
+                    A runner count and timestamp will be appended to this prefix.
+            speed (str):
+                    How fast to run inference. Should be one of: ["slow", "medium", "fast"].
+                    Defaults to "fast".
+        """
+        super().__init__(name=name, prefix="pluginref-runner")
+        self._graph = graph
+
+        self.speed = util.default(speed, "fast")
+
+        VALID_SPEEDS = ["slow", "medium", "fast"]
+        if self.speed not in VALID_SPEEDS:
+            # Like Polygraphy, extension modules should use `G_LOGGER.critical()` for any unrecoverable errors.
+            G_LOGGER.critical(
+                f"Invalid speed: {self.speed}. Note: Valid speeds are: {VALID_SPEEDS}"
+            )
+
+    @util.check_called_by("activate")
+    def activate_impl(self):
+        # As with the loader, the `graph` argument could be either a `gs.Graph` or a callable that
+        # returns one, such as a loader, so we try to call it.
+        self.graph, _ = util.invoke_if_callable(self._graph)
+
+    #
+    # All the methods from this point forward are guaranteed to be called only after `activate()`,
+    # so we can assume that `self.graph` will be available.
+    #
+
+    @util.check_called_by("get_input_metadata")
+    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.
+        meta = TensorMetadata()
+        for tensor in self.graph.inputs:
+            meta.add(tensor.name, tensor.dtype, tensor.shape)
+        return meta
+
+    @util.check_called_by("infer")
+    def infer_impl(self, feed_dict):
+        start = time.time()
+
+        # Since our runner only supports Identity, all we need to do for inference is bind node outputs to their inputs.
+        # We'll begin with a copy of the input tensors:
+        tensor_values = copy.copy(feed_dict)
+
+        for node in self.graph.nodes:
+            # We don't support non-Identity nodes, so we'll report an error if we see one
+            if node.op != "Identity":
+                G_LOGGER.critical(
+                    f"Encountered an unsupported type of node: {node.op}."
+                    "Note: This runner only supports Identity nodes!"
+                )
+
+            inp_tensor = node.inputs[0]
+            out_tensor = node.outputs[0]
+            # The output of an Identity node should be identical to its input.
+            tensor_values[out_tensor.name] = tensor_values[inp_tensor.name]
+
+        # Find the output tensors based on `self.graph.outputs` and create a dictionary that we can return:
+        outputs = OrderedDict()
+        for out in self.graph.outputs:
+            outputs[out.name] = tensor_values[out.name]
+
+        # Next we'll implement our artifical delay so that we can see amazing performance gains in "fast" mode!
+        delay = {"slow": 1.0, "medium": 0.5, "fast": 0.0}[self.speed]
+        time.sleep(delay)
+
+        end = time.time()
+
+        # In order to allow Polygraphy to report inference times accurately, runners are responsible for reporting
+        # their own inference time. This is done by setting the `self.inference_time` attribute.
+        self.inference_time = end - start
+        return outputs
+
+    @util.check_called_by("deactivate")
+    def deactivate_impl(self):
+        del self.graph

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/polygraphy_reshape_destroyer/export.py

@@ -0,0 +1,34 @@
+#
+# 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.
+#
+
+"""
+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_reshape_destroyer.args import ReplaceReshapeArgs, IdentityOnlyRunnerArgs
+
+
+# The entry point is expected to take no arguments and return a list of argument group instances.
+#
+# NOTE: Argument groups will be parsed in the order in which they are provided,
+#       and after all of Polygraphy's built-in argument groups.
+def export_argument_groups():
+    return [
+        ReplaceReshapeArgs(),
+        IdentityOnlyRunnerArgs(),
+    ]

tools/Polygraphy/examples/dev/02_extending_polygraphy_run/extension_module/setup.py

@@ -0,0 +1,68 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+from setuptools import find_packages, setup
+
+import polygraphy_reshape_destroyer
+
+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(
+        # NOTE: You will want to edit most of these fields for your custom extension module.
+        name="polygraphy_reshape_destroyer",
+        version=polygraphy_reshape_destroyer.__version__,
+        description="Polygraphy Reshape Destroyer: Destroyer Of Reshapes",
+        long_description=open("README.md", "r", encoding="utf-8").read(),
+        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",
+            # Our included loader needs ONNX-GraphSurgeon to modify the model.
+            "onnx_graphsurgeon",
+            "numpy<2",
+        ],
+        packages=find_packages(exclude=("tests", "tests.*")),
+        # The format of the entry_points is:
+        # {
+        #   "polygraphy.run.plugins"    # Polygraphy run entrypoint (this should not be changed)
+        #      : ["<name-of-plugin>=module.<submodule...>:object"]
+        # }
+        entry_points={
+            "polygraphy.run.plugins": [
+                "reshape-destroyer=polygraphy_reshape_destroyer.export:export_argument_groups",
+            ]
+        },
+        zip_safe=True,
+        python_requires=">=3.6",
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/Polygraphy/examples/dev/README.md

@@ -0,0 +1,5 @@
+# Polygraphy Development Examples
+
+This directory includes examples related to Polygraphy development.
+This covers topics such as creating new command-line tools and adding
+new features to Polygraphy.