chore: import upstream snapshot with attribution

samples/python/quickly_deployable_plugins/qdp_runner.py

@@ -0,0 +1,363 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 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 tensorrt as trt
+import torch
+import numpy as np
+
+from polygraphy.backend.trt import (
+    CreateConfig,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+    network_from_onnx_path,
+    bytes_from_engine,
+    engine_from_bytes,
+)
+
+from polygraphy.backend.common import bytes_from_path
+from polygraphy import cuda
+
+import onnx_graphsurgeon as gs
+import onnx
+import os
+import argparse
+
+import tensorrt.plugin as trtp
+
+import qdp_defs
+import logging
+
+def run_add(enable_autotune=False):
+
+    if enable_autotune:
+        qdp_defs.register_autotune()
+
+    BLOCK_SIZE = 256
+
+    builder, network = create_network(strongly_typed=True)
+    x = torch.randint(10, (10, 3, 32, 32), dtype=torch.float32, device="cuda")
+
+    # Populate network
+    i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+
+    out = network.add_plugin(
+        trtp.op.sample.elemwise_add_plugin(i_x, block_size=BLOCK_SIZE)
+    )
+    out.get_output(0).name = "y"
+    network.mark_output(tensor=out.get_output(0))
+
+    builder.create_builder_config()
+
+    engine = engine_from_network(
+        (builder, network),
+        CreateConfig(),
+    )
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer(
+            {
+                "x": x,
+            },
+            copy_outputs_to_host=False,
+        )
+
+    if torch.allclose(x + 1, outputs["y"]):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+
+
+def run_inplace_add():
+    builder, network = create_network(strongly_typed=True)
+    x = torch.ones((10, 3, 32, 32), dtype=torch.float32, device="cuda")
+
+    x_clone = x.clone()
+
+    i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+
+    # Amounts to elementwise-add in the first and second plugins
+    deltas = (2, 4)
+
+    out0 = network.add_plugin(trtp.op.sample.elemwise_add_plugin_(i_x, delta=deltas[0]))
+    out1 = network.add_plugin(
+        trtp.op.sample.elemwise_add_plugin_(out0.get_output(0), delta=deltas[1])
+    )
+    out1.get_output(0).name = "y"
+    network.mark_output(tensor=out1.get_output(0))
+
+    builder.create_builder_config()
+
+    # Enable preview feature for aliasing plugin I/O
+    config = CreateConfig(
+        preview_features=[trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
+    )
+
+    engine = engine_from_network(
+        (builder, network),
+        config,
+    )
+
+    context = engine.create_execution_context()
+
+    stream = cuda.Stream()
+
+    context.set_tensor_address("x", x.data_ptr())
+    context.set_tensor_address("y", x.data_ptr())
+    context.execute_async_v3(stream.ptr)
+    stream.synchronize()
+
+    if torch.allclose(x, x_clone + sum(deltas), atol=1e-2):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+        print(x[0][0][0][:10])
+        print(x_clone[0][0][0][:10])
+
+
+def run_non_zero():
+    builder, network = create_network(strongly_typed=True)
+    inp_shape = (128, 128)
+
+    X = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
+
+    # Zero out some random indices
+    indices = np.random.choice(
+        np.prod(inp_shape),
+        replace=False,
+        size=np.random.randint(0, np.prod(inp_shape) + 1),
+    )
+    X[np.unravel_index(indices, inp_shape)] = 0
+
+    # Populate network
+    i_x = network.add_input(name="X", dtype=trt.DataType.FLOAT, shape=inp_shape)
+
+    out = network.add_plugin(trtp.op.sample.non_zero_plugin(i_x))
+    out.get_output(0).name = "Y"
+    network.mark_output(tensor=out.get_output(0))
+
+    builder.create_builder_config()
+
+    engine = engine_from_network(
+        (builder, network),
+        config=CreateConfig(),
+    )
+
+    Y_ref = np.transpose(np.nonzero(X))
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+        Y = Y[np.lexsort(np.fliplr(Y).T)]
+
+    if np.allclose(Y, Y_ref, atol=1e-3):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+
+
+def check_artifacts_dir_exists(artifacts_dir):
+    if not os.path.exists(artifacts_dir):
+        raise ValueError(f"artifacts_dir '{artifacts_dir}' does not exist")
+
+
+def run_circ_pad(
+    enable_multi_tactic=False, mode="onnx", artifacts_dir=None, save_or_load_engine=None, aot=False
+):
+
+    if enable_multi_tactic and aot:
+        qdp_defs.enable_multi_tactic_aot_circ_pad()
+    elif enable_multi_tactic:
+        qdp_defs.enable_multi_tactic_circ_pad()
+    else:
+        qdp_defs.enable_single_tactic_circ_pad()
+
+    inp_shape = (10, 3, 32, 32)
+    x = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
+
+    pads = np.array((1, 1, 1, 1), dtype=np.int32)
+
+    if save_or_load_engine is not None and save_or_load_engine is False:
+        check_artifacts_dir_exists(artifacts_dir)
+        engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
+        engine = engine_from_bytes(bytes_from_path(engine_path))
+    else:
+        if mode == "inetdef":
+            builder, network = create_network(strongly_typed=True)
+            i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+            out = network.add_plugin(trtp.op.sample.circ_pad_plugin(i_x, pads=pads), aot = aot)
+            out.get_output(0).name = "y"
+            network.mark_output(tensor=out.get_output(0))
+
+            engine = engine_from_network(
+                (builder, network),
+                CreateConfig(),
+            )
+        elif mode == "onnx":
+            if artifacts_dir is None:
+                raise ValueError("'artifacts_dir' must be specified in onnx mode")
+
+            check_artifacts_dir_exists(artifacts_dir)
+
+            onnx_path = os.path.join(artifacts_dir, "circ_pad.onnx")
+            var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
+            var_y = gs.Variable(name="y", dtype=np.float32)
+            circ_pad_node = gs.Node(
+                name="circ_pad_plugin 0",
+                op="circ_pad_plugin",
+                inputs=[var_x],
+                outputs=[var_y],
+                attrs={"pads": pads, "plugin_namespace": "sample", "aot": aot},
+            )
+            graph = gs.Graph(
+                nodes=[circ_pad_node], inputs=[var_x], outputs=[var_y], opset=16
+            )
+            onnx.save(gs.export_onnx(graph), onnx_path)
+
+            engine = engine_from_network(
+                network_from_onnx_path(onnx_path, strongly_typed=True), CreateConfig()
+            )
+        else:
+            raise ValueError(f"Unknown mode {mode}")
+
+        if save_or_load_engine is not None and save_or_load_engine is True:
+            check_artifacts_dir_exists(artifacts_dir)
+            engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
+            with open(engine_path, "wb") as f:
+                f.write(bytes_from_engine(engine))
+
+    Y_ref = np.pad(x, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"x": x})
+        Y = outputs["y"]
+
+        if np.allclose(Y, Y_ref, atol=1e-2):
+            print("Inference result is correct!")
+        else:
+            print("Inference result is incorrect!")
+
+
+def setup_add_sample(subparsers):
+    subparser = subparsers.add_parser("add", help="'add' sample help")
+    subparser.add_argument("--autotune", action="store_true", help="Enable autotuning")
+    subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_inplace_add_sample(subparsers):
+    subparser = subparsers.add_parser("inplace_add", help="inplace_add sample help")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_non_zero_sample(subparsers):
+    subparser = subparsers.add_parser("non_zero", help="non_zero sample help")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_circ_pad_sample(subparsers):
+    subparser = subparsers.add_parser("circ_pad", help="circ_pad sample help.")
+    subparser.add_argument(
+        "--multi_tactic", action="store_true", help="Enable multiple tactics."
+    )
+    subparser.add_argument(
+        "--save_engine", action="store_true", help="Save engine to the artifacts_dir."
+    )
+    subparser.add_argument(
+        "--load_engine",
+        action="store_true",
+        help="Load engine from the artifacts_dir. Ignores all other options.",
+    )
+    subparser.add_argument(
+        "--artifacts_dir",
+        type=str,
+        help="Whether to store (or retrieve) artifacts.",
+    )
+    subparser.add_argument(
+        "--mode",
+        type=str,
+        choices=["onnx", "inetdef"],
+        help="Whether to use ONNX parser or INetworkDefinition APIs to construct the network.",
+    )
+    subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin.")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose log output."
+    )
+
+    return subparser
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser = argparse.ArgumentParser(description="Main script help")
+    subparsers = parser.add_subparsers(dest="sample", help="Mode help", required=True)
+
+    setup_add_sample(subparsers)
+    setup_inplace_add_sample(subparsers)
+    circ_pad_subparser = setup_circ_pad_sample(subparsers)
+    setup_non_zero_sample(subparsers)
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.DEBUG)
+
+    if args.sample == "add":
+        run_add(args.autotune)
+    if args.sample == "inplace_add":
+        run_inplace_add()
+    if args.sample == "non_zero":
+        run_non_zero()
+    if args.sample == "circ_pad":
+        if args.mode == "onnx":
+            if args.artifacts_dir is None:
+                parser.error(
+                    "circ_pad: argument --mode: When mode is 'onnx', artifacts_dir is required"
+                )
+
+        save_or_load_engine = None
+
+        if args.load_engine is True:
+            if args.save_engine is True:
+                parser.error(
+                    "circ_pad: save_engine and load_engine cannot be specified at the same time. First save_engine and load_engine separately."
+                )
+            else:
+                if args.multi_tactic is True or args.mode is not None:
+                    print(
+                        "warning circ_pad: when load_engine is specified, all other options except 'artifacts_dir' is ignored."
+                    )
+
+            save_or_load_engine = False
+        else:
+            if args.mode is None:
+                circ_pad_subparser.print_help()
+                parser.error(
+                    "circ_pad: '--mode' option is required."
+                )
+
+        if args.save_engine is True:
+            save_or_load_engine = True
+
+        run_circ_pad(args.multi_tactic, args.mode, args.artifacts_dir, save_or_load_engine, args.aot)