From 782b57dc3937ea68678396c9b33a029d810e8e69 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 26 Aug 2025 08:22:34 -0700 Subject: [PATCH 1/2] refactor to better see import times --- .../cudf/cudf/core/_internals/aggregation.py | 2 +- .../cudf/cudf/core/column/numerical_base.py | 2 +- python/cudf/cudf/core/udf/__init__.py | 28 +++++-- python/cudf/cudf/core/udf/groupby_lowering.py | 37 ++++---- python/cudf/cudf/core/udf/groupby_typing.py | 71 ++++++++-------- python/cudf/cudf/core/udf/masked_lowering.py | 29 ++++--- python/cudf/cudf/core/udf/masked_typing.py | 35 ++++---- python/cudf/cudf/core/udf/strings_lowering.py | 58 ++++++------- python/cudf/cudf/core/udf/strings_typing.py | 41 +++++---- python/cudf/cudf/core/udf/utils.py | 75 +---------------- python/cudf/cudf/utils/_numba.py | 84 +++++++++++++++++++ 11 files changed, 252 insertions(+), 210 deletions(-) diff --git a/python/cudf/cudf/core/_internals/aggregation.py b/python/cudf/cudf/core/_internals/aggregation.py index 60e5823fc8d4..ae237d06b82c 100644 --- a/python/cudf/cudf/core/_internals/aggregation.py +++ b/python/cudf/cudf/core/_internals/aggregation.py @@ -9,7 +9,7 @@ import pylibcudf as plc from cudf.api.types import is_scalar -from cudf.core.udf.utils import compile_udf +from cudf.utils._numba import compile_udf from cudf.utils.dtypes import SUPPORTED_NUMPY_TO_PYLIBCUDF_TYPES if TYPE_CHECKING: diff --git a/python/cudf/cudf/core/column/numerical_base.py b/python/cudf/cudf/core/column/numerical_base.py index f25e1c6e4e35..07eac625c930 100644 --- a/python/cudf/cudf/core/column/numerical_base.py +++ b/python/cudf/cudf/core/column/numerical_base.py @@ -15,7 +15,7 @@ from cudf.core.column.column import ColumnBase, column_empty from cudf.core.missing import NA from cudf.core.mixins import Scannable -from cudf.core.udf.utils import compile_udf +from cudf.utils._numba import compile_udf from cudf.utils.dtypes import _get_nan_for_dtype if TYPE_CHECKING: diff --git a/python/cudf/cudf/core/udf/__init__.py b/python/cudf/cudf/core/udf/__init__.py index 85d454652b71..af4e5576c08c 100644 --- a/python/cudf/cudf/core/udf/__init__.py +++ b/python/cudf/cudf/core/udf/__init__.py @@ -1,9 +1,21 @@ # Copyright (c) 2022-2023, NVIDIA CORPORATION. -from . import ( - groupby_lowering, - groupby_typing, - masked_lowering, - masked_typing, - strings_lowering, - strings_typing, -) +#from . import ( +# groupby_lowering, +# groupby_typing, +# masked_lowering, +# masked_typing, +# strings_lowering, +# strings_typing, +#) + +from . import strings_typing, strings_lowering +strings_typing.register_strings_typing() +strings_lowering.register_strings_lowering() + +from . import masked_typing, masked_lowering +masked_typing.register_masked_typing() +masked_lowering.register_masked_lowering() + +from . import groupby_typing, groupby_lowering +groupby_typing.register_groupby_typing() +groupby_lowering.register_groupby_lowering() diff --git a/python/cudf/cudf/core/udf/groupby_lowering.py b/python/cudf/cudf/core/udf/groupby_lowering.py index fe0637cfaefc..49d017e946ca 100644 --- a/python/cudf/cudf/core/udf/groupby_lowering.py +++ b/python/cudf/cudf/core/udf/groupby_lowering.py @@ -91,7 +91,6 @@ def group_corr(context, builder, sig, args): return result -@lower_builtin(Group, types.Array, group_size_type, types.Array) def group_constructor(context, builder, sig, args): """ Instruction boilerplate used for instantiating a Group @@ -171,20 +170,24 @@ def cuda_Group_size(context, builder, sig, args): cuda_Group_count = cuda_Group_size +def register_groupby_lowering(): + + lower_builtin(Group, types.Array, group_size_type, types.Array)(group_constructor) + + for ty in SUPPORTED_GROUPBY_NUMBA_TYPES: + cuda_lower("GroupType.max", GroupType(ty))(cuda_Group_max) + cuda_lower("GroupType.min", GroupType(ty))(cuda_Group_min) + cuda_lower("GroupType.sum", GroupType(ty))(cuda_Group_sum) + cuda_lower("GroupType.count", GroupType(ty))(cuda_Group_count) + cuda_lower("GroupType.size", GroupType(ty))(cuda_Group_size) + cuda_lower("GroupType.mean", GroupType(ty))(cuda_Group_mean) + cuda_lower("GroupType.std", GroupType(ty))(cuda_Group_std) + cuda_lower("GroupType.var", GroupType(ty))(cuda_Group_var) + cuda_lower("GroupType.idxmax", GroupType(ty, types.int64))( + cuda_Group_idxmax + ) + cuda_lower("GroupType.idxmin", GroupType(ty, types.int64))( + cuda_Group_idxmin + ) + cuda_lower("GroupType.corr", GroupType(ty), GroupType(ty))(group_corr) -for ty in SUPPORTED_GROUPBY_NUMBA_TYPES: - cuda_lower("GroupType.max", GroupType(ty))(cuda_Group_max) - cuda_lower("GroupType.min", GroupType(ty))(cuda_Group_min) - cuda_lower("GroupType.sum", GroupType(ty))(cuda_Group_sum) - cuda_lower("GroupType.count", GroupType(ty))(cuda_Group_count) - cuda_lower("GroupType.size", GroupType(ty))(cuda_Group_size) - cuda_lower("GroupType.mean", GroupType(ty))(cuda_Group_mean) - cuda_lower("GroupType.std", GroupType(ty))(cuda_Group_std) - cuda_lower("GroupType.var", GroupType(ty))(cuda_Group_var) - cuda_lower("GroupType.idxmax", GroupType(ty, types.int64))( - cuda_Group_idxmax - ) - cuda_lower("GroupType.idxmin", GroupType(ty, types.int64))( - cuda_Group_idxmin - ) - cuda_lower("GroupType.corr", GroupType(ty), GroupType(ty))(group_corr) diff --git a/python/cudf/cudf/core/udf/groupby_typing.py b/python/cudf/cudf/core/udf/groupby_typing.py index eaa1d4c76b07..a6b8bfb44054 100644 --- a/python/cudf/cudf/core/udf/groupby_typing.py +++ b/python/cudf/cudf/core/udf/groupby_typing.py @@ -79,10 +79,7 @@ class GroupByJITDataFrame(Row): pass -register_model(GroupByJITDataFrame)(models.RecordModel) - -@typeof_impl.register(Group) def typeof_group(val, c): """ Tie Group and GroupType together such that when Numba @@ -97,7 +94,6 @@ def typeof_group(val, c): # The typing of the python "function" Group.__init__ # as it appears in python code -@type_callable(Group) def type_group(context): def typer(group_data, size, index): if ( @@ -110,7 +106,6 @@ def typer(group_data, size, index): return typer -@register_model(GroupType) class GroupModel(models.StructModel): """ Model backing GroupType instances. See the link below for details. @@ -315,12 +310,10 @@ def resolve(self, value, attr): ) -@cuda_registry.register_attr class DataFrameAttr(DataFrameAttributeTemplate): key = GroupByJITDataFrame -@cuda_registry.register_attr class GroupAttr(AttributeTemplate): key = GroupType @@ -355,41 +348,51 @@ def resolve_corr(self, mod): ) -for ty in SUPPORTED_GROUPBY_NUMBA_TYPES: - _register_cuda_unary_reduction_caller("Max", ty, ty) - _register_cuda_unary_reduction_caller("Min", ty, ty) - _register_cuda_idx_reduction_caller("IdxMax", ty) - _register_cuda_idx_reduction_caller("IdxMin", ty) - if ty in types.integer_domain: - _register_cuda_binary_reduction_caller("Corr", ty, ty, types.float64) +def register_groupby_typing(): + typeof_impl.register(Group)(typeof_group) + type_callable(Group)(type_group) + register_model(GroupType)(GroupModel) + cuda_registry.register_attr(DataFrameAttr) + cuda_registry.register_attr(GroupAttr) + + register_model(GroupByJITDataFrame)(models.RecordModel) + + for op in arith_ops + comparison_ops + unary_ops: + cuda_registry.register_global(op)(GroupOpBase) + + for attr in ("group_data", "index", "size"): + make_attribute_wrapper(GroupType, attr, attr) -_register_cuda_unary_reduction_caller("Sum", types.int32, types.int64) -_register_cuda_unary_reduction_caller("Sum", types.int64, types.int64) -_register_cuda_unary_reduction_caller("Sum", types.float32, types.float32) -_register_cuda_unary_reduction_caller("Sum", types.float64, types.float64) + _register_cuda_unary_reduction_caller("Sum", types.int32, types.int64) + _register_cuda_unary_reduction_caller("Sum", types.int64, types.int64) + _register_cuda_unary_reduction_caller("Sum", types.float32, types.float32) + _register_cuda_unary_reduction_caller("Sum", types.float64, types.float64) -_register_cuda_unary_reduction_caller("Mean", types.int32, types.float64) -_register_cuda_unary_reduction_caller("Mean", types.int64, types.float64) -_register_cuda_unary_reduction_caller("Mean", types.float32, types.float32) -_register_cuda_unary_reduction_caller("Mean", types.float64, types.float64) + _register_cuda_unary_reduction_caller("Mean", types.int32, types.float64) + _register_cuda_unary_reduction_caller("Mean", types.int64, types.float64) + _register_cuda_unary_reduction_caller("Mean", types.float32, types.float32) + _register_cuda_unary_reduction_caller("Mean", types.float64, types.float64) -_register_cuda_unary_reduction_caller("Std", types.int32, types.float64) -_register_cuda_unary_reduction_caller("Std", types.int64, types.float64) -_register_cuda_unary_reduction_caller("Std", types.float32, types.float32) -_register_cuda_unary_reduction_caller("Std", types.float64, types.float64) + _register_cuda_unary_reduction_caller("Std", types.int32, types.float64) + _register_cuda_unary_reduction_caller("Std", types.int64, types.float64) + _register_cuda_unary_reduction_caller("Std", types.float32, types.float32) + _register_cuda_unary_reduction_caller("Std", types.float64, types.float64) -_register_cuda_unary_reduction_caller("Var", types.int32, types.float64) -_register_cuda_unary_reduction_caller("Var", types.int64, types.float64) -_register_cuda_unary_reduction_caller("Var", types.float32, types.float32) -_register_cuda_unary_reduction_caller("Var", types.float64, types.float64) + _register_cuda_unary_reduction_caller("Var", types.int32, types.float64) + _register_cuda_unary_reduction_caller("Var", types.int64, types.float64) + _register_cuda_unary_reduction_caller("Var", types.float32, types.float32) + _register_cuda_unary_reduction_caller("Var", types.float64, types.float64) + for ty in SUPPORTED_GROUPBY_NUMBA_TYPES: + _register_cuda_unary_reduction_caller("Max", ty, ty) + _register_cuda_unary_reduction_caller("Min", ty, ty) + _register_cuda_idx_reduction_caller("IdxMax", ty) + _register_cuda_idx_reduction_caller("IdxMin", ty) -for attr in ("group_data", "index", "size"): - make_attribute_wrapper(GroupType, attr, attr) + if ty in types.integer_domain: + _register_cuda_binary_reduction_caller("Corr", ty, ty, types.float64) -for op in arith_ops + comparison_ops + unary_ops: - cuda_registry.register_global(op)(GroupOpBase) diff --git a/python/cudf/cudf/core/udf/masked_lowering.py b/python/cudf/cudf/core/udf/masked_lowering.py index fb561cdc306b..f33bef671805 100644 --- a/python/cudf/cudf/core/udf/masked_lowering.py +++ b/python/cudf/cudf/core/udf/masked_lowering.py @@ -233,19 +233,6 @@ def register_const_op(op): cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op) -# register all lowering at init -for binary_op in arith_ops + bitwise_ops + comparison_ops: - register_arithmetic_op(binary_op) - register_const_op(binary_op) - # null op impl can be shared between all ops - cuda_lower(binary_op, MaskedType, NAType)(masked_scalar_null_op_impl) - cuda_lower(binary_op, NAType, MaskedType)(masked_scalar_null_op_impl) - -# register all lowering at init -for unary_op in unary_ops: - register_unary_op(unary_op) -register_unary_op(abs) - @cuda_lower(operator.is_, MaskedType, NAType) @cuda_lower(operator.is_, NAType, MaskedType) @@ -413,3 +400,19 @@ def lower_constant_masked(context, builder, ty, val): masked.value = context.get_constant(ty.value_type, val.value) masked.valid = context.get_constant(types.boolean, val.valid) return masked._getvalue() + +def register_masked_lowering(): + # register all lowering at init + for binary_op in arith_ops + bitwise_ops + comparison_ops: + register_arithmetic_op(binary_op) + register_const_op(binary_op) + # null op impl can be shared between all ops + cuda_lower(binary_op, MaskedType, NAType)(masked_scalar_null_op_impl) + cuda_lower(binary_op, NAType, MaskedType)(masked_scalar_null_op_impl) + + # register all lowering at init + for unary_op in unary_ops: + register_unary_op(unary_op) + register_unary_op(abs) + + diff --git a/python/cudf/cudf/core/udf/masked_typing.py b/python/cudf/cudf/core/udf/masked_typing.py index 220fea7b04f1..a58f88f9d09f 100644 --- a/python/cudf/cudf/core/udf/masked_typing.py +++ b/python/cudf/cudf/core/udf/masked_typing.py @@ -467,16 +467,6 @@ def generic(self, args, kws): return nb_signature(return_type, args[0]) -for binary_op in arith_ops + bitwise_ops + comparison_ops: - # Every op shares the same typing class - cuda_decl_registry.register_global(binary_op)(MaskedScalarArithOp) - cuda_decl_registry.register_global(binary_op)(MaskedScalarNullOp) - cuda_decl_registry.register_global(binary_op)(MaskedScalarScalarOp) - -for unary_op in unary_ops: - cuda_decl_registry.register_global(unary_op)(MaskedScalarUnaryOp) - - # Strings functions and utilities def _is_valid_string_arg(ty): return ( @@ -550,10 +540,6 @@ def generic(self, args, kws): ) -for op in comparison_ops: - cuda_decl_registry.register_global(op)(MaskedStringViewCmpOp) - - def create_masked_binary_attr(attrname, retty): """ Helper function wrapping numba's low level extension API. Provides @@ -685,5 +671,22 @@ def resolve_value(self, mod): return managed_udf_string -cuda_decl_registry.register_attr(MaskedStringViewAttrs) -cuda_decl_registry.register_attr(MaskedManagedUDFStringAttrs) +def register_masked_typing(): + for binary_op in arith_ops + bitwise_ops + comparison_ops: + # Every op shares the same typing class + cuda_decl_registry.register_global(binary_op)(MaskedScalarArithOp) + cuda_decl_registry.register_global(binary_op)(MaskedScalarNullOp) + cuda_decl_registry.register_global(binary_op)(MaskedScalarScalarOp) + + for unary_op in unary_ops: + cuda_decl_registry.register_global(unary_op)(MaskedScalarUnaryOp) + + + for op in comparison_ops: + cuda_decl_registry.register_global(op)(MaskedStringViewCmpOp) + + + + cuda_decl_registry.register_attr(MaskedStringViewAttrs) + cuda_decl_registry.register_attr(MaskedManagedUDFStringAttrs) + diff --git a/python/cudf/cudf/core/udf/strings_lowering.py b/python/cudf/cudf/core/udf/strings_lowering.py index 61f69cb8c711..1fb26b0ab785 100644 --- a/python/cudf/cudf/core/udf/strings_lowering.py +++ b/python/cudf/cudf/core/udf/strings_lowering.py @@ -794,32 +794,34 @@ def upper_or_lower_impl(context, builder, sig, args): cuda_lower(op, MaskedType(string_view))(upper_or_lower_impl) -create_masked_binary_string_func("MaskedType.strip", strip_impl, udf_string) -create_masked_binary_string_func("MaskedType.lstrip", lstrip_impl, udf_string) -create_masked_binary_string_func("MaskedType.rstrip", rstrip_impl, udf_string) -create_masked_binary_string_func( - "MaskedType.startswith", - startswith_impl, - types.boolean, -) -create_masked_binary_string_func( - "MaskedType.endswith", endswith_impl, types.boolean -) -create_masked_binary_string_func("MaskedType.find", find_impl, size_type) -create_masked_binary_string_func("MaskedType.rfind", rfind_impl, size_type) -create_masked_binary_string_func("MaskedType.count", count_impl, size_type) -create_masked_binary_string_func( - operator.contains, contains_impl, types.boolean -) - +def register_strings_lowering(): + create_masked_binary_string_func("MaskedType.strip", strip_impl, udf_string) + create_masked_binary_string_func("MaskedType.lstrip", lstrip_impl, udf_string) + create_masked_binary_string_func("MaskedType.rstrip", rstrip_impl, udf_string) + create_masked_binary_string_func( + "MaskedType.startswith", + startswith_impl, + types.boolean, + ) + create_masked_binary_string_func( + "MaskedType.endswith", endswith_impl, types.boolean + ) + create_masked_binary_string_func("MaskedType.find", find_impl, size_type) + create_masked_binary_string_func("MaskedType.rfind", rfind_impl, size_type) + create_masked_binary_string_func("MaskedType.count", count_impl, size_type) + create_masked_binary_string_func( + operator.contains, contains_impl, types.boolean + ) + + + create_masked_unary_identifier_func("MaskedType.isalnum", isalnum_impl) + create_masked_unary_identifier_func("MaskedType.isalpha", isalpha_impl) + create_masked_unary_identifier_func("MaskedType.isdigit", isdigit_impl) + create_masked_unary_identifier_func("MaskedType.isupper", isupper_impl) + create_masked_unary_identifier_func("MaskedType.islower", islower_impl) + create_masked_unary_identifier_func("MaskedType.isspace", isspace_impl) + create_masked_unary_identifier_func("MaskedType.isdecimal", isdecimal_impl) + create_masked_unary_identifier_func("MaskedType.istitle", istitle_impl) + create_masked_upper_or_lower("MaskedType.upper", upper_impl) + create_masked_upper_or_lower("MaskedType.lower", lower_impl) -create_masked_unary_identifier_func("MaskedType.isalnum", isalnum_impl) -create_masked_unary_identifier_func("MaskedType.isalpha", isalpha_impl) -create_masked_unary_identifier_func("MaskedType.isdigit", isdigit_impl) -create_masked_unary_identifier_func("MaskedType.isupper", isupper_impl) -create_masked_unary_identifier_func("MaskedType.islower", islower_impl) -create_masked_unary_identifier_func("MaskedType.isspace", isspace_impl) -create_masked_unary_identifier_func("MaskedType.isdecimal", isdecimal_impl) -create_masked_unary_identifier_func("MaskedType.istitle", istitle_impl) -create_masked_upper_or_lower("MaskedType.upper", upper_impl) -create_masked_upper_or_lower("MaskedType.lower", lower_impl) diff --git a/python/cudf/cudf/core/udf/strings_typing.py b/python/cudf/cudf/core/udf/strings_typing.py index e065c0ad75d8..2bc55593931a 100644 --- a/python/cudf/cudf/core/udf/strings_typing.py +++ b/python/cudf/cudf/core/udf/strings_typing.py @@ -51,7 +51,6 @@ def return_as(self): return ManagedUDFString() -@register_model(StringView) class stringview_model(models.StructModel): # from string_view.hpp: _members = ( @@ -70,7 +69,6 @@ def __init__(self, dmm, fe_type): super().__init__(dmm, fe_type, self._members) -@register_model(UDFString) class udf_string_model(models.StructModel): # from udf_string.hpp: # private: @@ -91,7 +89,6 @@ def __init__(self, dmm, fe_type): udf_string = UDFString() -@register_model(ManagedUDFString) class managed_udf_string_model(models.StructModel): _members = (("meminfo", types.voidptr), ("udf_string", udf_string)) @@ -142,7 +139,6 @@ def prepare_args(self, ty, val, **kwargs): # String functions -@cuda_decl_registry.register_global(len) class StringLength(AbstractTemplate): """ provide the length of a cudf::string_view like struct @@ -161,7 +157,6 @@ def NRT_decref(st): pass -@cuda_decl_registry.register_global(NRT_decref) class NRT_decref_typing(AbstractTemplate): def generic(self, args, kws): if isinstance(args[0], ManagedUDFString): @@ -300,23 +295,33 @@ def resolve_replace(self, mod): ) -@cuda_decl_registry.register_attr class ManagedUDFStringAttrs(StringViewAttrs): key = managed_udf_string -cuda_decl_registry.register_attr(StringViewAttrs) -cuda_decl_registry.register_attr(ManagedUDFStringAttrs) +def register_strings_typing(): -register_stringview_binaryop(operator.eq, types.boolean) -register_stringview_binaryop(operator.ne, types.boolean) -register_stringview_binaryop(operator.lt, types.boolean) -register_stringview_binaryop(operator.gt, types.boolean) -register_stringview_binaryop(operator.le, types.boolean) -register_stringview_binaryop(operator.ge, types.boolean) + register_model(StringView)(stringview_model) + register_model(UDFString)(udf_string_model) + register_model(ManagedUDFString)(managed_udf_string_model) -# st in other -register_stringview_binaryop(operator.contains, types.boolean) + cuda_decl_registry.register_global(NRT_decref)(NRT_decref_typing) + + cuda_decl_registry.register_global(len)(StringLength) + + cuda_decl_registry.register_attr(StringViewAttrs) + cuda_decl_registry.register_attr(ManagedUDFStringAttrs) + + register_stringview_binaryop(operator.eq, types.boolean) + register_stringview_binaryop(operator.ne, types.boolean) + register_stringview_binaryop(operator.lt, types.boolean) + register_stringview_binaryop(operator.gt, types.boolean) + register_stringview_binaryop(operator.le, types.boolean) + register_stringview_binaryop(operator.ge, types.boolean) + + # st in other + register_stringview_binaryop(operator.contains, types.boolean) + + # st + other + register_stringview_binaryop(operator.add, managed_udf_string) -# st + other -register_stringview_binaryop(operator.add, managed_udf_string) diff --git a/python/cudf/cudf/core/udf/utils.py b/python/cudf/cudf/core/udf/utils.py index 682f5267bebc..8cc8ed01b9f0 100644 --- a/python/cudf/cudf/core/udf/utils.py +++ b/python/cudf/cudf/core/udf/utils.py @@ -38,6 +38,7 @@ STRING_TYPES, TIMEDELTA_TYPES, ) +from cudf.core.udf._compile_udf import compile_udf, make_cache_key if TYPE_CHECKING: from collections.abc import Callable @@ -63,11 +64,6 @@ precompiled: cachetools.LRUCache = cachetools.LRUCache(maxsize=32) -# This cache is keyed on the (signature, code, closure variables) of UDFs, so -# it can hit for distinct functions that are similar. The lru_cache wrapping -# compile_udf misses for these similar functions, but doesn't need to serialize -# closure variables to check for a hit. -_udf_code_cache: cachetools.LRUCache = cachetools.LRUCache(maxsize=32) UDF_SHIM_FILE = os.path.join( @@ -144,75 +140,6 @@ def _mask_get(mask, pos): return (mask[pos // MASK_BITSIZE] >> (pos % MASK_BITSIZE)) & 1 -def make_cache_key(udf, sig): - """ - Build a cache key for a user defined function. Used to avoid - recompiling the same function for the same set of types - """ - codebytes = udf.__code__.co_code - constants = udf.__code__.co_consts - names = udf.__code__.co_names - - if udf.__closure__ is not None: - cvars = tuple(x.cell_contents for x in udf.__closure__) - cvarbytes = dumps(cvars) - else: - cvarbytes = b"" - - return names, constants, codebytes, cvarbytes, sig - - -def compile_udf(udf, type_signature): - """Compile ``udf`` with `numba` - - Compile a python callable function ``udf`` with - `numba.cuda.compile_ptx_for_current_device(device=True)` using - ``type_signature`` into CUDA PTX together with the generated output type. - - The output is expected to be passed to the PTX parser in `libcudf` - to generate a CUDA device function to be inlined into CUDA kernels, - compiled at runtime and launched. - - Parameters - ---------- - udf: - a python callable function - - type_signature: - a tuple that specifies types of each of the input parameters of ``udf``. - The types should be one in `numba.types` and could be converted from - numpy types with `numba.numpy_support.from_dtype(...)`. - - Returns - ------- - ptx_code: - The compiled CUDA PTX - - output_type: - An numpy type - - """ - key = make_cache_key(udf, type_signature) - res = _udf_code_cache.get(key) - if res: - return res - - # We haven't compiled a function like this before, so need to fall back to - # compilation with Numba - ptx_code, return_type = cuda.compile_ptx_for_current_device( - udf, type_signature, device=True - ) - if not isinstance(return_type, MaskedType): - output_type = numpy_support.as_dtype(return_type).type - else: - output_type = return_type - - # Populate the cache for this function - res = (ptx_code, output_type) - _udf_code_cache[key] = res - - return res - def _generate_cache_key(frame, func: Callable, args, suffix="__APPLY_UDF"): """Create a cache key that uniquely identifies a compilation. diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py index 66b98ea513fb..240702d423ed 100644 --- a/python/cudf/cudf/utils/_numba.py +++ b/python/cudf/cudf/utils/_numba.py @@ -4,6 +4,15 @@ import numba from numba import config as numba_config from packaging import version +from pickle import dumps +import cachetools + + +# This cache is keyed on the (signature, code, closure variables) of UDFs, so +# it can hit for distinct functions that are similar. The lru_cache wrapping +# compile_udf misses for these similar functions, but doesn't need to serialize +# closure variables to check for a hit. +_udf_code_cache: cachetools.LRUCache = cachetools.LRUCache(maxsize=32) # Avoids using contextlib.contextmanager due to additional overhead @@ -28,3 +37,78 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: ) if self.is_numba_lt_061: numba_config.CAPTURED_ERRORS = self.CAPTURED_ERRORS + + + + +def make_cache_key(udf, sig): + """ + Build a cache key for a user defined function. Used to avoid + recompiling the same function for the same set of types + """ + codebytes = udf.__code__.co_code + constants = udf.__code__.co_consts + names = udf.__code__.co_names + + if udf.__closure__ is not None: + cvars = tuple(x.cell_contents for x in udf.__closure__) + cvarbytes = dumps(cvars) + else: + cvarbytes = b"" + + return names, constants, codebytes, cvarbytes, sig + + + +def compile_udf(udf, type_signature): + """Compile ``udf`` with `numba` + + Compile a python callable function ``udf`` with + `numba.cuda.compile_ptx_for_current_device(device=True)` using + ``type_signature`` into CUDA PTX together with the generated output type. + + The output is expected to be passed to the PTX parser in `libcudf` + to generate a CUDA device function to be inlined into CUDA kernels, + compiled at runtime and launched. + + Parameters + ---------- + udf: + a python callable function + + type_signature: + a tuple that specifies types of each of the input parameters of ``udf``. + The types should be one in `numba.types` and could be converted from + numpy types with `numba.numpy_support.from_dtype(...)`. + + Returns + ------- + ptx_code: + The compiled CUDA PTX + + output_type: + An numpy type + + """ + key = make_cache_key(udf, type_signature) + res = _udf_code_cache.get(key) + if res: + return res + + # We haven't compiled a function like this before, so need to fall back to + # compilation with Numba + ptx_code, return_type = cuda.compile_ptx_for_current_device( + udf, type_signature, device=True + ) + breakpoint() + if not isinstance(return_type, MaskedType): + output_type = numpy_support.as_dtype(return_type).type + else: + output_type = return_type + + # Populate the cache for this function + res = (ptx_code, output_type) + _udf_code_cache[key] = res + + return res + From c8cacfdd9ab39990ed3baa6ca2d84b58216ed07c Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 9 Sep 2025 06:28:16 -0700 Subject: [PATCH 2/2] importing, running --- python/cudf/cudf/core/udf/masked_lowering.py | 200 ++++++- python/cudf/cudf/core/udf/masked_typing.py | 97 ++- python/cudf/cudf/core/udf/strings_lowering.py | 565 +++++++----------- python/cudf/cudf/core/udf/udf_kernel_base.py | 3 +- python/cudf/cudf/core/udf/utils.py | 6 +- python/cudf/cudf/utils/_numba.py | 16 +- 6 files changed, 453 insertions(+), 434 deletions(-) diff --git a/python/cudf/cudf/core/udf/masked_lowering.py b/python/cudf/cudf/core/udf/masked_lowering.py index f33bef671805..db2c75fc09d8 100644 --- a/python/cudf/cudf/core/udf/masked_lowering.py +++ b/python/cudf/cudf/core/udf/masked_lowering.py @@ -23,10 +23,37 @@ NAType, _supported_masked_types, ) -from cudf.core.udf.strings_typing import managed_udf_string +from cudf.core.udf.strings_lowering import ( + contains_impl, + count_impl, + endswith_impl, + find_impl, + isalnum_impl, + isalpha_impl, + isdecimal_impl, + isdigit_impl, + islower_impl, + isspace_impl, + istitle_impl, + isupper_impl, + len_impl, + lower_impl, + lstrip_impl, + replace_impl, + rfind_impl, + rstrip_impl, + startswith_impl, + strip_impl, + upper_impl, +) +from cudf.core.udf.strings_typing import ( + managed_udf_string, + size_type, + string_view, + udf_string, +) -@cuda_lowering_registry.lower_constant(NAType) def constant_na(context, builder, ty, pyval): # This handles None, etc. return context.get_dummy_value() @@ -233,7 +260,6 @@ def register_const_op(op): cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op) - @cuda_lower(operator.is_, MaskedType, NAType) @cuda_lower(operator.is_, NAType, MaskedType) def masked_scalar_is_null_impl(context, builder, sig, args): @@ -401,7 +427,133 @@ def lower_constant_masked(context, builder, ty, val): masked.valid = context.get_constant(types.boolean, val.valid) return masked._getvalue() + +def masked_len_impl(context, builder, sig, args): + ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + masked_sv_ty = sig.args[0] + masked_sv = cgutils.create_struct_proxy(masked_sv_ty)( + context, builder, value=args[0] + ) + result = len_impl( + context, builder, size_type(string_view), (masked_sv.value,) + ) + ret.value = result + ret.valid = masked_sv.valid + + return ret._getvalue() + + +def masked_string_view_replace_impl(context, builder, sig, args): + ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + src_masked, to_replace_masked, replacement_masked = _masked_proxies( + context, builder, MaskedType(string_view), *args + ) + result = replace_impl( + context, + builder, + nb_signature(udf_string, string_view, string_view, string_view), + (src_masked.value, to_replace_masked.value, replacement_masked.value), + ) + + ret.value = result + ret.valid = builder.and_( + builder.and_(src_masked.valid, to_replace_masked.valid), + replacement_masked.valid, + ) + + return ret._getvalue() + + +def create_masked_binary_string_func(op, cuda_func, retty): + """ + Provide a wrapper around numba's low-level extension API which + produces the boilerplate needed to implement a binary function + of two masked strings. + """ + + def masked_binary_func_impl(context, builder, sig, args): + ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + + lhs_masked = cgutils.create_struct_proxy(sig.args[0])( + context, builder, value=args[0] + ) + rhs_masked = cgutils.create_struct_proxy(sig.args[0])( + context, builder, value=args[1] + ) + + result = cuda_func( + context, + builder, + nb_signature(retty, string_view, string_view), + (lhs_masked.value, rhs_masked.value), + ) + + ret.value = result + ret.valid = builder.and_(lhs_masked.valid, rhs_masked.valid) + + return ret._getvalue() + + cuda_lower(op, MaskedType(string_view), MaskedType(string_view))( + masked_binary_func_impl + ) + + +def create_masked_unary_identifier_func(op, cuda_func): + """ + Provide a wrapper around numba's low-level extension API which + produces the boilerplate needed to implement a unary function + of a masked string. + """ + + def masked_unary_func_impl(context, builder, sig, args): + ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + masked_str = cgutils.create_struct_proxy(sig.args[0])( + context, builder, value=args[0] + ) + + result = cuda_func( + context, + builder, + types.boolean(string_view, string_view), + (masked_str.value,), + ) + ret.value = result + ret.valid = masked_str.valid + return ret._getvalue() + + cuda_lower(op, MaskedType(string_view))(masked_unary_func_impl) + + +def create_masked_upper_or_lower(op, cuda_func): + def upper_or_lower_impl(context, builder, sig, args): + ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + masked_str = cgutils.create_struct_proxy(sig.args[0])( + context, builder, value=args[0] + ) + + result = cuda_func( + context, + builder, + udf_string(string_view), + (masked_str.value,), + ) + ret.value = result + ret.valid = masked_str.valid + return ret._getvalue() + + cuda_lower(op, MaskedType(string_view))(upper_or_lower_impl) + + +def _masked_proxies(context, builder, maskedty, *args): + return tuple( + cgutils.create_struct_proxy(maskedty)(context, builder, value=arg) + for arg in args + ) + + def register_masked_lowering(): + cuda_lowering_registry.lower_constant(NAType)(constant_na) + # register all lowering at init for binary_op in arith_ops + bitwise_ops + comparison_ops: register_arithmetic_op(binary_op) @@ -415,4 +567,46 @@ def register_masked_lowering(): register_unary_op(unary_op) register_unary_op(abs) + cuda_lower(len, MaskedType(string_view))(masked_len_impl) + cuda_lower(len, MaskedType(udf_string))(masked_len_impl) + cuda_lower( + "MaskedType.replace", + MaskedType(string_view), + MaskedType(string_view), + MaskedType(string_view), + )(masked_string_view_replace_impl) + + create_masked_binary_string_func( + "MaskedType.strip", strip_impl, udf_string + ) + create_masked_binary_string_func( + "MaskedType.lstrip", lstrip_impl, udf_string + ) + create_masked_binary_string_func( + "MaskedType.rstrip", rstrip_impl, udf_string + ) + create_masked_binary_string_func( + "MaskedType.startswith", + startswith_impl, + types.boolean, + ) + create_masked_binary_string_func( + "MaskedType.endswith", endswith_impl, types.boolean + ) + create_masked_binary_string_func("MaskedType.find", find_impl, size_type) + create_masked_binary_string_func("MaskedType.rfind", rfind_impl, size_type) + create_masked_binary_string_func("MaskedType.count", count_impl, size_type) + create_masked_binary_string_func( + operator.contains, contains_impl, types.boolean + ) + create_masked_unary_identifier_func("MaskedType.isalnum", isalnum_impl) + create_masked_unary_identifier_func("MaskedType.isalpha", isalpha_impl) + create_masked_unary_identifier_func("MaskedType.isdigit", isdigit_impl) + create_masked_unary_identifier_func("MaskedType.isupper", isupper_impl) + create_masked_unary_identifier_func("MaskedType.islower", islower_impl) + create_masked_unary_identifier_func("MaskedType.isspace", isspace_impl) + create_masked_unary_identifier_func("MaskedType.isdecimal", isdecimal_impl) + create_masked_unary_identifier_func("MaskedType.istitle", istitle_impl) + create_masked_upper_or_lower("MaskedType.upper", upper_impl) + create_masked_upper_or_lower("MaskedType.lower", lower_impl) diff --git a/python/cudf/cudf/core/udf/masked_typing.py b/python/cudf/cudf/core/udf/masked_typing.py index a58f88f9d09f..e9ac260f2508 100644 --- a/python/cudf/cudf/core/udf/masked_typing.py +++ b/python/cudf/cudf/core/udf/masked_typing.py @@ -187,47 +187,12 @@ def __eq__(self, other): # For typing a Masked constant value defined outside a kernel (e.g. captured in # a closure). -@typeof_impl.register(api.Masked) def typeof_masked(val, c): return MaskedType(typeof(val.value)) -# Implemented typing for Masked(value, valid) - the construction of a Masked -# type in a kernel. -@cuda_decl_registry.register -class MaskedConstructor(ConcreteTemplate): - key = api.Masked - cases = [ - nb_signature(MaskedType(t), t, types.boolean) - for t in _supported_masked_types - ] - - -# Typing for `api.Masked` -@cuda_decl_registry.register_attr -class ClassesTemplate(AttributeTemplate): - key = types.Module(api) - - def resolve_Masked(self, mod): - return types.Function(MaskedConstructor) - - -# Registration of the global is also needed for Numba to type api.Masked -cuda_decl_registry.register_global(api, types.Module(api)) -# For typing bare Masked (as in `from .api import Masked` -cuda_decl_registry.register_global( - api.Masked, types.Function(MaskedConstructor) -) - - -# Provide access to `m.value` and `m.valid` in a kernel for a Masked `m`. -make_attribute_wrapper(MaskedType, "value", "value") -make_attribute_wrapper(MaskedType, "valid", "valid") - - # Tell numba how `MaskedType` is constructed on the backend in terms # of primitive things that exist at the LLVM level -@register_model(MaskedType) class MaskedModel(models.StructModel): def __init__(self, dmm, fe_type): # This struct has two members, a value and a validity @@ -268,7 +233,6 @@ def unify(self, context, other): na_type = NAType() -@typeof_impl.register(type(NA)) def typeof_na(val, c): """ Tie instances of _NAType (cudf.NA) to our NAType. @@ -278,9 +242,6 @@ def typeof_na(val, c): return na_type -register_model(NAType)(models.OpaqueModel) - - # Ultimately, we want numba to produce PTX code that specifies how to implement # an operation on two singular `Masked` structs together, which is defined # as producing a new `Masked` with the right validity and if valid, @@ -366,7 +327,6 @@ def generic(self, args, kws): ) -@cuda_decl_registry.register_global(operator.is_) class MaskedScalarIsNull(AbstractTemplate): """ Typing for `Masked is cudf.NA` @@ -379,7 +339,6 @@ def generic(self, args, kws): return nb_signature(types.boolean, na_type, args[1]) -@cuda_decl_registry.register_global(operator.truth) class MaskedScalarTruth(AbstractTemplate): """ Typing for `if Masked` @@ -393,7 +352,6 @@ def generic(self, args, kws): return nb_signature(types.boolean, MaskedType(types.boolean)) -@cuda_decl_registry.register_global(float) class MaskedScalarFloatCast(AbstractTemplate): """ Typing for float(Masked) @@ -408,7 +366,6 @@ def generic(self, args, kws): return nb_signature(MaskedType(types.float64), args[0]) -@cuda_decl_registry.register_global(int) class MaskedScalarIntCast(AbstractTemplate): """ Typing for int(Masked) @@ -423,7 +380,6 @@ def generic(self, args, kws): return nb_signature(MaskedType(types.int64), args[0]) -@cuda_decl_registry.register_global(abs) class MaskedScalarAbsoluteValue(AbstractTemplate): """ Typing for the builtin function abs. Returns the same @@ -450,7 +406,6 @@ def generic(self, args, kws): return nb_signature(MaskedType(return_type), args[0]) -@cuda_decl_registry.register_global(api.pack_return) class UnpackReturnToMasked(AbstractTemplate): """ Turn a returned MaskedType into its value and validity @@ -504,7 +459,6 @@ def len_typing(self, args, kws): return nb_signature(size_type, args[0]) -@register_masked_string_function(operator.add) def concat_typing(self, args, kws): if _is_valid_string_arg(args[0]) and _is_valid_string_arg(args[1]): return nb_signature( @@ -514,7 +468,6 @@ def concat_typing(self, args, kws): ) -@register_masked_string_function(operator.contains) def contains_typing(self, args, kws): if _is_valid_string_arg(args[0]) and _is_valid_string_arg(args[1]): return nb_signature( @@ -672,6 +625,43 @@ def resolve_value(self, mod): def register_masked_typing(): + typeof_impl.register(api.Masked)(typeof_masked) + + # Implemented typing for Masked(value, valid) - the construction of a Masked + # type in a kernel. + class MaskedConstructor(ConcreteTemplate): + key = api.Masked + cases = [ + nb_signature(MaskedType(t), t, types.boolean) + for t in _supported_masked_types + ] + + cuda_decl_registry.register(MaskedConstructor) + + class ClassesTemplate(AttributeTemplate): + key = types.Module(api) + + def resolve_Masked(self, mod): + return types.Function(MaskedConstructor) + + cuda_decl_registry.register_attr(ClassesTemplate) + + # Registration of the global is also needed for Numba to type api.Masked + cuda_decl_registry.register_global(api, types.Module(api)) + # For typing bare Masked (as in `from .api import Masked` + cuda_decl_registry.register_global( + api.Masked, types.Function(MaskedConstructor) + ) + + # Provide access to `m.value` and `m.valid` in a kernel for a Masked `m`. + make_attribute_wrapper(MaskedType, "value", "value") + make_attribute_wrapper(MaskedType, "valid", "valid") + + register_model(MaskedType)(MaskedModel) + + typeof_impl.register(type(NA))(typeof_na) + register_model(NAType)(models.OpaqueModel) + for binary_op in arith_ops + bitwise_ops + comparison_ops: # Every op shares the same typing class cuda_decl_registry.register_global(binary_op)(MaskedScalarArithOp) @@ -681,12 +671,19 @@ def register_masked_typing(): for unary_op in unary_ops: cuda_decl_registry.register_global(unary_op)(MaskedScalarUnaryOp) - for op in comparison_ops: cuda_decl_registry.register_global(op)(MaskedStringViewCmpOp) - - cuda_decl_registry.register_attr(MaskedStringViewAttrs) cuda_decl_registry.register_attr(MaskedManagedUDFStringAttrs) + cuda_decl_registry.register_global(operator.is_)(MaskedScalarIsNull) + cuda_decl_registry.register_global(operator.truth)(MaskedScalarTruth) + cuda_decl_registry.register_global(float)(MaskedScalarFloatCast) + cuda_decl_registry.register_global(int)(MaskedScalarIntCast) + cuda_decl_registry.register_global(abs)(MaskedScalarAbsoluteValue) + cuda_decl_registry.register_global(api.pack_return)(UnpackReturnToMasked) + + register_masked_string_function(len)(len_typing) + register_masked_string_function(operator.add)(concat_typing) + register_masked_string_function(operator.contains)(contains_typing) diff --git a/python/cudf/cudf/core/udf/strings_lowering.py b/python/cudf/cudf/core/udf/strings_lowering.py index 1fb26b0ab785..54eb3d8a8147 100644 --- a/python/cudf/cudf/core/udf/strings_lowering.py +++ b/python/cudf/cudf/core/udf/strings_lowering.py @@ -17,7 +17,6 @@ get_character_flags_table_ptr, get_special_case_mapping_table_ptr, ) -from cudf.core.udf.masked_typing import MaskedType from cudf.core.udf.strings_typing import ( NRT_decref, managed_udf_string, @@ -126,7 +125,6 @@ def _declare_upper_or_lower(func): # casts -@cuda_lowering_registry.lower_cast(types.StringLiteral, string_view) def cast_string_literal_to_string_view(context, builder, fromty, toty, val): """ Cast a literal to a string_view @@ -146,7 +144,6 @@ def cast_string_literal_to_string_view(context, builder, fromty, toty, val): return sv._getvalue() -@cuda_lowering_registry.lower_cast(string_view, managed_udf_string) def cast_string_view_to_managed_udf_string( context, builder, fromty, toty, val ): @@ -178,7 +175,6 @@ def cast_string_view_to_managed_udf_string( return managed._getvalue() -@cuda_lowering_registry.lower_cast(managed_udf_string, string_view) def cast_managed_udf_string_to_string_view( context, builder, fromty, toty, val ): @@ -225,12 +221,6 @@ def call_create_string_view_from_udf_string(udf_str, sv): # Return string setitem impl with an extra incref -@cuda_lower( - operator.setitem, - types.CPointer(managed_udf_string), - types.Integer, - types.Any, -) def setitem_cpointer_managed_udf_string(context, builder, sig, args): base_ptr, idx, val = args elem_ptr = builder.gep(base_ptr, [idx]) @@ -246,7 +236,6 @@ def call_len_string_view(st): return _string_view_len(st) -@cuda_lower(len, string_view) def len_impl(context, builder, sig, args): sv_ptr = builder.alloca(args[0].type) builder.store(args[0], sv_ptr) @@ -260,7 +249,6 @@ def len_impl(context, builder, sig, args): return result -@cuda_lower(NRT_decref, managed_udf_string) def decref_managed_udf_string(context, builder, sig, args): managed_ptr = args[0] managed = cgutils.create_struct_proxy(managed_udf_string)( @@ -279,7 +267,6 @@ def call_concat_string_view(result, lhs, rhs): return _concat_string_view(result, lhs, rhs) -@cuda_lower(operator.add, string_view, string_view) def concat_impl(context, builder, sig, args): lhs_ptr = builder.alloca(args[0].type) rhs_ptr = builder.alloca(args[1].type) @@ -315,7 +302,6 @@ def call_string_view_replace(result, src, to_replace, replacement): return _string_view_replace(result, src, to_replace, replacement) -@cuda_lower("StringView.replace", string_view, string_view, string_view) def replace_impl(context, builder, sig, args): src_ptr = builder.alloca(args[0].type) to_replace_ptr = builder.alloca(args[1].type) @@ -349,479 +335,330 @@ def replace_impl(context, builder, sig, args): return managed._getvalue() -def create_binary_string_func(binary_func, retty): +def register_binary_string_func(binary_func, cuda_func, retty): """ Provide a wrapper around numba's low-level extension API which produces the boilerplate needed to implement a binary function of two strings. """ - def deco(cuda_func): - @cuda_lower(binary_func, string_view, string_view) - def binary_func_impl(context, builder, sig, args): - lhs_ptr = builder.alloca(args[0].type) - rhs_ptr = builder.alloca(args[1].type) - builder.store(args[0], lhs_ptr) - builder.store(args[1], rhs_ptr) - - # these conditional statements should compile out - if retty != udf_string: - # binary function of two strings yielding a fixed-width type - # example: str.startswith(other) -> bool - # shim functions can return the value through nb_retval - result = context.compile_internal( - builder, - cuda_func, - nb_signature(retty, _STR_VIEW_PTR, _STR_VIEW_PTR), - (lhs_ptr, rhs_ptr), - ) - return result - else: - # binary function of two strings yielding a new string - # example: str.strip(other) -> str - # shim functions can not return a struct due to C linkage - # so we create a new udf_string and pass a pointer to it - # for the shim function to write the output to. The return - # value of compile_internal is therefore discarded (although - # this may change in the future if we need to return error - # codes, for instance). - - managed_ptr = builder.alloca( - context.data_model_manager[ - managed_udf_string - ].get_value_type() - ) - udf_str_ptr = builder.gep( - managed_ptr, [ir.IntType(32)(0), ir.IntType(32)(1)] - ) - - meminfo = context.compile_internal( - builder, - cuda_func, - types.voidptr( - _UDF_STRING_PTR, _STR_VIEW_PTR, _STR_VIEW_PTR - ), - (udf_str_ptr, lhs_ptr, rhs_ptr), - ) - managed = cgutils.create_struct_proxy(managed_udf_string)( - context, - builder, - value=builder.load( - managed_ptr - ), # {i8*, {i8*, i32, i32}}* -> {i8*, {i8*, i32, i32}} - ) - managed.meminfo = meminfo - - return managed._getvalue() - - # binary_func can be attribute-like: str.binary_func - # or operator-like: binary_func(str, other) - if isinstance(binary_func, str): - binary_func_impl = cuda_lower( - f"StringView.{binary_func}", string_view, string_view - )(binary_func_impl) - binary_func_impl = cuda_lower( - f"ManagedUDFString.{binary_func}", string_view, string_view - )(binary_func_impl) + @cuda_lower(binary_func, string_view, string_view) + def binary_func_impl(context, builder, sig, args): + lhs_ptr = builder.alloca(args[0].type) + rhs_ptr = builder.alloca(args[1].type) + builder.store(args[0], lhs_ptr) + builder.store(args[1], rhs_ptr) + + # these conditional statements should compile out + if retty != udf_string: + # binary function of two strings yielding a fixed-width type + # example: str.startswith(other) -> bool + # shim functions can return the value through nb_retval + result = context.compile_internal( + builder, + cuda_func, + nb_signature(retty, _STR_VIEW_PTR, _STR_VIEW_PTR), + (lhs_ptr, rhs_ptr), + ) + return result else: - binary_func_impl = cuda_lower( - binary_func, string_view, string_view - )(binary_func_impl) + # binary function of two strings yielding a new string + # example: str.strip(other) -> str + # shim functions can not return a struct due to C linkage + # so we create a new udf_string and pass a pointer to it + # for the shim function to write the output to. The return + # value of compile_internal is therefore discarded (although + # this may change in the future if we need to return error + # codes, for instance). + + managed_ptr = builder.alloca( + context.data_model_manager[managed_udf_string].get_value_type() + ) + udf_str_ptr = builder.gep( + managed_ptr, [ir.IntType(32)(0), ir.IntType(32)(1)] + ) + + meminfo = context.compile_internal( + builder, + cuda_func, + types.voidptr(_UDF_STRING_PTR, _STR_VIEW_PTR, _STR_VIEW_PTR), + (udf_str_ptr, lhs_ptr, rhs_ptr), + ) + managed = cgutils.create_struct_proxy(managed_udf_string)( + context, + builder, + value=builder.load( + managed_ptr + ), # {i8*, {i8*, i32, i32}}* -> {i8*, {i8*, i32, i32}} + ) + managed.meminfo = meminfo + + return managed._getvalue() + + # binary_func can be attribute-like: str.binary_func + # or operator-like: binary_func(str, other) + if isinstance(binary_func, str): + cuda_lower(f"StringView.{binary_func}", string_view, string_view)( + binary_func_impl + ) + cuda_lower( + f"ManagedUDFString.{binary_func}", string_view, string_view + )(binary_func_impl) + else: + cuda_lower(binary_func, string_view, string_view)(binary_func_impl) + + +def create_unary_identifier_func(id_func, cuda_func): + """ + Provide a wrapper around numba's low-level extension API which + produces the boilerplate needed to implement a unary function + of a string. + """ + + def id_func_impl(context, builder, sig, args): + str_ptr = builder.alloca(args[0].type) + builder.store(args[0], str_ptr) + + # Lookup table required for conversion functions + # must be resolved at runtime after context initialization, + # therefore cannot be a global variable + tbl_ptr = context.get_constant( + types.uintp, get_character_flags_table_ptr() + ) + result = context.compile_internal( + builder, + cuda_func, + nb_signature(types.boolean, _STR_VIEW_PTR, types.uintp), + (str_ptr, tbl_ptr), + ) - return binary_func_impl + return result - return deco + cuda_lower(f"StringView.{id_func}", string_view)(id_func_impl) + cuda_lower(f"UDFString.{id_func}", string_view)(id_func_impl) + + +def create_upper_or_lower(id_func, cuda_func): + """ + Provide a wrapper around numba's low-level extension API which + produces the boilerplate needed to implement either the upper + or lower attrs of a string view. + """ + + def id_func_impl(context, builder, sig, args): + str_ptr = builder.alloca(args[0].type) + builder.store(args[0], str_ptr) + + # Lookup table required for conversion functions + # must be resolved at runtime after context initialization, + # therefore cannot be a global variable + flags_tbl_ptr = context.get_constant( + types.uintp, get_character_flags_table_ptr() + ) + cases_tbl_ptr = context.get_constant( + types.uintp, get_character_cases_table_ptr() + ) + special_tbl_ptr = context.get_constant( + types.uintp, get_special_case_mapping_table_ptr() + ) + + managed_ptr = builder.alloca( + context.data_model_manager[managed_udf_string].get_value_type() + ) + udf_str_ptr = builder.gep( + managed_ptr, [ir.IntType(32)(0), ir.IntType(32)(1)] + ) + meminfo = context.compile_internal( + builder, + cuda_func, + types.voidptr( + _UDF_STRING_PTR, + _STR_VIEW_PTR, + types.uintp, + types.uintp, + types.uintp, + ), + ( + udf_str_ptr, + str_ptr, + flags_tbl_ptr, + cases_tbl_ptr, + special_tbl_ptr, + ), + ) + managed = cgutils.create_struct_proxy(managed_udf_string)( + context, + builder, + value=builder.load( + managed_ptr + ), # {i8*, {i8*, i32, i32}}* -> {i8*, {i8*, i32, i32}} + ) + managed.meminfo = meminfo + return managed._getvalue() + + cuda_lower(f"StringView.{id_func}", string_view)(id_func_impl) + cuda_lower(f"UDFString.{id_func}", string_view)(id_func_impl) -@create_binary_string_func(operator.contains, types.boolean) def contains_impl(st, substr): return _string_view_contains(st, substr) -@create_binary_string_func(operator.eq, types.boolean) def eq_impl(st, rhs): return _string_view_eq(st, rhs) -@create_binary_string_func(operator.ne, types.boolean) def ne_impl(st, rhs): return _string_view_ne(st, rhs) -@create_binary_string_func(operator.ge, types.boolean) def ge_impl(st, rhs): return _string_view_ge(st, rhs) -@create_binary_string_func(operator.le, types.boolean) def le_impl(st, rhs): return _string_view_le(st, rhs) -@create_binary_string_func(operator.gt, types.boolean) def gt_impl(st, rhs): return _string_view_gt(st, rhs) -@create_binary_string_func(operator.lt, types.boolean) def lt_impl(st, rhs): return _string_view_lt(st, rhs) -@create_binary_string_func("strip", udf_string) def strip_impl(result, to_strip, strip_char): return _string_view_strip(result, to_strip, strip_char) -@create_binary_string_func("lstrip", udf_string) def lstrip_impl(result, to_strip, strip_char): return _string_view_lstrip(result, to_strip, strip_char) -@create_binary_string_func("rstrip", udf_string) def rstrip_impl(result, to_strip, strip_char): return _string_view_rstrip(result, to_strip, strip_char) -@create_binary_string_func("startswith", types.boolean) def startswith_impl(sv, substr): return _string_view_startswith(sv, substr) -@create_binary_string_func("endswith", types.boolean) def endswith_impl(sv, substr): return _string_view_endswith(sv, substr) -@create_binary_string_func("count", size_type) def count_impl(st, substr): return _string_view_count(st, substr) -@create_binary_string_func("find", size_type) def find_impl(sv, substr): return _string_view_find(sv, substr) -@create_binary_string_func("rfind", size_type) def rfind_impl(sv, substr): return _string_view_rfind(sv, substr) -def create_unary_identifier_func(id_func): - """ - Provide a wrapper around numba's low-level extension API which - produces the boilerplate needed to implement a unary function - of a string. - """ - - def deco(cuda_func): - @cuda_lower(f"StringView.{id_func}", string_view) - @cuda_lower(f"UDFString.{id_func}", string_view) - def id_func_impl(context, builder, sig, args): - str_ptr = builder.alloca(args[0].type) - builder.store(args[0], str_ptr) - - # Lookup table required for conversion functions - # must be resolved at runtime after context initialization, - # therefore cannot be a global variable - tbl_ptr = context.get_constant( - types.uintp, get_character_flags_table_ptr() - ) - result = context.compile_internal( - builder, - cuda_func, - nb_signature(types.boolean, _STR_VIEW_PTR, types.uintp), - (str_ptr, tbl_ptr), - ) - - return result - - return id_func_impl - - return deco - - -def create_upper_or_lower(id_func): - """ - Provide a wrapper around numba's low-level extension API which - produces the boilerplate needed to implement either the upper - or lower attrs of a string view. - """ - - def deco(cuda_func): - @cuda_lower(f"StringView.{id_func}", string_view) - @cuda_lower(f"UDFString.{id_func}", string_view) - def id_func_impl(context, builder, sig, args): - str_ptr = builder.alloca(args[0].type) - builder.store(args[0], str_ptr) - - # Lookup table required for conversion functions - # must be resolved at runtime after context initialization, - # therefore cannot be a global variable - flags_tbl_ptr = context.get_constant( - types.uintp, get_character_flags_table_ptr() - ) - cases_tbl_ptr = context.get_constant( - types.uintp, get_character_cases_table_ptr() - ) - special_tbl_ptr = context.get_constant( - types.uintp, get_special_case_mapping_table_ptr() - ) - - managed_ptr = builder.alloca( - context.data_model_manager[managed_udf_string].get_value_type() - ) - udf_str_ptr = builder.gep( - managed_ptr, [ir.IntType(32)(0), ir.IntType(32)(1)] - ) - meminfo = context.compile_internal( - builder, - cuda_func, - types.voidptr( - _UDF_STRING_PTR, - _STR_VIEW_PTR, - types.uintp, - types.uintp, - types.uintp, - ), - ( - udf_str_ptr, - str_ptr, - flags_tbl_ptr, - cases_tbl_ptr, - special_tbl_ptr, - ), - ) - managed = cgutils.create_struct_proxy(managed_udf_string)( - context, - builder, - value=builder.load( - managed_ptr - ), # {i8*, {i8*, i32, i32}}* -> {i8*, {i8*, i32, i32}} - ) - managed.meminfo = meminfo - return managed._getvalue() - - return id_func_impl - - return deco - - -@create_upper_or_lower("upper") def upper_impl(result, st, flags, cases, special): return _string_view_upper(result, st, flags, cases, special) -@create_upper_or_lower("lower") def lower_impl(result, st, flags, cases, special): return _string_view_lower(result, st, flags, cases, special) -@create_unary_identifier_func("isdigit") def isdigit_impl(st, tbl): return _string_view_isdigit(st, tbl) -@create_unary_identifier_func("isalnum") def isalnum_impl(st, tbl): return _string_view_isalnum(st, tbl) -@create_unary_identifier_func("isalpha") def isalpha_impl(st, tbl): return _string_view_isalpha(st, tbl) -@create_unary_identifier_func("isnumeric") def isnumeric_impl(st, tbl): return _string_view_isnumeric(st, tbl) -@create_unary_identifier_func("isdecimal") def isdecimal_impl(st, tbl): return _string_view_isdecimal(st, tbl) -@create_unary_identifier_func("isspace") def isspace_impl(st, tbl): return _string_view_isspace(st, tbl) -@create_unary_identifier_func("isupper") def isupper_impl(st, tbl): return _string_view_isupper(st, tbl) -@create_unary_identifier_func("islower") def islower_impl(st, tbl): return _string_view_islower(st, tbl) -@create_unary_identifier_func("istitle") def istitle_impl(st, tbl): return _string_view_istitle(st, tbl) -@cuda_lower(len, MaskedType(string_view)) -@cuda_lower(len, MaskedType(udf_string)) -def masked_len_impl(context, builder, sig, args): - ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) - masked_sv_ty = sig.args[0] - masked_sv = cgutils.create_struct_proxy(masked_sv_ty)( - context, builder, value=args[0] - ) - result = len_impl( - context, builder, size_type(string_view), (masked_sv.value,) - ) - ret.value = result - ret.valid = masked_sv.valid - - return ret._getvalue() - - -def _masked_proxies(context, builder, maskedty, *args): - return tuple( - cgutils.create_struct_proxy(maskedty)(context, builder, value=arg) - for arg in args - ) - - -@cuda_lower( - "MaskedType.replace", - MaskedType(string_view), - MaskedType(string_view), - MaskedType(string_view), -) -def masked_string_view_replace_impl(context, builder, sig, args): - ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) - src_masked, to_replace_masked, replacement_masked = _masked_proxies( - context, builder, MaskedType(string_view), *args +def register_strings_lowering(): + # casts + cuda_lowering_registry.lower_cast(types.StringLiteral, string_view)( + cast_string_literal_to_string_view ) - result = replace_impl( - context, - builder, - nb_signature(udf_string, string_view, string_view, string_view), - (src_masked.value, to_replace_masked.value, replacement_masked.value), + cuda_lowering_registry.lower_cast(string_view, managed_udf_string)( + cast_string_view_to_managed_udf_string ) - - ret.value = result - ret.valid = builder.and_( - builder.and_(src_masked.valid, to_replace_masked.valid), - replacement_masked.valid, + cuda_lowering_registry.lower_cast(managed_udf_string, string_view)( + cast_managed_udf_string_to_string_view ) - return ret._getvalue() + cuda_lower( + operator.setitem, + types.CPointer(managed_udf_string), + types.Integer, + types.Any, + )(setitem_cpointer_managed_udf_string) + cuda_lower(len, string_view)(len_impl) + cuda_lower(NRT_decref, managed_udf_string)(decref_managed_udf_string) -def create_masked_binary_string_func(op, cuda_func, retty): - """ - Provide a wrapper around numba's low-level extension API which - produces the boilerplate needed to implement a binary function - of two masked strings. - """ - - def masked_binary_func_impl(context, builder, sig, args): - ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) + cuda_lower(operator.add, string_view, string_view)(concat_impl) - lhs_masked = cgutils.create_struct_proxy(sig.args[0])( - context, builder, value=args[0] - ) - rhs_masked = cgutils.create_struct_proxy(sig.args[0])( - context, builder, value=args[1] - ) - - result = cuda_func( - context, - builder, - nb_signature(retty, string_view, string_view), - (lhs_masked.value, rhs_masked.value), - ) - - ret.value = result - ret.valid = builder.and_(lhs_masked.valid, rhs_masked.valid) - - return ret._getvalue() - - cuda_lower(op, MaskedType(string_view), MaskedType(string_view))( - masked_binary_func_impl + cuda_lower("StringView.replace", string_view, string_view, string_view)( + replace_impl ) - -def create_masked_unary_identifier_func(op, cuda_func): - """ - Provide a wrapper around numba's low-level extension API which - produces the boilerplate needed to implement a unary function - of a masked string. - """ - - def masked_unary_func_impl(context, builder, sig, args): - ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) - masked_str = cgutils.create_struct_proxy(sig.args[0])( - context, builder, value=args[0] - ) - - result = cuda_func( - context, - builder, - types.boolean(string_view, string_view), - (masked_str.value,), - ) - ret.value = result - ret.valid = masked_str.valid - return ret._getvalue() - - cuda_lower(op, MaskedType(string_view))(masked_unary_func_impl) - - -def create_masked_upper_or_lower(op, cuda_func): - def upper_or_lower_impl(context, builder, sig, args): - ret = cgutils.create_struct_proxy(sig.return_type)(context, builder) - masked_str = cgutils.create_struct_proxy(sig.args[0])( - context, builder, value=args[0] - ) - - result = cuda_func( - context, - builder, - udf_string(string_view), - (masked_str.value,), - ) - ret.value = result - ret.valid = masked_str.valid - return ret._getvalue() - - cuda_lower(op, MaskedType(string_view))(upper_or_lower_impl) - - -def register_strings_lowering(): - create_masked_binary_string_func("MaskedType.strip", strip_impl, udf_string) - create_masked_binary_string_func("MaskedType.lstrip", lstrip_impl, udf_string) - create_masked_binary_string_func("MaskedType.rstrip", rstrip_impl, udf_string) - create_masked_binary_string_func( - "MaskedType.startswith", - startswith_impl, - types.boolean, - ) - create_masked_binary_string_func( - "MaskedType.endswith", endswith_impl, types.boolean - ) - create_masked_binary_string_func("MaskedType.find", find_impl, size_type) - create_masked_binary_string_func("MaskedType.rfind", rfind_impl, size_type) - create_masked_binary_string_func("MaskedType.count", count_impl, size_type) - create_masked_binary_string_func( + register_binary_string_func( operator.contains, contains_impl, types.boolean ) - - - create_masked_unary_identifier_func("MaskedType.isalnum", isalnum_impl) - create_masked_unary_identifier_func("MaskedType.isalpha", isalpha_impl) - create_masked_unary_identifier_func("MaskedType.isdigit", isdigit_impl) - create_masked_unary_identifier_func("MaskedType.isupper", isupper_impl) - create_masked_unary_identifier_func("MaskedType.islower", islower_impl) - create_masked_unary_identifier_func("MaskedType.isspace", isspace_impl) - create_masked_unary_identifier_func("MaskedType.isdecimal", isdecimal_impl) - create_masked_unary_identifier_func("MaskedType.istitle", istitle_impl) - create_masked_upper_or_lower("MaskedType.upper", upper_impl) - create_masked_upper_or_lower("MaskedType.lower", lower_impl) - + register_binary_string_func(operator.eq, eq_impl, types.boolean) + register_binary_string_func(operator.ne, ne_impl, types.boolean) + register_binary_string_func(operator.ge, ge_impl, types.boolean) + register_binary_string_func(operator.le, le_impl, types.boolean) + register_binary_string_func(operator.gt, gt_impl, types.boolean) + register_binary_string_func(operator.lt, lt_impl, types.boolean) + register_binary_string_func("strip", strip_impl, udf_string) + register_binary_string_func("lstrip", lstrip_impl, udf_string) + register_binary_string_func("rstrip", rstrip_impl, udf_string) + register_binary_string_func("startswith", startswith_impl, types.boolean) + register_binary_string_func("endswith", endswith_impl, types.boolean) + register_binary_string_func("count", count_impl, size_type) + register_binary_string_func("find", find_impl, size_type) + register_binary_string_func("rfind", rfind_impl, size_type) + + create_upper_or_lower("upper", upper_impl) + create_upper_or_lower("lower", lower_impl) + + create_unary_identifier_func("isdigit", isdigit_impl) + create_unary_identifier_func("isalnum", isalnum_impl) + create_unary_identifier_func("isalpha", isalpha_impl) + create_unary_identifier_func("isnumeric", isnumeric_impl) + create_unary_identifier_func("isdecimal", isdecimal_impl) + create_unary_identifier_func("isspace", isspace_impl) + create_unary_identifier_func("isupper", isupper_impl) + create_unary_identifier_func("islower", islower_impl) + create_unary_identifier_func("istitle", istitle_impl) diff --git a/python/cudf/cudf/core/udf/udf_kernel_base.py b/python/cudf/cudf/core/udf/udf_kernel_base.py index d0c248bfb127..136d644b3255 100644 --- a/python/cudf/cudf/core/udf/udf_kernel_base.py +++ b/python/cudf/cudf/core/udf/udf_kernel_base.py @@ -18,10 +18,9 @@ _generate_cache_key, _masked_array_type_from_col, _supported_cols_from_frame, - compile_udf, precompiled as kernel_cache, ) -from cudf.utils._numba import _CUDFNumbaConfig +from cudf.utils._numba import _CUDFNumbaConfig, compile_udf from cudf.utils.performance_tracking import _performance_tracking diff --git a/python/cudf/cudf/core/udf/utils.py b/python/cudf/cudf/core/udf/utils.py index 8cc8ed01b9f0..63d618b667cc 100644 --- a/python/cudf/cudf/core/udf/utils.py +++ b/python/cudf/cudf/core/udf/utils.py @@ -3,7 +3,6 @@ import functools import os -from pickle import dumps from typing import TYPE_CHECKING import cachetools @@ -21,7 +20,6 @@ from cudf._lib import strings_udf from cudf.core.buffer import as_buffer -from cudf.core.udf.masked_typing import MaskedType from cudf.core.udf.nrt_utils import nrt_enabled from cudf.core.udf.strings_typing import ( NRT_decref, @@ -29,6 +27,7 @@ str_view_arg_handler, string_view, ) +from cudf.utils._numba import make_cache_key from cudf.utils.dtypes import ( BOOL_TYPES, CUDF_STRING_DTYPE, @@ -38,7 +37,6 @@ STRING_TYPES, TIMEDELTA_TYPES, ) -from cudf.core.udf._compile_udf import compile_udf, make_cache_key if TYPE_CHECKING: from collections.abc import Callable @@ -65,7 +63,6 @@ precompiled: cachetools.LRUCache = cachetools.LRUCache(maxsize=32) - UDF_SHIM_FILE = os.path.join( os.path.dirname(strings_udf.__file__), "..", "core", "udf", "shim.fatbin" ) @@ -140,7 +137,6 @@ def _mask_get(mask, pos): return (mask[pos // MASK_BITSIZE] >> (pos % MASK_BITSIZE)) & 1 - def _generate_cache_key(frame, func: Callable, args, suffix="__APPLY_UDF"): """Create a cache key that uniquely identifies a compilation. diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py index 240702d423ed..6064d14a0654 100644 --- a/python/cudf/cudf/utils/_numba.py +++ b/python/cudf/cudf/utils/_numba.py @@ -1,12 +1,13 @@ # Copyright (c) 2023-2025, NVIDIA CORPORATION. from __future__ import annotations -import numba -from numba import config as numba_config -from packaging import version from pickle import dumps -import cachetools +import cachetools +import numba +from numba import config as numba_config, cuda +from numba.np import numpy_support +from packaging import version # This cache is keyed on the (signature, code, closure variables) of UDFs, so # it can hit for distinct functions that are similar. The lru_cache wrapping @@ -39,8 +40,6 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: numba_config.CAPTURED_ERRORS = self.CAPTURED_ERRORS - - def make_cache_key(udf, sig): """ Build a cache key for a user defined function. Used to avoid @@ -59,7 +58,6 @@ def make_cache_key(udf, sig): return names, constants, codebytes, cvarbytes, sig - def compile_udf(udf, type_signature): """Compile ``udf`` with `numba` @@ -100,8 +98,7 @@ def compile_udf(udf, type_signature): ptx_code, return_type = cuda.compile_ptx_for_current_device( udf, type_signature, device=True ) - breakpoint() - if not isinstance(return_type, MaskedType): + if return_type.is_internal: output_type = numpy_support.as_dtype(return_type).type else: output_type = return_type @@ -111,4 +108,3 @@ def compile_udf(udf, type_signature): _udf_code_cache[key] = res return res -