[Story][MLIR] Migrate the cuDF UDF backend from `numba-cuda` to `numba-cuda-mlir`
- 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
Assessment
This issue has not been assessed yet.