diff --git a/python/cudf/cudf/tests/test_mvc.py b/python/cudf/cudf/tests/test_mvc.py new file mode 100644 index 00000000000..7dd25ebc500 --- /dev/null +++ b/python/cudf/cudf/tests/test_mvc.py @@ -0,0 +1,99 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. +import subprocess +import sys + +import pytest + +IS_CUDA_11 = False +IS_CUDA_12 = False +try: + from ptxcompiler.patch import safe_get_versions +except ModuleNotFoundError: + from cudf.utils._ptxcompiler import safe_get_versions + +# do not test cuda 12 if pynvjitlink isn't present +HAVE_PYNVJITLINK = False +try: + import pynvjitlink # noqa: F401 + + HAVE_PYNVJITLINK = True +except ModuleNotFoundError: + pass + + +versions = safe_get_versions() +driver_version, runtime_version = versions + +if (11, 0) <= driver_version < (12, 0): + IS_CUDA_11 = True +if (12, 0) <= driver_version < (13, 0): + IS_CUDA_12 = True + + +TEST_BODY = """ +@numba.cuda.jit +def test_kernel(x): + id = numba.cuda.grid(1) + if id < len(x): + x[id] += 1 + +s = cudf.Series([1, 2, 3]) +with _CUDFNumbaConfig(): + test_kernel.forall(len(s))(s) +""" + +CUDA_11_TEST = ( + """ +import numba.cuda +import cudf +from cudf.utils._numba import _CUDFNumbaConfig, patch_numba_linker_cuda_11 + + +patch_numba_linker_cuda_11() +""" + + TEST_BODY +) + + +CUDA_12_TEST = ( + """ +import numba.cuda +import cudf +from cudf.utils._numba import _CUDFNumbaConfig +from pynvjitlink.patch import ( + patch_numba_linker as patch_numba_linker_pynvjitlink, +) + +patch_numba_linker_pynvjitlink() +""" + + TEST_BODY +) + + +@pytest.mark.parametrize( + "test", + [ + pytest.param( + CUDA_11_TEST, + marks=pytest.mark.skipif( + not IS_CUDA_11, + reason="Minor Version Compatibility test for CUDA 11", + ), + ), + pytest.param( + CUDA_12_TEST, + marks=pytest.mark.skipif( + not IS_CUDA_12 or not HAVE_PYNVJITLINK, + reason="Minor Version Compatibility test for CUDA 12", + ), + ), + ], +) +def test_numba_mvc(test): + cp = subprocess.run( + [sys.executable, "-c", test], + capture_output=True, + cwd="/", + ) + + assert cp.returncode == 0 diff --git a/python/cudf/cudf/tests/test_numba_import.py b/python/cudf/cudf/tests/test_numba_import.py deleted file mode 100644 index 238a32a94fa..00000000000 --- a/python/cudf/cudf/tests/test_numba_import.py +++ /dev/null @@ -1,48 +0,0 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. -import subprocess -import sys - -import pytest - -IS_CUDA_11 = False -try: - from ptxcompiler.patch import NO_DRIVER, safe_get_versions - - versions = safe_get_versions() - if versions != NO_DRIVER: - driver_version, runtime_version = versions - if driver_version < (12, 0): - IS_CUDA_11 = True -except ModuleNotFoundError: - pass - -TEST_NUMBA_MVC_ENABLED = """ -import numba.cuda -import cudf -from cudf.utils._numba import _CUDFNumbaConfig, _patch_numba_mvc - - -_patch_numba_mvc() - -@numba.cuda.jit -def test_kernel(x): - id = numba.cuda.grid(1) - if id < len(x): - x[id] += 1 - -s = cudf.Series([1, 2, 3]) -with _CUDFNumbaConfig(): - test_kernel.forall(len(s))(s) -""" - - -@pytest.mark.skipif( - not IS_CUDA_11, reason="Minor Version Compatibility test for CUDA 11" -) -def test_numba_mvc_enabled_cuda_11(): - cp = subprocess.run( - [sys.executable, "-c", TEST_NUMBA_MVC_ENABLED], - capture_output=True, - cwd="/", - ) - assert cp.returncode == 0 diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py index 09afb5680bd..bc0d6f37d89 100644 --- a/python/cudf/cudf/utils/_numba.py +++ b/python/cudf/cudf/utils/_numba.py @@ -7,6 +7,19 @@ from numba import config as numba_config +try: + from pynvjitlink.patch import ( + patch_numba_linker as patch_numba_linker_pynvjitlink, + ) +except ImportError: + + def patch_numba_linker_pynvjitlink(): + warnings.warn( + "CUDA Toolkit is newer than CUDA driver. " + "Numba features will not work in this configuration. " + ) + + CC_60_PTX_FILE = os.path.join( os.path.dirname(__file__), "../core/udf/shim_60.ptx" ) @@ -65,7 +78,7 @@ def _get_ptx_file(path, prefix): return regular_result[1] -def _patch_numba_mvc(): +def patch_numba_linker_cuda_11(): # Enable the config option for minor version compatibility numba_config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY = 1 @@ -106,29 +119,19 @@ def _setup_numba(): versions = safe_get_versions() if versions != NO_DRIVER: driver_version, runtime_version = versions - if driver_version >= (12, 0) and runtime_version > driver_version: - warnings.warn( - f"Using CUDA toolkit version {runtime_version} with CUDA " - f"driver version {driver_version} requires minor version " - "compatibility, which is not yet supported for CUDA " - "driver versions 12.0 and above. It is likely that many " - "cuDF operations will not work in this state. Please " - f"install CUDA toolkit version {driver_version} to " - "continue using cuDF." - ) - else: - # Support MVC for all CUDA versions in the 11.x range - ptx_toolkit_version = _get_cuda_version_from_ptx_file( - CC_60_PTX_FILE - ) - # Numba thinks cubinlinker is only needed if the driver is older - # than the CUDA runtime, but when PTX files are present, it might - # also need to patch because those PTX files may be compiled by - # a CUDA version that is newer than the driver as well - if (driver_version < ptx_toolkit_version) or ( - driver_version < runtime_version - ): - _patch_numba_mvc() + ptx_toolkit_version = _get_cuda_version_from_ptx_file(CC_60_PTX_FILE) + + # MVC is required whenever any PTX is newer than the driver + # This could be the shipped PTX file or the PTX emitted by + # the version of NVVM on the user system, the latter aligning + # with the runtime version + if (driver_version < ptx_toolkit_version) or ( + driver_version < runtime_version + ): + if driver_version < (12, 0): + patch_numba_linker_cuda_11() + else: + patch_numba_linker_pynvjitlink() def _get_cuda_version_from_ptx_file(path): @@ -171,6 +174,8 @@ def _get_cuda_version_from_ptx_file(path): "7.8": (11, 8), "8.0": (12, 0), "8.1": (12, 1), + "8.2": (12, 2), + "8.3": (12, 3), } cuda_ver = ver_map.get(version)