Numba 0.57 compatibility fixes (#13271)

For compatibility with Numba 0.57 (required for full CUDA 12.0 support) the following changes are necessary:

- We can no longer use `Record` types directly, because Numba 0.57 assumes that it can convert a `Record` back to a NumPy dtype using its internal mechanism. The internal mechanism doesn't recognize cuDF extension types and has no provision for adding support for new types. To work around this, we directly encode a correspondence between our `Row` type (directly derived from a `Record`) and the NumPy `object` dtype. An `object` dtype is acceptable for this use because we use our row types for type checking and inference, but not code generation. Because code generation is not needed, we don't need to have an accurate representation of the low-level layout of a `Row` in the typing.
- A deprecated address space conversion function was removed from Numba. We switch to using another method that performs the address space conversion itself.
- The data layout moved from the `nvvm` module on to an `NVVM` object. Rather than attempting to get this from Numba, we just define it in cuDF, since there is only one possible data layout in CUDA anyway. This allows the same code path to work with all Numba versions (and removes a previous workaround from it having once moved before).

These changes should be compatible with both Numba 0.56 (the minimum required version at present) and 0.57 (can be used, will become a requirement soon).

Tested locally with Numba 0.57.0rc1 (which doesn't differ from the recently-tagged release in any way material to CUDA).

Authors:
  - Graham Markall (https://github.com/gmarkall)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Bradley Dice (https://github.com/bdice)

URL: #13271

python/cudf/cudf/core/udf/groupby_utils.py

@@ -6,7 +6,6 @@
 from numba import cuda, types
 from numba.cuda.cudadrv.devices import get_context
 from numba.np import numpy_support
-from numba.types import Record
 
 import cudf.core.udf.utils
 from cudf.core.udf.groupby_typing import (
@@ -19,6 +18,7 @@
     groupby_apply_kernel_template,
 )
 from cudf.core.udf.utils import (
+    Row,
     _generate_cache_key,
     _get_extensionty_size,
     _get_kernel,
@@ -32,10 +32,9 @@
 
 def _get_frame_groupby_type(dtype, index_dtype):
     """
-    Get the numba `Record` type corresponding to a frame.
-    Models the column as a dictionary like data structure
-    containing GroupTypes.
-    See numba.np.numpy_support.from_struct_dtype for details.
+    Get the Numba type corresponding to a row of grouped data. Models the
+    column as a Record-like data structure containing GroupTypes. See
+    numba.np.numpy_support.from_struct_dtype for details.
 
     Parameters
     ----------
@@ -74,7 +73,7 @@ def _get_frame_groupby_type(dtype, index_dtype):
 
     # Numba requires that structures are aligned for the CUDA target
     _is_aligned_struct = True
-    return Record(fields, offset, _is_aligned_struct)
+    return Row(fields, offset, _is_aligned_struct)
 
 
 def _groupby_apply_kernel_string_from_template(frame, args):

python/cudf/cudf/core/udf/row_function.py

@@ -4,7 +4,6 @@
 import numpy as np
 from numba import cuda
 from numba.np import numpy_support
-from numba.types import Record
 
 from cudf.core.udf.api import Masked, pack_return
 from cudf.core.udf.masked_typing import MaskedType
@@ -16,6 +15,7 @@
     unmasked_input_initializer_template,
 )
 from cudf.core.udf.utils import (
+    Row,
     _all_dtypes_from_frame,
     _construct_signature,
     _get_extensionty_size,
@@ -29,15 +29,12 @@
 
 def _get_frame_row_type(dtype):
     """
-    Get the numba `Record` type corresponding to a frame.
-    Models each column and its mask as a MaskedType and
-    models the row as a dictionary like data structure
-    containing these MaskedTypes.
-    Large parts of this function are copied with comments
-    from the Numba internals and slightly modified to
-    account for validity bools to be present in the final
-    struct.
-    See numba.np.numpy_support.from_struct_dtype for details.
+    Get the Numba type of a row in a frame. Models each column and its mask as
+    a MaskedType and models the row as a dictionary like data structure
+    containing these MaskedTypes. Large parts of this function are copied with
+    comments from the Numba internals and slightly modified to account for
+    validity bools to be present in the final struct. See
+    numba.np.numpy_support.from_struct_dtype for details.
     """
 
     # Create the numpy structured type corresponding to the numpy dtype.
@@ -89,7 +86,7 @@ def _get_frame_row_type(dtype):
 
     # Numba requires that structures are aligned for the CUDA target
     _is_aligned_struct = True
-    return Record(fields, offset, _is_aligned_struct)
+    return Row(fields, offset, _is_aligned_struct)
 
 
 def _row_kernel_string_from_template(frame, row_type, args):

python/cudf/cudf/core/udf/strings_lowering.py

@@ -7,7 +7,6 @@
 from numba.core import cgutils
 from numba.core.datamodel import default_manager
 from numba.core.typing import signature as nb_signature
-from numba.cuda.cudadrv import nvvm
 from numba.cuda.cudaimpl import (
     lower as cuda_lower,
     registry as cuda_lowering_registry,
@@ -126,9 +125,8 @@ def cast_string_literal_to_string_view(context, builder, fromty, toty, val):
     sv = cgutils.create_struct_proxy(string_view)(context, builder)
 
     # set the empty strview data pointer to point to the literal value
-    s = context.insert_const_string(builder.module, fromty.literal_value)
-    sv.data = context.insert_addrspace_conv(
-        builder, s, nvvm.ADDRSPACE_CONSTANT
+    sv.data = context.insert_string_const_addrspace(
+        builder, fromty.literal_value
     )
     sv.length = context.get_constant(size_type, len(fromty.literal_value))
     sv.bytes = context.get_constant(

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

@@ -11,12 +11,12 @@
 from cubinlinker.patch import _numba_version_ok, get_logger, new_patched_linker
 from cuda import cudart
 from numba import cuda, typeof
-from numba.core.datamodel import default_manager
+from numba.core.datamodel import default_manager, models
 from numba.core.errors import TypingError
-from numba.cuda.cudadrv import nvvm
+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, Tuple, boolean, int64, void
+from numba.types import CPointer, Poison, Record, Tuple, boolean, int64, void
 
 import rmm
 
@@ -240,6 +240,26 @@ def _construct_signature(frame, return_type, args):
     return sig
 
 
+class Row(Record):
+    # Numba's Record type provides a convenient abstraction for representing a
+    # row, in that it provides a mapping from strings (column / field names) to
+    # types. However, it cannot be used directly since it assumes that all its
+    # fields can be converted to NumPy types by Numba's internal conversion
+    # mechanism (`numba.np_support.as_dtype). This is not the case for cuDF
+    # extension types that might be the column types (e.g. masked types, string
+    # types or group types).
+    #
+    # We use this type for type inference and type checking, but not in code
+    # generation. For this use case, it is sufficient to provide a dtype for a
+    # row that corresponds to any Python object.
+    @property
+    def dtype(self):
+        return np.dtype("object")
+
+
+register_model(Row)(models.RecordModel)
+
+
 @cuda.jit(device=True)
 def _mask_get(mask, pos):
     """Return the validity of mask[pos] as a word."""
@@ -356,14 +376,20 @@ def _post_process_output_col(col, retty):
     return as_column(col, retty)
 
 
+# The only supported data layout in NVVM.
+# See: https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html?#data-layout
+_nvvm_data_layout = (
+    "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-"
+    "i128:128:128-f32:32:32-f64:64:64-v16:16:16-v32:32:32-"
+    "v64:64:64-v128:128:128-n16:32:64"
+)
+
+
 def _get_extensionty_size(ty):
     """
     Return the size of an extension type in bytes
     """
-    data_layout = nvvm.data_layout
-    if isinstance(data_layout, dict):
-        data_layout = data_layout[64]
-    target_data = ll.create_target_data(data_layout)
+    target_data = ll.create_target_data(_nvvm_data_layout)
     llty = default_manager[ty].get_value_type()
     return llty.get_abi_size(target_data)