From 5589be837e9e0ad06dd9f8edec83ba057d8a5582 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 11 Jun 2026 06:03:04 -0700 Subject: [PATCH 1/6] [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 26b2d5190e6a58179b2304494b5c9083bf785e65 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 29 Jun 2026 07:04:05 -0700 Subject: [PATCH 2/6] 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 74cbb400b5750c50a29c630f714a1e35a338337f Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 31 Aug 2026 12:35:14 -0700 Subject: [PATCH 3/6] Fold prior-stack review feedback into unary ops Apply recurring feedback from PRs #22884/#22885/#22886 to the unary-op code: add type annotations to the new typing templates and lowering functions, convert leading comments into class/function docstrings (and fill the TODO docstring stubs), and drop the redundant ``ref`` parameter from the unary tests so assertions compute the expected value via the op itself. Annotate the shared op lists in _ops.py so the newly annotated factory signatures type-check. --- python/cudf/cudf/core/udf/_ops.py | 16 ++-- .../core/udf/mlir_backend/masked_lowering.py | 85 +++++++++++-------- .../core/udf/mlir_backend/masked_typing.py | 42 ++++++--- .../mlir_backend/test_masked_lowering.py | 76 ++++++++++------- 4 files changed, 139 insertions(+), 80 deletions(-) diff --git a/python/cudf/cudf/core/udf/_ops.py b/python/cudf/cudf/core/udf/_ops.py index 4a1b400ce406..d00423602d05 100644 --- a/python/cudf/cudf/core/udf/_ops.py +++ b/python/cudf/cudf/core/udf/_ops.py @@ -1,10 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021-2022, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import math import operator +from collections.abc import Callable +from typing import Any -arith_ops = [ +arith_ops: list[Callable[..., Any]] = [ operator.add, operator.sub, operator.mul, @@ -21,9 +23,13 @@ operator.imod, ] -bitwise_ops = [operator.and_, operator.or_, operator.xor] +bitwise_ops: list[Callable[..., Any]] = [ + operator.and_, + operator.or_, + operator.xor, +] -unary_ops = [ +unary_ops: list[Callable[..., Any]] = [ math.acos, math.acosh, math.asin, @@ -57,7 +63,7 @@ operator.invert, ] -comparison_ops = [ +comparison_ops: list[Callable[..., Any]] = [ operator.eq, operator.ne, operator.lt, 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..8b6a485dfba3 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -313,8 +313,19 @@ def _lower_masked_binary_null( builder.store_var(target, packed) -def _make_temp_var(builder, base_var, name_suffix, numba_type): - """TODO: write docstring.""" +def _make_temp_var( + builder: MLIRLower, + base_var: Var, + name_suffix: str, + numba_type: types.Type, +) -> Var: + """Create a fresh typed IR ``Var`` for staging an intermediate value. + + Used to feed the masked payload into a registered scalar lowering (which + operates on plain IR vars) and to receive its result. The name is derived + from ``base_var`` plus ``name_suffix`` so typemap keys stay unique when one + operand feeds several unary calls in a single expression. + """ scope = getattr(base_var, "scope", None) loc = getattr(base_var, "loc", None) name = f"$masked_uop_{base_var.name}_{name_suffix}" @@ -323,22 +334,23 @@ def _make_temp_var(builder, base_var, name_suffix, numba_type): return temp -# 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): +def _make_lower_masked_unary(op: Callable) -> Callable: + """``(Masked)``: 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 bit. + """ + + def _lower( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list + ) -> None: 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 + 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, 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) @@ -365,19 +377,21 @@ def _lower(builder, target, args, kwargs): builder.load_var(out_var), builder.get_mlir_type(result_inner_ty), ) - packed = _pack_masked( - builder, target_type, result_val, m_valid - ) + 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< None: + """``operator.invert`` (bitwise ~) on Masked integers. + + There is no scalar ``@lower`` for invert, so compute ``xori(x, -1)``. The + all-ones mask is the signed constant -1 (two's complement); ``(1< None: + """``bool(m)`` / ``operator.truth(m)``: ``m.valid and bool(m.value)``.""" 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]) @@ -409,22 +425,21 @@ def _lower_masked_truth(builder, target, args, kwargs): builder.store_var(target, result) -# int(m) -> Masked(int64); float(m) -> Masked(float64). -def _make_lower_masked_numeric_cast(): - def _lower(builder, target, args, kwargs): +def _make_lower_masked_numeric_cast() -> Callable: + """``int(m)`` -> ``Masked(int64)`` / ``float(m)`` -> ``Masked(float64)``: + cast the payload to the target value type, preserving the validity bit. + """ + + def _lower( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list + ) -> None: target_type = builder.get_numba_type(target.name) - target_value_mlir_ty = builder.get_mlir_type( - target_type.value_type - ) + 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] - ) + 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 - ) + packed = _pack_masked(builder, target_type, casted, m_valid) builder.store_var(target, packed) 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..205c5e4f693d 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -218,9 +218,14 @@ def generic( return None -# Resolve the underlying scalar op on the value type, wrap the result. class MaskedScalarUnaryOp(AbstractTemplate): - def generic(self, args, kws): + """``(Masked)``: resolve the underlying scalar op on the value type + and wrap the result back in a ``MaskedType``. + """ + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: if len(args) == 1 and isinstance(args[0], MaskedType): return_type = self.context.resolve_function_type( self.key, (args[0].value_type,), kws @@ -229,34 +234,49 @@ def generic(self, args, kws): 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): + """``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). + """ + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: 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): + """``float(m)`` -> ``Masked(float64)``.""" + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: 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): + """``int(m)`` -> ``Masked(int64)``.""" + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: 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): + """``abs(m)`` -> ``Masked(result)``.""" + + def generic( + self, args: tuple[types.Type, ...], kws: dict + ) -> Signature | None: if len(args) == 1 and isinstance(args[0], MaskedType): return_type = self.context.resolve_function_type( self.key, (args[0].value_type,), 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 a270e174cb47..7a9f95f52e53 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,11 +464,9 @@ def k(out_valid, a, av): assert bool(out_valid.get()[0]) is False -@pytest.mark.parametrize( - "op,ref", [(operator.neg, lambda x: -x), (operator.pos, lambda x: +x)] -) -def test_masked_unary_sign(op, ref): - """TODO: write docstring.""" +@pytest.mark.parametrize("op", [operator.neg, operator.pos]) +def test_masked_unary_sign(op): + """Sign unary ops on a Masked carry the operand's validity.""" @cuda.jit( types.void( @@ -492,50 +490,59 @@ def k(out_v, out_valid, a, av): cp.array([7], dtype=np.int64), cp.array([False], dtype=np.bool_), ) - assert int(out_v.get()[0]) == ref(7) + assert int(out_v.get()[0]) == op(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): - """TODO: write docstring.""" + """``~m`` on a Masked integer inverts the payload bits.""" - @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) + @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_)) + _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", + "fn", [ - (math.sin, math.sin), - (math.cos, math.cos), - (math.sqrt, math.sqrt), - (math.exp, math.exp), + math.sin, + math.cos, + math.sqrt, + math.exp, ], ) -def test_masked_unary_math(fn, ref): - """TODO: write docstring.""" +def test_masked_unary_math(fn): + """math.* unary functions lower on a Masked's payload.""" - @cuda.jit(types.void(types.float64[::1], types.float64[::1], types.boolean[::1])) + @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_) + 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) + np.testing.assert_allclose(float(out.get()[0]), fn(1.5), rtol=1e-12) @pytest.mark.parametrize("x", [-9, 0, 12]) def test_masked_abs(x): - """TODO: write docstring.""" + """``abs(m)`` returns a Masked with the absolute value of the payload.""" @cuda.jit( types.void( @@ -566,16 +573,18 @@ def k(out_v, out_valid, a, av): @pytest.mark.parametrize( "value,valid,expected", [ - (5, True, True), # valid & truthy + (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): - """TODO: write docstring.""" + """``bool(m)`` is ``m.valid and bool(m.value)``.""" - @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): out[0] = bool(Masked(a[0], av[0])) @@ -590,9 +599,11 @@ def k(out, a, av): def test_masked_bool_in_if_condition(): - """TODO: write docstring.""" + """A Masked used directly in an ``if`` uses its truth value.""" - @cuda.jit(types.void(types.int64[::1], types.int64[::1], types.boolean[::1])) + @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: @@ -601,14 +612,21 @@ def k(out, a, av): 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_)) + _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_)) + _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(): - """TODO: write docstring.""" + """``float(m)`` casts the payload to float64, preserving validity.""" @cuda.jit( types.void( @@ -637,7 +655,7 @@ def k(out_v, out_valid, a, av): def test_masked_int_cast(): - """TODO: write docstring.""" + """``int(m)`` casts the payload to int64, preserving validity.""" @cuda.jit( types.void( From cad42fd9f38aea9a35f04a88955b1b0c1031ce91 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Mon, 31 Aug 2026 12:52:48 -0700 Subject: [PATCH 4/6] Fix bool(Masked) for float payloads convert(float -> i1) lowers via arith.fptoui, which truncates instead of testing truthiness, so bool(Masked(1.0)) came out False. Take a dedicated float path computing (payload != 0) via arith.cmpf UNE, which also yields the Python-correct bool(nan) is True. Add float-payload truth tests. --- .../core/udf/mlir_backend/masked_lowering.py | 18 ++++++++++-- .../mlir_backend/test_masked_lowering.py | 29 +++++++++++++++++++ 2 files changed, 44 insertions(+), 3 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 8b6a485dfba3..b8e2e8e27903 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -415,12 +415,24 @@ def _lower_masked_invert( def _lower_masked_truth( builder: MLIRLower, target: Var, args: list[Var], kwargs: list ) -> None: - """``bool(m)`` / ``operator.truth(m)``: ``m.valid and bool(m.value)``.""" + """``bool(m)`` / ``operator.truth(m)``: ``m.valid and bool(m.value)``. + + Float payloads take a dedicated path: ``convert(float -> i1)`` lowers via + ``arith.fptoui`` to a 1-bit integer, which truncates rather than testing + truthiness (e.g. ``bool(1.0)`` would come out ``False``). Compute + ``payload != 0`` directly instead; the unordered ``UNE`` predicate also + gives the Python-correct ``bool(nan) is True``. + """ 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]) - bool_mlir_ty = builder.get_mlir_type(types.boolean) - payload_as_bool = bool_of(convert(m_val, bool_mlir_ty)) + inner_ty = builder.get_numba_type(args[0].name).value_type + if isinstance(inner_ty, types.Float): + zero = arith.constant(m_val.type, 0.0) + payload_as_bool = arith.cmpf(arith.CmpFPredicate.UNE, m_val, zero) + else: + bool_mlir_ty = builder.get_mlir_type(types.boolean) + payload_as_bool = bool_of(convert(m_val, bool_mlir_ty)) result = arith.select(m_valid, payload_as_bool, false()) builder.store_var(target, result) 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 7a9f95f52e53..c3c3027ed4d2 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 @@ -598,6 +598,35 @@ def k(out, a, av): assert bool(out.get()[0]) is expected +@pytest.mark.parametrize( + "value,valid,expected", + [ + (1.0, True, True), # valid & truthy + (-2.5, True, True), # valid & truthy (negative) + (0.0, True, False), # valid & falsy + (float("nan"), True, True), # nan is truthy, like Python + (1.0, False, False), # invalid -> False regardless of payload + ], +) +def test_masked_bool_truth_float(value, valid, expected): + """``bool(m)`` on a float payload tests ``payload != 0`` (not fptoui).""" + + @cuda.jit( + types.void(types.boolean[::1], types.float64[::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.float64), + cp.array([valid], dtype=np.bool_), + ) + assert bool(out.get()[0]) is expected + + def test_masked_bool_in_if_condition(): """A Masked used directly in an ``if`` uses its truth value.""" From c9cab650a8fbb00ec8e38cf0b035df1637f76c87 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 3 Sep 2026 05:51:21 -0700 Subject: [PATCH 5/6] Parametrize unary op tests over operand validity Per review feedback, sweep valid=True/False on the masked unary tests (neg/pos, invert, math, abs, float/int cast) and assert the result validity tracks the operand, rather than hard-coding a single mask value. Extend the invert and math kernels to emit the validity bit so it can be checked. --- .../mlir_backend/test_masked_lowering.py | 83 +++++++++++++------ 1 file changed, 56 insertions(+), 27 deletions(-) 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 c3c3027ed4d2..3d5386d548ca 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,8 +464,9 @@ def k(out_valid, a, av): assert bool(out_valid.get()[0]) is False +@pytest.mark.parametrize("valid", [True, False]) @pytest.mark.parametrize("op", [operator.neg, operator.pos]) -def test_masked_unary_sign(op): +def test_masked_unary_sign(op, valid): """Sign unary ops on a Masked carry the operand's validity.""" @cuda.jit( @@ -488,30 +489,45 @@ def k(out_v, out_valid, a, av): out_v, out_valid, cp.array([7], dtype=np.int64), - cp.array([False], dtype=np.bool_), + cp.array([valid], dtype=np.bool_), ) assert int(out_v.get()[0]) == op(7) # validity carried from the operand - assert bool(out_valid.get()[0]) is False + assert bool(out_valid.get()[0]) is valid +@pytest.mark.parametrize("valid", [True, False]) @pytest.mark.parametrize("x", [5, 0, -6, 255]) -def test_masked_invert(x): +def test_masked_invert(x, valid): """``~m`` on a Masked integer inverts the payload bits.""" @cuda.jit( - types.void(types.int64[::1], types.int64[::1], types.boolean[::1]) + types.void( + types.int64[::1], + types.boolean[::1], + types.int64[::1], + types.boolean[::1], + ) ) - def k(out, a, av): - out[0] = (~Masked(a[0], av[0])).value + def k(out_v, out_valid, a, av): + m = ~Masked(a[0], av[0]) + out_v[0] = m.value + out_valid[0] = m.valid - out = cp.zeros(1, dtype=np.int64) + out_v = cp.zeros(1, dtype=np.int64) + out_valid = cp.zeros(1, dtype=np.bool_) _launch( - k, out, cp.array([x], dtype=np.int64), cp.array([True], dtype=np.bool_) + k, + out_v, + out_valid, + cp.array([x], dtype=np.int64), + cp.array([valid], dtype=np.bool_), ) - assert int(out.get()[0]) == ~x + assert int(out_v.get()[0]) == ~x + assert bool(out_valid.get()[0]) is valid +@pytest.mark.parametrize("valid", [True, False]) @pytest.mark.parametrize( "fn", [ @@ -521,27 +537,38 @@ def k(out, a, av): math.exp, ], ) -def test_masked_unary_math(fn): +def test_masked_unary_math(fn, valid): """math.* unary functions lower on a Masked's payload.""" @cuda.jit( - types.void(types.float64[::1], types.float64[::1], types.boolean[::1]) + types.void( + types.float64[::1], + types.boolean[::1], + types.float64[::1], + types.boolean[::1], + ) ) - def k(out, a, av): - out[0] = fn(Masked(a[0], av[0])).value + def k(out_v, out_valid, a, av): + m = fn(Masked(a[0], av[0])) + out_v[0] = m.value + out_valid[0] = m.valid - out = cp.zeros(1, dtype=np.float64) + out_v = cp.zeros(1, dtype=np.float64) + out_valid = cp.zeros(1, dtype=np.bool_) _launch( k, - out, + out_v, + out_valid, cp.array([1.5], dtype=np.float64), - cp.array([True], dtype=np.bool_), + cp.array([valid], dtype=np.bool_), ) - np.testing.assert_allclose(float(out.get()[0]), fn(1.5), rtol=1e-12) + np.testing.assert_allclose(float(out_v.get()[0]), fn(1.5), rtol=1e-12) + assert bool(out_valid.get()[0]) is valid +@pytest.mark.parametrize("valid", [True, False]) @pytest.mark.parametrize("x", [-9, 0, 12]) -def test_masked_abs(x): +def test_masked_abs(x, valid): """``abs(m)`` returns a Masked with the absolute value of the payload.""" @cuda.jit( @@ -564,10 +591,10 @@ def k(out_v, out_valid, a, av): out_v, out_valid, cp.array([x], dtype=np.int64), - cp.array([True], dtype=np.bool_), + cp.array([valid], dtype=np.bool_), ) assert int(out_v.get()[0]) == abs(x) - assert bool(out_valid.get()[0]) is True + assert bool(out_valid.get()[0]) is valid @pytest.mark.parametrize( @@ -654,7 +681,8 @@ def k(out, a, av): assert int(out.get()[0]) == 0 -def test_masked_float_cast(): +@pytest.mark.parametrize("valid", [True, False]) +def test_masked_float_cast(valid): """``float(m)`` casts the payload to float64, preserving validity.""" @cuda.jit( @@ -677,13 +705,14 @@ def k(out_v, out_valid, a, av): out_v, out_valid, cp.array([3], dtype=np.int64), - cp.array([True], dtype=np.bool_), + cp.array([valid], dtype=np.bool_), ) assert float(out_v.get()[0]) == 3.0 - assert bool(out_valid.get()[0]) is True + assert bool(out_valid.get()[0]) is valid -def test_masked_int_cast(): +@pytest.mark.parametrize("valid", [True, False]) +def test_masked_int_cast(valid): """``int(m)`` casts the payload to int64, preserving validity.""" @cuda.jit( @@ -706,7 +735,7 @@ def k(out_v, out_valid, a, av): out_v, out_valid, cp.array([3.9], dtype=np.float64), - cp.array([False], dtype=np.bool_), + cp.array([valid], dtype=np.bool_), ) assert int(out_v.get()[0]) == 3 - assert bool(out_valid.get()[0]) is False + assert bool(out_valid.get()[0]) is valid From 6090a263276ecadc25ca246cbff398f60ab2dbb0 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Thu, 3 Sep 2026 06:17:07 -0700 Subject: [PATCH 6/6] Fix integer truth narrowing and not_ on Masked Two correctness issues from review: - bool(Masked(int)) went through convert(payload -> i1), which is arith.trunci and keeps only the low bit, so even values like Masked(2) tested as False. Compare against zero with arith.cmpi ne at full width instead (mirroring the existing float path). Factor the shared truth logic into _masked_truth_value. - operator.not_ was in the generic unary loops, producing a Masked(bool) that preserved validity; but "not m" must be plain logical negation, so an invalid (falsy) operand yields True, not an invalid masked value. Exclude not_ from the generic typing/lowering loops, type it as a plain boolean, and lower it as not (m.valid and bool(m.value)). Add even-value truth cases and valid/invalid regressions for both "not m" and operator.not_. --- .../core/udf/mlir_backend/masked_lowering.py | 58 +++++++++++++------ .../core/udf/mlir_backend/masked_typing.py | 12 +++- .../mlir_backend/test_masked_lowering.py | 48 ++++++++++++++- 3 files changed, 97 insertions(+), 21 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 b8e2e8e27903..4bdb9493a26e 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -12,10 +12,10 @@ 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, + true, ) from numba_cuda_mlir.models import PrimitiveModel, register_model from numba_cuda_mlir.numba_cuda.core import ir as numba_ir @@ -412,29 +412,50 @@ def _lower_masked_invert( builder.store_var(target, packed) -def _lower_masked_truth( - builder: MLIRLower, target: Var, args: list[Var], kwargs: list -) -> None: - """``bool(m)`` / ``operator.truth(m)``: ``m.valid and bool(m.value)``. +def _masked_truth_value(builder: MLIRLower, arg: Var) -> mlir_ir.Value: + """Return the ``i1`` truth value of a Masked: ``m.valid and bool(m.value)``. + + The payload truthiness cannot go through ``convert(payload -> i1)`` because + that narrows numerically rather than testing truth: + + * ``convert(int -> i1)`` is ``arith.trunci``, keeping only the low bit, so + ``bool(Masked(2))`` would come out ``False``. + * ``convert(float -> i1)`` is ``arith.fptoui``, which truncates similarly, + so ``bool(Masked(1.0))`` would come out ``False``. - Float payloads take a dedicated path: ``convert(float -> i1)`` lowers via - ``arith.fptoui`` to a 1-bit integer, which truncates rather than testing - truthiness (e.g. ``bool(1.0)`` would come out ``False``). Compute - ``payload != 0`` directly instead; the unordered ``UNE`` predicate also - gives the Python-correct ``bool(nan) is True``. + Compare against zero directly instead. For floats the unordered ``UNE`` + predicate also gives the Python-correct ``bool(nan) is True``. """ - m = builder.load_var(args[0]) + m = builder.load_var(arg) st = llvm.StructType(m.type) m_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1]) - inner_ty = builder.get_numba_type(args[0].name).value_type + inner_ty = builder.get_numba_type(arg.name).value_type if isinstance(inner_ty, types.Float): zero = arith.constant(m_val.type, 0.0) payload_as_bool = arith.cmpf(arith.CmpFPredicate.UNE, m_val, zero) else: - bool_mlir_ty = builder.get_mlir_type(types.boolean) - payload_as_bool = bool_of(convert(m_val, bool_mlir_ty)) - result = arith.select(m_valid, payload_as_bool, false()) - builder.store_var(target, result) + zero = arith.constant(m_val.type, 0) + payload_as_bool = arith.cmpi(arith.CmpIPredicate.ne, m_val, zero) + return arith.select(m_valid, payload_as_bool, false()) + + +def _lower_masked_truth( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list +) -> None: + """``bool(m)`` / ``operator.truth(m)``: ``m.valid and bool(m.value)``.""" + builder.store_var(target, _masked_truth_value(builder, args[0])) + + +def _lower_masked_not( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list +) -> None: + """``not m`` / ``operator.not_(m)``: ``not (m.valid and bool(m.value))``. + + The result is a plain boolean, not a Masked: an invalid operand is falsy, + so ``not`` of it is ``True`` rather than an invalid masked value. + """ + truth = _masked_truth_value(builder, args[0]) + builder.store_var(target, arith.xori(truth, true())) def _make_lower_masked_numeric_cast() -> Callable: @@ -507,7 +528,9 @@ def _register() -> None: lower(binary_op, NAType, MaskedType)(_lower_masked_binary_null) for unary_op in unary_ops: - if unary_op is operator.invert: + # invert has no scalar lowering to delegate to (handled below), and + # not_ is logical negation returning a plain bool rather than a Masked. + if unary_op in (operator.invert, operator.not_): continue lower(unary_op, MaskedType)(_make_lower_masked_unary(unary_op)) lower(abs, MaskedType)(_make_lower_masked_unary(abs)) @@ -515,6 +538,7 @@ def _register() -> None: lower(operator.truth, MaskedType)(_lower_masked_truth) lower(bool, MaskedType)(_lower_masked_truth) + lower(operator.not_, MaskedType)(_lower_masked_not) lower(float, MaskedType)(_make_lower_masked_numeric_cast()) lower(int, MaskedType)(_make_lower_masked_numeric_cast()) 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 205c5e4f693d..501adcbfd0bf 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -235,10 +235,12 @@ def generic( class MaskedScalarTruth(AbstractTemplate): - """``bool(m)`` / ``operator.truth(m)`` -> boolean. + """Masked -> plain boolean for the truth-testing unary ops. - The runtime result is ``m.valid and bool(m.value)``; the *type* is a - plain boolean (used directly in ``if`` conditions). + Covers ``bool(m)`` / ``operator.truth(m)`` and ``not m`` / + ``operator.not_(m)``. The *type* is always a plain boolean (used directly + in ``if`` conditions); the runtime result differs per op and is handled in + lowering. """ def generic( @@ -300,9 +302,13 @@ def _register() -> None: typing_registry.register_global(binary_op)(MaskedScalarScalarOp) for unary_op in unary_ops: + # not_ is logical negation: it returns a plain boolean, not a Masked. + if unary_op is operator.not_: + continue typing_registry.register_global(unary_op)(MaskedScalarUnaryOp) typing_registry.register_global(operator.truth)(MaskedScalarTruth) typing_registry.register_global(bool)(MaskedScalarTruth) + typing_registry.register_global(operator.not_)(MaskedScalarTruth) typing_registry.register_global(float)(MaskedScalarFloatCast) typing_registry.register_global(int)(MaskedScalarIntCast) typing_registry.register_global(abs)(MaskedScalarAbsoluteValue) 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 3d5386d548ca..430909767eae 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 @@ -600,9 +600,12 @@ def k(out_v, out_valid, a, av): @pytest.mark.parametrize( "value,valid,expected", [ - (5, True, True), # valid & truthy + (5, True, True), # valid & odd truthy + (2, True, True), # even; trunc->i1 would drop low bit + (-2, True, True), # valid & negative even truthy (0, True, False), # valid & falsy (5, False, False), # invalid -> False regardless of payload + (2, False, False), (0, False, False), ], ) @@ -625,6 +628,49 @@ def k(out, a, av): assert bool(out.get()[0]) is expected +@pytest.mark.parametrize("use_operator", [False, True]) +@pytest.mark.parametrize( + "value,valid,expected", + [ + (5, True, False), # valid & truthy -> not True + (0, True, True), # valid & falsy -> not False + (5, False, True), # invalid is falsy -> not False + (0, False, True), + ], +) +def test_masked_not(value, valid, expected, use_operator): + """``not m`` / ``operator.not_`` -> plain bool; invalid Masked is falsy.""" + + if use_operator: + + @cuda.jit( + types.void( + types.boolean[::1], types.int64[::1], types.boolean[::1] + ) + ) + def k(out, a, av): + out[0] = operator.not_(Masked(a[0], av[0])) + + else: + + @cuda.jit( + types.void( + types.boolean[::1], types.int64[::1], types.boolean[::1] + ) + ) + def k(out, a, av): + out[0] = not 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 + + @pytest.mark.parametrize( "value,valid,expected", [