wehub-resource-sync
364 lines
12 KiB
Python
364 lines
12 KiB
Python
#
|
|
# 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)
|