cuDF numba cuda 12 updates

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -48,7 +48,7 @@ dependencies:
 - nbsphinx
 - ninja
 - notebook
-- numba>=0.56.4,<0.57
+- numba>=0.57
 - numpy>=1.21,<1.24
 - numpydoc
 - nvcc_linux-64=11.8

conda/recipes/cudf/meta.yaml

@@ -53,7 +53,7 @@ requirements:
     - cython >=0.29,<0.30
     - scikit-build >=0.13.1
     - setuptools
-    - numba >=0.56.4,<0.57
+    - numba >=0.57
     - dlpack >=0.5,<0.6.0a0
     - pyarrow =11
     - libcudf ={{ version }}
@@ -65,7 +65,7 @@ requirements:
     - typing_extensions
     - pandas >=1.3,<1.6.0dev0
     - cupy >=12.0.0
-    - numba >=0.56.4,<0.57
+    - numba >=0.57
     - numpy >=1.21,<1.24  # Temporarily upper bound numpy to avoid overflow deprecations
     - {{ pin_compatible('pyarrow', max_pin='x.x.x') }}
     - libcudf {{ version }}

dependencies.yaml

@@ -350,7 +350,7 @@ dependencies:
         packages:
           - cachetools
           - cuda-python>=11.7.1,<12.0
-          - &numba numba>=0.56.4,<0.57
+          - &numba numba>=0.57
           - nvtx>=0.2.1
           - packaging
           - rmm==23.6.*

python/cudf/cudf/__init__.py

@@ -5,7 +5,7 @@
 validate_setup()
 
 import cupy
-from numba import config as numba_config, cuda
+from numba import config as numba_config
 
 import rmm
 from rmm.allocators.cupy import rmm_cupy_allocator
@@ -80,35 +80,19 @@
     read_text,
 )
 from cudf.options import describe_option, get_option, set_option
+from cudf.utils._numba_setup import _setup_numba
 from cudf.utils.dtypes import _NA_REP
 from cudf.utils.utils import clear_cache, set_allocator
 
-try:
-    from cubinlinker.patch import patch_numba_linker_if_needed
-except ImportError:
-    pass
-else:
-    # Patch Numba to support CUDA enhanced compatibility.
-    # cuDF requires a stronger set of conditions than what is
-    # checked by patch_numba_linker_if_needed due to the PTX
-    # files needed for JIT Groupby Apply and string UDFs
-    from cudf.core.udf.utils import _PTX_FILE, _setup_numba_linker
+_setup_numba()
 
-    _setup_numba_linker(_PTX_FILE)
-
-    del patch_numba_linker_if_needed
+# This must be imported after _setup_numba is called and the numba
+# config is modified otherwise the config options will have no effect
+from numba import cuda
 
 cuda.set_memory_manager(RMMNumbaManager)
 cupy.cuda.set_allocator(rmm_cupy_allocator)
 
-try:
-    # Numba 0.54: Disable low occupancy warnings
-    numba_config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
-except AttributeError:
-    # Numba < 0.54: No occupancy warnings
-    pass
-del numba_config
-
 
 rmm.register_reinitialize_hook(clear_cache)
 

python/cudf/cudf/core/udf/utils.py

@@ -1,6 +1,5 @@
 # Copyright (c) 2020-2023, NVIDIA CORPORATION.
 
-import glob
 import os
 from typing import Any, Callable, Dict
 
@@ -13,7 +12,6 @@
 from numba.core.datamodel import default_manager, models
 from numba.core.errors import TypingError
 from numba.core.extending import register_model
-from numba.cuda.cudadrv.driver import Linker
 from numba.np import numpy_support
 from numba.types import CPointer, Poison, Record, Tuple, boolean, int64, void
 
@@ -33,6 +31,7 @@
     udf_string,
 )
 from cudf.utils import cudautils
