diff --git a/python/cudf/cudf/core/udf/groupby_utils.py b/python/cudf/cudf/core/udf/groupby_utils.py index f3270b93406..ae09dd1d704 100644 --- a/python/cudf/cudf/core/udf/groupby_utils.py +++ b/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): diff --git a/python/cudf/cudf/core/udf/row_function.py b/python/cudf/cudf/core/udf/row_function.py index bfb04e38e7d..e040836f97d 100644 --- a/python/cudf/cudf/core/udf/row_function.py +++ b/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): diff --git a/python/cudf/cudf/core/udf/strings_lowering.py b/python/cudf/cudf/core/udf/strings_lowering.py index a53722f505d..fdce404d887 100644 --- a/python/cudf/cudf/core/udf/strings_lowering.py +++ b/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( diff --git a/python/cudf/cudf/core/udf/utils.py b/python/cudf/cudf/core/udf/utils.py index b3b846658f6..0b7544752e2 100644 --- a/python/cudf/cudf/core/udf/utils.py +++ b/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)