From dc13b5048066697b2e5309c71d2ce1ac003dc3f8 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 11 Jun 2026 06:01:43 -0700 Subject: [PATCH 01/14] [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 3be1a3bf6026..3da9c650c12d 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 @@ -13,6 +15,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 @@ -231,3 +234,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 9db5a9e1eb9429581eb1d3338ee98fd052269364 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 06:53:52 -0700 Subject: [PATCH 02/14] 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 3da9c650c12d..b4e45705804e 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 @@ -236,9 +236,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 da8c675ab18c295cf64d961f885de918360d0457 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 8 Jul 2026 06:33:42 -0700 Subject: [PATCH 03/14] 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 b4e45705804e..fd0d2feab6d6 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 @@ -236,15 +236,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( @@ -266,7 +262,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( @@ -309,18 +305,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( @@ -346,11 +342,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( @@ -376,13 +372,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( @@ -408,7 +404,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 bb45b42fec5f16615a9498bf9ba37c4931f88e57 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Tue, 14 Jul 2026 09:22:58 -0700 Subject: [PATCH 04/14] 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 fd0d2feab6d6..e4ef4a4e0508 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 @@ -410,9 +410,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 @@ -432,7 +435,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]) @@ -440,7 +445,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 dca111791a73d9d124391f7bad8e115a8032e632 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 15 Jul 2026 06:14:11 -0700 Subject: [PATCH 05/14] 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 99360499a6467aef981a4dfe99534a8350d3f9a8 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 15 Jul 2026 06:15:03 -0700 Subject: [PATCH 06/14] 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 ( From deef12e06cb7ed21b6127c0a62b51a5ca7a28931 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 11 Jun 2026 06:03:04 -0700 Subject: [PATCH 07/14] [MLIR] PR 5: Masked unary ops + coercions Adds unary operator typing/lowering and scalar coercions over numeric/boolean MaskedType values: * unary ops (-x, math.sin(x), etc.) -> Masked(result), validity carried; delegate to the registered scalar lowering * operator.invert (~x) on integer payloads via arith.xori(x, -1) * abs(m) -> Masked(result) * bool(m) / truth -> m.valid and bool(m.value) * int(m) -> Masked(int64); float(m) -> Masked(float64) Tests: +21 kernel tests (sign, invert, math.* delegation, abs, truth across valid/invalid/falsy, bool-in-if, int/float coercion). --- .../core/udf/mlir_backend/masked_lowering.py | 146 ++++++++++++- .../core/udf/mlir_backend/masked_typing.py | 63 +++++- .../mlir_backend/test_masked_lowering.py | 205 ++++++++++++++++++ 3 files changed, 412 insertions(+), 2 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 21cd3b171b1c..c85270d20916 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -12,13 +12,21 @@ 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 ( + bool_of, coerce_numpy_scalars_for_binary_op, convert, + false, ) from numba_cuda_mlir.models import PrimitiveModel, register_model +from numba_cuda_mlir.numba_cuda.core import ir as numba_ir 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._ops import ( + arith_ops, + bitwise_ops, + comparison_ops, + unary_ops, +) from cudf.core.udf.api import Masked from cudf.core.udf.mlir_backend.masked_typing import ( MaskedType, @@ -305,6 +313,129 @@ def _lower_masked_binary_null( builder.store_var(target, packed) +def _make_temp_var(builder, base_var, name_suffix, numba_type): + """Create a synthetic IR var + typemap entry. + + Used to feed a scalar value into a registered numba_cuda_mlir scalar + lowering (e.g. ``math.sin``) and read its result back, since those + lowerings operate on IR vars rather than raw SSA values. + """ + scope = getattr(base_var, "scope", None) + loc = getattr(base_var, "loc", None) + name = f"$masked_uop_{base_var.name}_{name_suffix}" + temp = numba_ir.Var(scope=scope, name=name, loc=loc) + builder.fndesc.typemap[temp.name] = numba_type + return temp + + +# --- Unary ops ------------------------------------------------------------ +# Generic unary: delegate the scalar op to the registered numba_cuda_mlir +# scalar lowering (``math.sin`` -> math dialect, ``operator.neg`` -> arith, +# etc.), then re-wrap with the operand's validity. +def _make_lower_masked_unary(op): + def _lower(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + result_inner_ty = target_type.value_type + operand_inner_ty = builder.get_numba_type( + args[0].name + ).value_type + + m = builder.load_var(args[0]) + st = llvm.StructType(m.type) + m_val, m_valid = _extract_masked_value_valid( + m, st.body[0], st.body[1] + ) + m_val = convert(m_val, builder.get_mlir_type(operand_inner_ty)) + + sig = result_inner_ty(operand_inner_ty) + cg = builder.get_registered_builder(op, sig) + if cg is None: + raise NotImplementedError( + "No MLIR lowering for unary " + f"{getattr(op, '__name__', op)!r} on {operand_inner_ty}; " + f"signature {sig}" + ) + # The same operand var can feed multiple unary calls in one + # expression (e.g. ``sin(x) + lgamma(x)``); suffix the temp var + # name by op so typemap keys stay unique. + op_tag = getattr(op, "__name__", "op") + op_var = _make_temp_var( + builder, args[0], f"{op_tag}_in", operand_inner_ty + ) + out_var = _make_temp_var( + builder, args[0], f"{op_tag}_out", result_inner_ty + ) + builder.store_var(op_var, m_val) + cg(builder, out_var, [op_var], []) + result_val = convert( + builder.load_var(out_var), + builder.get_mlir_type(result_inner_ty), + ) + packed = _pack_masked( + builder, target_type, result_val, m_valid + ) + builder.store_var(target, packed) + + return _lower + + +# ``operator.invert`` (bitwise ~) on Masked integers: there is no scalar +# @lower for invert, so do ``xori(x, -1)``. The all-ones mask is the +# signed constant -1 (two's complement); ``(1< Masked(int64); float(m) -> Masked(float64). +def _make_lower_masked_numeric_cast(): + def _lower(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + target_value_mlir_ty = builder.get_mlir_type( + target_type.value_type + ) + m = builder.load_var(args[0]) + st = llvm.StructType(m.type) + m_val, m_valid = _extract_masked_value_valid( + m, st.body[0], st.body[1] + ) + casted = builder.mlir_convert(m_val, target_value_mlir_ty) + packed = _pack_masked( + builder, target_type, casted, m_valid + ) + builder.store_var(target, packed) + + return _lower + + def _register() -> None: """Register the data model and lowerings with ``numba_cuda_mlir``. @@ -354,5 +485,18 @@ def _register() -> None: lower(binary_op, MaskedType, NAType)(_lower_masked_binary_null) lower(binary_op, NAType, MaskedType)(_lower_masked_binary_null) + for unary_op in unary_ops: + if unary_op is operator.invert: + continue + lower(unary_op, MaskedType)(_make_lower_masked_unary(unary_op)) + lower(abs, MaskedType)(_make_lower_masked_unary(abs)) + lower(operator.invert, MaskedType)(_lower_masked_invert) + + lower(operator.truth, MaskedType)(_lower_masked_truth) + lower(bool, MaskedType)(_lower_masked_truth) + + lower(float, MaskedType)(_make_lower_masked_numeric_cast()) + lower(int, MaskedType)(_make_lower_masked_numeric_cast()) + _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 14a7b9f77a6e..4344040f89fa 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -19,7 +19,12 @@ 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._ops import ( + arith_ops, + bitwise_ops, + comparison_ops, + unary_ops, +) from cudf.core.udf.api import Masked _SUPPORTED_MASKED_VALUE_TYPE_CLASSES = ( @@ -213,6 +218,54 @@ def generic( return None +# --- Unary ops: `` Masked`` -> ``Masked(result)`` --- +# Resolve the underlying scalar op on the value type, wrap the result. +class MaskedScalarUnaryOp(AbstractTemplate): + def generic(self, args, kws): + if len(args) == 1 and isinstance(args[0], MaskedType): + return_type = self.context.resolve_function_type( + self.key, (args[0].value_type,), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0]) + return None + + +# ``bool(m)`` / ``operator.truth(m)`` -> boolean. The runtime result is +# ``m.valid and bool(m.value)``; the *type* is a plain boolean (used +# directly in ``if`` conditions). +class MaskedScalarTruth(AbstractTemplate): + def generic(self, args, kws): + if len(args) == 1 and isinstance(args[0], MaskedType): + return nb_signature(types.boolean, args[0]) + return None + + +# ``float(m)`` -> Masked(float64); ``int(m)`` -> Masked(int64). +class MaskedScalarFloatCast(AbstractTemplate): + def generic(self, args, kws): + if len(args) == 1 and isinstance(args[0], MaskedType): + return nb_signature(MaskedType(types.float64), args[0]) + return None + + +class MaskedScalarIntCast(AbstractTemplate): + def generic(self, args, kws): + if len(args) == 1 and isinstance(args[0], MaskedType): + return nb_signature(MaskedType(types.int64), args[0]) + return None + + +# ``abs(m)`` -> Masked(result). +class MaskedScalarAbsoluteValue(AbstractTemplate): + def generic(self, args, kws): + if len(args) == 1 and isinstance(args[0], MaskedType): + return_type = self.context.resolve_function_type( + self.key, (args[0].value_type,), kws + ).return_type + return nb_signature(MaskedType(return_type), args[0]) + return None + + def _register() -> None: """Register typing for ``Masked`` and ``MaskedType`` attributes with ``numba_cuda_mlir``. Called once at module import. @@ -227,5 +280,13 @@ def _register() -> None: typing_registry.register_global(binary_op)(MaskedScalarNullOp) typing_registry.register_global(binary_op)(MaskedScalarScalarOp) + for unary_op in unary_ops: + typing_registry.register_global(unary_op)(MaskedScalarUnaryOp) + typing_registry.register_global(operator.truth)(MaskedScalarTruth) + typing_registry.register_global(bool)(MaskedScalarTruth) + typing_registry.register_global(float)(MaskedScalarFloatCast) + typing_registry.register_global(int)(MaskedScalarIntCast) + typing_registry.register_global(abs)(MaskedScalarAbsoluteValue) + _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 e4ef4a4e0508..b2e5e76d818a 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,7 @@ from __future__ import annotations +import math import operator import cupy as cp @@ -461,3 +462,207 @@ def k(out_valid, a, av): cp.array([True], dtype=np.bool_), # valid operand; NA still poisons ) assert bool(out_valid.get()[0]) is False + + +# --- unary ops + coercions ------------------------------------------------- + + +@pytest.mark.parametrize( + "op,ref", [(operator.neg, lambda x: -x), (operator.pos, lambda x: +x)] +) +def test_masked_unary_sign(op, ref): + """``-m`` / ``+m`` apply the scalar op and carry 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])) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([7], dtype=np.int64), + cp.array([False], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == ref(7) + # validity carried from the operand + assert bool(out_valid.get()[0]) is False + + +@pytest.mark.parametrize("x", [5, 0, -6, 255]) +def test_masked_invert(x): + """``~m`` is bitwise-not on the integer payload.""" + + @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) + def k(out, a, av): + out[0] = (~Masked(a[0], av[0])).value + + out = cp.zeros(1, dtype=np.int64) + _launch(k, out, cp.array([x], dtype=np.int64), cp.array([True], dtype=np.bool_)) + assert int(out.get()[0]) == ~x + + +@pytest.mark.parametrize( + "fn,ref", + [ + (math.sin, math.sin), + (math.cos, math.cos), + (math.sqrt, math.sqrt), + (math.exp, math.exp), + ], +) +def test_masked_unary_math(fn, ref): + """``math.*`` unary ops delegate to the scalar lowering.""" + + @cuda.jit(types.void(types.float64[::1], types.float64[::1], types.boolean[::1])) + def k(out, a, av): + out[0] = fn(Masked(a[0], av[0])).value + + out = cp.zeros(1, dtype=np.float64) + _launch( + k, out, cp.array([1.5], dtype=np.float64), cp.array([True], dtype=np.bool_) + ) + np.testing.assert_allclose(float(out.get()[0]), ref(1.5), rtol=1e-12) + + +@pytest.mark.parametrize("x", [-9, 0, 12]) +def test_masked_abs(x): + """``abs(m)`` -> Masked with the absolute value.""" + + @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 = abs(Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([x], dtype=np.int64), + cp.array([True], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == abs(x) + assert bool(out_valid.get()[0]) is True + + +@pytest.mark.parametrize( + "value,valid,expected", + [ + (5, True, True), # valid & truthy + (0, True, False), # valid & falsy + (5, False, False), # invalid -> False regardless of payload + (0, False, False), + ], +) +def test_masked_bool_truth(value, valid, expected): + """``bool(m)`` is ``m.valid and bool(m.value)``.""" + + @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) + def k(out, a, av): + out[0] = bool(Masked(a[0], av[0])) + + out = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out, + cp.array([value], dtype=np.int64), + cp.array([valid], dtype=np.bool_), + ) + assert bool(out.get()[0]) is expected + + +def test_masked_bool_in_if_condition(): + """``bool(m)`` works as an ``if`` predicate inside a UDF.""" + + @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) + def k(out, a, av): + m = Masked(a[0], av[0]) + if m: + out[0] = 1 + else: + out[0] = 0 + + out = cp.zeros(1, dtype=np.int64) + _launch(k, out, cp.array([5], dtype=np.int64), cp.array([True], dtype=np.bool_)) + assert int(out.get()[0]) == 1 + _launch(k, out, cp.array([5], dtype=np.int64), cp.array([False], dtype=np.bool_)) + assert int(out.get()[0]) == 0 + + +def test_masked_float_cast(): + """``float(m)`` -> Masked(float64), validity carried.""" + + @cuda.jit( + types.void( + types.float64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av): + m = float(Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.float64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([3], dtype=np.int64), + cp.array([True], dtype=np.bool_), + ) + assert float(out_v.get()[0]) == 3.0 + assert bool(out_valid.get()[0]) is True + + +def test_masked_int_cast(): + """``int(m)`` -> Masked(int64), truncating the float payload.""" + + @cuda.jit( + types.void( + types.int64[::1], + types.boolean[::1], + types.float64[::1], + types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av): + m = int(Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([3.9], dtype=np.float64), + cp.array([False], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == 3 + assert bool(out_valid.get()[0]) is False From b7b6e462becc6b6da4ab6f465d1b265d7c9a08f7 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 07:04:05 -0700 Subject: [PATCH 08/14] Remove section dividers; stub docstrings as TODO --- .../core/udf/mlir_backend/masked_lowering.py | 8 +------- .../core/udf/mlir_backend/masked_typing.py | 1 - .../mlir_backend/test_masked_lowering.py | 19 ++++++++----------- 3 files changed, 9 insertions(+), 19 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 c85270d20916..1a90704e9969 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -314,12 +314,7 @@ def _lower_masked_binary_null( def _make_temp_var(builder, base_var, name_suffix, numba_type): - """Create a synthetic IR var + typemap entry. - - Used to feed a scalar value into a registered numba_cuda_mlir scalar - lowering (e.g. ``math.sin``) and read its result back, since those - lowerings operate on IR vars rather than raw SSA values. - """ + """TODO: write docstring.""" scope = getattr(base_var, "scope", None) loc = getattr(base_var, "loc", None) name = f"$masked_uop_{base_var.name}_{name_suffix}" @@ -328,7 +323,6 @@ def _make_temp_var(builder, base_var, name_suffix, numba_type): return temp -# --- Unary ops ------------------------------------------------------------ # Generic unary: delegate the scalar op to the registered numba_cuda_mlir # scalar lowering (``math.sin`` -> math dialect, ``operator.neg`` -> arith, # etc.), then re-wrap with the operand's validity. 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 4344040f89fa..8af841083698 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -218,7 +218,6 @@ def generic( return None -# --- Unary ops: `` Masked`` -> ``Masked(result)`` --- # Resolve the underlying scalar op on the value type, wrap the result. class MaskedScalarUnaryOp(AbstractTemplate): def generic(self, args, kws): 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 b2e5e76d818a..a270e174cb47 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 @@ -464,14 +464,11 @@ def k(out_valid, a, av): assert bool(out_valid.get()[0]) is False -# --- unary ops + coercions ------------------------------------------------- - - @pytest.mark.parametrize( "op,ref", [(operator.neg, lambda x: -x), (operator.pos, lambda x: +x)] ) def test_masked_unary_sign(op, ref): - """``-m`` / ``+m`` apply the scalar op and carry validity.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -502,7 +499,7 @@ def k(out_v, out_valid, a, av): @pytest.mark.parametrize("x", [5, 0, -6, 255]) def test_masked_invert(x): - """``~m`` is bitwise-not on the integer payload.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) def k(out, a, av): @@ -523,7 +520,7 @@ def k(out, a, av): ], ) def test_masked_unary_math(fn, ref): - """``math.*`` unary ops delegate to the scalar lowering.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.float64[::1], types.float64[::1], types.boolean[::1])) def k(out, a, av): @@ -538,7 +535,7 @@ def k(out, a, av): @pytest.mark.parametrize("x", [-9, 0, 12]) def test_masked_abs(x): - """``abs(m)`` -> Masked with the absolute value.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -576,7 +573,7 @@ def k(out_v, out_valid, a, av): ], ) def test_masked_bool_truth(value, valid, expected): - """``bool(m)`` is ``m.valid and bool(m.value)``.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) def k(out, a, av): @@ -593,7 +590,7 @@ def k(out, a, av): def test_masked_bool_in_if_condition(): - """``bool(m)`` works as an ``if`` predicate inside a UDF.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) def k(out, a, av): @@ -611,7 +608,7 @@ def k(out, a, av): def test_masked_float_cast(): - """``float(m)`` -> Masked(float64), validity carried.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -640,7 +637,7 @@ def k(out_v, out_valid, a, av): def test_masked_int_cast(): - """``int(m)`` -> Masked(int64), truncating the float payload.""" + """TODO: write docstring.""" @cuda.jit( types.void( From 81481a4e8d7d5d09ae7789cff074e38500b47177 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Fri, 12 Jun 2026 04:40:44 -0700 Subject: [PATCH 09/14] [MLIR] PR 6: datetime64 / timedelta64 value types Extends MaskedType's supported value types to NPDatetime / NPTimedelta (units ns/us/ms/s) so UDFs can operate on datetime64 / timedelta64 columns: * _SUPPORTED_MASKED_VALUE_TYPE_CLASSES and _supported_value_type_instances gain datetime/timedelta; the Masked constructor and .value/.valid work generically. * binary-op lowering gains a datetime delegation path: dt/td +/- dt/td routes through the registered numba_cuda_mlir scalar datetime lowering (unit scaling) instead of a raw i64 op; comparisons use the existing numeric path. _apply_masked_binary_op now takes inner_ty1/inner_ty2/ ref_var so the Masked-Masked and Masked-scalar callers can opt in. Tests: +2 typing (supported-set spot checks, no-poison) and +5 kernel (dt-dt->td, dt+td->dt, td+td->td, comparison, validity propagation). Device arrays use the int64-view trick since cupy rejects temporal dtypes. --- .../core/udf/mlir_backend/masked_lowering.py | 74 +++++++++- .../core/udf/mlir_backend/masked_typing.py | 12 +- .../mlir_backend/test_masked_lowering.py | 132 ++++++++++++++++++ .../mlir_backend/test_masked_typing.py | 12 ++ 4 files changed, 226 insertions(+), 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 1a90704e9969..c207f7f3426e 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -18,6 +18,7 @@ false, ) from numba_cuda_mlir.models import PrimitiveModel, register_model +from numba_cuda_mlir.numba_cuda import typing as nb_typing from numba_cuda_mlir.numba_cuda.core import ir as numba_ir from numba_cuda_mlir.numba_cuda.types.misc import unliteral @@ -199,6 +200,45 @@ def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null): builder.store_var(target, valid) +# datetime64 / timedelta64 ``+``/``-`` need numba_cuda_mlir's ``datetime`` +# lowering (which scales by unit), not a raw i64 op. +def _needs_datetimelike_delegate(op, ty1, ty2): + if op not in (operator.add, operator.sub): + return False + return isinstance( + ty1, (types.NPDatetime, types.NPTimedelta) + ) or isinstance(ty2, (types.NPDatetime, types.NPTimedelta)) + + +def _apply_masked_datetimelike_binary( + builder, target, target_type, v1, v2, result_valid, op, ty1, ty2, + ref_var, +): + """Delegate ``dt/td +/- dt/td`` to the registered scalar datetime + lowering, then pack with the (already-computed) validity.""" + ret_ty = target_type.value_type + nb_sig = nb_typing.signature(ret_ty, ty1, ty2) + cg = builder.get_registered_builder(op, nb_sig) + if cg is None: + raise NotImplementedError( + f"No MLIR lowering for masked {op!r} with {ty1}, {ty2}; " + f"signature {nb_sig}" + ) + in1 = _make_temp_var(builder, ref_var, "mdt_l", ty1) + in2 = _make_temp_var(builder, ref_var, "mdt_r", ty2) + outv = _make_temp_var(builder, ref_var, "mdt_o", ret_ty) + builder.store_var(in1, convert(v1, builder.get_mlir_type(ty1))) + builder.store_var(in2, convert(v2, builder.get_mlir_type(ty2))) + cg(builder, outv, [in1, in2], ()) + result_val = convert( + builder.load_var(outv), builder.get_mlir_type(ret_ty) + ) + packed = _pack_masked( + builder, target_type, result_val, result_valid + ) + builder.store_var(target, packed) + + def _apply_masked_binary_op( builder: MLIRLower, target: Var, @@ -207,11 +247,29 @@ def _apply_masked_binary_op( v2: mlir_ir.Value, result_valid: mlir_ir.Value, op: Callable, + *, + inner_ty1: types.Type | None = None, + inner_ty2: types.Type | None = None, + ref_var: Var | None = None, ) -> 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. """ + # datetime/timedelta add/sub: delegate to the unit-aware scalar + # lowering when we know the operand inner types. + if ( + inner_ty1 is not None + and inner_ty2 is not None + and ref_var is not None + and _needs_datetimelike_delegate(op, inner_ty1, inner_ty2) + ): + _apply_masked_datetimelike_binary( + builder, target, target_type, v1, v2, result_valid, op, + inner_ty1, inner_ty2, ref_var, + ) + return + 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 @@ -239,8 +297,11 @@ def _lower( 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) + ty1 = builder.get_numba_type(args[0].name).value_type + ty2 = builder.get_numba_type(args[1].name).value_type _apply_masked_binary_op( - builder, target, target_type, v1, v2, result_valid, op + builder, target, target_type, v1, v2, result_valid, op, + inner_ty1=ty1, inner_ty2=ty2, ref_var=args[0], ) return _lower @@ -287,13 +348,20 @@ def _lower( 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_inner_ty = builder.get_numba_type(m_var.name).value_type + s_ty = builder.get_numba_type(s_var.name) + s_inner_ty = ( + unliteral(s_ty) if isinstance(s_ty, types.Literal) else s_ty + ) if masked_first: _apply_masked_binary_op( - builder, target, target_type, m_val, s_val, m_valid, op + builder, target, target_type, m_val, s_val, m_valid, op, + inner_ty1=m_inner_ty, inner_ty2=s_inner_ty, ref_var=m_var, ) else: _apply_masked_binary_op( - builder, target, target_type, s_val, m_val, m_valid, op + builder, target, target_type, s_val, m_val, m_valid, op, + inner_ty1=s_inner_ty, inner_ty2=m_inner_ty, ref_var=m_var, ) return _lower 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 8af841083698..837919884204 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -27,14 +27,24 @@ ) from cudf.core.udf.api import Masked +# Datetime / timedelta resolutions cudf UDFs support. Mirrors the units +# used by the column dtypes flowing into the kernels. +_units = ("ns", "us", "ms", "s") + _SUPPORTED_MASKED_VALUE_TYPE_CLASSES = ( types.Number, types.Boolean, + types.NPDatetime, + types.NPTimedelta, ) _supported_value_type_instances = ( - nb_types.integer_domain | nb_types.real_domain | {nb_types.boolean} + nb_types.integer_domain + | nb_types.real_domain + | {nb_types.boolean} + | {nb_types.NPDatetime(u) for u in _units} + | {nb_types.NPTimedelta(u) for u in _units} ) 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 a270e174cb47..3f45eaa18e2c 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 @@ -663,3 +663,135 @@ def k(out_v, out_valid, a, av): ) assert int(out_v.get()[0]) == 3 assert bool(out_valid.get()[0]) is False + + +# --- datetime / timedelta ------------------------------------------------- +# +# cupy rejects datetime64/timedelta64 dtypes directly, so device arrays are +# allocated as int64 and viewed as the temporal dtype (the same trick +# ``cudf.core.udf.utils._return_arr_from_dtype`` uses). + +_DT = types.NPDatetime("ns") +_TD = types.NPTimedelta("ns") + + +def _dt_in(values): + return cp.asarray(np.array(values, dtype="int64")).view("datetime64[ns]") + + +def _td_in(values): + return cp.asarray(np.array(values, dtype="int64")).view("timedelta64[ns]") + + +def _dt_out(): + return cp.zeros(1, dtype=np.int64).view("datetime64[ns]") + + +def _td_out(): + return cp.zeros(1, dtype=np.int64).view("timedelta64[ns]") + + +def _bool(v): + return cp.array([v], dtype=np.bool_) + + +def test_masked_datetime_minus_datetime_is_timedelta(): + """``Masked(dt) - Masked(dt)`` -> Masked(timedelta), unit-aware.""" + + @cuda.jit( + types.void( + _TD[::1], types.boolean[::1], + _DT[::1], types.boolean[::1], + _DT[::1], types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av, b, bv): + m = Masked(a[0], av[0]) - Masked(b[0], bv[0]) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = _td_out() + out_valid = cp.zeros(1, dtype=np.bool_) + _launch(k, out_v, out_valid, _dt_in([1000]), _bool(True), + _dt_in([400]), _bool(True)) + assert int(out_v.get().view("int64")[0]) == 600 + assert bool(out_valid.get()[0]) is True + + +def test_masked_datetime_plus_timedelta_is_datetime(): + """``Masked(dt) + Masked(td)`` -> Masked(datetime).""" + + @cuda.jit( + types.void( + _DT[::1], + _DT[::1], types.boolean[::1], + _TD[::1], types.boolean[::1], + ) + ) + def k(out_v, a, av, t, tv): + out_v[0] = (Masked(a[0], av[0]) + Masked(t[0], tv[0])).value + + out_v = _dt_out() + _launch(k, out_v, _dt_in([1000]), _bool(True), _td_in([250]), _bool(True)) + assert int(out_v.get().view("int64")[0]) == 1250 + + +def test_masked_timedelta_plus_timedelta_is_timedelta(): + """``Masked(td) + Masked(td)`` -> Masked(timedelta).""" + + @cuda.jit( + types.void( + _TD[::1], + _TD[::1], types.boolean[::1], + _TD[::1], types.boolean[::1], + ) + ) + def k(out_v, a, av, b, bv): + out_v[0] = (Masked(a[0], av[0]) + Masked(b[0], bv[0])).value + + out_v = _td_out() + _launch(k, out_v, _td_in([300]), _bool(True), _td_in([120]), _bool(True)) + assert int(out_v.get().view("int64")[0]) == 420 + + +@pytest.mark.parametrize( + "op,ref", [(operator.lt, lambda a, b: a < b), (operator.gt, lambda a, b: a > b)] +) +def test_masked_datetime_comparison(op, ref): + """Comparisons on datetime payloads flow through the numeric path.""" + + @cuda.jit( + types.void( + types.boolean[::1], + _DT[::1], types.boolean[::1], + _DT[::1], types.boolean[::1], + ) + ) + def k(out, a, av, b, bv): + out[0] = op(Masked(a[0], av[0]), Masked(b[0], bv[0])).value + + out = cp.zeros(1, dtype=np.bool_) + _launch(k, out, _dt_in([400]), _bool(True), _dt_in([1000]), _bool(True)) + assert bool(out.get()[0]) == ref(400, 1000) + + +def test_masked_datetime_arith_validity_propagates(): + """One invalid datetime operand -> invalid timedelta result.""" + + @cuda.jit( + types.void( + _TD[::1], types.boolean[::1], + _DT[::1], types.boolean[::1], + _DT[::1], types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av, b, bv): + m = Masked(a[0], av[0]) - Masked(b[0], bv[0]) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = _td_out() + out_valid = cp.ones(1, dtype=np.bool_) + _launch(k, out_v, out_valid, _dt_in([1000]), _bool(True), + _dt_in([400]), _bool(False)) + assert bool(out_valid.get()[0]) is False diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py index 20372a64d34f..2397d7723370 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import pytest from numba_cuda_mlir import types from cudf.core.udf.mlir_backend.masked_typing import ( @@ -52,6 +53,17 @@ def test_masked_type_unsupported_value_becomes_poison(): assert isinstance(masked.value_type, types.Poison) +@pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"]) +def test_masked_datetime_timedelta_not_poisoned(unit): + """``MaskedType(datetime64/timedelta64)`` keeps its value type (not Poison).""" + dt = MaskedType(types.NPDatetime(unit)) + td = MaskedType(types.NPTimedelta(unit)) + assert dt.value_type == types.NPDatetime(unit) + assert td.value_type == types.NPTimedelta(unit) + assert not isinstance(dt.value_type, types.Poison) + assert not isinstance(td.value_type, types.Poison) + + def test_na_type_singleton_repr(): """``NAType`` repr is ``"NA"``.""" from cudf.core.udf.mlir_backend.masked_typing import NAType, na_type From 95514facb38cd63232b3d6736af082f80fae6fc5 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 07:06:54 -0700 Subject: [PATCH 10/14] Remove section dividers; stub docstrings as TODO --- .../cudf/core/udf/mlir_backend/masked_lowering.py | 3 +-- .../mlir_backend/test_masked_lowering.py | 11 +++++------ .../mlir_backend/test_masked_typing.py | 2 +- 3 files changed, 7 insertions(+), 9 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 c207f7f3426e..0a37b216ecdb 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -214,8 +214,7 @@ def _apply_masked_datetimelike_binary( builder, target, target_type, v1, v2, result_valid, op, ty1, ty2, ref_var, ): - """Delegate ``dt/td +/- dt/td`` to the registered scalar datetime - lowering, then pack with the (already-computed) validity.""" + """TODO: write docstring.""" ret_ty = target_type.value_type nb_sig = nb_typing.signature(ret_ty, ty1, ty2) cg = builder.get_registered_builder(op, nb_sig) 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 3f45eaa18e2c..50739b5c28d2 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 @@ -665,7 +665,6 @@ def k(out_v, out_valid, a, av): assert bool(out_valid.get()[0]) is False -# --- datetime / timedelta ------------------------------------------------- # # cupy rejects datetime64/timedelta64 dtypes directly, so device arrays are # allocated as int64 and viewed as the temporal dtype (the same trick @@ -696,7 +695,7 @@ def _bool(v): def test_masked_datetime_minus_datetime_is_timedelta(): - """``Masked(dt) - Masked(dt)`` -> Masked(timedelta), unit-aware.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -719,7 +718,7 @@ def k(out_v, out_valid, a, av, b, bv): def test_masked_datetime_plus_timedelta_is_datetime(): - """``Masked(dt) + Masked(td)`` -> Masked(datetime).""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -737,7 +736,7 @@ def k(out_v, a, av, t, tv): def test_masked_timedelta_plus_timedelta_is_timedelta(): - """``Masked(td) + Masked(td)`` -> Masked(timedelta).""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -758,7 +757,7 @@ def k(out_v, a, av, b, bv): "op,ref", [(operator.lt, lambda a, b: a < b), (operator.gt, lambda a, b: a > b)] ) def test_masked_datetime_comparison(op, ref): - """Comparisons on datetime payloads flow through the numeric path.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -776,7 +775,7 @@ def k(out, a, av, b, bv): def test_masked_datetime_arith_validity_propagates(): - """One invalid datetime operand -> invalid timedelta result.""" + """TODO: write docstring.""" @cuda.jit( types.void( diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py index 2397d7723370..357ab6f10b9c 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py @@ -55,7 +55,7 @@ def test_masked_type_unsupported_value_becomes_poison(): @pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"]) def test_masked_datetime_timedelta_not_poisoned(unit): - """``MaskedType(datetime64/timedelta64)`` keeps its value type (not Poison).""" + """TODO: write docstring.""" dt = MaskedType(types.NPDatetime(unit)) td = MaskedType(types.NPTimedelta(unit)) assert dt.value_type == types.NPDatetime(unit) From 7c58d026e96e3f52d5153461d403b4e9bcb6a1de Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 15 Jun 2026 10:54:14 -0700 Subject: [PATCH 11/14] [MLIR] PR 7: Masked membership (value in tuple) Adds `operator.contains` typing/lowering for `value in (...)` where the value is a Masked numeric/boolean/temporal scalar: * literal tuple of constants -> OR of equality vs each constant * homogeneous UniTuple -> reduce equality across the tuple-as-tensor via linalg.ReduceOp Both produce Masked(boolean) carrying the Masked operand's validity. (String membership -- substr in str -- arrives with the string value type in a later PR; the typing template intentionally accepts any non-Poison Masked value type.) Tests: +13 kernel tests (literal-tuple int/float hit+miss, runtime UniTuple hit+miss, invalid-operand validity propagation). --- .../core/udf/mlir_backend/masked_lowering.py | 99 ++++++++++++++++++- .../core/udf/mlir_backend/masked_typing.py | 30 ++++++ .../mlir_backend/test_masked_lowering.py | 77 +++++++++++++++ 3 files changed, 205 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 0a37b216ecdb..063a0fdd4530 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -9,13 +9,19 @@ 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._mlir.dialects import arith, linalg, llvm, tensor +from numba_cuda_mlir._mlir.extras import types as T from numba_cuda_mlir.extending import lower_cast, lowering_registry from numba_cuda_mlir.lowering_utilities import ( bool_of, coerce_numpy_scalars_for_binary_op, + concretize_tuple_to_tensor, convert, + equal, false, + float_of, + int_of, + try_extract_constant, ) from numba_cuda_mlir.models import PrimitiveModel, register_model from numba_cuda_mlir.numba_cuda import typing as nb_typing @@ -497,6 +503,90 @@ def _lower(builder, target, args, kwargs): return _lower +# --- membership: ``value in container`` -> Masked(boolean) ---------------- +def _const_mlir_for_membership(py_const, mlir_ty): + if isinstance(py_const, float): + return float_of(py_const, mlir_ty) + if isinstance(py_const, bool): + return int_of(int(py_const), mlir_ty) + return int_of(py_const, mlir_ty) + + +# ``value in (c0, c1, ...)`` literal tuple: OR of equality vs each const. +def _lower_masked_literal_tuple_contains(builder, target, args, kwargs): + tup = builder.load_var(args[0]) + m = builder.load_var(args[1]) + st = llvm.StructType(m.type) + m_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1]) + + constant_values = [] + for x in tup: + cv = try_extract_constant(x) + if cv is None: + raise NotImplementedError( + "Masked membership in a tuple is only implemented for " + f"constant tuple elements, got {x!r}" + ) + constant_values.append(cv) + + result = false() + for const_val in constant_values: + c = _const_mlir_for_membership(const_val, m_val.type) + m_v, c_v = coerce_numpy_scalars_for_binary_op(m_val, c) + result = arith.ori(result, equal(m_v, c_v)) + + bool_mlir_ty = builder.get_mlir_type(types.boolean) + undef_bool = llvm.UndefOp(bool_mlir_ty) + final_bool = arith.select(m_valid, result, undef_bool) + target_type = builder.get_numba_type(target.name) + packed = _pack_masked(builder, target_type, final_bool, m_valid) + builder.store_var(target, packed) + + +# ``value in homogeneous_tuple``: reduce equality across the tuple. +def _lower_masked_unittuple_contains(builder, target, args, kwargs): + tup = builder.load_var(args[0]) + if not isinstance(tup, tuple): + raise NotImplementedError( + f"UniTuple contains expects a lowered tuple, got {type(tup)}" + ) + tup_t = concretize_tuple_to_tensor(tup) + + m = builder.load_var(args[1]) + st = llvm.StructType(m.type) + m_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1]) + elem_ty = tup_t.type.element_type + m_cmp = convert(m_val, elem_ty) + + def body(_op, element, accumulator): + found = equal(element, m_cmp) + found = arith.ori(found, accumulator) + linalg.yield_([found]) + + result_type = mlir_ir.RankedTensorType.get((), T.bool()) + init = tensor.splat(result_type, false(), []) + dims_attr = mlir_ir.DenseI64ArrayAttr.get([0]) + reduce_op = linalg.ReduceOp( + result=[result_type], + inputs=[tup_t], + inits=[init], + dimensions=dims_attr, + ) + block = reduce_op.combiner.blocks.append( + tup_t.type.element_type, result_type.element_type + ) + with mlir_ir.InsertionPoint(block): + body(reduce_op, *block.arguments) + combined = bool_of(tensor.extract(reduce_op.results[0], [])) + + bool_mlir_ty = builder.get_mlir_type(types.boolean) + undef_bool = llvm.UndefOp(bool_mlir_ty) + final_bool = arith.select(m_valid, combined, undef_bool) + target_type = builder.get_numba_type(target.name) + packed = _pack_masked(builder, target_type, final_bool, m_valid) + builder.store_var(target, packed) + + def _register() -> None: """Register the data model and lowerings with ``numba_cuda_mlir``. @@ -559,5 +649,12 @@ def _register() -> None: lower(float, MaskedType)(_make_lower_masked_numeric_cast()) lower(int, MaskedType)(_make_lower_masked_numeric_cast()) + lower(operator.contains, types.Tuple, MaskedType)( + _lower_masked_literal_tuple_contains + ) + lower(operator.contains, types.UniTuple, MaskedType)( + _lower_masked_unittuple_contains + ) + _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 837919884204..bf8626e54767 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -275,6 +275,32 @@ def generic(self, args, kws): return None +# ``value in (a, b, ...)`` membership -> Masked(boolean). ``value`` is a +# Masked scalar; the container is a literal Tuple of constants or a +# homogeneous UniTuple. (String membership -- substr in str -- arrives with +# the string value type in a later PR.) +def _is_masked_membership_item(ty): + return isinstance(ty, MaskedType) and not isinstance( + ty.value_type, types.Poison + ) + + +class MaskedSequenceContainsTemplate(AbstractTemplate): + def generic(self, args, kws): + if len(args) != 2 or kws: + return None + container, item = args + if not _is_masked_membership_item(item): + return None + if isinstance(container, types.Tuple) and all( + isinstance(x, types.Literal) for x in container.types + ): + return nb_signature(MaskedType(types.boolean), container, item) + if isinstance(container, types.UniTuple): + return nb_signature(MaskedType(types.boolean), container, item) + return None + + def _register() -> None: """Register typing for ``Masked`` and ``MaskedType`` attributes with ``numba_cuda_mlir``. Called once at module import. @@ -297,5 +323,9 @@ def _register() -> None: typing_registry.register_global(int)(MaskedScalarIntCast) typing_registry.register_global(abs)(MaskedScalarAbsoluteValue) + typing_registry.register_global(operator.contains)( + MaskedSequenceContainsTemplate + ) + _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 50739b5c28d2..3633bd264a04 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 @@ -794,3 +794,80 @@ def k(out_v, out_valid, a, av, b, bv): _launch(k, out_v, out_valid, _dt_in([1000]), _bool(True), _dt_in([400]), _bool(False)) assert bool(out_valid.get()[0]) is False + + +# --- membership: value in (...) ------------------------------------------- + + +@pytest.mark.parametrize("x,expected", [(1, True), (3, True), (5, True), (4, False), (0, False)]) +def test_masked_in_literal_tuple_int(x, expected): + """``Masked(int) in (1, 3, 5)`` (literal tuple) -> Masked(boolean).""" + + @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) + def k(out, a, av): + out[0] = (Masked(a[0], av[0]) in (1, 3, 5)).value + + out = cp.zeros(1, dtype=np.bool_) + _launch(k, out, cp.array([x], dtype=np.int64), cp.array([True], dtype=np.bool_)) + assert bool(out.get()[0]) is expected + + +@pytest.mark.parametrize("x,expected", [(1.5, True), (2.5, True), (9.9, False)]) +def test_masked_in_literal_tuple_float(x, expected): + """Literal-tuple membership works for float payloads too.""" + + @cuda.jit(types.void(types.boolean[::1], types.float64[::1], types.boolean[::1])) + def k(out, a, av): + out[0] = (Masked(a[0], av[0]) in (1.5, 2.5)).value + + out = cp.zeros(1, dtype=np.bool_) + _launch(k, out, cp.array([x], dtype=np.float64), cp.array([True], dtype=np.bool_)) + assert bool(out.get()[0]) is expected + + +@pytest.mark.parametrize("x,expected", [(2, True), (7, True), (9, True), (8, False)]) +def test_masked_in_unittuple_int(x, expected): + """``Masked(int) in (a, b, c)`` (runtime homogeneous UniTuple).""" + + @cuda.jit( + types.void( + types.boolean[::1], types.int64[::1], types.boolean[::1], + types.int64[::1], types.int64[::1], types.int64[::1], + ) + ) + def k(out, a, av, p, q, r): + out[0] = (Masked(a[0], av[0]) in (p[0], q[0], r[0])).value + + out = cp.zeros(1, dtype=np.bool_) + _launch( + k, out, + cp.array([x], dtype=np.int64), cp.array([True], dtype=np.bool_), + cp.array([2], dtype=np.int64), + cp.array([7], dtype=np.int64), + cp.array([9], dtype=np.int64), + ) + assert bool(out.get()[0]) is expected + + +def test_masked_in_tuple_invalid_propagates(): + """An invalid Masked operand makes the membership result invalid.""" + + @cuda.jit( + types.void( + types.boolean[::1], types.boolean[::1], + types.int64[::1], types.boolean[::1], + ) + ) + def k(out_v, out_valid, a, av): + m = Masked(a[0], av[0]) in (1, 3, 5) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.bool_) + out_valid = cp.ones(1, dtype=np.bool_) + _launch( + k, out_v, out_valid, + cp.array([3], dtype=np.int64), # would be a hit, but operand invalid + cp.array([False], dtype=np.bool_), + ) + assert bool(out_valid.get()[0]) is False From 3be7aecc8d767a97430fb7b7dfe190dc91efc814 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 07:16:39 -0700 Subject: [PATCH 12/14] Remove section dividers; stub docstrings as TODO --- .../cudf/core/udf/mlir_backend/masked_lowering.py | 1 - .../mlir_backend/test_masked_lowering.py | 11 ++++------- 2 files changed, 4 insertions(+), 8 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 063a0fdd4530..44ddfc4d6ff3 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -503,7 +503,6 @@ def _lower(builder, target, args, kwargs): return _lower -# --- membership: ``value in container`` -> Masked(boolean) ---------------- def _const_mlir_for_membership(py_const, mlir_ty): if isinstance(py_const, float): return float_of(py_const, mlir_ty) 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 3633bd264a04..d0a44cb2af75 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 @@ -796,12 +796,9 @@ def k(out_v, out_valid, a, av, b, bv): assert bool(out_valid.get()[0]) is False -# --- membership: value in (...) ------------------------------------------- - - @pytest.mark.parametrize("x,expected", [(1, True), (3, True), (5, True), (4, False), (0, False)]) def test_masked_in_literal_tuple_int(x, expected): - """``Masked(int) in (1, 3, 5)`` (literal tuple) -> Masked(boolean).""" + """TODO: write docstring.""" @cuda.jit(types.void(types.boolean[::1], types.int64[::1], types.boolean[::1])) def k(out, a, av): @@ -814,7 +811,7 @@ def k(out, a, av): @pytest.mark.parametrize("x,expected", [(1.5, True), (2.5, True), (9.9, False)]) def test_masked_in_literal_tuple_float(x, expected): - """Literal-tuple membership works for float payloads too.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.boolean[::1], types.float64[::1], types.boolean[::1])) def k(out, a, av): @@ -827,7 +824,7 @@ def k(out, a, av): @pytest.mark.parametrize("x,expected", [(2, True), (7, True), (9, True), (8, False)]) def test_masked_in_unittuple_int(x, expected): - """``Masked(int) in (a, b, c)`` (runtime homogeneous UniTuple).""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -850,7 +847,7 @@ def k(out, a, av, p, q, r): def test_masked_in_tuple_invalid_propagates(): - """An invalid Masked operand makes the membership result invalid.""" + """TODO: write docstring.""" @cuda.jit( types.void( From 924a332d913ffa37349ce92f2b03e84af8008fb7 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 15 Jun 2026 12:52:35 -0700 Subject: [PATCH 13/14] pack_return --- .../core/udf/mlir_backend/masked_lowering.py | 42 ++++++++++- .../core/udf/mlir_backend/masked_typing.py | 16 ++++- .../mlir_backend/test_masked_lowering.py | 69 +++++++++++++++++++ 3 files changed, 125 insertions(+), 2 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 44ddfc4d6ff3..29cde6e0b5f6 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -34,7 +34,7 @@ comparison_ops, unary_ops, ) -from cudf.core.udf.api import Masked +from cudf.core.udf.api import Masked, pack_return from cudf.core.udf.mlir_backend.masked_typing import ( MaskedType, NAType, @@ -586,6 +586,26 @@ def body(_op, element, accumulator): builder.store_var(target, packed) +# --- pack_return ---------------------------------------------------------- +# ``pack_return(masked)`` -> identity. +def _lower_pack_return_masked(builder, target, args, kwargs): + builder.store_var(target, builder.load_var(args[0])) + + +# ``pack_return(scalar)`` -> Masked(scalar, True). +def _lower_pack_return_scalar(builder, target, args, kwargs): + target_type = builder.get_numba_type(target.name) + value_mlir_ty = builder.get_mlir_type(target_type.value_type) + scalar_val = convert(builder.load_var(args[0]), value_mlir_ty) + valid_one = arith.constant( + result=builder.get_mlir_type(types.boolean), value=1 + ) + packed = _pack_masked( + builder, target_type, scalar_val, valid_one + ) + builder.store_var(target, packed) + + def _register() -> None: """Register the data model and lowerings with ``numba_cuda_mlir``. @@ -655,5 +675,25 @@ def _register() -> None: _lower_masked_unittuple_contains ) + lower(pack_return, MaskedType)(_lower_pack_return_masked) + # Register per concrete scalar shape so the dispatcher matches exactly + # rather than falling through ``types.Number`` (which would shadow + # Boolean). + for scalar_ty in ( + types.Integer, + types.int8, + types.int16, + types.int32, + types.int64, + types.uint8, + types.uint16, + types.uint32, + types.uint64, + types.float32, + types.float64, + types.boolean, + ): + lower(pack_return, scalar_ty)(_lower_pack_return_scalar) + _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 bf8626e54767..697714758181 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -25,7 +25,7 @@ comparison_ops, unary_ops, ) -from cudf.core.udf.api import Masked +from cudf.core.udf.api import Masked, pack_return # Datetime / timedelta resolutions cudf UDFs support. Mirrors the units # used by the column dtypes flowing into the kernels. @@ -301,6 +301,18 @@ def generic(self, args, kws): return None +# ``pack_return(x)`` -> Masked. Identity for a Masked input; wrap a bare +# numeric/boolean scalar with valid=True. Used by the apply-kernel templates +# to normalize a UDF's return value (which may be a Masked or a plain scalar). +class PackReturnTemplate(AbstractTemplate): + def generic(self, args, kws): + if isinstance(args[0], MaskedType): + return nb_signature(args[0], args[0]) + if isinstance(args[0], (types.Number, types.Boolean)): + return nb_signature(MaskedType(args[0]), args[0]) + return None + + def _register() -> None: """Register typing for ``Masked`` and ``MaskedType`` attributes with ``numba_cuda_mlir``. Called once at module import. @@ -327,5 +339,7 @@ def _register() -> None: MaskedSequenceContainsTemplate ) + typing_registry.register_global(pack_return)(PackReturnTemplate) + _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 d0a44cb2af75..159c4150f09b 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 @@ -868,3 +868,72 @@ def k(out_v, out_valid, a, av): cp.array([False], dtype=np.bool_), ) assert bool(out_valid.get()[0]) is False + + +# --- pack_return ----------------------------------------------------------- +# +# ``pack_return`` is the bridge the apply-kernel templates call on a UDF's +# return value. Calling it directly here exercises the two lowering paths in +# isolation, well before the kernel templates that use it exist in the stack. + + +from cudf.core.udf.api import pack_return # noqa: E402 + + +@pytest.mark.parametrize("valid", [True, False]) +def test_pack_return_masked_is_identity(valid): + """``pack_return(masked)`` passes the struct through unchanged.""" + + @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 = pack_return(Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, out_v, out_valid, + cp.array([42], dtype=np.int64), + cp.array([valid], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == 42 + assert bool(out_valid.get()[0]) is valid + + +@pytest.mark.parametrize("nb_ty,np_dtype,sample", _DTYPE_SAMPLES) +def test_pack_return_scalar_wraps_valid(nb_ty, np_dtype, sample): + """``pack_return(scalar)`` -> Masked(scalar, valid=True) for every dtype.""" + + @cuda.jit(types.void(nb_ty[::1], types.boolean[::1], nb_ty[::1])) + def k(out_v, out_valid, a): + m = pack_return(a[0]) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np_dtype) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch(k, out_v, out_valid, cp.array([sample], dtype=np_dtype)) + assert out_v.get()[0] == sample + assert bool(out_valid.get()[0]) is True + + +def test_pack_return_scalar_literal_constant(): + """A bare integer literal returned from the UDF wraps to valid Masked.""" + + @cuda.jit(types.void(types.int64[::1], types.boolean[::1])) + def k(out_v, out_valid): + m = pack_return(7) + out_v[0] = m.value + out_valid[0] = m.valid + + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch(k, out_v, out_valid) + assert int(out_v.get()[0]) == 7 + assert bool(out_valid.get()[0]) is True From bf3cd36ba12fd8d19442c67c1639253b2a39d37a Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 07:29:08 -0700 Subject: [PATCH 14/14] Remove section dividers; stub docstrings as TODO --- python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py | 1 - .../private_objects/mlir_backend/test_masked_lowering.py | 7 +++---- 2 files changed, 3 insertions(+), 5 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 29cde6e0b5f6..8d98eba2b38a 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -586,7 +586,6 @@ def body(_op, element, accumulator): builder.store_var(target, packed) -# --- pack_return ---------------------------------------------------------- # ``pack_return(masked)`` -> identity. def _lower_pack_return_masked(builder, target, args, kwargs): builder.store_var(target, builder.load_var(args[0])) 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 159c4150f09b..98b7301da7fa 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 @@ -870,7 +870,6 @@ def k(out_v, out_valid, a, av): assert bool(out_valid.get()[0]) is False -# --- pack_return ----------------------------------------------------------- # # ``pack_return`` is the bridge the apply-kernel templates call on a UDF's # return value. Calling it directly here exercises the two lowering paths in @@ -882,7 +881,7 @@ def k(out_v, out_valid, a, av): @pytest.mark.parametrize("valid", [True, False]) def test_pack_return_masked_is_identity(valid): - """``pack_return(masked)`` passes the struct through unchanged.""" + """TODO: write docstring.""" @cuda.jit( types.void( @@ -908,7 +907,7 @@ def k(out_v, out_valid, a, av): @pytest.mark.parametrize("nb_ty,np_dtype,sample", _DTYPE_SAMPLES) def test_pack_return_scalar_wraps_valid(nb_ty, np_dtype, sample): - """``pack_return(scalar)`` -> Masked(scalar, valid=True) for every dtype.""" + """TODO: write docstring.""" @cuda.jit(types.void(nb_ty[::1], types.boolean[::1], nb_ty[::1])) def k(out_v, out_valid, a): @@ -924,7 +923,7 @@ def k(out_v, out_valid, a): def test_pack_return_scalar_literal_constant(): - """A bare integer literal returned from the UDF wraps to valid Masked.""" + """TODO: write docstring.""" @cuda.jit(types.void(types.int64[::1], types.boolean[::1])) def k(out_v, out_valid):