[MLIR] MaskedType: core type, constructor, accessors by brandon-b-miller · Pull Request #22884 · NVIDIA/cudf

brandon-b-miller · 2026-06-14T17:07:21Z

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/.

copy-pr-bot · 2026-06-14T17:07:24Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-06-22T18:37:05Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds MLIR backend support for MaskedType in cuDF UDFs through new typing and lowering modules, plus pytest coverage for constructor and attribute behavior.

Changes

MaskedType MLIR backend

Layer / File(s)	Summary
MaskedType definition and typing registration `python/cudf/cudf/core/udf/mlir_backend/masked_typing.py`, `python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py`	`MaskedType` is defined as a Numba type parameterized by an inner `value_type`, with `types.Literal` unwrapping and `types.Poison` fallback for unsupported inner types. `_register()` installs a `MaskedConstructor` `ConcreteTemplate` for `Masked(value, valid)` signatures and a `MaskedTypeAttrs` resolver that maps `.value` to the inner type and `.valid` to `types.boolean`. Tests cover repr, equality, hash, unliteralization, poison-wrapping, and supported-type set contents.
LLVM data model and constructor/getattr lowering `python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py`, `python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py`	`_pack_masked` builds the `{value, valid}` LLVM struct via an undefined value and two field inserts. `_register()` attaches an identified `StructType` data model to `MaskedType`, lowers the `Masked(value, valid)` constructor by coercing arguments and calling `_pack_masked`, and lowers `.value`/`.valid` getattr by extracting the corresponding struct field and converting to the destination MLIR type. Tests exercise literal-`True`, literal-`False`, and runtime valid flags across integer, float, and boolean dtypes.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

numba

Suggested reviewers

gforsyth
wence-
vyasr

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.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 clearly summarizes the main MLIR change: MaskedType plus its constructor and accessors.
Description check	✅ Passed	The description is directly related to the changeset and accurately summarizes the new MaskedType support and tests.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch mlir/pr2-masked

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

vyasr

Generally makes sense. I still don't fully grasp what goes in the lowering file and what goes in the typing file, it seems to me like some aspects of type definitions in the typing file are still relevant for lowering (e.g. the constructor signatures), could you explain the division a bit more?

brandon-b-miller · 2026-06-24T02:37:04Z

Generally makes sense. I still don't fully grasp what goes in the lowering file and what goes in the typing file, it seems to me like some aspects of type definitions in the typing file are still relevant for lowering (e.g. the constructor signatures), could you explain the division a bit more?

@vyasr the split maps to the two phase compilation pipeline used by numba (typing/lowering). The typing phase's only job is to determine, for every expression in the UDF, what the resuling type is. This phase emits no code, it just annotates every variable and result with a type. MaskedType is a new type, and at start theres no expressions at all known to numba that result in an instance of it. Much of the code in typing.py adds the very first expression one can write that results in such an instance - in this case, that expression is when you write Masked(val, valid) in the code under cuda.jit, this shall yield a MaskedType whose value_type is the same as the scalar value passed to its ctor. Stuff that defines typing for expressions like this belongs in the typing file.

The lowering phase only triggers when typing for the whole blob completes successfully. Its job is to actually emit the code that computes the value. Since this might be dependent on what the input types for that expression actually are, the author needs a registration in the lowering file that associates a procedure for computing that value with a valid set of input types. For lower_masked_constructor, the process is not much more than "load both operands, then put them into a struct that has a member for each". These implementation functions that use the MLIR python bindings to write MLIR belong in the lowering file.

mroeschke

Overall looks good. Just a few last comments

mroeschke

