Enable CEC for strings_udf (#11884)

This PR removes the runtime checks for CEC in `strings_udf`.

Authors:
  - https://github.com/brandon-b-miller

Approvers:
  - Graham Markall (https://github.com/gmarkall)
  - Lawrence Mitchell (https://github.com/wence-)
  - Ray Douglass (https://github.com/raydouglass)

URL: #11884

ci/gpu/build.sh

@@ -295,22 +295,15 @@ py.test -n 8 --cache-clear --basetemp="$WORKSPACE/custreamz-cuda-tmp" --junitxml
 gpuci_logger "Installing strings_udf"
 gpuci_mamba_retry install strings_udf -c "${CONDA_BLD_DIR}" -c "${CONDA_ARTIFACT_PATH}"
 
-# only install strings_udf after cuDF is finished testing without its presence
 cd "$WORKSPACE/python/strings_udf/strings_udf"
 gpuci_logger "Python py.test for strings_udf"
+py.test -n 8 --cache-clear --basetemp="$WORKSPACE/strings-udf-cuda-tmp" --junitxml="$WORKSPACE/junit-strings-udf.xml" -v --cov-config=.coveragerc --cov=strings_udf --cov-report=xml:"$WORKSPACE/python/strings_udf/strings-udf-coverage.xml" --cov-report term tests
 
-STRINGS_UDF_PYTEST_RETCODE=0
-py.test -n 8 --cache-clear --basetemp="$WORKSPACE/strings-udf-cuda-tmp" --junitxml="$WORKSPACE/junit-strings-udf.xml" -v --cov-config=.coveragerc --cov=strings_udf --cov-report=xml:"$WORKSPACE/python/strings_udf/strings-udf-coverage.xml" --cov-report term tests || STRINGS_UDF_PYTEST_RETCODE=$?
+# retest cuDF UDFs
+cd "$WORKSPACE/python/cudf/cudf"
+gpuci_logger "Python py.test retest cuDF UDFs"
+py.test -n 8 --cache-clear --basetemp="$WORKSPACE/cudf-cuda-strings-udf-tmp" --ignore="$WORKSPACE/python/cudf/cudf/benchmarks" --junitxml="$WORKSPACE/junit-cudf-strings-udf.xml" -v --cov-config="$WORKSPACE/python/cudf/.coveragerc" --cov=cudf --cov-report=xml:"$WORKSPACE/python/cudf/cudf-strings-udf-coverage.xml" --cov-report term --dist=loadscope tests/test_udf_masked_ops.py
 
-if [ ${STRINGS_UDF_PYTEST_RETCODE} -eq 5 ]; then
-    echo "No strings UDF tests were run, but this script will continue to execute."
-elif [ ${STRINGS_UDF_PYTEST_RETCODE} -ne 0 ]; then
-    exit ${STRINGS_UDF_PYTEST_RETCODE}
-else
-    cd "$WORKSPACE/python/cudf/cudf"
-    gpuci_logger "Python py.test retest cuDF UDFs"
-    py.test -n 8 --cache-clear --basetemp="$WORKSPACE/cudf-cuda-strings-udf-tmp" --ignore="$WORKSPACE/python/cudf/cudf/benchmarks" --junitxml="$WORKSPACE/junit-cudf-strings-udf.xml" -v --cov-config="$WORKSPACE/python/cudf/.coveragerc" --cov=cudf --cov-report=xml:"$WORKSPACE/python/cudf/cudf-strings-udf-coverage.xml" --cov-report term --dist=loadscope tests
-fi
 
 # Run benchmarks with both cudf and pandas to ensure compatibility is maintained.
 # Benchmarks are run in DEBUG_ONLY mode, meaning that only small data sizes are used.

conda/recipes/strings_udf/meta.yaml

@@ -50,7 +50,7 @@ requirements:
     - cudf ={{ version }}
     - {{ pin_compatible('cudatoolkit', max_pin='x', min_pin='x') }}
     - cachetools
-    - ptxcompiler  # [linux64]  # CUDA enhanced compatibility. See https://github.com/rapidsai/ptxcompiler
+    - ptxcompiler >=0.7.0  # [linux64]  # CUDA enhanced compatibility. See https://github.com/rapidsai/ptxcompiler
 test:                                   # [linux64]
   requires:                             # [linux64]
     - cudatoolkit {{ cuda_version }}.*  # [linux64]

python/cudf/cudf/core/udf/__init__.py

@@ -23,17 +23,20 @@
     | {types.boolean}
 )
 _STRING_UDFS_ENABLED = False
+
 try:
     import strings_udf
+    from strings_udf import ptxpath
+
+    if ptxpath:
+        utils.ptx_files.append(ptxpath)
 
-    if strings_udf.ENABLED:
-        from . import strings_typing  # isort: skip
-        from . import strings_lowering  # isort: skip
-        from strings_udf import ptxpath
         from strings_udf._lib.cudf_jit_udf import to_string_view_array
         from strings_udf._typing import str_view_arg_handler, string_view
 
-        # add an overload of MaskedType.__init__(string_view, bool)
+        from . import strings_typing  # isort: skip
+        from . import strings_lowering  # isort: skip
+
         cuda_lower(api.Masked, strings_typing.string_view, types.boolean)(
             masked_lowering.masked_constructor
         )
@@ -47,13 +50,11 @@
         utils.launch_arg_getters[dtype("O")] = to_string_view_array
         utils.masked_array_types[dtype("O")] = string_view
         utils.JIT_SUPPORTED_TYPES |= STRING_TYPES
-        utils.ptx_files.append(ptxpath)
+
         utils.arg_handlers.append(str_view_arg_handler)
         row_function.itemsizes[dtype("O")] = string_view.size_bytes
 
         _STRING_UDFS_ENABLED = True
-    else:
-        del strings_udf
 
 except ImportError as e:
     # allow cuDF to work without strings_udf

python/strings_udf/strings_udf/__init__.py

@@ -1,36 +1,20 @@
 # Copyright (c) 2022, NVIDIA CORPORATION.
 import glob
 import os
-import re
-import subprocess
-import sys
 
+from cubinlinker.patch import _numba_version_ok, get_logger, new_patched_linker
 from numba import cuda
-from ptxcompiler.patch import CMD
+from numba.cuda.cudadrv.driver import Linker
+from ptxcompiler.patch import NO_DRIVER, safe_get_versions
 
 from . import _version
 
 __version__ = _version.get_versions()["version"]
 
-ENABLED = False
+logger = get_logger()
 
-
-def compiler_from_ptx_file(path):
-    """Parse a PTX file header and extract the CUDA version used to compile it.
-    Here is an example PTX header that this function should parse:
-    // Generated by NVIDIA NVVM Compiler
-    //
-    // Compiler Build ID: CL-30672275
-    // Cuda compilation tools, release 11.5, V11.5.119
-    // Based on NVVM 7
-    """
-    file = open(path).read()
-    major, minor = (
-        re.search(r"Cuda compilation tools, release ([0-9\.]+)", file)
-        .group(1)
-        .split(".")
-    )
-    return int(major), int(minor)
+# tracks the version of CUDA used to build the c++ and PTX components
+STRINGS_UDF_PTX_VERSION = (11, 5)
 
 
 def _get_appropriate_file(sms, cc):
@@ -41,59 +25,70 @@ def _get_appropriate_file(sms, cc):
         return None
 
 
-# adapted from PTXCompiler
-cp = subprocess.run([sys.executable, "-c", CMD], capture_output=True)
-if cp.returncode == 0:
-    # must have a driver to proceed
-    versions = [int(s) for s in cp.stdout.strip().split()]
-    driver_version = tuple(versions[:2])
-    runtime_version = tuple(versions[2:])
+def maybe_patch_numba_linker(driver_version):
+    # Numba thinks cubinlinker is only needed if the driver is older than the ctk
+    # but when strings_udf is present, it might also need to patch because the PTX
+    # file strings_udf relies on may be newer than the driver as well
+    if driver_version < STRINGS_UDF_PTX_VERSION:
+        logger.debug(
+            "Driver version %s.%s needs patching due to strings_udf"
+            % 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")
+
+
+def _get_ptx_file():
+    if "RAPIDS_NO_INITIALIZE" in os.environ:
+        cc = int(os.environ.get("STRINGS_UDF_CC", "52"))
+    else:
+        dev = cuda.get_current_device()
 
-    # CUDA enhanced compatibility not yet enabled
-    if driver_version >= runtime_version:
         # Load the highest compute capability file available that is less than
         # the current device's.
-        dev = cuda.get_current_device()
         cc = int("".join(str(x) for x in dev.compute_capability))
-        files = glob.glob(
-            os.path.join(os.path.dirname(__file__), "shim_*.ptx")
+    files = glob.glob(os.path.join(os.path.dirname(__file__), "shim_*.ptx"))
+    if len(files) == 0:
+        raise RuntimeError(
+            "This strings_udf installation is missing the necessary PTX "
+            f"files for compute capability {cc}. "
+            "Please file an issue reporting this error and how you "
+            "installed cudf and strings_udf."
+            "https://github.com/rapidsai/cudf/issues"
         )
-        if len(files) == 0:
-            raise RuntimeError(
-                "This strings_udf installation is missing the necessary PTX "
-                "files. Please file an issue reporting this error and how you "
-                "installed cudf and strings_udf."
-            )
-
-        suffix_a_sm = None
-        regular_sms = []
-
-        for f in files:
-            file_name = os.path.basename(f)
-            sm_number = file_name.rstrip(".ptx").lstrip("shim_")
-            if sm_number.endswith("a"):
-                processed_sm_number = int(sm_number.rstrip("a"))
-                if processed_sm_number == cc:
-                    suffix_a_sm = (processed_sm_number, f)
-            else:
-                regular_sms.append((int(sm_number), f))
-
-        regular_result = None
-
-        if regular_sms:
-            regular_result = _get_appropriate_file(regular_sms, cc)
-
-        if suffix_a_sm is None and regular_result is None:
-            raise RuntimeError(
-                "This strings_udf installation is missing the necessary PTX "
-                f"files that are <={cc}."
-            )
-        elif suffix_a_sm is not None:
-            ptxpath = suffix_a_sm[1]
-        else:
-            ptxpath = regular_result[1]
 
-        if driver_version >= compiler_from_ptx_file(ptxpath):
-            ENABLED = True
+    regular_sms = []
+
+    for f in files:
+        file_name = os.path.basename(f)
+        sm_number = file_name.rstrip(".ptx").lstrip("shim_")
+        if sm_number.endswith("a"):
+            processed_sm_number = int(sm_number.rstrip("a"))
+            if processed_sm_number == cc:
+                return f
         else:
-            del ptxpath
+            regular_sms.append((int(sm_number), f))
+
+    regular_result = None
+
+    if regular_sms:
+        regular_result = _get_appropriate_file(regular_sms, cc)
+
+    if regular_result is None:
+        raise RuntimeError(
+            "This strings_udf installation is missing the necessary PTX "
+            f"files that are <={cc}."
+        )
+    else:
+        return regular_result[1]
+
+
+ptxpath = None
+versions = safe_get_versions()
+if versions != NO_DRIVER:
+    driver_version, runtime_version = versions
+    maybe_patch_numba_linker(driver_version)
+    ptxpath = _get_ptx_file()

python/strings_udf/strings_udf/tests/test_string_udfs.py

@@ -15,9 +15,6 @@
 from strings_udf._lib.cudf_jit_udf import to_string_view_array
 from strings_udf._typing import str_view_arg_handler, string_view
 
-if not strings_udf.ENABLED:
-    pytest.skip("Strings UDF not enabled.", allow_module_level=True)
-
 
 def get_kernel(func, dtype):
     """