wehub-resource-sync
136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-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.
|
|
#
|
|
|
|
from cuda.bindings import driver as cuda, runtime as cudart, nvrtc
|
|
|
|
import numpy as np
|
|
import os
|
|
from common_runtime import cuda_call, create_cuda_context, cuda_init, cuda_get_device, cuda_memcpy_htod
|
|
import argparse
|
|
import threading
|
|
|
|
import tensorrt as trt
|
|
import cupy as cp
|
|
|
|
|
|
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",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def volume(d):
|
|
return np.prod(d)
|
|
|
|
|
|
def getComputeCapacity(devID):
|
|
major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
|
|
minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
|
|
return (major, minor)
|
|
|
|
|
|
# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/common.py
|
|
class KernelHelper:
|
|
def __init__(self, code, devID):
|
|
prog = cuda_call(
|
|
nvrtc.nvrtcCreateProgram(str.encode(code), b"sourceCode.cu", 0, [], [])
|
|
)
|
|
cuda_root = None
|
|
for env_name in ("CUDA_PATH", "CUDA_HOME"):
|
|
cand = os.getenv(env_name)
|
|
if cand and os.path.isfile(os.path.join(cand, "include", "cuda_fp16.h")):
|
|
cuda_root = cand
|
|
break
|
|
if cuda_root is None:
|
|
raise RuntimeError(
|
|
"Neither CUDA_PATH nor CUDA_HOME points at a CUDA install containing include/cuda_fp16.h"
|
|
)
|
|
include_dirs = os.path.join(cuda_root, "include")
|
|
|
|
# Initialize CUDA
|
|
cuda_call(cudart.cudaFree(0))
|
|
|
|
major, minor = getComputeCapacity(devID)
|
|
_, nvrtc_minor = cuda_call(nvrtc.nvrtcVersion())
|
|
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,
|
|
('-I' + include_dirs).encode("UTF-8"),
|
|
b"--std=c++11",
|
|
b"-default-device",
|
|
]
|
|
cuda_call(nvrtc.nvrtcCompileProgram(prog, len(opts), opts))
|
|
except RuntimeError as err:
|
|
logSize = cuda_call(nvrtc.nvrtcGetProgramLogSize(prog))
|
|
log = b" " * logSize
|
|
cuda_call(nvrtc.nvrtcGetProgramLog(prog, log))
|
|
print(log.decode())
|
|
print(err)
|
|
exit(-1)
|
|
|
|
if use_cubin:
|
|
dataSize = cuda_call(nvrtc.nvrtcGetCUBINSize(prog))
|
|
data = b" " * dataSize
|
|
cuda_call(nvrtc.nvrtcGetCUBIN(prog, data))
|
|
else:
|
|
dataSize = cuda_call(nvrtc.nvrtcGetPTXSize(prog))
|
|
data = b" " * dataSize
|
|
cuda_call(nvrtc.nvrtcGetPTX(prog, data))
|
|
|
|
self.module = cuda_call(cuda.cuModuleLoadData(np.char.array(data)))
|
|
|
|
def getFunction(self, name):
|
|
return cuda_call(cuda.cuModuleGetFunction(self.module, name))
|
|
|
|
|
|
class CudaCtxManager(trt.IPluginResource):
|
|
def __init__(self, device=None):
|
|
trt.IPluginResource.__init__(self)
|
|
self.device = device
|
|
self.cuda_ctx = None
|
|
|
|
def clone(self):
|
|
cloned = CudaCtxManager()
|
|
cloned.__dict__.update(self.__dict__)
|
|
# Delay the CUDA ctx creation until clone()
|
|
# since only a cloned resource is registered by TRT
|
|
cloned.cuda_ctx = create_cuda_context(self.device)
|
|
return cloned
|
|
|
|
def release(self):
|
|
cuda_call(cuda.cuCtxDestroy(self.cuda_ctx))
|
|
|
|
class UnownedMemory:
|
|
def __init__(self, ptr, shape, dtype):
|
|
mem = cp.cuda.UnownedMemory(ptr, volume(shape) * cp.dtype(dtype).itemsize, self)
|
|
cupy_ptr = cp.cuda.MemoryPointer(mem, 0)
|
|
self.d = cp.ndarray(shape, dtype=dtype, memptr=cupy_ptr)
|