[MLIR] Scaffold MLIR based UDF backend by brandon-b-miller · Pull Request #22766 · NVIDIA/cudf

brandon-b-miller · 2026-06-03T14:28:26Z

This PR adds initial scaffolding for a rewrite of the UDF numba-cuda extension backend using numba-cuda-mlir.

There's no conda packages yet, so this PR leaves updating the dep on the conda side as a todo. Xref NVIDIA/numba-cuda-mlir#97

coderabbitai · 2026-06-03T14:34:56Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b2d3c1a1-4c11-42f6-aed5-537035a24b29

📥 Commits

Reviewing files that changed from the base of the PR and between 7f73bb4 and a5f1f75.

📒 Files selected for processing (2)

dependencies.yaml
python/cudf/pyproject.toml

✅ Files skipped from review due to trivial changes (1)

python/cudf/pyproject.toml

🚧 Files skipped from review as they are similar to previous changes (1)

dependencies.yaml

📝 Walkthrough

Summary by CodeRabbit

Chores
- Added numba-cuda-mlir>=0.3.0 to runtime dependencies, including CUDA-version-specific variants.
- Updated packaging notes around how this requirement is handled.
Tests
- Added pytest coverage to ensure the MLIR backend namespace is importable when available.
Documentation
- Introduced module-level docstring/headers for the MLIR backend namespace.

Walkthrough

This PR adds numba-cuda-mlir>=0.3.0 as a runtime dependency across build configuration and project metadata, establishes the cudf.core.udf.mlir_backend package namespace, and introduces a basic plumbing test to verify the module is correctly structured and available when the dependency is present.

Changes

numba-cuda-mlir dependency and MLIR backend package

Layer / File(s)	Summary
Dependency infrastructure for numba-cuda-mlir `python/cudf/pyproject.toml`, `dependencies.yaml`, `conda/recipes/cudf/recipe.yaml`	Declare `numba-cuda-mlir>=0.3.0` in pyproject dependencies; define `depends_on_numba_cuda_mlir` in dependencies.yaml with CUDA 12/13 suffixed variant selectors and wire into global and target-specific build matrix includes; document conda recipe transition with TODO comment.
MLIR backend package namespace and plumbing test `python/cudf/cudf/core/udf/mlir_backend/__init__.py`, `python/cudf/cudf/tests/private_objects/mlir_backend/__init__.py`, `python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py`	Create the `cudf.core.udf.mlir_backend` package with SPDX header and module docstring; add test package structure; introduce test_plumbing.py that conditionally runs when `numba_cuda_mlir` is importable, verifying the MLIR backend module is accessible and intentionally empty.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested labels

improvement

Suggested reviewers

gforsyth

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[MLIR] Scaffold MLIR based UDF backend' accurately and specifically describes the main change—adding initial scaffolding for a MLIR-based UDF backend.
Description check	✅ Passed	The description directly explains the PR's purpose: adding scaffolding for a UDF numba-cuda extension backend rewrite using numba-cuda-mlir, with context about conda package availability.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

bdice

Can we start work on a conda recipe? I'd like to land that before merging.

bdice · 2026-06-05T17:37:15Z

