NVIDIA / NVIDIA/cudf

[Story][MLIR] Migrate the cuDF UDF backend from `numba-cuda` to `numba-cuda-mlir`

Open
#23,144 0 comments 0 reactions 1 assignee Claimed by @brandon-b-miller View on GitHub
feature request Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

## Summary

Track the incremental migration of cuDF's UDF compilation pipeline (`Series.apply`, `DataFrame.apply`, `groupby.apply`) from the legacy `numba` / `numba-cuda` extension API to the new **`numba-cuda-mlir`** backend, which registers types/lowerings via `numba_cuda_mlir.extending.{typing_registry, lowering_registry}` and lowers through MLIR.

This is a large rewrite (~4–5k lines) delivered as a series of small, independently reviewable PRs rather than one cutover.

## Motivation
Migrate cuDF to the new MLIR based python compiler, as it will be via this package that new cuda versions, etc are supported. This requires a rewrite of some of our LLVMlite based lowerings to MLIR based ones as well as a few other changes to fix compatibility with the new package's registration systems.

## Approach & principles

- **Incremental, atomic PRs**, each with, self contained tests
(kernel level unit tests under
`python/cudf/cudf/tests/private_objects/mlir_backend/` that exercise just the
typing+lowering being added no full `df.apply` needed).
- **Exactly one cutover PR at the very end.** Intermediate PRs are "lights on":
the new backend is registered/tested in isolation while the production path
stays on the legacy backend; where e2e behavior needs exercising early,
tests monkeypatch the specific pipeline pieces.
- `numba-cuda-mlir` is currently a **pip only** dependency (conda packaging in
progress) tests gate with `pytest.importorskip("numba_cuda_mlir")` so conda
CI stays green.

## Task breakdown

### Phase 0 — Plumbing & packaging
- [x] Add `numba-cuda-mlir` dependency + backend scaffold; default `CUDA_ENABLE_NRT=False` #22766
- [ ] Conda packaging for `numba-cuda-mlir` (staged-recipes) — prerequisite for conda CI coverage

### Phase 1 — MaskedType: all non-string operations
- [x] Core `MaskedType`: type, `Masked(value, valid)` constructor, `.value`/`.valid`, data model #22884
- [x] NA handling: `NAType` / `cudf.NA`, unify, `is` / `is not` NA, NA/scalar/Masked→Masked casts #22885
- [ ] Numeric binary operators #22886
- [ ] Unary operators, truth, numeric casts, abs #23119
- [ ] Datetime64 / timedelta64 value types & arithmetic
- [ ] Containment (`value in (...)`)
- [ ] `pack_return` (normalize UDF return to `MaskedType`)

### Phase 2 — Full elementwise pipeline running through MaskedType
- [ ] Wire the scalar / `apply` kernel path so a real UDF compiles + runs
end-to-end on the MLIR backend through `MaskedType` (validated via
monkeypatched pipeline tests; no production cutover yet)

### Phase 3 — MLIRString scalar
- [ ] Introduce the string scalar type + string ops (MLIR-side) with tests

### Phase 4 — Strings → MaskedType
- [ ] Plumb string value types into `MaskedType` (nullable string UDFs)

### Phase 5 — Cooperative array (groupby)
- [ ] `CooperativeArray` type + groupby-reduction pipeline running on the new
backend (no NRT-managed groupby yet)

### Phase 6 — Cutover
- [ ] Single PR flipping the production UDF path (`utils.py` /
`scalar_function.py` / `groupby_function.py`) to the MLIR backend and removing
the legacy path

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.