chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,135 @@
+#
+# 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)