From 11e84bdb86833b8481d23010118f963ab982414a Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 11 Jun 2026 06:01:43 -0700 Subject: [PATCH 1/6] [MLIR] PR 4: Masked binary ops (numeric) Adds binary arithmetic / bitwise / comparison operator typing and lowering over numeric/boolean MaskedType values: * Masked Masked (validity = AND of operand validities) * Masked scalar / scalar Masked (validity carried) * Masked NA / NA Masked (result invalid) Lowering is numeric-only; unit-aware datetime delegation is deferred to the datetime PR. Tests: +20 kernel tests covering arith value/validity, all six comparisons, scalar-on-either-side, the row['a']<1 literal regression guard, and NA poisoning. --- .../core/udf/mlir_backend/masked_lowering.py | 137 ++++++++++- .../core/udf/mlir_backend/masked_typing.py | 68 ++++++ .../mlir_backend/test_masked_lowering.py | 230 ++++++++++++++++++ 3 files changed, 434 insertions(+), 1 deletion(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index 7625d36916c0..0ced7f935e19 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -6,13 +6,19 @@ from functools import partial from typing import TYPE_CHECKING +import numpy as np from numba_cuda_mlir import types from numba_cuda_mlir._mlir import ir as mlir_ir from numba_cuda_mlir._mlir.dialects import arith, llvm from numba_cuda_mlir.extending import lower_cast, lowering_registry -from numba_cuda_mlir.lowering_utilities import convert +from numba_cuda_mlir.lowering_utilities import ( + coerce_numpy_scalars_for_binary_op, + convert, +) from numba_cuda_mlir.models import PrimitiveModel, register_model +from numba_cuda_mlir.numba_cuda.types.misc import unliteral +from cudf.core.udf._ops import arith_ops, bitwise_ops, comparison_ops from cudf.core.udf.api import Masked from cudf.core.udf.mlir_backend.masked_typing import ( MaskedType, @@ -183,6 +189,116 @@ def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null): builder.store_var(target, valid) +# --- Binary ops ----------------------------------------------------------- +# Shared helper: apply ``op(v1, v2)`` to two scalar MLIR values, convert the +# result to the target Masked's value type, and pack it with the given +# validity bit. Numeric/boolean only at this layer. +def _apply_masked_binary_op( + builder, target, target_type, v1, v2, result_valid, op +): + target_value_mlir_ty = builder.get_mlir_type(target_type.value_type) + v1, v2 = coerce_numpy_scalars_for_binary_op(v1, v2) + # Comparisons compute on the (already coerced) operand type and + # produce i1; arithmetic/bitwise compute on the target value type. + is_cmp = op in comparison_ops + operand_ty = v1.type if is_cmp else target_value_mlir_ty + v1 = convert(v1, operand_ty) + v2 = convert(v2, operand_ty) + result_val = convert(op(v1, v2), target_value_mlir_ty) + packed = _pack_masked( + builder, target_type, result_val, result_valid + ) + builder.store_var(target, packed) + + +# ``Masked Masked``: AND the validity bits. +def _make_lower_masked_binary(op): + def _lower(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + m1 = builder.load_var(args[0]) + m2 = builder.load_var(args[1]) + st1 = llvm.StructType(m1.type) + st2 = llvm.StructType(m2.type) + v1, valid1 = _extract_masked_value_valid( + m1, st1.body[0], st1.body[1] + ) + v2, valid2 = _extract_masked_value_valid( + m2, st2.body[0], st2.body[1] + ) + result_valid = arith.andi(valid1, valid2) + _apply_masked_binary_op( + builder, target, target_type, v1, v2, result_valid, op + ) + + return _lower + + +def _scalar_value_from_var(builder, s_var, m_var, masked_value_mlir_ty): + """Resolve the scalar operand for the Masked-vs-scalar path. + + Prefer a materialized constant when the scalar is a Literal so we + never mistake the masked operand for the scalar (e.g. + ``row['a'] < 1`` must not become ``row['a'] < row['a']``). + """ + s_ty = builder.get_numba_type(s_var.name) + if isinstance(s_ty, types.Literal): + py_val = s_ty.literal_value + base_ty = unliteral(s_ty) + mlir_ty = builder.get_mlir_type(base_ty) + if isinstance(py_val, (bool, np.bool_)) or ( + hasattr(mlir_ty, "width") and mlir_ty.width == 1 + ): + py_val = 1 if py_val else 0 + return arith.constant(mlir_ty, py_val) + s_raw = builder.load_var(s_var) + if getattr(s_var, "name", None) == getattr(m_var, "name", None): + raise RuntimeError( + "Masked vs scalar lowering: scalar variable is the same as " + "the masked variable; cannot extract a distinct scalar." + ) + return s_raw + + +# ``Masked scalar`` and ``scalar Masked``: carry the Masked +# operand's validity. +def _make_lower_masked_binary_scalar(op, masked_first): + def _lower(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + m_var, s_var = ( + (args[0], args[1]) if masked_first else (args[1], args[0]) + ) + m = builder.load_var(m_var) + st = llvm.StructType(m.type) + m_val, m_valid = _extract_masked_value_valid( + m, st.body[0], st.body[1] + ) + s_val = _scalar_value_from_var(builder, s_var, m_var, st.body[0]) + if masked_first: + _apply_masked_binary_op( + builder, target, target_type, m_val, s_val, m_valid, op + ) + else: + _apply_masked_binary_op( + builder, target, target_type, s_val, m_val, m_valid, op + ) + + return _lower + + +# ``Masked NA`` / ``NA Masked``: result is invalid. +def _lower_masked_binary_null(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + value_mlir_ty = builder.get_mlir_type(target_type.value_type) + undef_val = llvm.UndefOp(value_mlir_ty) + valid_zero = arith.constant( + result=builder.get_mlir_type(types.boolean), value=0 + ) + packed = _pack_masked( + builder, target_type, undef_val, valid_zero + ) + builder.store_var(target, packed) + + def _register() -> None: """Register the data model and lowerings with ``numba_cuda_mlir``. @@ -213,5 +329,24 @@ def _register() -> None: lower(operator.is_not, MaskedType, NAType)(is_not_na) lower(operator.is_not, NAType, MaskedType)(is_not_na) + for binary_op in arith_ops + bitwise_ops + comparison_ops: + lower(binary_op, MaskedType, MaskedType)( + _make_lower_masked_binary(binary_op) + ) + lower(binary_op, MaskedType, types.Number)( + _make_lower_masked_binary_scalar(binary_op, True) + ) + lower(binary_op, types.Number, MaskedType)( + _make_lower_masked_binary_scalar(binary_op, False) + ) + lower(binary_op, MaskedType, types.Boolean)( + _make_lower_masked_binary_scalar(binary_op, True) + ) + lower(binary_op, types.Boolean, MaskedType)( + _make_lower_masked_binary_scalar(binary_op, False) + ) + lower(binary_op, MaskedType, NAType)(_lower_masked_binary_null) + lower(binary_op, NAType, MaskedType)(_lower_masked_binary_null) + _register() diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py index 7a8beec820ca..c7a389a47c45 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -18,6 +18,7 @@ from numba_cuda_mlir.typing import signature as nb_signature from cudf.core.missing import NA +from cudf.core.udf._ops import arith_ops, bitwise_ops, comparison_ops from cudf.core.udf.api import Masked _SUPPORTED_MASKED_VALUE_TYPE_CLASSES = ( @@ -139,6 +140,68 @@ def generic(self, args, kws): return None +# ``Masked Masked``: resolve the underlying scalar op on the two +# value types, then wrap the result back in a MaskedType. +class MaskedScalarArithOp(AbstractTemplate): + def generic(self, args, kws): + if isinstance(args[0], MaskedType) and isinstance( + args[1], MaskedType + ): + return_type = self.context.resolve_function_type( + self.key, (args[0].value_type, args[1].value_type), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0], args[1]) + return None + + +# ``Masked scalar`` and ``scalar Masked`` (scalar may be a +# Literal, e.g. ``row['a'] == 1``). +class MaskedScalarScalarOp(AbstractTemplate): + def generic(self, args, kws): + if isinstance(args[0], MaskedType) and isinstance( + args[1], (types.Number, types.Boolean) + ): + return_type = self.context.resolve_function_type( + self.key, (args[0].value_type, args[1]), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0], args[1]) + if isinstance(args[0], MaskedType) and isinstance( + args[1], types.Literal + ): + scalar_ty = unliteral(args[1]) + return_type = self.context.resolve_function_type( + self.key, (args[0].value_type, scalar_ty), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0], args[1]) + if isinstance( + args[0], (types.Number, types.Boolean) + ) and isinstance(args[1], MaskedType): + return_type = self.context.resolve_function_type( + self.key, (args[0], args[1].value_type), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0], args[1]) + if isinstance(args[0], types.Literal) and isinstance( + args[1], MaskedType + ): + scalar_ty = unliteral(args[0]) + return_type = self.context.resolve_function_type( + self.key, (scalar_ty, args[1].value_type), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0], args[1]) + return None + + +# ``Masked NA`` / ``NA Masked``: result type is the Masked +# operand's type; the lowering produces an invalid (poisoned) value. +class MaskedScalarNullOp(AbstractTemplate): + def generic(self, args, kws): + if isinstance(args[0], MaskedType) and isinstance(args[1], NAType): + return nb_signature(args[0], args[0], na_type) + if isinstance(args[0], NAType) and isinstance(args[1], MaskedType): + return nb_signature(args[1], na_type, args[1]) + return None + + def _register() -> None: """Register typing for ``Masked`` and ``MaskedType`` attributes with ``numba_cuda_mlir``. Called once at module import. @@ -148,5 +211,10 @@ def _register() -> None: typing_registry.register_global(operator.is_)(MaskedNAComparison) typing_registry.register_global(operator.is_not)(MaskedNAComparison) + for binary_op in arith_ops + bitwise_ops + comparison_ops: + typing_registry.register_global(binary_op)(MaskedScalarArithOp) + typing_registry.register_global(binary_op)(MaskedScalarNullOp) + typing_registry.register_global(binary_op)(MaskedScalarScalarOp) + _register() diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py index 7c2bc5433f79..f633f4b72608 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py @@ -3,6 +3,8 @@ from __future__ import annotations +import operator + import cupy as cp import numpy as np import pytest @@ -16,6 +18,7 @@ import cudf.core.udf.mlir_backend.masked_lowering import cudf.core.udf.mlir_backend.masked_typing # noqa: F401 +from cudf.core.missing import NA from cudf.core.udf.api import Masked from cudf.core.udf.utils import DEPRECATED_SM_REGEX @@ -234,3 +237,230 @@ def k(out, v, valid): result = out.get() assert bool(result[0]) is valid_in assert bool(result[1]) is valid_in + + +# --- binary ops ------------------------------------------------------------ + + +_ARITH = [ + (operator.add, lambda a, b: a + b), + (operator.sub, lambda a, b: a - b), + (operator.mul, lambda a, b: a * b), +] + + +@pytest.mark.parametrize("op,ref", _ARITH) +def test_masked_masked_arith_value(op, ref): + """``Masked(a) Masked(b)`` computes ``op(a, b)`` in the value field.""" + + @cuda.jit( + types.void( + types.int64[::1], + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out, a, av, b, bv): + m = op(Masked(a[0], av[0]), Masked(b[0], bv[0])) + out[0] = m.value + + a, b = 12, 5 + in_a = cp.array([a], dtype=np.int64) + in_b = cp.array([b], dtype=np.int64) + true_ = cp.array([True], dtype=np.bool_) + out = cp.zeros(1, dtype=np.int64) + _launch(k, out, in_a, true_, in_b, true_) + assert int(out.get()[0]) == ref(a, b) + + +@pytest.mark.parametrize( + "av,bv,expected", + [ + (True, True, True), + (True, False, False), + (False, True, False), + (False, False, False), + ], +) +def test_masked_masked_validity_is_anded(av, bv, expected): + """``Masked op Masked`` validity is the AND of the operand validities.""" + + @cuda.jit( + types.void( + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out_valid, a, a_valid, b, b_valid): + m = Masked(a[0], a_valid[0]) + Masked(b[0], b_valid[0]) + out_valid[0] = m.valid + + in_a = cp.array([1], dtype=np.int64) + in_b = cp.array([2], dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_valid, + in_a, + cp.array([av], dtype=np.bool_), + in_b, + cp.array([bv], dtype=np.bool_), + ) + assert bool(out_valid.get()[0]) is expected + + +_CMP = [ + (operator.lt, lambda a, b: a < b), + (operator.le, lambda a, b: a <= b), + (operator.gt, lambda a, b: a > b), + (operator.ge, lambda a, b: a >= b), + (operator.eq, lambda a, b: a == b), + (operator.ne, lambda a, b: a != b), +] + + +@pytest.mark.parametrize("op,ref", _CMP) +@pytest.mark.parametrize("a,b", [(3, 5), (5, 5), (8, 5)]) +def test_masked_masked_comparison(op, ref, a, b): + """Comparison of two Masked values yields a Masked(boolean).""" + + @cuda.jit( + types.void( + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out, x, xv, y, yv): + m = op(Masked(x[0], xv[0]), Masked(y[0], yv[0])) + out[0] = m.value + + true_ = cp.array([True], dtype=np.bool_) + out = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out, + cp.array([a], dtype=np.int64), + true_, + cp.array([b], dtype=np.int64), + true_, + ) + assert bool(out.get()[0]) == ref(a, b) + + +@pytest.mark.parametrize("op,ref", _ARITH) +def test_masked_scalar_arith(op, ref): + """``Masked(a) literal`` carries the Masked operand's validity.""" + + @cuda.jit( + types.void( + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av): + m = op(Masked(a[0], av[0]), 4) + out_v[0] = m.value + out_valid[0] = m.valid + + a = 10 + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([a], dtype=np.int64), + cp.array([False], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == ref(a, 4) + # validity is carried from the (invalid) Masked operand + assert bool(out_valid.get()[0]) is False + + +@pytest.mark.parametrize("op,ref", _ARITH) +def test_scalar_masked_arith(op, ref): + """``literal Masked(a)`` puts the scalar on the left.""" + + @cuda.jit( + types.void( + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av): + m = op(100, Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid + + a = 30 + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([a], dtype=np.int64), + cp.array([True], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == ref(100, a) + assert bool(out_valid.get()[0]) is True + + +def test_masked_scalar_comparison_against_literal(): + """``Masked(a) < literal`` -- the scalar literal must not be confused + with the masked operand (regression guard for ``row['a'] < 1``).""" + + @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) + def k(out, a, av): + m = Masked(a[0], av[0]) < 7 + out[0] = m.value + + true_ = cp.array([True], dtype=np.bool_) + out = cp.zeros(1, dtype=np.bool_) + _launch(k, out, cp.array([3], dtype=np.int64), true_) + assert bool(out.get()[0]) is True + out = cp.zeros(1, dtype=np.bool_) + _launch(k, out, cp.array([9], dtype=np.int64), true_) + assert bool(out.get()[0]) is False + + +@pytest.mark.parametrize("na_first", [True, False]) +def test_masked_binary_with_na_is_invalid(na_first): + """``Masked NA`` (and ``NA Masked``) produce an invalid result.""" + if na_first: + + @cuda.jit( + types.void(types.boolean[::1], types.int64[::1], types.boolean[::1]) + ) + def k(out_valid, a, av): + m = NA + Masked(a[0], av[0]) + out_valid[0] = m.valid + else: + + @cuda.jit( + types.void(types.boolean[::1], types.int64[::1], types.boolean[::1]) + ) + def k(out_valid, a, av): + m = Masked(a[0], av[0]) + NA + out_valid[0] = m.valid + + out_valid = cp.ones(1, dtype=np.bool_) + _launch( + k, + out_valid, + cp.array([5], dtype=np.int64), + cp.array([True], dtype=np.bool_), # valid operand; NA still poisons + ) + assert bool(out_valid.get()[0]) is False From 9c4b5aeeb98e01651ecac02615be1080c8905590 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 06:53:52 -0700 Subject: [PATCH 2/6] Remove section divider comments --- python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py | 1 - .../tests/private_objects/mlir_backend/test_masked_lowering.py | 3 --- 2 files changed, 4 deletions(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index 0ced7f935e19..ed7cbd814201 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -189,7 +189,6 @@ def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null): builder.store_var(target, valid) -# --- Binary ops ----------------------------------------------------------- # Shared helper: apply ``op(v1, v2)`` to two scalar MLIR values, convert the # result to the target Masked's value type, and pack it with the given # validity bit. Numeric/boolean only at this layer. diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py index f633f4b72608..6eef0fa0738a 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py @@ -239,9 +239,6 @@ def k(out, v, valid): assert bool(result[1]) is valid_in -# --- binary ops ------------------------------------------------------------ - - _ARITH = [ (operator.add, lambda a, b: a + b), (operator.sub, lambda a, b: a - b), From 259bf2a5491ce492bba4cbd1d21627b057731c71 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 8 Jul 2026 06:33:42 -0700 Subject: [PATCH 3/6] address reviews --- .../core/udf/mlir_backend/masked_lowering.py | 54 +++++++++++++------ .../core/udf/mlir_backend/masked_typing.py | 31 +++++++---- .../mlir_backend/test_masked_lowering.py | 42 +++++++-------- 3 files changed, 80 insertions(+), 47 deletions(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index ed7cbd814201..7dde5c0ba721 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -3,6 +3,7 @@ from __future__ import annotations import operator +from collections.abc import Callable from functools import partial from typing import TYPE_CHECKING @@ -189,12 +190,19 @@ def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null): builder.store_var(target, valid) -# Shared helper: apply ``op(v1, v2)`` to two scalar MLIR values, convert the -# result to the target Masked's value type, and pack it with the given -# validity bit. Numeric/boolean only at this layer. def _apply_masked_binary_op( - builder, target, target_type, v1, v2, result_valid, op -): + builder: MLIRLower, + target: Var, + target_type: MaskedType, + v1: mlir_ir.Value, + v2: mlir_ir.Value, + result_valid: mlir_ir.Value, + op: Callable, +) -> None: + """Apply ``op(v1, v2)`` to two scalar MLIR values, convert the result to + the target Masked's value type, and pack it with the given validity bit. + Numeric/boolean only at this layer. + """ target_value_mlir_ty = builder.get_mlir_type(target_type.value_type) v1, v2 = coerce_numpy_scalars_for_binary_op(v1, v2) # Comparisons compute on the (already coerced) operand type and @@ -210,9 +218,12 @@ def _apply_masked_binary_op( builder.store_var(target, packed) -# ``Masked Masked``: AND the validity bits. -def _make_lower_masked_binary(op): - def _lower(builder, target, args, kwargs): +def _make_lower_masked_binary(op: Callable) -> Callable: + """``Masked Masked``: AND the validity bits.""" + + def _lower( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list + ) -> None: target_type = builder.get_numba_type(target.name) m1 = builder.load_var(args[0]) m2 = builder.load_var(args[1]) @@ -232,7 +243,12 @@ def _lower(builder, target, args, kwargs): return _lower -def _scalar_value_from_var(builder, s_var, m_var, masked_value_mlir_ty): +def _scalar_value_from_var( + builder: MLIRLower, + s_var: Var, + m_var: Var, + masked_value_mlir_ty: mlir_ir.Type, +) -> mlir_ir.Value: """Resolve the scalar operand for the Masked-vs-scalar path. Prefer a materialized constant when the scalar is a Literal so we @@ -258,10 +274,16 @@ def _scalar_value_from_var(builder, s_var, m_var, masked_value_mlir_ty): return s_raw -# ``Masked scalar`` and ``scalar Masked``: carry the Masked -# operand's validity. -def _make_lower_masked_binary_scalar(op, masked_first): - def _lower(builder, target, args, kwargs): +def _make_lower_masked_binary_scalar( + op: Callable, masked_first: bool +) -> Callable: + """``Masked scalar`` and ``scalar Masked``: carry the Masked + operand's validity. + """ + + def _lower( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list + ) -> None: target_type = builder.get_numba_type(target.name) m_var, s_var = ( (args[0], args[1]) if masked_first else (args[1], args[0]) @@ -284,8 +306,10 @@ def _lower(builder, target, args, kwargs): return _lower -# ``Masked NA`` / ``NA Masked``: result is invalid. -def _lower_masked_binary_null(builder, target, args, kwargs): +def _lower_masked_binary_null( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list +) -> None: + """``Masked NA`` / ``NA Masked``: result is invalid.""" target_type = builder.get_numba_type(target.name) value_mlir_ty = builder.get_mlir_type(target_type.value_type) undef_val = llvm.UndefOp(value_mlir_ty) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py index c7a389a47c45..c58d91c6fe9d 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -14,6 +14,7 @@ AbstractTemplate, AttributeTemplate, ConcreteTemplate, + Signature, ) from numba_cuda_mlir.typing import signature as nb_signature @@ -140,10 +141,14 @@ def generic(self, args, kws): return None -# ``Masked Masked``: resolve the underlying scalar op on the two -# value types, then wrap the result back in a MaskedType. class MaskedScalarArithOp(AbstractTemplate): - def generic(self, args, kws): + """``Masked Masked``: resolve the underlying scalar op on the two + value types, then wrap the result back in a ``MaskedType``. + """ + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: if isinstance(args[0], MaskedType) and isinstance( args[1], MaskedType ): @@ -154,10 +159,14 @@ def generic(self, args, kws): return None -# ``Masked scalar`` and ``scalar Masked`` (scalar may be a -# Literal, e.g. ``row['a'] == 1``). class MaskedScalarScalarOp(AbstractTemplate): - def generic(self, args, kws): + """``Masked scalar`` and ``scalar Masked`` (scalar may be a + ``Literal``, e.g. ``row['a'] == 1``). + """ + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: if isinstance(args[0], MaskedType) and isinstance( args[1], (types.Number, types.Boolean) ): @@ -191,10 +200,14 @@ def generic(self, args, kws): return None -# ``Masked NA`` / ``NA Masked``: result type is the Masked -# operand's type; the lowering produces an invalid (poisoned) value. class MaskedScalarNullOp(AbstractTemplate): - def generic(self, args, kws): + """``Masked NA`` / ``NA Masked``: result type is the Masked + operand's type; the lowering produces an invalid (poisoned) value. + """ + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: if isinstance(args[0], MaskedType) and isinstance(args[1], NAType): return nb_signature(args[0], args[0], na_type) if isinstance(args[0], NAType) and isinstance(args[1], MaskedType): diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py index 6eef0fa0738a..9d0fd4ad5cee 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py @@ -239,15 +239,11 @@ def k(out, v, valid): assert bool(result[1]) is valid_in -_ARITH = [ - (operator.add, lambda a, b: a + b), - (operator.sub, lambda a, b: a - b), - (operator.mul, lambda a, b: a * b), -] +_ARITH = [operator.add, operator.sub, operator.mul] -@pytest.mark.parametrize("op,ref", _ARITH) -def test_masked_masked_arith_value(op, ref): +@pytest.mark.parametrize("op", _ARITH) +def test_masked_masked_arith_value(op): """``Masked(a) Masked(b)`` computes ``op(a, b)`` in the value field.""" @cuda.jit( @@ -269,7 +265,7 @@ def k(out, a, av, b, bv): true_ = cp.array([True], dtype=np.bool_) out = cp.zeros(1, dtype=np.int64) _launch(k, out, in_a, true_, in_b, true_) - assert int(out.get()[0]) == ref(a, b) + assert int(out.get()[0]) == op(a, b) @pytest.mark.parametrize( @@ -312,18 +308,18 @@ def k(out_valid, a, a_valid, b, b_valid): _CMP = [ - (operator.lt, lambda a, b: a < b), - (operator.le, lambda a, b: a <= b), - (operator.gt, lambda a, b: a > b), - (operator.ge, lambda a, b: a >= b), - (operator.eq, lambda a, b: a == b), - (operator.ne, lambda a, b: a != b), + operator.lt, + operator.le, + operator.gt, + operator.ge, + operator.eq, + operator.ne, ] -@pytest.mark.parametrize("op,ref", _CMP) +@pytest.mark.parametrize("op", _CMP) @pytest.mark.parametrize("a,b", [(3, 5), (5, 5), (8, 5)]) -def test_masked_masked_comparison(op, ref, a, b): +def test_masked_masked_comparison(op, a, b): """Comparison of two Masked values yields a Masked(boolean).""" @cuda.jit( @@ -349,11 +345,11 @@ def k(out, x, xv, y, yv): cp.array([b], dtype=np.int64), true_, ) - assert bool(out.get()[0]) == ref(a, b) + assert bool(out.get()[0]) == op(a, b) -@pytest.mark.parametrize("op,ref", _ARITH) -def test_masked_scalar_arith(op, ref): +@pytest.mark.parametrize("op", _ARITH) +def test_masked_scalar_arith(op): """``Masked(a) literal`` carries the Masked operand's validity.""" @cuda.jit( @@ -379,13 +375,13 @@ def k(out_v, out_valid, a, av): cp.array([a], dtype=np.int64), cp.array([False], dtype=np.bool_), ) - assert int(out_v.get()[0]) == ref(a, 4) + assert int(out_v.get()[0]) == op(a, 4) # validity is carried from the (invalid) Masked operand assert bool(out_valid.get()[0]) is False -@pytest.mark.parametrize("op,ref", _ARITH) -def test_scalar_masked_arith(op, ref): +@pytest.mark.parametrize("op", _ARITH) +def test_scalar_masked_arith(op): """``literal Masked(a)`` puts the scalar on the left.""" @cuda.jit( @@ -411,7 +407,7 @@ def k(out_v, out_valid, a, av): cp.array([a], dtype=np.int64), cp.array([True], dtype=np.bool_), ) - assert int(out_v.get()[0]) == ref(100, a) + assert int(out_v.get()[0]) == op(100, a) assert bool(out_valid.get()[0]) is True From 72201ee0ce079d4f9fe9080941ef120d0da8c610 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 14 Jul 2026 09:22:58 -0700 Subject: [PATCH 4/6] lint --- .../core/udf/mlir_backend/masked_lowering.py | 20 +++++-------------- .../core/udf/mlir_backend/masked_typing.py | 10 ++++------ .../mlir_backend/test_masked_lowering.py | 15 ++++++++++---- 3 files changed, 20 insertions(+), 25 deletions(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index 7dde5c0ba721..459649a4215a 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -212,9 +212,7 @@ def _apply_masked_binary_op( v1 = convert(v1, operand_ty) v2 = convert(v2, operand_ty) result_val = convert(op(v1, v2), target_value_mlir_ty) - packed = _pack_masked( - builder, target_type, result_val, result_valid - ) + packed = _pack_masked(builder, target_type, result_val, result_valid) builder.store_var(target, packed) @@ -229,12 +227,8 @@ def _lower( m2 = builder.load_var(args[1]) st1 = llvm.StructType(m1.type) st2 = llvm.StructType(m2.type) - v1, valid1 = _extract_masked_value_valid( - m1, st1.body[0], st1.body[1] - ) - v2, valid2 = _extract_masked_value_valid( - m2, st2.body[0], st2.body[1] - ) + v1, valid1 = _extract_masked_value_valid(m1, st1.body[0], st1.body[1]) + v2, valid2 = _extract_masked_value_valid(m2, st2.body[0], st2.body[1]) result_valid = arith.andi(valid1, valid2) _apply_masked_binary_op( builder, target, target_type, v1, v2, result_valid, op @@ -290,9 +284,7 @@ def _lower( ) m = builder.load_var(m_var) st = llvm.StructType(m.type) - m_val, m_valid = _extract_masked_value_valid( - m, st.body[0], st.body[1] - ) + m_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1]) s_val = _scalar_value_from_var(builder, s_var, m_var, st.body[0]) if masked_first: _apply_masked_binary_op( @@ -316,9 +308,7 @@ def _lower_masked_binary_null( valid_zero = arith.constant( result=builder.get_mlir_type(types.boolean), value=0 ) - packed = _pack_masked( - builder, target_type, undef_val, valid_zero - ) + packed = _pack_masked(builder, target_type, undef_val, valid_zero) builder.store_var(target, packed) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py index c58d91c6fe9d..14a7b9f77a6e 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -149,9 +149,7 @@ class MaskedScalarArithOp(AbstractTemplate): def generic( self, args: tuple[types.Type, ...], kws: dict ) -> Signature | None: - if isinstance(args[0], MaskedType) and isinstance( - args[1], MaskedType - ): + if isinstance(args[0], MaskedType) and isinstance(args[1], MaskedType): return_type = self.context.resolve_function_type( self.key, (args[0].value_type, args[1].value_type), kws ).return_type @@ -182,9 +180,9 @@ def generic( self.key, (args[0].value_type, scalar_ty), kws ).return_type return nb_signature(MaskedType(return_type), args[0], args[1]) - if isinstance( - args[0], (types.Number, types.Boolean) - ) and isinstance(args[1], MaskedType): + if isinstance(args[0], (types.Number, types.Boolean)) and isinstance( + args[1], MaskedType + ): return_type = self.context.resolve_function_type( self.key, (args[0], args[1].value_type), kws ).return_type diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py index 9d0fd4ad5cee..4977f0d5cc8a 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py @@ -413,9 +413,12 @@ def k(out_v, out_valid, a, av): def test_masked_scalar_comparison_against_literal(): """``Masked(a) < literal`` -- the scalar literal must not be confused - with the masked operand (regression guard for ``row['a'] < 1``).""" + with the masked operand (regression guard for ``row['a'] < 1``). + """ - @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) + @cuda.jit( + types.void(types.boolean[::1], types.int64[::1], types.boolean[::1]) + ) def k(out, a, av): m = Masked(a[0], av[0]) < 7 out[0] = m.value @@ -435,7 +438,9 @@ def test_masked_binary_with_na_is_invalid(na_first): if na_first: @cuda.jit( - types.void(types.boolean[::1], types.int64[::1], types.boolean[::1]) + types.void( + types.boolean[::1], types.int64[::1], types.boolean[::1] + ) ) def k(out_valid, a, av): m = NA + Masked(a[0], av[0]) @@ -443,7 +448,9 @@ def k(out_valid, a, av): else: @cuda.jit( - types.void(types.boolean[::1], types.int64[::1], types.boolean[::1]) + types.void( + types.boolean[::1], types.int64[::1], types.boolean[::1] + ) ) def k(out_valid, a, av): m = Masked(a[0], av[0]) + NA From e10bfbed4fb7dd3e85e935916eaea07343966ce4 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 15 Jul 2026 06:14:11 -0700 Subject: [PATCH 5/6] simplify --- .../core/udf/mlir_backend/masked_lowering.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index 459649a4215a..16b8080b252b 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -240,14 +240,12 @@ def _lower( def _scalar_value_from_var( builder: MLIRLower, s_var: Var, - m_var: Var, - masked_value_mlir_ty: mlir_ir.Type, ) -> mlir_ir.Value: """Resolve the scalar operand for the Masked-vs-scalar path. - Prefer a materialized constant when the scalar is a Literal so we - never mistake the masked operand for the scalar (e.g. - ``row['a'] < 1`` must not become ``row['a'] < row['a']``). + A ``Literal`` operand carries its value in the type rather than as a + distinct runtime register, so materialize it directly as a constant; + genuine runtime scalars are loaded from their variable. """ s_ty = builder.get_numba_type(s_var.name) if isinstance(s_ty, types.Literal): @@ -259,13 +257,7 @@ def _scalar_value_from_var( ): py_val = 1 if py_val else 0 return arith.constant(mlir_ty, py_val) - s_raw = builder.load_var(s_var) - if getattr(s_var, "name", None) == getattr(m_var, "name", None): - raise RuntimeError( - "Masked vs scalar lowering: scalar variable is the same as " - "the masked variable; cannot extract a distinct scalar." - ) - return s_raw + return builder.load_var(s_var) def _make_lower_masked_binary_scalar( @@ -285,7 +277,7 @@ def _lower( m = builder.load_var(m_var) st = llvm.StructType(m.type) m_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1]) - s_val = _scalar_value_from_var(builder, s_var, m_var, st.body[0]) + s_val = _scalar_value_from_var(builder, s_var) if masked_first: _apply_masked_binary_op( builder, target, target_type, m_val, s_val, m_valid, op From 39b35fbdfc73cc5ddbe71fe6926983d078988def Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 15 Jul 2026 06:15:03 -0700 Subject: [PATCH 6/6] lint --- python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py index 16b8080b252b..21cd3b171b1c 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -3,7 +3,6 @@ from __future__ import annotations import operator -from collections.abc import Callable from functools import partial from typing import TYPE_CHECKING @@ -28,6 +27,8 @@ ) if TYPE_CHECKING: + from collections.abc import Callable + from numba_cuda_mlir.mlir_lowering import MLIRLower from numba_cuda_mlir.numba_cuda.core.ir import Var from numba_cuda_mlir.numba_cuda.datamodel.manager import (