Skip to content

Consolidate extension type registration #19938

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
brandon-b-miller wants to merge 3 commits into from
Closed

Consolidate extension type registration #19938

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
782b57d
refactor to better see import times
brandon-b-miller Aug 26, 2025
8979d2b
Merge branch 'branch-25.10' into refactor-cudf-import-times
brandon-b-miller Sep 8, 2025
c8cacfd
importing, running
brandon-b-miller Sep 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/cudf/cudf/core/_internals/aggregation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion python/cudf/cudf/core/column/numerical_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
28 changes: 20 additions & 8 deletions python/cudf/cudf/core/udf/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()
37 changes: 20 additions & 17 deletions python/cudf/cudf/core/udf/groupby_lowering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
71 changes: 37 additions & 34 deletions python/cudf/cudf/core/udf/groupby_typing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 (
Expand All @@ -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.
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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)
Loading