From 799c3f9dcceaf5cf5ab7843664919e8dc7b60ab1 Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Wed, 24 May 2023 20:07:44 -0500
Subject: [PATCH] Patch numba if it is imported first to ensure minor version
 compatibility works. (#13433)

`cudf` conditionally sets the Numba configuration option `CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY` if it is needed for the current driver/runtime. Using numba 0.57.0, if `numba.cuda` is imported before `cudf`, this option is read too late, and the necessary imports of ptxcompiler/cubinlinker are not available in `numba.cuda.cudadrv.driver`. This PR patches Numba by adding those imports late if (and only if) MVC is needed and `numba.cuda` has already been imported.

See also: https://github.com/numba/numba/issues/8977

Authors:
  - Bradley Dice (https://github.com/bdice)
  - https://github.com/brandon-b-miller

Approvers:
  - https://github.com/brandon-b-miller
  - Graham Markall (https://github.com/gmarkall)

URL: https://github.com/rapidsai/cudf/pull/13433
---
 python/cudf/cudf/tests/test_numba_import.py | 46 +++++++++++++++++++++
 python/cudf/cudf/utils/_numba.py            | 30 +++++++++++---
 2 files changed, 71 insertions(+), 5 deletions(-)
 create mode 100644 python/cudf/cudf/tests/test_numba_import.py

diff --git a/python/cudf/cudf/tests/test_numba_import.py b/python/cudf/cudf/tests/test_numba_import.py
new file mode 100644
index 00000000000..dcde0f68aa2
--- /dev/null
+++ b/python/cudf/cudf/tests/test_numba_import.py
@@ -0,0 +1,46 @@
+# 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
+    )
+    assert cp.returncode == 0
diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py
index 194db9c90a6..09afb5680bd 100644
--- a/python/cudf/cudf/utils/_numba.py
+++ b/python/cudf/cudf/utils/_numba.py
@@ -2,9 +2,10 @@
 
 import glob
 import os
+import sys
 import warnings
 
-from numba import config
+from numba import config as numba_config
 
 CC_60_PTX_FILE = os.path.join(
     os.path.dirname(__file__), "../core/udf/shim_60.ptx"
@@ -64,6 +65,25 @@ def _get_ptx_file(path, prefix):
         return regular_result[1]
 
 
+def _patch_numba_mvc():
+    # Enable the config option for minor version compatibility
+    numba_config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY = 1
+
+    if "numba.cuda" in sys.modules:
+        # Patch numba for version 0.57.0 MVC support, which must know the
+        # config value at import time. We cannot guarantee the order of imports
+        # between cudf and numba.cuda so we patch numba to ensure it has these
+        # names available.
+        # See https://github.com/numba/numba/issues/8977 for details.
+        import numba.cuda
+        from cubinlinker import CubinLinker, CubinLinkerError
+        from ptxcompiler import compile_ptx
+
+        numba.cuda.cudadrv.driver.compile_ptx = compile_ptx
+        numba.cuda.cudadrv.driver.CubinLinker = CubinLinker
+        numba.cuda.cudadrv.driver.CubinLinkerError = CubinLinkerError
+
+
 def _setup_numba():
     """
     Configure the numba linker for use with cuDF. This consists of
@@ -108,7 +128,7 @@ def _setup_numba():
             if (driver_version < ptx_toolkit_version) or (
                 driver_version < runtime_version
             ):
-                config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY = 1
+                _patch_numba_mvc()
 
 
 def _get_cuda_version_from_ptx_file(path):
@@ -164,8 +184,8 @@ def _get_cuda_version_from_ptx_file(path):
 
 class _CUDFNumbaConfig:
     def __enter__(self):
-        self.enter_val = config.CUDA_LOW_OCCUPANCY_WARNINGS
-        config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
+        self.enter_val = numba_config.CUDA_LOW_OCCUPANCY_WARNINGS
+        numba_config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
 
     def __exit__(self, exc_type, exc_value, traceback):
-        config.CUDA_LOW_OCCUPANCY_WARNINGS = self.enter_val
+        numba_config.CUDA_LOW_OCCUPANCY_WARNINGS = self.enter_val