One comment and the failing typing checks otherwise LGTM. (Feel free to use # type: ignores if the typing is difficult)

vyasr

Very close. I have a few smaller requests then we can merge.

…d_typing.py Co-authored-by: Vyas Ramasubramani <vyasr@nvidia.com>

numba_cuda_mlir/types.py imports ``bool`` unconditionally, but the vendored numba_cuda.types only defines it for numpy>=2.0, so collection fails on the oldest-deps CI matrix (numpy 1.26). Define the alias at conftest import time. TODO: remove once a numba-cuda-mlir release >0.4.0 fixes the import.

mroeschke · 2026-06-30T17:52:13Z

+# AttributeError is unconditional -- so skip the MLIR backend tests on numpy < 2
+# instead of erroring at collection.
+# TODO: remove once a numba-cuda-mlir release > 0.4.0 supports numpy < 2.0.
+if np.lib.NumpyVersion(np.__version__) < "2.0.0":


Just FYI, #22962 appears close to merging and will raise the NumPy minimum version to 2.0 which should allow you to remove this (likely in a follow up)

Let's do this in #22885

vyasr · 2026-06-30T20:23:22Z

/merge

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

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

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

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

brandon-b-miller force-pushed the mlir/pr2-masked branch from f33dcd1 to 3ab1d60 Compare June 16, 2026 15:36

brandon-b-miller marked this pull request as ready for review June 22, 2026 18:32

brandon-b-miller requested a review from a team as a code owner June 22, 2026 18:32

brandon-b-miller requested review from bdice and vyasr June 22, 2026 18:32

mroeschke reviewed Jun 23, 2026

View reviewed changes

Comment thread python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py

brandon-b-miller commented Jun 23, 2026

View reviewed changes

Comment thread python/cudf/cudf/core/udf/mlir_backend/masked_typing.py

brandon-b-miller commented Jun 23, 2026

View reviewed changes

Comment thread python/cudf/cudf/core/udf/mlir_backend/masked_typing.py

vyasr reviewed Jun 24, 2026

View reviewed changes

vyasr mentioned this pull request Jun 24, 2026

Consolidate extension type registration #19938

Closed

brandon-b-miller added feature request New feature or request non-breaking Non-breaking change labels Jun 24, 2026

mroeschke reviewed Jun 25, 2026

View reviewed changes

mroeschke approved these changes Jun 26, 2026

View reviewed changes

Comment thread python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py Outdated

brandon-b-miller force-pushed the mlir/pr2-masked branch from 0baeca1 to b248bf3 Compare June 26, 2026 22:07

vyasr requested changes Jun 27, 2026

View reviewed changes

brandon-b-miller force-pushed the mlir/pr2-masked branch from b248bf3 to 756652e Compare June 29, 2026 12:59

vyasr approved these changes Jun 29, 2026

View reviewed changes

brandon-b-miller force-pushed the mlir/pr2-masked branch from 26171e1 to d111062 Compare June 29, 2026 20:57

brandon-b-miller added 4 commits June 30, 2026 04:51

minimum maskedtype

5c7acef

excize pack_return for now

74f0c09

clean

da0dec6

Drop test module docstrings and narration comments

e14cccc

brandon-b-miller and others added 14 commits June 30, 2026 04:51

refactor

49cf166

Use resolve_value/resolve_valid for MaskedType attribute typing

5a7931f

Reject keyword args in the Masked constructor lowering

3794bbf

Use descriptive names in MaskedType data model

631accb

Rename _pack_masked_result -> _pack_masked (value/valid params)

a301e30

address reviews

47cd423

Inline cuda.jit in tests, drop _jit helper

ba5616a

Use a numba-cuda-mlir config context manager for kernel launches

c0aead8

Parametrize Masked constructor tests over validity (literal + runtime)

d68bff1

Update python/cudf/cudf/tests/private_objects/mlir_backend/test_maske…

0eff97d

…d_typing.py Co-authored-by: Vyas Ramasubramani <vyasr@nvidia.com>

Don't assert on the payload of an invalid Masked

b63b3ea

precommit

8a67d0e

lint

bc6878c

brandon-b-miller force-pushed the mlir/pr2-masked branch from d111062 to bc6878c Compare June 30, 2026 11:51

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

further compat fixes

b648b3c

skip tests for numpy <2.0

1e999bb

mroeschke reviewed Jun 30, 2026

View reviewed changes

rapids-bot Bot merged commit fcb3913 into main Jun 30, 2026
125 of 126 checks passed

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

brandon-b-miller deleted the mlir/pr2-masked branch July 1, 2026 00:50

vyasr mentioned this pull request Jul 1, 2026

[MLIR] MaskedType: NA handling #22885

Merged

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

coderabbitai Bot mentioned this pull request Jul 8, 2026

[MLIR] MaskedType: numeric binary operators #22886

Open

Uh oh!

Conversation

brandon-b-miller commented Jun 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jun 14, 2026

Uh oh!

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

Uh oh!

Reviews paused

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vyasr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brandon-b-miller commented Jun 24, 2026

Uh oh!

mroeschke left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mroeschke left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vyasr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mroeschke Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

vyasr Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

vyasr commented Jun 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

brandon-b-miller commented Jun 14, 2026 •

edited

Loading

coderabbitai Bot commented Jun 22, 2026 •

edited

Loading