From 02c1f3c75fae3ffa8acc95f444483ef6a4706929 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:28:17 +0000 Subject: [PATCH 1/2] Replace rolling.apply implementation with numba cuda mlir implementation --- .../cudf/cudf/core/_internals/aggregation.py | 42 +----- python/cudf/cudf/core/udf/rolling_utils.py | 120 ++++++++++++++++++ python/cudf/cudf/core/window/rolling.py | 47 +++++-- python/cudf/cudf/tests/window/test_rolling.py | 42 ++++++ 4 files changed, 204 insertions(+), 47 deletions(-) create mode 100644 python/cudf/cudf/core/udf/rolling_utils.py diff --git a/python/cudf/cudf/core/_internals/aggregation.py b/python/cudf/cudf/core/_internals/aggregation.py index 54260a571bb5..799875b4249d 100644 --- a/python/cudf/cudf/core/_internals/aggregation.py +++ b/python/cudf/cudf/core/_internals/aggregation.py @@ -4,14 +4,9 @@ from typing import TYPE_CHECKING, Literal -import numpy as np -from numba.np import numpy_support - import pylibcudf as plc from cudf.api.types import is_scalar -from cudf.core.udf.utils import compile_udf -from cudf.utils.dtypes import SUPPORTED_NUMPY_TO_PYLIBCUDF_TYPES if TYPE_CHECKING: from collections.abc import Callable @@ -238,28 +233,6 @@ def any(cls) -> Self: def all(cls) -> Self: return cls(plc.aggregation.all()) - # Rolling aggregations - @classmethod - def from_udf(cls, op, *args, **kwargs) -> Self: - # Handling UDF type - nb_type = numpy_support.from_dtype(kwargs["dtype"]) - type_signature = (nb_type[:],) - ptx_code, output_dtype = compile_udf(op, type_signature) - output_np_dtype = np.dtype(output_dtype) - if output_np_dtype not in SUPPORTED_NUMPY_TO_PYLIBCUDF_TYPES: - raise TypeError( - f"Result of window function has unsupported dtype {op[1]}" - ) - - return cls( - plc.aggregation.udf( - ptx_code, - plc.DataType( - SUPPORTED_NUMPY_TO_PYLIBCUDF_TYPES[output_np_dtype] - ), - ) - ) - def make_aggregation( op: str | Callable, kwargs: dict | None = None @@ -268,15 +241,10 @@ def make_aggregation( Parameters ---------- op : str or callable - If callable, must meet one of the following requirements: - - * Is of the form lambda x: x.agg(*args, **kwargs), where - `agg` is the name of a supported aggregation. Used to - to specify aggregations that take arguments, e.g., - `lambda x: x.quantile(0.5)`. - * Is a user defined aggregation function that operates on - group values. In this case, the output dtype must be - specified in the `kwargs` dictionary. + If callable, must be of the form lambda x: x.agg(*args, **kwargs), + where `agg` is the name of a supported aggregation. Used to + specify aggregations that take arguments, e.g., + `lambda x: x.quantile(0.5)`. \*\*kwargs : dict, optional Any keyword arguments to be passed to the op. @@ -292,8 +260,6 @@ def make_aggregation( elif callable(op): if op is list: return Aggregation.collect() - elif "dtype" in kwargs: - return Aggregation.from_udf(op, **kwargs) else: return op(Aggregation) raise TypeError(f"Unknown aggregation {op}") diff --git a/python/cudf/cudf/core/udf/rolling_utils.py b/python/cudf/cudf/core/udf/rolling_utils.py new file mode 100644 index 000000000000..18c5ca9f7e6e --- /dev/null +++ b/python/cudf/cudf/core/udf/rolling_utils.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +import cupy as cp +import numpy as np +from numba_cuda_mlir import cuda, compiler +from numba_cuda_mlir.numba_cuda.core import config as _mlir_config +from numba_cuda_mlir.numba_cuda.np import numpy_support + +from cudf.core.column.column import ColumnBase, as_column +from cudf.core.udf.utils import UDFError +from cudf.utils.performance_tracking import _performance_tracking + +if TYPE_CHECKING: + from collections.abc import Callable + + +class _MLIRNumbaCudaConfig: + """Silence numba_cuda_mlir low-occupancy warnings during launch.""" + + def __enter__(self) -> None: + self._low_occupancy_warnings = _mlir_config.CUDA_LOW_OCCUPANCY_WARNINGS + _mlir_config.CUDA_LOW_OCCUPANCY_WARNINGS = 0 + + def __exit__(self, exc_type, exc_value, traceback) -> None: + _mlir_config.CUDA_LOW_OCCUPANCY_WARNINGS = self._low_occupancy_warnings + + +def _get_udf_return_type(func: Callable, value_dtype: np.dtype) -> np.dtype: + """Compile ``func`` for a 1D window of ``value_dtype`` to infer its + output dtype.""" + nb_value_type = numpy_support.from_dtype(value_dtype) + signature = (nb_value_type[::1],) + try: + _, return_type = compiler.compile( + func, signature, device=True, output="ptx" + ) + except Exception as e: + raise UDFError(str(e)) from e + return np.dtype(numpy_support.as_dtype(return_type)) + + +def _make_rolling_kernel(device_func): + @cuda.jit + def _kernel(data, start, end, out, valid, min_periods): + i = cuda.grid(1) + if i < out.size: + begin = start[i] + stop = end[i] + count = stop - begin + if count >= min_periods and count > 0: + out[i] = device_func(data[begin:stop]) + valid[i] = True + else: + valid[i] = False + + return _kernel + + +@functools.lru_cache(maxsize=32) +def _compile_or_get_kernel(func: Callable, value_dtype: np.dtype): + return_dtype = _get_udf_return_type(func, value_dtype) + device_func = cuda.jit(device=True)(func) + kernel = _make_rolling_kernel(device_func) + return kernel, return_dtype + + +@_performance_tracking +def jit_rolling_apply( + source_column: ColumnBase, + start: cp.ndarray, + end: cp.ndarray, + min_periods: int, + func: Callable, +) -> ColumnBase: + """Apply a user-defined function to each rolling window using a custom + CUDA kernel compiled with ``numba_cuda_mlir``. + + Parameters + ---------- + source_column : ColumnBase + The (non-null) numeric column the windows are drawn from. + start, end : cupy.ndarray + ``size_type`` arrays giving the absolute ``[start, end)`` row + indices of each row's window. + min_periods : int + Minimum number of observations in a window required to produce a + non-null result. + func : callable + The user-defined function. Receives a 1D array (the window) and + returns a scalar. + """ + value_dtype = source_column.dtype + kernel, return_dtype = _compile_or_get_kernel(func, value_dtype) + + n = len(source_column) + if n == 0: + return as_column(cp.empty(0, dtype=return_dtype)) + + data = source_column.values + out = cp.empty(n, dtype=return_dtype) + valid = cp.zeros(n, dtype=np.bool_) + + threads_per_block = 128 + blocks = (n + threads_per_block - 1) // threads_per_block + + with _MLIRNumbaCudaConfig(): + kernel[blocks, threads_per_block]( + data, start, end, out, valid, min_periods + ) + cuda.synchronize() + + result = as_column(out) + valid_col = as_column(valid) + mask, null_count = valid_col.as_mask() + return result.set_mask(mask, null_count) diff --git a/python/cudf/cudf/core/window/rolling.py b/python/cudf/cudf/core/window/rolling.py index 73f932110439..7bc593471fcf 100644 --- a/python/cudf/cudf/core/window/rolling.py +++ b/python/cudf/cudf/core/window/rolling.py @@ -356,6 +356,30 @@ def _plc_windows(self) -> WindowTypePair: f"not {type(self.window).__name__}" ) + def _window_start_end(self) -> tuple[cupy.ndarray, cupy.ndarray]: + """ + Return the absolute ``[start, end)`` row indices of each row's window + as ``size_type`` cupy arrays, used by the UDF (``apply``) kernel path. + """ + n = len(self.obj) + idx = cupy.arange(n, dtype=SIZE_TYPE_DTYPE) + pre, fwd = self._plc_windows + if isinstance(pre, int): + start = idx - (pre - 1) + end = idx + (fwd + 1) + else: + preceding = cupy.asarray( + ColumnBase.from_pylibcudf(pre).astype(SIZE_TYPE_DTYPE).values + ) + following = cupy.asarray( + ColumnBase.from_pylibcudf(fwd).astype(SIZE_TYPE_DTYPE).values + ) + start = idx - preceding + np.int32(1) + end = idx + following + np.int32(1) + start = cupy.clip(start, 0, n).astype(SIZE_TYPE_DTYPE) + end = cupy.clip(end, 0, n).astype(SIZE_TYPE_DTYPE) + return start, end + def _apply_agg_column( self, source_column: ColumnBase, agg_name: str | Callable, **agg_kwargs ) -> ColumnBase: @@ -363,17 +387,24 @@ def _apply_agg_column( # pandas window aggregations operate on the category values, # not the codes source_column = source_column._get_decategorized_column() # type: ignore[attr-defined] + + min_periods = 1 if self.min_periods is None else self.min_periods + + if callable(agg_name): + from cudf.core.udf.rolling_utils import jit_rolling_apply + + start, end = self._window_start_end() + return jit_rolling_apply( + source_column, start, end, min_periods, agg_name + ) + pre, fwd = self._plc_windows rolling_agg = aggregation.make_aggregation( - agg_name, - {"dtype": source_column.dtype} - if callable(agg_name) - else agg_kwargs, + agg_name, agg_kwargs ).plc_obj - min_periods = 1 if self.min_periods is None else self.min_periods - if self.min_periods == 0 and isinstance(agg_name, str): + if self.min_periods == 0: # libcudf supports min_periods=0 and returns identity values for windows with # insufficient observations: SUM and COUNT return 0, MIN returns the maximum # value for the type, MAX returns the minimum value for the type. Only SUM and @@ -394,9 +425,7 @@ def _apply_agg_column( plc_result, dtype_from_pylibcudf_column(plc_result) ) - if isinstance(agg_name, str): - return col.astype(np.dtype("float64")) - return col + return col.astype(np.dtype("float64")) def _reduce( self, diff --git a/python/cudf/cudf/tests/window/test_rolling.py b/python/cudf/cudf/tests/window/test_rolling.py index 4f1af902a5a0..61c4a963fd4f 100644 --- a/python/cudf/cudf/tests/window/test_rolling.py +++ b/python/cudf/cudf/tests/window/test_rolling.py @@ -370,6 +370,48 @@ def some_func(A): ) +@pytest.mark.parametrize("window_size", [1, 2, 3]) +@pytest.mark.parametrize("min_periods", [1, 2, 3]) +def test_rolling_groupby_numba_udf(window_size, min_periods): + if min_periods > window_size: + pytest.skip("min_periods cannot exceed window_size") + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 2, 2, 2, 2], + "b": [1.0, 2.0, 4.0, 8.0, 9.0, 4.0, 2.0], + } + ) + gdf = cudf.from_pandas(pdf) + + def some_func(A): + b = 0 + for a in A: + b = b + a**2 + return b / len(A) + + assert_eq( + pdf.groupby("a").rolling(window_size, min_periods).apply(some_func), + gdf.groupby("a").rolling(window_size, min_periods).apply(some_func), + ) + + +def test_rolling_numba_udf_base_indexer(): + indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=3) + pdf = pd.DataFrame({"a": [1.0, 2.0, 4.0, 9.0, 9.0, 4.0]}) + gdf = cudf.from_pandas(pdf) + + def some_func(A): + b = 0 + for a in A: + b = b + a + return b / len(A) + + assert_eq( + pdf.rolling(window=indexer, min_periods=1).apply(some_func), + gdf.rolling(window=indexer, min_periods=1).apply(some_func), + ) + + def test_rolling_groupby_simple(supported_rolling_reductions): pdf = pd.DataFrame( { From 6f80d191202d5916fd5134a051a29a5e4a040aaa Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:31:31 +0000 Subject: [PATCH 2/2] pre-commit --- python/cudf/cudf/core/udf/rolling_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/core/udf/rolling_utils.py b/python/cudf/cudf/core/udf/rolling_utils.py index 18c5ca9f7e6e..c465faa6a506 100644 --- a/python/cudf/cudf/core/udf/rolling_utils.py +++ b/python/cudf/cudf/core/udf/rolling_utils.py @@ -7,7 +7,7 @@ import cupy as cp import numpy as np -from numba_cuda_mlir import cuda, compiler +from numba_cuda_mlir import compiler, cuda from numba_cuda_mlir.numba_cuda.core import config as _mlir_config from numba_cuda_mlir.numba_cuda.np import numpy_support @@ -32,7 +32,8 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: def _get_udf_return_type(func: Callable, value_dtype: np.dtype) -> np.dtype: """Compile ``func`` for a 1D window of ``value_dtype`` to infer its - output dtype.""" + output dtype. + """ nb_value_type = numpy_support.from_dtype(value_dtype) signature = (nb_value_type[::1],) try: