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 21cd3b171b1c..4bdb9493a26e 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -14,11 +14,19 @@ from numba_cuda_mlir.lowering_utilities import ( 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 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,171 @@ def _lower_masked_binary_null( builder.store_var(target, packed) +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}" + temp = numba_ir.Var(scope=scope, name=name, loc=loc) + builder.fndesc.typemap[temp.name] = numba_type + return temp + + +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 + + 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 + + +def _lower_masked_invert( + builder: MLIRLower, target: Var, args: list[Var], kwargs: list +) -> 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< 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``. + + Compare against zero directly instead. For floats the unordered ``UNE`` + predicate also gives the Python-correct ``bool(nan) is True``. + """ + 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(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: + 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: + """``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) + 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 +527,21 @@ 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: + # 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)) + lower(operator.invert, MaskedType)(_lower_masked_invert) + + 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()) + _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..501adcbfd0bf 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,75 @@ def generic( return None +class MaskedScalarUnaryOp(AbstractTemplate): + """``(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 + ).return_type + return nb_signature(MaskedType(return_type), args[0]) + return None + + +class MaskedScalarTruth(AbstractTemplate): + """Masked -> plain boolean for the truth-testing unary ops. + + 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( + 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 + + +class MaskedScalarFloatCast(AbstractTemplate): + """``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): + """``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 + + +class MaskedScalarAbsoluteValue(AbstractTemplate): + """``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 + ).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 +301,17 @@ def _register() -> None: typing_registry.register_global(binary_op)(MaskedScalarNullOp) 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) + _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..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 @@ -3,6 +3,7 @@ from __future__ import annotations +import math import operator import cupy as cp @@ -461,3 +462,326 @@ 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 + + +@pytest.mark.parametrize("valid", [True, False]) +@pytest.mark.parametrize("op", [operator.neg, operator.pos]) +def test_masked_unary_sign(op, valid): + """Sign unary ops on a Masked carry the 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])) + 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([valid], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == op(7) + # validity carried from the operand + 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, valid): + """``~m`` on a Masked integer inverts the payload bits.""" + + @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 = ~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([valid], dtype=np.bool_), + ) + 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", + [ + math.sin, + math.cos, + math.sqrt, + math.exp, + ], +) +def test_masked_unary_math(fn, valid): + """math.* unary functions lower on a Masked's payload.""" + + @cuda.jit( + types.void( + types.float64[::1], + types.boolean[::1], + types.float64[::1], + types.boolean[::1], + ) + ) + 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_v = cp.zeros(1, dtype=np.float64) + out_valid = cp.zeros(1, dtype=np.bool_) + _launch( + k, + out_v, + out_valid, + cp.array([1.5], dtype=np.float64), + cp.array([valid], dtype=np.bool_), + ) + 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, valid): + """``abs(m)`` returns a Masked with the absolute value of the payload.""" + + @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([valid], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == abs(x) + assert bool(out_valid.get()[0]) is valid + + +@pytest.mark.parametrize( + "value,valid,expected", + [ + (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), + ], +) +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 + + +@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", + [ + (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.""" + + @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 + + +@pytest.mark.parametrize("valid", [True, False]) +def test_masked_float_cast(valid): + """``float(m)`` casts the payload to float64, preserving validity.""" + + @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([valid], dtype=np.bool_), + ) + assert float(out_v.get()[0]) == 3.0 + assert bool(out_valid.get()[0]) is valid + + +@pytest.mark.parametrize("valid", [True, False]) +def test_masked_int_cast(valid): + """``int(m)`` casts the payload to int64, preserving validity.""" + + @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([valid], dtype=np.bool_), + ) + assert int(out_v.get()[0]) == 3 + assert bool(out_valid.get()[0]) is valid