[MLIR] MaskedType: NA handling (#22885)

Part of the MLIR UDF backend stack. Depends on #22766 (plumbing). Adds
NA handling for MaskedType: NAType + cudf.NA typeof, MaskedType / NAType
unify, `m is NA` / `m is not NA`, and the NA / scalar / Masked->Masked
casts used for branch unification.

Unification covers expressions in the target code like the following
where `x` is `MaskedType`.
```
if x is cudf.NA:
    return x
else:
    return 42
```

42 is a scalar here so the unification code kicks in and types the
overall expression as returning MaskedType, and upcasts the 42 to
Masked(42, True)

Author: vyasr

python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py

@@ -2,17 +2,23 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import operator
+from functools import partial
 from typing import TYPE_CHECKING
 
 from numba_cuda_mlir import types
 from numba_cuda_mlir._mlir import ir as mlir_ir
-from numba_cuda_mlir._mlir.dialects import llvm
-from numba_cuda_mlir.extending import lowering_registry
+from numba_cuda_mlir._mlir.dialects import arith, llvm
+from numba_cuda_mlir.extending import lower_cast, lowering_registry
 from numba_cuda_mlir.lowering_utilities import convert
 from numba_cuda_mlir.models import PrimitiveModel, register_model
 
 from cudf.core.udf.api import Masked
-from cudf.core.udf.mlir_backend.masked_typing import MaskedType
+from cudf.core.udf.mlir_backend.masked_typing import (
+    MaskedType,
+    NAType,
+    na_type,
+)
 
 if TYPE_CHECKING:
     from numba_cuda_mlir.mlir_lowering import MLIRLower
@@ -62,6 +68,21 @@ def __init__(
         super().__init__(data_model_manager, masked_type, struct_type)
 
 
+def _extract_masked_value_valid(struct_val, value_mlir_ty, valid_ty):
+    """Pull the ``(value, valid)`` SSA values out of a ``Masked`` struct."""
+    v = llvm.extractvalue(
+        res=value_mlir_ty,
+        container=struct_val,
+        position=mlir_ir.DenseI64ArrayAttr.get([0]),
+    )
+    valid = llvm.extractvalue(
+        res=valid_ty,
+        container=struct_val,
+        position=mlir_ir.DenseI64ArrayAttr.get([1]),
+    )
+    return v, valid
+
+
 def _lower_masked_constructor(
     builder: MLIRLower, target: Var, args: list[Var], kwargs: list
 ) -> None:
@@ -105,6 +126,63 @@ def _lower_masked_getattr(
     builder.store_var(target, convert(field_value, target_mlir_ty))
 
 
+# ``cast(NA -> Masked)`` and ``cast(scalar -> Masked)``. Both build a Masked
+# struct for the target's value type; they differ only in the payload (NA has
+# none, so use undef) and the validity bit (NA -> invalid, scalar -> valid).
+# Triggered by branch unification, e.g. ``return x if cond else cudf.NA`` or
+# ``return 5``.
+def _cast_to_masked(context, builder, from_ty, to_ty, val):
+    value_mlir_ty = builder.get_mlir_type(to_ty.value_type)
+    if isinstance(from_ty, NAType):
+        value = llvm.UndefOp(value_mlir_ty)
+        valid = 0
+    else:
+        value = convert(val, value_mlir_ty)
+        valid = 1
+    valid_const = arith.constant(
+        result=builder.get_mlir_type(types.boolean), value=valid
+    )
+    return _pack_masked(builder, to_ty, value, valid_const)
+
+
+# ``cast(Masked -> Masked)``: branch unification across different
+# inner widths (e.g. one branch returns Masked(int32), another
+# returns Masked(float64); Numba unifies to Masked(float64)).
+# Promote the payload, preserve the validity bit.
+def _cast_masked_to_masked(context, builder, from_ty, to_ty, val):
+    if from_ty.value_type == to_ty.value_type:
+        return val
+    st = llvm.StructType(val.type)
+    m_val, m_valid = _extract_masked_value_valid(val, st.body[0], st.body[1])
+    value_mlir_ty = builder.get_mlir_type(to_ty.value_type)
+    casted = convert(m_val, value_mlir_ty)
+    return _pack_masked(builder, to_ty, casted, m_valid)
+
+
+# ``is``/``is not`` against NA are registered for both operand orders
+# (``m is NA`` and ``NA is m``), so find the MaskedType operand rather than
+# assuming which position it is in.
+def _masked_operand(builder, args):
+    """Return the ``MaskedType`` operand of a ``Masked``/``NA`` comparison."""
+    for arg in args:
+        if isinstance(builder.get_numba_type(arg.name), MaskedType):
+            return arg
+    raise TypeError("expected a MaskedType operand")
+
+
+# ``is``/``is not`` against NA both reduce to the validity bit: ``m is NA`` ->
+# ``not m.valid`` and ``m is not NA`` -> ``m.valid``. Registered for both
+# operators (and operand orders) via partials below.
+def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null):
+    m = builder.load_var(_masked_operand(builder, args))
+    st = llvm.StructType(m.type)
+    _, valid = _extract_masked_value_valid(m, st.body[0], st.body[1])
+    if is_null:
+        one = arith.constant(valid.type, 1)
+        valid = arith.xori(valid, one)
+    builder.store_var(target, valid)
+
+
 def _register() -> None:
     """Register the data model and lowerings with ``numba_cuda_mlir``.
 
@@ -123,5 +201,17 @@ def _register() -> None:
 
     lowering_registry.lower_getattr_generic(MaskedType)(_lower_masked_getattr)
 
+    lower_cast(na_type, MaskedType)(_cast_to_masked)
+    for _scalar_cls in (types.Integer, types.Float, types.Boolean):
+        lower_cast(_scalar_cls, MaskedType)(_cast_to_masked)
+    lower_cast(MaskedType, MaskedType)(_cast_masked_to_masked)
+
+    is_na = partial(_lower_masked_na_compare, is_null=True)
+    is_not_na = partial(_lower_masked_na_compare, is_null=False)
+    lower(operator.is_, MaskedType, NAType)(is_na)
+    lower(operator.is_, NAType, MaskedType)(is_na)
+    lower(operator.is_not, MaskedType, NAType)(is_not_na)
+    lower(operator.is_not, NAType, MaskedType)(is_not_na)
+
 
 _register()

python/cudf/cudf/core/udf/mlir_backend/masked_typing.py

@@ -2,16 +2,22 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from numba_cuda_mlir import types
+import operator
+
+from numba_cuda_mlir import models, types
 from numba_cuda_mlir.extending import typing_registry
+from numba_cuda_mlir.models import register_model
 from numba_cuda_mlir.numba_cuda import types as nb_types
+from numba_cuda_mlir.numba_cuda.extending import typeof_impl
 from numba_cuda_mlir.numba_cuda.types.misc import unliteral
 from numba_cuda_mlir.numba_cuda.typing.templates import (
+    AbstractTemplate,
     AttributeTemplate,
     ConcreteTemplate,
 )
 from numba_cuda_mlir.typing import signature as nb_signature
 
+from cudf.core.missing import NA
 from cudf.core.udf.api import Masked
 
 _SUPPORTED_MASKED_VALUE_TYPE_CLASSES = (
@@ -49,6 +55,55 @@ def __hash__(self) -> int:
         # parameter ``value_type`` matches, so numba can cache them by repr.
         return hash(repr(self))
 
+    def unify(self, context, other):
+        """Pick a common type when branches return different shapes.
+
+        ``return x if cond else cudf.NA`` unifies ``MaskedType`` and
+        ``NAType``; ``return x if cond else 5`` unifies ``MaskedType``
+        and a scalar; two branches returning different ``Masked``
+        widths unify their inner value types. Returning ``None``
+        signals "no unifier" and lets numba fall through.
+        See https://numba.pydata.org/numba-doc/dev/user/troubleshoot.html#my-code-has-a-type-unification-problem
+        """
+        if isinstance(other, NAType):
+            return self
+        other_value_type = (
+            other.value_type if isinstance(other, MaskedType) else other
+        )
+        unified = context.unify_pairs(self.value_type, other_value_type)
+        return MaskedType(unified) if unified is not None else None
+
+
+class NAType(types.Type):
+    """Type for ``cudf.NA`` -- the missing-value sentinel that can flow
+    through a UDF branch and unify into a ``MaskedType``.
+    """
+
+    def __init__(self):
+        super().__init__(name="NA")
+
+    def unify(self, context, other):
+        # See https://numba.pydata.org/numba-doc/dev/user/troubleshoot.html#my-code-has-a-type-unification-problem
+        # NA + Masked is delegated to MaskedType.unify (see above) so we
+        # only need to handle NA + (NA | scalar) here.
+        if isinstance(other, MaskedType):
+            return None
+        if isinstance(other, NAType):
+            return self
+        return MaskedType(other)
+
+
+na_type = NAType()
+
+
+@typeof_impl.register(type(NA))
+def _typeof_na(val, c):
+    return na_type
+
+
+# NAType has no payload; OpaqueModel is the right data model.
+register_model(NAType)(models.OpaqueModel)
+
 
 # ``Masked(value, valid)`` constructor: produces a ``Masked(value_ty)``.
 class MaskedConstructor(ConcreteTemplate):
@@ -71,12 +126,27 @@ def resolve_valid(self, typ: MaskedType) -> types.Type:
         return types.boolean
 
 
+# ``is``/``is not`` between a ``MaskedType`` and ``NA`` (either operand order)
+# both type to boolean; the is/is-not distinction is handled in lowering, so a
+# single template serves both operators.
+class MaskedNAComparison(AbstractTemplate):
+    def generic(self, args, kws):
+        if len(args) != 2 or kws:
+            return None
+        lhs, rhs = args
+        if {type(lhs), type(rhs)} == {MaskedType, NAType}:
+            return nb_signature(types.boolean, lhs, rhs)
+        return None
+
+
 def _register() -> None:
     """Register typing for ``Masked`` and ``MaskedType`` attributes with
     ``numba_cuda_mlir``. Called once at module import.
     """
     typing_registry.register_global(Masked, types.Function(MaskedConstructor))
     typing_registry.register_attr(MaskedTypeAttrs)
+    typing_registry.register_global(operator.is_)(MaskedNAComparison)
+    typing_registry.register_global(operator.is_not)(MaskedNAComparison)
 
 
 _register()