+from cudf.utils._numba_setup import _get_ptx_file
 from cudf.utils.dtypes import (
     BOOL_TYPES,
     DATETIME_TYPES,
@@ -63,58 +62,6 @@
 precompiled: cachetools.LRUCache = cachetools.LRUCache(maxsize=32)
 launch_arg_getters: Dict[Any, Any] = {}
 
-
-def _get_best_ptx_file(archs, max_compute_capability):
-    """
-    Determine of the available PTX files which one is
-    the most recent up to and including the device cc
-    """
-    filtered_archs = [x for x in archs if x[0] <= max_compute_capability]
-    if filtered_archs:
-        return max(filtered_archs, key=lambda y: y[0])
-    else:
-        return None
-
-
-def _get_ptx_file(path, prefix):
-    if "RAPIDS_NO_INITIALIZE" in os.environ:
-        # cc=60 ptx is always built
-        cc = int(os.environ.get("STRINGS_UDF_CC", "60"))
-    else:
-        dev = cuda.get_current_device()
-
-        # Load the highest compute capability file available that is less than
-        # the current device's.
-        cc = int("".join(str(x) for x in dev.compute_capability))
-    files = glob.glob(os.path.join(path, f"{prefix}*.ptx"))
-    if len(files) == 0:
-        raise RuntimeError(f"Missing PTX files for cc={cc}")
-    regular_sms = []
-
-    for f in files:
-        file_name = os.path.basename(f)
-        sm_number = file_name.rstrip(".ptx").lstrip(prefix)
-        if sm_number.endswith("a"):
-            processed_sm_number = int(sm_number.rstrip("a"))
-            if processed_sm_number == cc:
-                return f
-        else:
-            regular_sms.append((int(sm_number), f))
-
-    regular_result = None
-
-    if regular_sms:
-        regular_result = _get_best_ptx_file(regular_sms, cc)
-
-    if regular_result is None:
-        raise RuntimeError(
-            "This cuDF installation is missing the necessary PTX "
-            f"files that are <={cc}."
-        )
-    else:
-        return regular_result[1]
-
-
 _PTX_FILE = _get_ptx_file(os.path.dirname(__file__), "shim_")
 
 
@@ -390,97 +337,6 @@ def _get_extensionty_size(ty):
     return llty.get_abi_size(target_data)
 
 
-def _get_cuda_version_from_ptx_file(path):
-    """
-    https://docs.nvidia.com/cuda/parallel-thread-execution/
-    Each PTX module must begin with a .version
-    directive specifying the PTX language version
-
-    example header:
-    //
-    // Generated by NVIDIA NVVM Compiler
-    //
-    // Compiler Build ID: CL-31057947
-    // Cuda compilation tools, release 11.6, V11.6.124
-    // Based on NVVM 7.0.1
-    //
-
-    .version 7.6
-    .target sm_52
-    .address_size 64
-
-    """
-    with open(path) as ptx_file:
-        for line in ptx_file:
-            if line.startswith(".version"):
-                ver_line = line
-                break
-        else:
-            raise ValueError("Could not read CUDA version from ptx file.")
-    version = ver_line.strip("\n").split(" ")[1]
-    # from ptx_docs/release_notes above:
-    ver_map = {
-        "7.5": (11, 5),
-        "7.6": (11, 6),
-        "7.7": (11, 7),
-        "7.8": (11, 8),
-        "8.0": (12, 0),
-    }
-
-    cuda_ver = ver_map.get(version)
-    if cuda_ver is None:
-        raise ValueError(
-            f"Could not map PTX version {version} to a CUDA version"
-        )
-
-    return cuda_ver
-
-
-def _setup_numba_linker(path):
-    from ptxcompiler.patch import NO_DRIVER, safe_get_versions
-
-    from cudf.core.udf.utils import (
-        _get_cuda_version_from_ptx_file,
-        maybe_patch_numba_linker,
-    )
-
-    versions = safe_get_versions()
-    if versions != NO_DRIVER:
-        driver_version, runtime_version = versions
-        ptx_toolkit_version = _get_cuda_version_from_ptx_file(path)
-        maybe_patch_numba_linker(
-            driver_version, runtime_version, ptx_toolkit_version
-        )
-
-
-def maybe_patch_numba_linker(
-    driver_version, runtime_version, ptx_toolkit_version
-):
-    from cubinlinker.patch import (
-        _numba_version_ok,
-        get_logger,
-        new_patched_linker,
-    )
-
-    # Numba thinks cubinlinker is only needed if the driver is older than
-    # the ctk, but when PTX files are present, it might also need to patch
-    # because those PTX files may newer than the driver as well
-    logger = get_logger()
-
-    if (driver_version < ptx_toolkit_version) or (
-        driver_version < runtime_version
-    ):
-        logger.debug(
-            "Driver version %s.%s needs patching due to PTX files"
-            % driver_version
-        )
-        if _numba_version_ok:
-            logger.debug("Patching Numba Linker")
-            Linker.new = new_patched_linker
-        else:
-            logger.debug("Cannot patch Numba Linker - unsupported version")
-
-
 @initfunc
 def set_malloc_heap_size(size=None):
     """