NVIDIA / NVIDIA/cudf

[FEA] Should `Rolling.apply` use pure numba rather than jitify?

Open
#18,033 0 comments 0 reactions 0 assignees View on GitHub
feature request numba Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**
cuDF supports [rolling.apply](https://docs.rapids.ai/api/cudf/latest/user_guide/api_docs/api/cudf.series.rolling/#cudf.Series.rolling) for executing a custom python function over the specified rolling windows. It works by taking a PTX string compiled in python by numba and handing it off to a rolling aggregation with `kind == PTX`, which is then used through `rolling_window`. Ultimately this invokes jitify and a series of parsing and compilation steps leading to the final [kernel](https://github.com/rapidsai/cudf/blob/81bb6f1ea46d633548b9bf534c93a28ba66909ea/cpp/src/rolling/jit/kernel.cu#L45) that computes the result.

Historically a similar process served to make APIs like `Series.applymap` work but over time we migrated to an approach that puts together the final kernel in numba rather than c++ for several reasons:

- The parsing approach through jitify trips up in several useful cases, such as when numba delivers PTX containing multiple function definitions
- The full numba approach allows for extension types to support null values
- The full numba approach benefits from features such as LTO through pynvjitlink

**Describe the solution you'd like**
A pure numba implementation of `Rolling.apply`, possibly with extension type support. For non nullable data, the implementation could look something like this:

```python
def count_if_gt_3(window):
count = 0
for i in window:
if i > 3:
count += 1
return count

devfunc = cuda.jit(device=True)(count_if_gt_3)
out = np.zeros(len(s))

@cuda.jit
def kernel(data, win_size, min_periods, out):
tid = cuda.grid(1)

start = max(0, tid - win_size + 1)
end = tid + 1

thread_win = data[start:end]

res = devfunc(thread_win)
out[tid] = res

```

The above seems to get me the correct result in a few test cases locally, of course the real impl would need more to account for `min_periods`, other dtypes, and more.

A _null sensitive_ implementation seems possible as well building on a lot of what we already have. The idea revolves around assembling the data into a `cuda.local.array` of `MaskedType` and then passing that array into the UDF as written. The existing implementations of operations between MaskedTypes should take care of the rest.

```python
@cuda.jit
def kernel(data, mask, win_size, min_periods, out_data, out_mask):
tid = cuda.grid(1)

start = max(0, tid - win_size + 1)
end = tid + 1

local = cuda.local.array(win_size, Masked(types.int64))

# place this window of data into thread local memory as an array of MaskedTypes
for i in range(0, end-start):
local[i] = Masked(data[start+i], mask_get(mask, start+i))

# the device function now iterates through the array of MaskedTypes
# any operations are resolved through MaskedType's overloads
res = devfunc(local)
out_data[tid] = res.value
out_mask[tid] = res.valid

```
Currently however creating a `cuda.local.array` of extension types needs a few changes (cc @gmarkall). The above would enable handling nulls explicitly within supported UDFs in conditional logic:

```python
def count_if_gt_3(window):
count = 0
for i in window:
if i != cudf.NA:
if i > 3:
count += 1
else:
return -1
return count
```

**Describe alternatives you've considered**
One disadvantage of this approach is that we'd need to reimplement logic around handling all of the keyword arguments currently supported for rolling without a custom aggregation, such as `center` and `min_periods`. This would have to be maintained separately from libcudf.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.