diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh index 113810665..9c3e9bfe8 100755 --- a/ci/test_wheel.sh +++ b/ci/test_wheel.sh @@ -13,6 +13,11 @@ python -m pip install \ "cuda-python==${CUDA_VER_MAJOR_MINOR%.*}.*" \ pytest +rapids-logger "Install wheel" +package=$(realpath wheel/numba_cuda*.whl) +echo "Wheel path: $package" +python -m pip install $package + GET_TEST_BINARY_DIR=" import numba_cuda @@ -34,11 +39,6 @@ else fi -rapids-logger "Install wheel" -package=$(realpath wheel/numba_cuda*.whl) -echo "Wheel path: $package" -python -m pip install $package - rapids-logger "Check GPU usage" nvidia-smi diff --git a/ci/test_wheel_deps_wheels.sh b/ci/test_wheel_deps_wheels.sh index 2518e7b07..59d6b8536 100755 --- a/ci/test_wheel_deps_wheels.sh +++ b/ci/test_wheel_deps_wheels.sh @@ -23,6 +23,12 @@ python -m pip install \ pytest +rapids-logger "Install wheel" +package=$(realpath wheel/numba_cuda*.whl) +echo "Package path: $package" +python -m pip install $package + + rapids-logger "Build tests" PY_SCRIPT=" import numba_cuda @@ -30,16 +36,11 @@ root = numba_cuda.__file__.rstrip('__init__.py') test_dir = root + \"numba/cuda/tests/test_binary_generation/\" print(test_dir) " - NUMBA_CUDA_TEST_BIN_DIR=$(python -c "$PY_SCRIPT") pushd $NUMBA_CUDA_TEST_BIN_DIR make popd -rapids-logger "Install wheel" -package=$(realpath wheel/numba_cuda*.whl) -echo "Package path: $package" -python -m pip install $package rapids-logger "Check GPU usage" nvidia-smi diff --git a/ci/test_wheel_pynvjitlink.sh b/ci/test_wheel_pynvjitlink.sh index ab586b6b1..961b70015 100755 --- a/ci/test_wheel_pynvjitlink.sh +++ b/ci/test_wheel_pynvjitlink.sh @@ -14,6 +14,12 @@ python -m pip install \ rapids-logger "Install pynvjitlink" python -m pip install pynvjitlink-cu12 + +rapids-logger "Install wheel" +package=$(realpath wheel/numba_cuda*.whl) +echo "Package path: $package" +python -m pip install $package + rapids-logger "Build tests" PY_SCRIPT=" import numba_cuda @@ -27,10 +33,6 @@ pushd $NUMBA_CUDA_TEST_BIN_DIR make popd -rapids-logger "Install wheel" -package=$(realpath wheel/numba_cuda*.whl) -echo "Package path: $package" -python -m pip install $package rapids-logger "Check GPU usage" nvidia-smi diff --git a/numba_cuda/numba/cuda/codegen.py b/numba_cuda/numba/cuda/codegen.py index 768d87c77..4030c6452 100644 --- a/numba_cuda/numba/cuda/codegen.py +++ b/numba_cuda/numba/cuda/codegen.py @@ -5,6 +5,7 @@ from .cudadrv import devices, driver, nvvm, runtime from numba.cuda.cudadrv.libs import get_cudalib from numba.cuda.cudadrv.linkable_code import LinkableCode +from numba.cuda.runtime.nrt import NRT_LIBRARY import os import subprocess @@ -362,9 +363,17 @@ def _reduce_states(self): but loaded functions are discarded. They are recreated when needed after deserialization. """ + nrt = False if self._linking_files: - msg = "Cannot pickle CUDACodeLibrary with linking files" - raise RuntimeError(msg) + if ( + len(self._linking_files) == 1 + and NRT_LIBRARY in self._linking_files + ): + nrt = True + else: + msg = "Cannot pickle CUDACodeLibrary with linking files" + raise RuntimeError(msg) + if not self._finalized: raise RuntimeError("Cannot pickle unfinalized CUDACodeLibrary") return dict( @@ -378,6 +387,7 @@ def _reduce_states(self): max_registers=self._max_registers, nvvm_options=self._nvvm_options, needs_cudadevrt=self.needs_cudadevrt, + nrt=nrt, ) @classmethod @@ -393,6 +403,7 @@ def _rebuild( max_registers, nvvm_options, needs_cudadevrt, + nrt, ): """ Rebuild an instance. @@ -409,6 +420,8 @@ def _rebuild( instance.needs_cudadevrt = needs_cudadevrt instance._finalized = True + if nrt: + instance._linking_files = {NRT_LIBRARY} return instance diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py index 452239d71..a8a4140de 100644 --- a/numba_cuda/numba/cuda/cudadrv/driver.py +++ b/numba_cuda/numba/cuda/cudadrv/driver.py @@ -49,7 +49,7 @@ from .drvapi import cu_occupancy_b2d_size, cu_stream_callback_pyobj, cu_uuid from .mappings import FILE_EXTENSION_MAP from .linkable_code import LinkableCode, LTOIR, Fatbin, Object -from numba.cuda.utils import _readenv +from numba.cuda.utils import _readenv, cached_file_read from numba.cuda.cudadrv import enums, drvapi, nvrtc try: @@ -2806,8 +2806,7 @@ def add_file(self, path, kind): """Add code from a file to the link""" def add_cu_file(self, path): - with open(path, "rb") as f: - cu = f.read() + cu = cached_file_read(path, how="rb") self.add_cu(cu, os.path.basename(path)) def add_file_guess_ext(self, path_or_code, ignore_nonlto=False): @@ -2963,8 +2962,7 @@ def add_file(self, path, kind): raise ImportError(_MVC_ERROR_MESSAGE) from err try: - with open(path, "rb") as f: - data = f.read() + data = cached_file_read(path, how="rb") except FileNotFoundError: raise LinkerError(f"{path} not found") @@ -3338,8 +3336,7 @@ def add_object(self, obj, name=""): def add_file(self, path, kind): try: - with open(path, "rb") as f: - data = f.read() + data = cached_file_read(path, "rb") except FileNotFoundError: raise LinkerError(f"{path} not found") diff --git a/numba_cuda/numba/cuda/cudadrv/linkable_code.py b/numba_cuda/numba/cuda/cudadrv/linkable_code.py index 244a51741..a9d09f1d0 100644 --- a/numba_cuda/numba/cuda/cudadrv/linkable_code.py +++ b/numba_cuda/numba/cuda/cudadrv/linkable_code.py @@ -16,16 +16,24 @@ class LinkableCode: :param teardown_callback: A function called just prior to the unloading of a module that has this code object linked into it. + :param nrt: If True, assume this object contains NRT function calls and + add NRT source code to the final link. """ def __init__( - self, data, name=None, setup_callback=None, teardown_callback=None + self, + data, + name=None, + setup_callback=None, + teardown_callback=None, + nrt=False, ): if setup_callback and not callable(setup_callback): raise TypeError("setup_callback must be callable") if teardown_callback and not callable(teardown_callback): raise TypeError("teardown_callback must be callable") + self.nrt = nrt self._name = name self._data = data self.setup_callback = setup_callback diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py index 72344aa22..6e40ea5eb 100644 --- a/numba_cuda/numba/cuda/dispatcher.py +++ b/numba_cuda/numba/cuda/dispatcher.py @@ -1,6 +1,5 @@ import numpy as np import os -import re import sys import ctypes import functools @@ -21,7 +20,9 @@ kernel_fixup, ExternFunction, ) +import re from numba.cuda.cudadrv import driver +from numba.cuda.cudadrv.linkable_code import LinkableCode from numba.cuda.cudadrv.devices import get_context from numba.cuda.descriptor import cuda_target from numba.cuda.errors import ( @@ -29,7 +30,7 @@ normalize_kernel_dimensions, ) from numba.cuda import types as cuda_types -from numba.cuda.runtime.nrt import rtsys +from numba.cuda.runtime.nrt import rtsys, NRT_LIBRARY from numba.cuda.locks import module_init_lock from numba import cuda @@ -262,6 +263,13 @@ def link_to_library_functions( self.reload_init = [] def maybe_link_nrt(self, link, tgt_ctx, asm): + """ + Add the NRT source code to the link if the neccesary conditions are met. + NRT must be enabled for the CUDATargetContext, and either NRT functions + must be detected in the kernel asm or an NRT enabled LinkableCode object + must be passed. + """ + if not tgt_ctx.enable_nrt: return @@ -271,13 +279,19 @@ def maybe_link_nrt(self, link, tgt_ctx, asm): + all_nrt + r")\s*\([^)]*\)\s*;" ) - + link_nrt = False nrt_in_asm = re.findall(pattern, asm) - - basedir = os.path.dirname(os.path.abspath(__file__)) - if nrt_in_asm: - nrt_path = os.path.join(basedir, "runtime", "nrt.cu") - link.append(nrt_path) + if len(nrt_in_asm) > 0: + link_nrt = True + if not link_nrt: + for file in link: + if isinstance(file, LinkableCode): + if file.nrt: + link_nrt = True + break + + if link_nrt: + link.append(NRT_LIBRARY) @property def library(self): diff --git a/numba_cuda/numba/cuda/runtime/nrt.cu b/numba_cuda/numba/cuda/runtime/nrt.cu index a318dd4dd..0e3e44071 100644 --- a/numba_cuda/numba/cuda/runtime/nrt.cu +++ b/numba_cuda/numba/cuda/runtime/nrt.cu @@ -4,30 +4,14 @@ #include #include "memsys.cuh" +#include "nrt.cuh" -typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info); -typedef void (*NRT_dealloc_func)(void* ptr, void* dealloc_info); - -typedef struct MemInfo NRT_MemInfo; - -extern "C" { -struct MemInfo { - cuda::atomic refct; - NRT_dtor_function dtor; - void* dtor_info; - void* data; - size_t size; -}; -} extern "C" __global__ void NRT_MemSys_set(NRT_MemSys *memsys_ptr) { TheMSys = memsys_ptr; } -static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi); -static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out); -extern "C" __device__ void* NRT_Allocate_External(size_t size); extern "C" __device__ void* NRT_Allocate(size_t size) { @@ -177,6 +161,7 @@ extern "C" __device__ void NRT_decref(NRT_MemInfo* mi) } } + #endif extern "C" __device__ void NRT_incref(NRT_MemInfo* mi) diff --git a/numba_cuda/numba/cuda/runtime/nrt.cuh b/numba_cuda/numba/cuda/runtime/nrt.cuh new file mode 100644 index 000000000..8944733f1 --- /dev/null +++ b/numba_cuda/numba/cuda/runtime/nrt.cuh @@ -0,0 +1,41 @@ +#include + +typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info); +typedef void (*NRT_dealloc_func)(void* ptr, void* dealloc_info); + +extern "C" +struct MemInfo { + cuda::atomic refct; + NRT_dtor_function dtor; + void* dtor_info; + void* data; + size_t size; +}; +typedef struct MemInfo NRT_MemInfo; + +extern "C" __device__ void* NRT_Allocate(size_t size); +extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi, + void* data, + size_t size, + NRT_dtor_function dtor, + void* dtor_info); +static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi); +static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out); +extern "C" __device__ void* NRT_Allocate_External(size_t size); +extern "C" __device__ void NRT_decref(NRT_MemInfo* mi); +extern "C" __device__ void NRT_incref(NRT_MemInfo* mi); +extern "C" __device__ void* NRT_Allocate_External(size_t size); +static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out); +static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi); +extern "C" __device__ NRT_MemInfo *NRT_MemInfo_alloc_aligned(size_t size, unsigned align); +extern "C" __device__ void* NRT_MemInfo_data_fast(NRT_MemInfo *mi); +extern "C" __device__ void NRT_MemInfo_call_dtor(NRT_MemInfo* mi); +extern "C" __device__ void NRT_MemInfo_destroy(NRT_MemInfo* mi); +extern "C" __device__ void NRT_dealloc(NRT_MemInfo* mi); +extern "C" __device__ void NRT_Free(void* ptr); +extern "C" __device__ NRT_MemInfo* NRT_MemInfo_new(void* data, size_t size, NRT_dtor_function dtor, void* dtor_info); +extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi, + void* data, + size_t size, + NRT_dtor_function dtor, + void* dtor_info); diff --git a/numba_cuda/numba/cuda/runtime/nrt.py b/numba_cuda/numba/cuda/runtime/nrt.py index 8b3be0c8e..c978ce023 100644 --- a/numba_cuda/numba/cuda/runtime/nrt.py +++ b/numba_cuda/numba/cuda/runtime/nrt.py @@ -13,7 +13,8 @@ ) from numba.cuda.cudadrv import devices from numba.cuda.api import get_current_device -from numba.cuda.utils import _readenv +from numba.cuda.utils import _readenv, cached_file_read +from numba.cuda.cudadrv.linkable_code import CUSource # Check environment variable or config for NRT statistics enablement @@ -32,6 +33,11 @@ config.CUDA_ENABLE_NRT = ENABLE_NRT +def get_include(): + """Return the include path for the NRT header""" + return os.path.dirname(os.path.abspath(__file__)) + + # Protect method to ensure NRT memory allocation and initialization def _alloc_init_guard(method): """ @@ -340,3 +346,9 @@ def print_memsys(self, stream=None): # Create an instance of the runtime rtsys = _Runtime() + + +basedir = os.path.dirname(os.path.abspath(__file__)) +nrt_path = os.path.join(basedir, "nrt.cu") +nrt_src = cached_file_read(nrt_path) +NRT_LIBRARY = CUSource(nrt_src, name="nrt.cu", nrt=True) diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py index a621fe625..9bd01b37d 100644 --- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py +++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py @@ -4,11 +4,86 @@ import numpy as np import unittest from numba.cuda.testing import CUDATestCase - from numba.tests.support import run_in_subprocess, override_config - +from numba.cuda import get_current_device +from numba.cuda.cudadrv.nvrtc import compile +from numba import types +from numba.cuda.cudadecl import registry as cuda_decl_registry +from numba.core.typing import signature +from numba.cuda.cudaimpl import lower as cuda_lower from numba import cuda -from numba.cuda.runtime.nrt import rtsys +from numba.cuda.runtime.nrt import rtsys, get_include +from numba.core.typing.templates import AbstractTemplate +from numba.cuda.cudadrv.linkable_code import ( + CUSource, + PTXSource, + Fatbin, + Cubin, + Archive, + Object, +) + + +TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR") + +if TEST_BIN_DIR: + + def make_linkable_code(name, kind, mode): + path = os.path.join(TEST_BIN_DIR, name) + with open(path, mode) as f: + contents = f.read() + return kind(contents, nrt=True) + + nrt_extern_a = make_linkable_code("nrt_extern.a", Archive, "rb") + nrt_extern_cubin = make_linkable_code("nrt_extern.cubin", Cubin, "rb") + nrt_extern_cu = make_linkable_code( + "nrt_extern.cu", + CUSource, + "rb", + ) + nrt_extern_fatbin = make_linkable_code("nrt_extern.fatbin", Fatbin, "rb") + nrt_extern_fatbin_multi = make_linkable_code( + "nrt_extern_multi.fatbin", Fatbin, "rb" + ) + nrt_extern_o = make_linkable_code("nrt_extern.o", Object, "rb") + nrt_extern_ptx = make_linkable_code("nrt_extern.ptx", PTXSource, "rb") + + +def allocate_deallocate_handle(): + """ + Handle to call NRT_Allocate and NRT_Free + """ + pass + + +@cuda_decl_registry.register_global(allocate_deallocate_handle) +class AllocateShimImpl(AbstractTemplate): + def generic(self, args, kws): + return signature(types.void) + + +device_fun_shim = cuda.declare_device( + "device_allocate_deallocate", types.int32() +) + + +# wrapper to turn the above into a python callable +def call_device_fun_shim(): + return device_fun_shim() + + +@cuda_lower(allocate_deallocate_handle) +def allocate_deallocate_impl(context, builder, sig, args): + sig_ = types.int32() + # call the external function, passing the pointer + result = context.compile_internal( + builder, + call_device_fun_shim, + sig_, + (), + ) + + return result class TestNrtBasic(CUDATestCase): @@ -77,6 +152,50 @@ def g(out_ary): self.assertEqual(out_ary[0], 1) +class TestNrtLinking(CUDATestCase): + def run(self, result=None): + with override_config("CUDA_ENABLE_NRT", True): + super(TestNrtLinking, self).run(result) + + def test_nrt_detect_linked_ptx_file(self): + src = f"#include <{get_include()}/nrt.cuh>" + src += """ + extern "C" __device__ int device_allocate_deallocate(int* nb_retval){ + auto ptr = NRT_Allocate(1); + NRT_Free(ptr); + return 0; + } + """ + cc = get_current_device().compute_capability + ptx, _ = compile(src, "external_nrt.cu", cc) + + @cuda.jit(link=[PTXSource(ptx.encode(), nrt=True)]) + def kernel(): + allocate_deallocate_handle() + + kernel[1, 1]() + + @unittest.skipIf(not TEST_BIN_DIR, "necessary binaries not generated.") + def test_nrt_detect_linkable_code(self): + codes = ( + nrt_extern_a, + nrt_extern_cubin, + nrt_extern_cu, + nrt_extern_fatbin, + nrt_extern_fatbin_multi, + nrt_extern_o, + nrt_extern_ptx, + ) + for code in codes: + with self.subTest(code=code): + + @cuda.jit(link=[code]) + def kernel(): + allocate_deallocate_handle() + + kernel[1, 1]() + + class TestNrtStatistics(CUDATestCase): def setUp(self): self._stream = cuda.default_stream() diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile index 145f8b024..7d36e5cba 100644 --- a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile +++ b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile @@ -40,6 +40,8 @@ LTOIR_FLAGS := $(LTOIR_GENCODE) -dc OUTPUT_DIR := ./ +NRT_INCLUDE_DIR := $(shell python -c "from numba.cuda.runtime.nrt import get_include; print(get_include())") + all: @echo "GPU CC: $(GPU_CC)" @echo "Alternative CC: $(ALT_CC)" @@ -52,7 +54,16 @@ all: nvcc $(NVCC_FLAGS) $(OBJECT_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.o test_device_functions.cu nvcc $(NVCC_FLAGS) $(LIBRARY_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.a test_device_functions.cu + nvcc $(NVCC_FLAGS) $(CUBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.cubin nrt_extern.cu -I$(NRT_INCLUDE_DIR) + nvcc $(NVCC_FLAGS) $(FATBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.fatbin nrt_extern.cu -I$(NRT_INCLUDE_DIR) + nvcc $(NVCC_FLAGS) $(MULTI_FATBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern_multi.fatbin nrt_extern.cu -I$(NRT_INCLUDE_DIR) + nvcc $(NVCC_FLAGS) $(PTX_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.ptx nrt_extern.cu -I$(NRT_INCLUDE_DIR) + nvcc $(NVCC_FLAGS) $(OBJECT_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.o nrt_extern.cu -I$(NRT_INCLUDE_DIR) + nvcc $(NVCC_FLAGS) $(LIBRARY_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.a nrt_extern.cu -I$(NRT_INCLUDE_DIR) + # Generate LTO-IR wrapped in a fatbin nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.ltoir.o test_device_functions.cu + nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.ltoir.o nrt_extern.cu -I$(NRT_INCLUDE_DIR) # Generate LTO-IR in a "raw" LTO-IR container python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/test_device_functions.ltoir test_device_functions.cu + python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/nrt_extern.ltoir nrt_extern.cu --nrt diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py index b4d32a34c..d5e50e5c8 100644 --- a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py +++ b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py @@ -7,6 +7,7 @@ import sys from cuda import nvrtc +from numba.cuda.runtime.nrt import get_include # Magic number found at the start of an LTO-IR file LTOIR_MAGIC = 0x7F4E43ED @@ -88,7 +89,9 @@ def get_ltoir(source, name, arch): nvrtc.nvrtcCreateProgram(source.encode(), name.encode(), 0, [], []) ) - cuda_include_flags = determine_include_flags() + cuda_include_flags = determine_include_flags() + ( + [f"-I{get_include()}"] if args.nrt else [] + ) if cuda_include_flags is None: print("Error determining CUDA include flags. Exiting.", file=sys.stderr) sys.exit(1) @@ -160,7 +163,7 @@ def main(sourcepath, outputpath, arch): help="compute arch to target (e.g. sm_87). Defaults to sm_50.", default="sm_50", ) - + parser.add_argument("--nrt", action="store_true") args = parser.parse_args() outputpath = args.output diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu new file mode 100644 index 000000000..c9f73908e --- /dev/null +++ b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu @@ -0,0 +1,7 @@ +#include + +extern "C" __device__ int device_allocate_deallocate(int* nb_retval){ + auto ptr = NRT_Allocate(1); + NRT_Free(ptr); + return 0; +} diff --git a/numba_cuda/numba/cuda/utils.py b/numba_cuda/numba/cuda/utils.py index a66989135..dc99b71bd 100644 --- a/numba_cuda/numba/cuda/utils.py +++ b/numba_cuda/numba/cuda/utils.py @@ -1,6 +1,7 @@ import os import warnings import traceback +import functools def _readenv(name, ctor, default): @@ -20,3 +21,9 @@ def _readenv(name, ctor, default): RuntimeWarning, ) return default + + +@functools.lru_cache(maxsize=None) +def cached_file_read(filepath, how="r"): + with open(filepath, how) as f: + return f.read()