From b8c03c4fc05bb16e083a4a33567c760776791e50 Mon Sep 17 00:00:00 2001 From: isVoid Date: Wed, 18 Dec 2024 01:44:18 -0800 Subject: [PATCH 1/5] initial --- numba_cuda/numba/cuda/cudadrv/driver.py | 30 +++++---- .../cuda/tests/cudadrv/test_nvjitlink.py | 64 ++++++++++++++++++- 2 files changed, 77 insertions(+), 17 deletions(-) diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py index 465d33d..cd5bd85 100644 --- a/numba_cuda/numba/cuda/cudadrv/driver.py +++ b/numba_cuda/numba/cuda/cudadrv/driver.py @@ -42,6 +42,11 @@ from .linkable_code import LinkableCode, LTOIR, Fatbin, Object from numba.cuda.cudadrv import enums, drvapi, nvrtc +try: + from pynvjitlink.api import NvJitLinker, NvJitLinkError +except ImportError: + NvJitLinker, NvJitLinkError = None, None + USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING if USE_NV_BINDING: @@ -92,20 +97,6 @@ def _readenv(name, ctor, default): if not hasattr(config, "CUDA_ENABLE_PYNVJITLINK"): config.CUDA_ENABLE_PYNVJITLINK = ENABLE_PYNVJITLINK -if ENABLE_PYNVJITLINK: - try: - from pynvjitlink.api import NvJitLinker, NvJitLinkError - except ImportError: - raise ImportError( - "Using pynvjitlink requires the pynvjitlink package to be available" - ) - - if config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY: - raise ValueError( - "Can't set CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY and " - "CUDA_ENABLE_PYNVJITLINK at the same time" - ) - def make_logger(): logger = logging.getLogger(__name__) @@ -3061,6 +3052,17 @@ def __init__( lto=False, additional_flags=None, ): + if NvJitLinker is None: + raise ImportError( + "Using pynvjitlink requires the pynvjitlink package to be " + "available" + ) + + if config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY: + raise ValueError( + "Can't set CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY and " + "CUDA_ENABLE_PYNVJITLINK at the same time" + ) if cc is None: raise RuntimeError("PyNvJitLinker requires CC to be specified") diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py index 044895c..dd6ba61 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py @@ -2,6 +2,11 @@ from numba.cuda.testing import skip_on_cudasim from numba.cuda.testing import CUDATestCase from numba.cuda.cudadrv.driver import PyNvJitLinker +from numba.cuda import get_current_device + +from numba import cuda +from numba import config +from numba.tests.support import run_in_subprocess, override_config import itertools import os @@ -9,9 +14,6 @@ import contextlib import warnings -from numba.cuda import get_current_device -from numba import cuda -from numba import config TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR") if TEST_BIN_DIR: @@ -251,5 +253,61 @@ def kernel(): pass +class TestLinkerUsage(CUDATestCase): + """Test that whether pynvjitlink can be enabled by both environment variable + and modification of config at runtime. + """ + def test_linker_enabled_envvar(self): + # Linkable code is only supported via pynvjitlink + src = """if 1: + import os + from numba import cuda + + TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR") + if TEST_BIN_DIR: + test_device_functions_cubin = os.path.join( + TEST_BIN_DIR, "test_device_functions.cubin" + ) + print(TEST_BIN_DIR) + files = ( + test_device_functions_cubin, + ) + for lto in [True, False]: + for file in files: + sig = "uint32(uint32, uint32)" + add_from_numba = cuda.declare_device("add_from_numba", sig) + + @cuda.jit(link=[file], lto=lto) + def kernel(result): + result[0] = add_from_numba(1, 2) + + result = cuda.device_array(1) + kernel[1, 1](result) + assert result[0] == 3 + """ + env = os.environ.copy() + env['NUMBA_CUDA_ENABLE_PYNVJITLINK'] = "1" + print(env['NUMBA_CUDA_TEST_BIN_DIR']) + run_in_subprocess(src, env=env) + + def test_linker_enabled_config(self): + with override_config("CUDA_ENABLE_PYNVJITLINK", True): + files = ( + test_device_functions_cubin, + ) + for lto in [True, False]: + for file in files: + sig = "uint32(uint32, uint32)" + add_from_numba = cuda.declare_device("add_from_numba", sig) + + @cuda.jit(link=[file], lto=lto) + def kernel(result): + result[0] = add_from_numba(1, 2) + + result = cuda.device_array(1) + kernel[1, 1](result) + assert result[0] == 3 + + if __name__ == "__main__": unittest.main() From 38238bc6bba2288064a6c30fb5e368b085554921 Mon Sep 17 00:00:00 2001 From: isVoid Date: Wed, 18 Dec 2024 02:01:53 -0800 Subject: [PATCH 2/5] enable test iff pynvjitlink is installed --- numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py index dd6ba61..83b4ed8 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py @@ -8,6 +8,13 @@ from numba import config from numba.tests.support import run_in_subprocess, override_config +try: + import pynvjitlink # noqa: F401 + PYNVJITLINK_INSTALLED = True +except ImportError: + PYNVJITLINK_INSTALLED = False + + import itertools import os import io @@ -253,6 +260,7 @@ def kernel(): pass +@unittest.skipIf(not PYNVJITLINK_INSTALLED) class TestLinkerUsage(CUDATestCase): """Test that whether pynvjitlink can be enabled by both environment variable and modification of config at runtime. From c9d148b781b9a2d941d4e64e1b8215cbc6d333a2 Mon Sep 17 00:00:00 2001 From: isVoid Date: Wed, 18 Dec 2024 02:30:55 -0800 Subject: [PATCH 3/5] add skipif reason --- numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py index 83b4ed8..5c4b9fb 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py @@ -260,7 +260,9 @@ def kernel(): pass -@unittest.skipIf(not PYNVJITLINK_INSTALLED) +@unittest.skipIf( + not PYNVJITLINK_INSTALLED, reason="Pynvjitlink is not installed" +) class TestLinkerUsage(CUDATestCase): """Test that whether pynvjitlink can be enabled by both environment variable and modification of config at runtime. From 6e1a5740bb4b6fcf151ab3e7a082e55a493aca1e Mon Sep 17 00:00:00 2001 From: isVoid Date: Wed, 18 Dec 2024 06:27:14 -0800 Subject: [PATCH 4/5] test proper exception is raised when nvjitlink is disabled --- .../cuda/tests/cudadrv/test_nvjitlink.py | 76 +++++++++++-------- 1 file changed, 43 insertions(+), 33 deletions(-) diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py index 5c4b9fb..7c72c6a 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py @@ -267,9 +267,8 @@ class TestLinkerUsage(CUDATestCase): """Test that whether pynvjitlink can be enabled by both environment variable and modification of config at runtime. """ - def test_linker_enabled_envvar(self): - # Linkable code is only supported via pynvjitlink - src = """if 1: + + src = """if 1: import os from numba import cuda @@ -278,45 +277,56 @@ def test_linker_enabled_envvar(self): test_device_functions_cubin = os.path.join( TEST_BIN_DIR, "test_device_functions.cubin" ) - print(TEST_BIN_DIR) - files = ( - test_device_functions_cubin, - ) - for lto in [True, False]: - for file in files: - sig = "uint32(uint32, uint32)" - add_from_numba = cuda.declare_device("add_from_numba", sig) - @cuda.jit(link=[file], lto=lto) - def kernel(result): - result[0] = add_from_numba(1, 2) + sig = "uint32(uint32, uint32)" + add_from_numba = cuda.declare_device("add_from_numba", sig) - result = cuda.device_array(1) - kernel[1, 1](result) - assert result[0] == 3 + @cuda.jit(link=[test_device_functions_cubin], lto=True) + def kernel(result): + result[0] = add_from_numba(1, 2) + + result = cuda.device_array(1) + kernel[1, 1](result) + assert result[0] == 3 """ + + def test_linker_enabled_envvar(self): env = os.environ.copy() env['NUMBA_CUDA_ENABLE_PYNVJITLINK'] = "1" - print(env['NUMBA_CUDA_TEST_BIN_DIR']) - run_in_subprocess(src, env=env) + run_in_subprocess(self.src, env=env) + + def test_linker_disabled_envvar(self): + env = os.environ.copy() + env.pop('NUMBA_CUDA_ENABLE_PYNVJITLINK', None) + with self.assertRaisesRegex( + AssertionError, "LTO and additional flags require PyNvJitLinker" + ): + # Actual error raised is `ValueError`, but `run_in_subprocess` + # reraises as AssertionError. + run_in_subprocess(self.src, env=env) + + def _lto_test(self): + sig = "uint32(uint32, uint32)" + add_from_numba = cuda.declare_device("add_from_numba", sig) + + @cuda.jit(link=[test_device_functions_cubin], lto=True) + def kernel(result): + result[0] = add_from_numba(1, 2) + + result = cuda.device_array(1) + kernel[1, 1](result) + assert result[0] == 3 def test_linker_enabled_config(self): with override_config("CUDA_ENABLE_PYNVJITLINK", True): - files = ( - test_device_functions_cubin, - ) - for lto in [True, False]: - for file in files: - sig = "uint32(uint32, uint32)" - add_from_numba = cuda.declare_device("add_from_numba", sig) - - @cuda.jit(link=[file], lto=lto) - def kernel(result): - result[0] = add_from_numba(1, 2) + self._lto_test() - result = cuda.device_array(1) - kernel[1, 1](result) - assert result[0] == 3 + def test_linker_disabled_config(self): + with override_config("CUDA_ENABLE_PYNVJITLINK", False): + with self.assertRaisesRegex( + ValueError, "LTO and additional flags require PyNvJitLinker" + ): + self._lto_test() if __name__ == "__main__": From 2064641f756da58c2c48a0f40652faec931b4f44 Mon Sep 17 00:00:00 2001 From: isVoid Date: Wed, 18 Dec 2024 19:53:18 -0800 Subject: [PATCH 5/5] run all tests in subprocess --- .../cuda/tests/cudadrv/test_nvjitlink.py | 33 ++++++++----------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py index 7c72c6a..4ba6a23 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py @@ -270,7 +270,9 @@ class TestLinkerUsage(CUDATestCase): src = """if 1: import os - from numba import cuda + from numba import cuda, config + + {config} TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR") if TEST_BIN_DIR: @@ -293,7 +295,7 @@ def kernel(result): def test_linker_enabled_envvar(self): env = os.environ.copy() env['NUMBA_CUDA_ENABLE_PYNVJITLINK'] = "1" - run_in_subprocess(self.src, env=env) + run_in_subprocess(self.src.format(config=""), env=env) def test_linker_disabled_envvar(self): env = os.environ.copy() @@ -303,30 +305,23 @@ def test_linker_disabled_envvar(self): ): # Actual error raised is `ValueError`, but `run_in_subprocess` # reraises as AssertionError. - run_in_subprocess(self.src, env=env) - - def _lto_test(self): - sig = "uint32(uint32, uint32)" - add_from_numba = cuda.declare_device("add_from_numba", sig) - - @cuda.jit(link=[test_device_functions_cubin], lto=True) - def kernel(result): - result[0] = add_from_numba(1, 2) - - result = cuda.device_array(1) - kernel[1, 1](result) - assert result[0] == 3 + run_in_subprocess(self.src.format(config=""), env=env) def test_linker_enabled_config(self): - with override_config("CUDA_ENABLE_PYNVJITLINK", True): - self._lto_test() + env = os.environ.copy() + env.pop('NUMBA_CUDA_ENABLE_PYNVJITLINK', None) + run_in_subprocess(self.src.format( + config="config.CUDA_ENABLE_PYNVJITLINK = True"), env=env) def test_linker_disabled_config(self): + env = os.environ.copy() + env.pop('NUMBA_CUDA_ENABLE_PYNVJITLINK', None) with override_config("CUDA_ENABLE_PYNVJITLINK", False): with self.assertRaisesRegex( - ValueError, "LTO and additional flags require PyNvJitLinker" + AssertionError, "LTO and additional flags require PyNvJitLinker" ): - self._lto_test() + run_in_subprocess(self.src.format( + config="config.CUDA_ENABLE_PYNVJITLINK = False"), env=env) if __name__ == "__main__":