Replace rolling.apply implementation with numba-cuda-mlir

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}")

python/cudf/cudf/core/udf/rolling_utils.py

@@ -0,0 +1,121 @@
+# 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 compiler, cuda
+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)

python/cudf/cudf/core/window/rolling.py

@@ -356,24 +356,55 @@ 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:
         if isinstance(source_column.dtype, CategoricalDtype):
             # 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,

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(
         {