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..8d98eba2b38a 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py @@ -6,14 +6,35 @@ 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._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 convert +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 cudf.core.udf.api import Masked +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 + +from cudf.core.udf._ops import ( + arith_ops, + bitwise_ops, + comparison_ops, + unary_ops, +) +from cudf.core.udf.api import Masked, pack_return from cudf.core.udf.mlir_backend.masked_typing import ( MaskedType, NAType, @@ -21,6 +42,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 ( @@ -183,6 +206,405 @@ 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, +): + """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) + 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, + target_type: MaskedType, + v1: mlir_ir.Value, + 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 + # 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) + + +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]) + 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) + 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, + inner_ty1=ty1, inner_ty2=ty2, ref_var=args[0], + ) + + return _lower + + +def _scalar_value_from_var( + builder: MLIRLower, + s_var: Var, +) -> mlir_ir.Value: + """Resolve the scalar operand for the Masked-vs-scalar path. + + 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): + 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) + return builder.load_var(s_var) + + +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]) + ) + 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_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, + 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, + inner_ty1=s_inner_ty, inner_ty2=m_inner_ty, ref_var=m_var, + ) + + return _lower + + +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) + 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 _make_temp_var(builder, base_var, name_suffix, numba_type): + """TODO: write docstring.""" + 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 + + +# 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 _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) + + +# ``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``. @@ -213,5 +635,64 @@ 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) + + 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()) + + lower(operator.contains, types.Tuple, MaskedType)( + _lower_masked_literal_tuple_contains + ) + lower(operator.contains, types.UniTuple, MaskedType)( + _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 7a8beec820ca..697714758181 100644 --- a/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py +++ b/python/cudf/cudf/core/udf/mlir_backend/masked_typing.py @@ -14,20 +14,37 @@ AbstractTemplate, AttributeTemplate, ConcreteTemplate, + Signature, ) from numba_cuda_mlir.typing import signature as nb_signature from cudf.core.missing import NA -from cudf.core.udf.api import Masked +from cudf.core.udf._ops import ( + arith_ops, + bitwise_ops, + comparison_ops, + unary_ops, +) +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. +_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} ) @@ -139,6 +156,163 @@ def generic(self, args, kws): return None +class MaskedScalarArithOp(AbstractTemplate): + """``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): + 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 + + +class MaskedScalarScalarOp(AbstractTemplate): + """``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) + ): + 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 + + +class MaskedScalarNullOp(AbstractTemplate): + """``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): + return nb_signature(args[1], na_type, args[1]) + return None + + +# 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 + + +# ``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 + + +# ``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. @@ -148,5 +322,24 @@ 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) + + 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) + + typing_registry.register_global(operator.contains)( + 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 3be1a3bf6026..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 @@ -3,6 +3,9 @@ from __future__ import annotations +import math +import operator + import cupy as cp import numpy as np import pytest @@ -13,6 +16,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 +235,704 @@ def k(out, v, valid): result = out.get() assert bool(result[0]) is valid_in assert bool(result[1]) is valid_in + + +_ARITH = [operator.add, operator.sub, operator.mul] + + +@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( + 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]) == op(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, + operator.le, + operator.gt, + operator.ge, + operator.eq, + operator.ne, +] + + +@pytest.mark.parametrize("op", _CMP) +@pytest.mark.parametrize("a,b", [(3, 5), (5, 5), (8, 5)]) +def test_masked_masked_comparison(op, 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]) == op(a, b) + + +@pytest.mark.parametrize("op", _ARITH) +def test_masked_scalar_arith(op): + """``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]) == op(a, 4) + # validity is carried from the (invalid) Masked operand + assert bool(out_valid.get()[0]) is False + + +@pytest.mark.parametrize("op", _ARITH) +def test_scalar_masked_arith(op): + """``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]) == op(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 + + +@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.""" + + @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): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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 + + +# +# 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(): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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 + + +@pytest.mark.parametrize("x,expected", [(1, True), (3, True), (5, True), (4, False), (0, False)]) +def test_masked_in_literal_tuple_int(x, expected): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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 + + +# +# ``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): + """TODO: write docstring.""" + + @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): + """TODO: write docstring.""" + + @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(): + """TODO: write docstring.""" + + @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 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..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 @@ -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): + """TODO: write docstring.""" + 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