+    """`cudf.utils._numba` should default `CUDA_ENABLE_NRT` to False on import."""
+    import cudf.utils._numba  # noqa: F401  -- import for its side effect
+
+    assert numba_config.CUDA_ENABLE_NRT is False, (


Suggested change

assert numba_config.CUDA_ENABLE_NRT is False, (

assert not numba_config.CUDA_ENABLE_NRT, (

bdice · 2026-06-05T17:37:43Z

+    assert numba_config.CUDA_ENABLE_NRT is False
+    with nrt_enabled():
+        assert numba_config.CUDA_ENABLE_NRT is True
+    assert numba_config.CUDA_ENABLE_NRT is False


Please apply in various places below, too. (This assumes numba_config.CUDA_ENABLE_NRT is boolean and does not take on states like None.)

Suggested change

assert numba_config.CUDA_ENABLE_NRT is False

with nrt_enabled():

assert numba_config.CUDA_ENABLE_NRT is True

assert numba_config.CUDA_ENABLE_NRT is False

assert not numba_config.CUDA_ENABLE_NRT

with nrt_enabled():

assert numba_config.CUDA_ENABLE_NRT

assert not numba_config.CUDA_ENABLE_NRT

bdice · 2026-06-05T17:41:39Z

+# per-kernel for UDFs whose data model declares an NRT meminfo (e.g. string
+# returns), and restores it afterwards. Setting the global default to False
+# avoids linking NRT into kernels that don't need it.
+numba_config.CUDA_ENABLE_NRT = False


I think importing cudf should not change the state of this global config. Instead, we should use a context manager to enable it as needed and restore to the prior global state (regardless of whether it was True or False) when outside of cuDF usage of NRT.

brandon-b-miller · 2026-06-11T12:18:28Z

Failing cuda 12.2 jobs require NVIDIA/numba-cuda-mlir#102 to be released. We can skip cuda 12.2 for now or wait.

vyasr

OK this is super minimal now, so just adding the MLIR dep. That's fine for the moment, we can add things as they become available. If we're going to get a CUDA 12.2 supporting release from numba-cuda-mlir then let's wait on that to resolve the currently failing test.

brandon-b-miller · 2026-06-16T15:05:45Z

/merge

Add a basic MLIR based `MaskedType` and the ability to express the creation of one inside a `numba-cuda-mlir` kernel. Part of the MLIR UDF backend stack. Depends on #22766 (plumbing). Adds the MaskedType extension type: parameterized value type, the Masked(value, valid) constructor, the .value / .valid accessors, and the LLVM struct data model. Numeric/boolean value types only. Stacked on: #22766. Tests under tests/private_objects/mlir_backend/. Authors: - https://github.com/brandon-b-miller Approvers: - Matthew Roeschke (https://github.com/mroeschke) - Vyas Ramasubramani (https://github.com/vyasr) URL: #22884

Part of the MLIR UDF backend stack. Depends on #22766 (plumbing). Adds NA handling for MaskedType: NAType + cudf.NA typeof, MaskedType / NAType unify, `m is NA` / `m is not NA`, and the NA / scalar / Masked->Masked casts used for branch unification. Unification covers expressions in the target code like the following where `x` is `MaskedType`. ``` if x is cudf.NA: return x else: return 42 ``` 42 is a scalar here so the unification code kicks in and types the overall expression as returning MaskedType, and upcasts the 42 to Masked(42, True)

brandon-b-miller added 2 commits June 3, 2026 07:07

plumbing

3cac86a

merge/resolve

6ea3204

brandon-b-miller requested review from a team as code owners June 3, 2026 14:28

brandon-b-miller requested review from gforsyth, vyasr and wence- June 3, 2026 14:28

github-actions Bot assigned brandon-b-miller Jun 3, 2026

github-actions Bot added the Python Affects Python cuDF API. label Jun 3, 2026

github-project-automation Bot added this to cuDF Python Jun 3, 2026

GPUtester moved this to In Progress in cuDF Python Jun 3, 2026

brandon-b-miller added feature request New feature or request non-breaking Non-breaking change numba Numba issue labels Jun 3, 2026

bdice reviewed Jun 5, 2026

View reviewed changes

brandon-b-miller added 3 commits June 8, 2026 05:48

Merge remote-tracking branch 'upstream/main' into mlir/pr1-plumbing

a4a08ec

Merge remote-tracking branch 'upstream/main' into mlir/pr1-plumbing

18c7298

reduce PR scope

7f73bb4

bdice approved these changes Jun 10, 2026

View reviewed changes

Merge branch 'main' into mlir/pr1-plumbing

39fbfcd

vyasr approved these changes Jun 13, 2026

View reviewed changes

This was referenced Jun 14, 2026

[MLIR] MaskedType: core type, constructor, accessors #22884

Merged

[MLIR] MaskedType: NA handling #22885

Merged

[MLIR] MaskedType: numeric binary operators #22886

Open

Merge branch 'main' into mlir/pr1-plumbing

a5f1f75

rapids-bot Bot merged commit ff34bdb into NVIDIA:main Jun 16, 2026
136 checks passed

github-project-automation Bot moved this from In Progress to Done in cuDF Python Jun 16, 2026

brandon-b-miller mentioned this pull request Jul 7, 2026

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

Open

14 tasks

Navigation Menu

Uh oh!

[MLIR] Scaffold MLIR based UDF backend - #22766

[MLIR] Scaffold MLIR based UDF backend#22766
rapids-bot[bot] merged 7 commits into
NVIDIA:mainfrom
brandon-b-miller:mlir/pr1-plumbing

brandon-b-miller commented Jun 3, 2026

Uh oh!

coderabbitai Bot commented Jun 3, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

bdice left a comment

Uh oh!

bdice Jun 5, 2026

Uh oh!

bdice Jun 5, 2026

Uh oh!

bdice Jun 5, 2026

Uh oh!

brandon-b-miller commented Jun 11, 2026

Uh oh!

vyasr left a comment

Uh oh!

brandon-b-miller commented Jun 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	assert numba_config.CUDA_ENABLE_NRT is False, (
	assert not numba_config.CUDA_ENABLE_NRT, (

Uh oh!

Conversation

brandon-b-miller commented Jun 3, 2026

Uh oh!

coderabbitai Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

bdice left a comment

Choose a reason for hiding this comment

Uh oh!

bdice Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

brandon-b-miller commented Jun 11, 2026

Uh oh!

vyasr left a comment

Choose a reason for hiding this comment

Uh oh!

brandon-b-miller commented Jun 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

coderabbitai Bot commented Jun 3, 2026 •

edited

Loading