github-featured/nvidia--tensorrt
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

TensorRT 10.0 GA Release

Browse Source 
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
This commit is contained in:
Asfiya Baig
2024-04-25 16:01:48 -07:00
parent 28733f0fdc
commit d3107c885a
853 changed files with 11215 additions and 12167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/python/python_plugin/CMakeLists.txt
+1 -1
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
samples/python/python_plugin/circ_pad_plugin_cpp.py
+23 -8
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -29,14 +29,29 @@ from polygraphy.backend.trt import (
TrtRunner,
)
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Circular Padding plugin C++ example")
parser.add_argument('--precision', type=str, default="fp32", choices=["fp32", "fp16"], help="Precision to use for plugin")
parser.add_argument('--plugin-lib', type=str, help="Path to the Circular Padding plugin lib", required=True)
def parseArgs():
parser = argparse.ArgumentParser(
description="Options for Circular Padding plugin C++ example"
)
parser.add_argument(
"--precision",
type=str,
default="fp32",
choices=["fp32", "fp16"],
help="Precision to use for plugin",
)
parser.add_argument(
"--plugin-lib",
type=str,
help="Path to the Circular Padding plugin lib",
required=True,
)
return parser.parse_args()
if __name__ == "__main__":
args = parseArgs()
@@ -67,15 +82,15 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
if np.allclose(Y, Y_ref):
print("Inference result correct!")
else:
samples/python/python_plugin/circ_pad_plugin_cuda_python.py
+74 -38
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -24,14 +24,14 @@ from polygraphy.backend.trt import (
CreateConfig,
EngineFromNetwork,
NetworkFromOnnxPath,
TrtRunner
TrtRunner,
)
from polygraphy.json import to_json, from_json
from utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager
from cuda import cuda
circ_pad_half_kernel = r'''
circ_pad_half_kernel = r"""
#include <cuda_fp16.h>
extern "C" __global__
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
@@ -58,9 +58,9 @@ void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, hal
];
}
}
'''
"""
circ_pad_float_kernel = r'''
circ_pad_float_kernel = r"""
extern "C" __global__
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
@@ -86,7 +86,8 @@ void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, f
];
}
}
'''
"""
class CircPadPlugin(trt.IPluginV2DynamicExt):
def __init__(self, fc=None):
@@ -107,7 +108,9 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
self.cuDevice = None
if fc is not None:
assert set([f.name for f in fc]) == set(["pads", "N"]), "Field collection invalid"
assert set([f.name for f in fc]) == set(
["pads", "N"]
), "Field collection invalid"
for f in fc:
if f.name == "pads":
self.pads = f.data
@@ -116,11 +119,17 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
def initialize(self):
err, self.cuDevice = cuda.cuDeviceGet(0)
trt.get_plugin_registry().acquire_plugin_resource("cuda_ctx", CudaCtxManager(self.cuDevice))
self.all_pads_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N * 2))
self.orig_dims_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
trt.get_plugin_registry().acquire_plugin_resource(
"cuda_ctx", CudaCtxManager(self.cuDevice)
)
self.all_pads_d = checkCudaErrors(
cuda.cuMemAlloc(np.int32().itemsize * self.N * 2)
)
self.orig_dims_d = checkCudaErrors(
cuda.cuMemAlloc(np.int32().itemsize * self.N)
)
self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
def get_output_datatype(self, index, input_types):
return input_types[0]
@@ -157,11 +166,17 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# Copy vectors from host memory to device memory
if self.all_pads_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes)
)
if self.orig_dims_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes)
)
if self.Y_shape_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes)
)
self.Y_len_d = np.prod(out_dims)
@@ -205,25 +220,43 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
if inp_dtype == np.float32:
kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
_circ_pad_float_kernel = kernelHelper.getFunction(b'circ_pad_float')
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_float_kernel,
numBlocks, 1, 1,
blockSize, 1, 1,
0,
stream_ptr,
kernelArgs, 0))
_circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float")
checkCudaErrors(
cuda.cuLaunchKernel(
_circ_pad_float_kernel,
numBlocks,
1,
1,
blockSize,
1,
1,
0,
stream_ptr,
kernelArgs,
0,
)
)
elif inp_dtype == np.float16:
kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
_circ_pad_half_kernel = kernelHelper.getFunction(b'circ_pad_half')
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_half_kernel,
numBlocks, 1, 1,
blockSize, 1, 1,
0,
stream_ptr,
kernelArgs, 0))
_circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half")
checkCudaErrors(
cuda.cuLaunchKernel(
_circ_pad_half_kernel,
numBlocks,
1,
1,
blockSize,
1,
1,
0,
stream_ptr,
kernelArgs,
0,
)
)
else:
raise ValueError("inp_dtype not valid")
def clone(self):
cloned_plugin = CircPadPlugin()
cloned_plugin.__dict__.update(self.__dict__)
@@ -239,7 +272,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
trt.get_plugin_registry().release_plugin_resource("cuda_ctx")
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -248,7 +281,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
@@ -259,10 +292,12 @@ class CircPadPluginCreator(trt.IPluginCreator):
self.name = "CircPadPlugin"
self.plugin_namespace = ""
self.plugin_version = "1"
self.field_names = trt.PluginFieldCollection([
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32)
])
self.field_names = trt.PluginFieldCollection(
[
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32),
]
)
def create_plugin(self, name, fc):
return CircPadPlugin(fc)
@@ -273,12 +308,13 @@ class CircPadPluginCreator(trt.IPluginCreator):
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
# Initialize CUDA Driver API
err, = cuda.cuInit(0)
(err,) = cuda.cuInit(0)
# Retrieve handle for device 0
err, cuDevice = cuda.cuDeviceGet(0)
@@ -319,12 +355,12 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
samples/python/python_plugin/circ_pad_plugin_cupy.py
+46 -15
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -27,14 +27,15 @@ from polygraphy.backend.trt import (
CreateConfig,
EngineFromNetwork,
NetworkFromOnnxPath,
TrtRunner
TrtRunner,
)
from polygraphy.json import to_json, from_json
from utils import volume, parseArgs
circ_pad_half_kernel = cp.RawKernel(r'''
circ_pad_half_kernel = cp.RawKernel(
r"""
#include <cuda_fp16.h>
extern "C" __global__
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int const* Y_len) {
@@ -61,9 +62,12 @@ void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, hal
];
}
}
''', 'circ_pad_half')
""",
"circ_pad_half",
)
circ_pad_float_kernel = cp.RawKernel(r'''
circ_pad_float_kernel = cp.RawKernel(
r"""
extern "C" __global__
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int const* Y_len) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
@@ -89,14 +93,17 @@ void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, f
];
}
}
''', 'circ_pad_float')
""",
"circ_pad_float",
)
class CircPadPlugin(trt.IPluginV2DynamicExt):
def __init__(self, fc=None):
trt.IPluginV2DynamicExt.__init__(self)
self.pads = []
self.X_shape = []
self.num_outputs = 1
self.plugin_namespace = ""
self.plugin_type = "CircPadPlugin"
@@ -190,9 +197,31 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
with cuda_stream:
if inp_dtype == np.float32:
circ_pad_float_kernel((numBlocks,), (blockSize,), (a, self.all_pads_d, self.orig_dims_d, c, self.Y_shape_d, self.Y_len_d))
circ_pad_float_kernel(
(numBlocks,),
(blockSize,),
(
a,
self.all_pads_d,
self.orig_dims_d,
c,
self.Y_shape_d,
self.Y_len_d,
),
)
elif inp_dtype == np.float16:
circ_pad_half_kernel((numBlocks,), (blockSize,), (a, self.all_pads_d, self.orig_dims_d, c, self.Y_shape_d, self.Y_len_d))
circ_pad_half_kernel(
(numBlocks,),
(blockSize,),
(
a,
self.all_pads_d,
self.orig_dims_d,
c,
self.Y_shape_d,
self.Y_len_d,
),
)
else:
raise ValueError("inp_dtype not valid")
@@ -201,7 +230,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
cloned_plugin.__dict__.update(self.__dict__)
return cloned_plugin
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -213,17 +242,18 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
# def terminate(self):
# pass
class CircPadPluginCreator(trt.IPluginCreator):
def __init__(self):
trt.IPluginCreator.__init__(self)
self.name = "CircPadPlugin"
self.plugin_namespace = ""
self.plugin_version = "1"
@@ -233,13 +263,14 @@ class CircPadPluginCreator(trt.IPluginCreator):
def create_plugin(self, name, fc):
return CircPadPlugin(fc)
def deserialize_plugin(self, name, data):
j = dict(from_json(data.decode("utf-8")))
deserialized = CircPadPlugin()
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
@@ -275,12 +306,12 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py
+85 -41
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -23,7 +23,7 @@ from polygraphy.backend.trt import (
CreateConfig,
TrtRunner,
create_network,
engine_from_network
engine_from_network,
)
from polygraphy.json import to_json, from_json
@@ -31,7 +31,7 @@ from polygraphy.json import to_json, from_json
from utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager
from cuda import cuda
circ_pad_half_kernel = r'''
circ_pad_half_kernel = r"""
#include <cuda_fp16.h>
extern "C" __global__
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
@@ -58,9 +58,9 @@ void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, hal
];
}
}
'''
"""
circ_pad_float_kernel = r'''
circ_pad_float_kernel = r"""
extern "C" __global__
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
@@ -86,7 +86,8 @@ void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, f
];
}
}
'''
"""
class CircPadPlugin(trt.IPluginV2DynamicExt):
def __init__(self, fc=None):
@@ -107,7 +108,9 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
self.cuDevice = None
if fc is not None:
assert set([f.name for f in fc]) == set(["pads", "N"]), "Field collection invalid"
assert set([f.name for f in fc]) == set(
["pads", "N"]
), "Field collection invalid"
for f in fc:
if f.name == "pads":
self.pads = f.data
@@ -116,11 +119,17 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
def initialize(self):
err, self.cuDevice = cuda.cuDeviceGet(0)
trt.get_plugin_registry().acquire_plugin_resource("cuda_ctx", CudaCtxManager(self.cuDevice))
self.all_pads_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N * 2))
self.orig_dims_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
trt.get_plugin_registry().acquire_plugin_resource(
"cuda_ctx", CudaCtxManager(self.cuDevice)
)
self.all_pads_d = checkCudaErrors(
cuda.cuMemAlloc(np.int32().itemsize * self.N * 2)
)
self.orig_dims_d = checkCudaErrors(
cuda.cuMemAlloc(np.int32().itemsize * self.N)
)
self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
def get_output_datatype(self, index, input_types):
return input_types[0]
@@ -157,11 +166,17 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# Copy vectors from host memory to device memory
if self.all_pads_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes)
)
if self.orig_dims_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes)
)
if self.Y_shape_d:
checkCudaErrors(cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes))
checkCudaErrors(
cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes)
)
self.Y_len_d = np.prod(out_dims)
@@ -205,25 +220,43 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
if inp_dtype == np.float32:
kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
_circ_pad_float_kernel = kernelHelper.getFunction(b'circ_pad_float')
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_float_kernel,
numBlocks, 1, 1,
blockSize, 1, 1,
0,
stream_ptr,
kernelArgs, 0))
_circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float")
checkCudaErrors(
cuda.cuLaunchKernel(
_circ_pad_float_kernel,
numBlocks,
1,
1,
blockSize,
1,
1,
0,
stream_ptr,
kernelArgs,
0,
)
)
elif inp_dtype == np.float16:
kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
_circ_pad_half_kernel = kernelHelper.getFunction(b'circ_pad_half')
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_half_kernel,
numBlocks, 1, 1,
blockSize, 1, 1,
0,
stream_ptr,
kernelArgs, 0))
_circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half")
checkCudaErrors(
cuda.cuLaunchKernel(
_circ_pad_half_kernel,
numBlocks,
1,
1,
blockSize,
1,
1,
0,
stream_ptr,
kernelArgs,
0,
)
)
else:
raise ValueError("inp_dtype not valid")
def clone(self):
cloned_plugin = CircPadPlugin()
cloned_plugin.__dict__.update(self.__dict__)
@@ -239,7 +272,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
plg_registry.release_plugin_resource("cuda_ctx")
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -248,7 +281,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
@@ -259,10 +292,12 @@ class CircPadPluginCreator(trt.IPluginCreator):
self.name = "CircPadPlugin"
self.plugin_namespace = ""
self.plugin_version = "1"
self.field_names = trt.PluginFieldCollection([
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32)
])
self.field_names = trt.PluginFieldCollection(
[
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32),
]
)
def create_plugin(self, name, fc):
return CircPadPlugin(fc)
@@ -273,13 +308,14 @@ class CircPadPluginCreator(trt.IPluginCreator):
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
precision = np.float32 if args.precision == "fp32" else np.float16
# Initialize CUDA Driver API
err, = cuda.cuInit(0)
(err,) = cuda.cuInit(0)
# Retrieve handle for device 0
err, cuDevice = cuda.cuDeviceGet(0)
@@ -306,28 +342,36 @@ if __name__ == "__main__":
builder, network = create_network()
plg_creator = plg_registry.get_plugin_creator("CircPadPlugin", "1", "")
plugin_fields_list = [
trt.PluginField("pads", np.array(pads, dtype=np.int32), trt.PluginFieldType.INT32),
trt.PluginField(
"pads", np.array(pads, dtype=np.int32), trt.PluginFieldType.INT32
),
trt.PluginField("N", np.array([4], dtype=np.int32), trt.PluginFieldType.INT32),
]
pfc = trt.PluginFieldCollection(plugin_fields_list)
plugin = plg_creator.create_plugin("CircPadPlugin", pfc)
# Populate network
input_X = network.add_input(name="X", dtype=trt.float32 if precision==np.float32 else trt.float16, shape=X.shape)
input_X = network.add_input(
name="X",
dtype=trt.float32 if precision == np.float32 else trt.float16,
shape=X.shape,
)
out = network.add_plugin_v2([input_X], plugin)
out.get_output(0).name = "Y"
network.mark_output(tensor=out.get_output(0))
# Build engine
config = builder.create_builder_config()
engine = engine_from_network((builder, network), CreateConfig(fp16=precision==trt.float16))
engine = engine_from_network(
(builder, network), CreateConfig(fp16=precision == trt.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(engine, "trt_runner")as runner:
with TrtRunner(engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
if np.allclose(Y, Y_ref):
print("Inference result correct!")
else:
samples/python/python_plugin/circ_pad_plugin_numba.py
+10 -7
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -32,6 +32,7 @@ from polygraphy.backend.trt import (
from polygraphy.json import to_json, from_json
from utils import volume, parseArgs
@cuda.jit
def circ_pad(X, all_pads, orig_dims, Y, Y_shape, Y_len):
index = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
@@ -57,6 +58,7 @@ def circ_pad(X, all_pads, orig_dims, Y, Y_shape, Y_len):
)
]
class CircPadPlugin(trt.IPluginV2DynamicExt):
def __init__(self, fc=None):
trt.IPluginV2DynamicExt.__init__(self)
@@ -76,7 +78,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
return input_types[0]
def get_output_dimensions(self, output_index, inputs, exprBuilder):
output_dims = trt.DimsExprs(inputs[0])
for i in range(np.size(self.pads) // 2):
@@ -163,8 +165,8 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
cloned_plugin = CircPadPlugin()
cloned_plugin.__dict__.update(self.__dict__)
return cloned_plugin
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -176,7 +178,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
@@ -203,6 +205,7 @@ class CircPadPluginCreator(trt.IPluginCreator):
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
@@ -234,12 +237,12 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
samples/python/python_plugin/circ_pad_plugin_torch.py
+10 -8
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -33,12 +33,13 @@ import torch
from utils import volume, parseArgs
class CircPadPlugin(trt.IPluginV2DynamicExt):
def __init__(self, fc=None):
trt.IPluginV2DynamicExt.__init__(self)
self.pads = []
self.X_shape = []
self.num_outputs = 1
self.plugin_namespace = ""
self.plugin_type = "CircPadPlugin"
@@ -110,10 +111,10 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
a_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=a_ptr)
c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
a_t = torch.as_tensor(a_d, device='cuda')
a_t = torch.as_tensor(a_d, device="cuda")
# Use PyTorch functional op - no need to write kernel
out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode='circular')
out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode="circular")
cp.copyto(c_d, cp.reshape(cp.asarray(out), (-1,)))
return 0
@@ -123,7 +124,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
cloned_plugin.__dict__.update(self.__dict__)
return cloned_plugin
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -135,7 +136,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
@@ -162,6 +163,7 @@ class CircPadPluginCreator(trt.IPluginCreator):
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
@@ -193,12 +195,12 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
samples/python/python_plugin/circ_pad_plugin_triton.py
+49 -19
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -36,13 +36,26 @@ import torch
from utils import volume, parseArgs
@triton.jit
def circ_pad(X,
all_pads_0, all_pads_2, all_pads_4, all_pads_6,
orig_dims_0, orig_dims_1, orig_dims_2, orig_dims_3,
Y,
Y_shape_1, Y_shape_2, Y_shape_3,
X_len, Y_len, BLOCK_SIZE: tl.constexpr,):
def circ_pad(
X,
all_pads_0,
all_pads_2,
all_pads_4,
all_pads_6,
orig_dims_0,
orig_dims_1,
orig_dims_2,
orig_dims_3,
Y,
Y_shape_1,
Y_shape_2,
Y_shape_3,
X_len,
Y_len,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
@@ -58,7 +71,12 @@ def circ_pad(X,
j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
load_idx = orig_dims_3 * orig_dims_2 * orig_dims_1 * j0 + orig_dims_3 * orig_dims_2 * j1 + orig_dims_3 * j2 + j3
load_idx = (
orig_dims_3 * orig_dims_2 * orig_dims_1 * j0
+ orig_dims_3 * orig_dims_2 * j1
+ orig_dims_3 * j2
+ j3
)
mask_x = load_idx < X_len
x = tl.load(X + load_idx, mask=mask_x)
@@ -143,8 +161,8 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
a_d = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
a_t = torch.as_tensor(a_d, device='cuda')
c_t = torch.as_tensor(c_d, device='cuda')
a_t = torch.as_tensor(a_d, device="cuda")
c_t = torch.as_tensor(c_d, device="cuda")
N = len(self.X_shape)
all_pads = np.zeros((N * 2,), dtype=np.int32)
@@ -163,12 +181,23 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
blockSize = 256
numBlocks = (int((np.prod(out_dims) + blockSize - 1) // blockSize),)
circ_pad[numBlocks](a_t,
all_pads[0], all_pads[2], all_pads[4], all_pads[6],
orig_dims[0], orig_dims[1], orig_dims[2], orig_dims[3],
circ_pad[numBlocks](
a_t,
all_pads[0],
all_pads[2],
all_pads[4],
all_pads[6],
orig_dims[0],
orig_dims[1],
orig_dims[2],
orig_dims[3],
c_t,
out_dims[1], out_dims[2], out_dims[3],
int(np.prod(orig_dims)), int(np.prod(out_dims)), BLOCK_SIZE=256
out_dims[1],
out_dims[2],
out_dims[3],
int(np.prod(orig_dims)),
int(np.prod(out_dims)),
BLOCK_SIZE=256,
)
return 0
@@ -178,7 +207,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
cloned_plugin.__dict__.update(self.__dict__)
return cloned_plugin
#
#
# The following defaults take effect since the respective methods are not overriden
#
@@ -190,7 +219,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt):
# def get_workspace_size(self, input_desc, output_desc):
# return 0
# def destroy(self):
# pass
@@ -217,6 +246,7 @@ class CircPadPluginCreator(trt.IPluginCreator):
deserialized.__dict__.update(j)
return deserialized
if __name__ == "__main__":
args = parseArgs()
@@ -248,12 +278,12 @@ if __name__ == "__main__":
# build engine
build_engine = EngineFromNetwork(
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision == np.float16)
)
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
# Run
with TrtRunner(build_engine, "trt_runner")as runner:
with TrtRunner(build_engine, "trt_runner") as runner:
outputs = runner.infer({"X": X})
Y = outputs["Y"]
samples/python/python_plugin/circ_plugin_cpp/circ_pad_plugin.cu
+2 -3
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* 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");
@@ -109,8 +109,7 @@ __global__ void circPadKernel(
int32_t j2 = (i2 - allPads[4] + origDims[2]) % origDims[2];
int32_t j3 = (i3 - allPads[6] + origDims[3]) % origDims[3];
y[i] = x[origDims[3] * origDims[2] * origDims[1] * j0 + origDims[3] * origDims[2] * j1 + origDims[3] * j2
+ j3];
y[i] = x[origDims[3] * origDims[2] * origDims[1] * j0 + origDims[3] * origDims[2] * j1 + origDims[3] * j2 + j3];
}
}
samples/python/python_plugin/utils.py
+54 -21
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -23,15 +23,26 @@ import threading
import tensorrt as trt
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Circular Padding plugin C++ example")
parser.add_argument('--precision', type=str, default="fp32", choices=["fp32", "fp16"], help="Precision to use for plugin")
parser = argparse.ArgumentParser(
description="Options for Circular Padding plugin C++ example"
)
parser.add_argument(
"--precision",
type=str,
default="fp32",
choices=["fp32", "fp16"],
help="Precision to use for plugin",
)
return parser.parse_args()
def volume(d):
return np.prod(d)
# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/helper_cuda.py
def checkCudaErrors(result):
def _cudaGetErrorEnum(error):
@@ -43,9 +54,14 @@ def checkCudaErrors(result):
elif isinstance(error, nvrtc.nvrtcResult):
return nvrtc.nvrtcGetErrorString(error)[1]
else:
raise RuntimeError('Unknown error type: {}'.format(error))
raise RuntimeError("Unknown error type: {}".format(error))
if result[0].value:
raise RuntimeError("CUDA error code={}({})".format(result[0].value, _cudaGetErrorEnum(result[0])))
raise RuntimeError(
"CUDA error code={}({})".format(
result[0].value, _cudaGetErrorEnum(result[0])
)
)
if len(result) == 1:
return None
elif len(result) == 2:
@@ -53,34 +69,50 @@ def checkCudaErrors(result):
else:
return result[1:]
# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/common.py
class KernelHelper:
def __init__(self, code, devID):
prog = checkCudaErrors(nvrtc.nvrtcCreateProgram(str.encode(code), b'sourceCode.cu', 0, [], []))
CUDA_HOME = os.getenv('CUDA_HOME')
prog = checkCudaErrors(
nvrtc.nvrtcCreateProgram(str.encode(code), b"sourceCode.cu", 0, [], [])
)
CUDA_HOME = os.getenv("CUDA_HOME")
if CUDA_HOME == None:
CUDA_HOME = os.getenv('CUDA_PATH')
CUDA_HOME = os.getenv("CUDA_PATH")
if CUDA_HOME == None:
raise RuntimeError('Environment variable CUDA_HOME or CUDA_PATH is not set')
include_dirs = os.path.join(CUDA_HOME, 'include')
raise RuntimeError("Environment variable CUDA_HOME or CUDA_PATH is not set")
include_dirs = os.path.join(CUDA_HOME, "include")
# Initialize CUDA
checkCudaErrors(cudart.cudaFree(0))
major = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
minor = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
major = checkCudaErrors(
cudart.cudaDeviceGetAttribute(
cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID
)
)
minor = checkCudaErrors(
cudart.cudaDeviceGetAttribute(
cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID
)
)
_, nvrtc_minor = checkCudaErrors(nvrtc.nvrtcVersion())
use_cubin = (nvrtc_minor >= 1)
prefix = 'sm' if use_cubin else 'compute'
arch_arg = bytes(f'--gpu-architecture={prefix}_{major}{minor}', 'ascii')
use_cubin = nvrtc_minor >= 1
prefix = "sm" if use_cubin else "compute"
arch_arg = bytes(f"--gpu-architecture={prefix}_{major}{minor}", "ascii")
try:
opts = [b'--fmad=true', arch_arg, '--include-path={}'.format(include_dirs).encode('UTF-8'),
b'--std=c++11', b'-default-device']
opts = [
b"--fmad=true",
arch_arg,
"--include-path={}".format(include_dirs).encode("UTF-8"),
b"--std=c++11",
b"-default-device",
]
checkCudaErrors(nvrtc.nvrtcCompileProgram(prog, len(opts), opts))
except RuntimeError as err:
logSize = checkCudaErrors(nvrtc.nvrtcGetProgramLogSize(prog))
log = b' ' * logSize
log = b" " * logSize
checkCudaErrors(nvrtc.nvrtcGetProgramLog(prog, log))
print(log.decode())
print(err)
@@ -88,11 +120,11 @@ class KernelHelper:
if use_cubin:
dataSize = checkCudaErrors(nvrtc.nvrtcGetCUBINSize(prog))
data = b' ' * dataSize
data = b" " * dataSize
checkCudaErrors(nvrtc.nvrtcGetCUBIN(prog, data))
else:
dataSize = checkCudaErrors(nvrtc.nvrtcGetPTXSize(prog))
data = b' ' * dataSize
data = b" " * dataSize
checkCudaErrors(nvrtc.nvrtcGetPTX(prog, data))
self.module = checkCudaErrors(cuda.cuModuleLoadData(np.char.array(data)))
@@ -100,8 +132,9 @@ class KernelHelper:
def getFunction(self, name):
return checkCudaErrors(cuda.cuModuleGetFunction(self.module, name))
class CudaCtxManager(trt.IPluginResource):
def __init__(self, device = None):
def __init__(self, device=None):
trt.IPluginResource.__init__(self)
self.device = device
self.cuda_ctx = None