Skip to content

Commit c311cff

Browse files
Address review feedback on datetime masked ops
- Fill in docstrings for the datetime helpers and tests (drop the TODO placeholders); annotate the new lowering functions. - Delegate temporal *comparisons* (not just +/-) to numba-cuda-mlir's unit-aware scalar lowering so mixed-unit compares are correct (e.g. 1s > 2ns is True, not a raw 1 > 2); add a mixed-unit test. (CodeRabbit NVIDIA#23309) - Key datetime staging temp vars off the unique result ``target`` so two temporal ops sharing a left operand in one expression can't clobber each other's typemap entries; drop the now-redundant ref_var plumbing and de-specialize the _make_temp_var name prefix. (CodeRabbit NVIDIA#23309) - Drop the redundant ``ref`` param from the datetime comparison test, asserting via the op directly (prior-PR review pattern).
1 parent b0c5eb2 commit c311cff

3 files changed

Lines changed: 88 additions & 43 deletions

File tree

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

Lines changed: 47 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -200,29 +200,49 @@ def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null):
200200
builder.store_var(target, valid)
201201

202202

203-
# datetime64 / timedelta64 ``+``/``-`` need numba_cuda_mlir's ``datetime``
204-
# lowering (which scales by unit), not a raw i64 op.
205-
def _needs_datetimelike_delegate(op, ty1, ty2):
206-
if op not in (operator.add, operator.sub):
203+
def _needs_datetimelike_delegate(
204+
op: Callable, ty1: types.Type, ty2: types.Type
205+
) -> bool:
206+
"""Whether a masked binary ``op`` over datetime64/timedelta64 operands must
207+
delegate to numba_cuda_mlir's unit-aware scalar lowering.
208+
209+
A raw i64 op on the payloads is wrong whenever units matter:
210+
211+
* ``+``/``-`` must scale the operands to a common unit before combining.
212+
* comparisons must likewise compare on a common unit (``1 s < 2 ns`` is
213+
``1_000_000_000 < 2``, not ``1 < 2``).
214+
215+
So delegate for add/sub and all comparison operators when either operand is
216+
temporal.
217+
"""
218+
if op not in (operator.add, operator.sub, *comparison_ops):
207219
return False
208220
return isinstance(
209221
ty1, (types.NPDatetime, types.NPTimedelta)
210222
) or isinstance(ty2, (types.NPDatetime, types.NPTimedelta))
211223

212224

213225
def _apply_masked_datetimelike_binary(
214-
builder,
215-
target,
216-
target_type,
217-
v1,
218-
v2,
219-
result_valid,
220-
op,
221-
ty1,
222-
ty2,
223-
ref_var,
224-
):
225-
"""TODO: write docstring."""
226+
builder: MLIRLower,
227+
target: Var,
228+
target_type: MaskedType,
229+
v1: mlir_ir.Value,
230+
v2: mlir_ir.Value,
231+
result_valid: mlir_ir.Value,
232+
op: Callable,
233+
ty1: types.Type,
234+
ty2: types.Type,
235+
) -> None:
236+
"""Lower a temporal masked binary op by delegating to the registered
237+
unit-aware scalar builder, then repacking with the operand validity.
238+
239+
The two payloads are staged into fresh typed IR vars (with their true
240+
datetime/timedelta types), handed to numba_cuda_mlir's scalar lowering for
241+
``op`` (which handles unit scaling), and the result is converted to the
242+
target Masked's value type and packed with ``result_valid``. Staging vars
243+
are keyed off the unique ``target`` so distinct temporal ops in one
244+
expression never share a typemap slot.
245+
"""
226246
ret_ty = target_type.value_type
227247
nb_sig = nb_typing.signature(ret_ty, ty1, ty2)
228248
cg = builder.get_registered_builder(op, nb_sig)
@@ -231,9 +251,10 @@ def _apply_masked_datetimelike_binary(
231251
f"No MLIR lowering for masked {op!r} with {ty1}, {ty2}; "
232252
f"signature {nb_sig}"
233253
)
234-
in1 = _make_temp_var(builder, ref_var, "mdt_l", ty1)
235-
in2 = _make_temp_var(builder, ref_var, "mdt_r", ty2)
236-
outv = _make_temp_var(builder, ref_var, "mdt_o", ret_ty)
254+
tag = getattr(op, "__name__", "op")
255+
in1 = _make_temp_var(builder, target, f"dt_{tag}_l", ty1)
256+
in2 = _make_temp_var(builder, target, f"dt_{tag}_r", ty2)
257+
outv = _make_temp_var(builder, target, f"dt_{tag}_o", ret_ty)
237258
builder.store_var(in1, convert(v1, builder.get_mlir_type(ty1)))
238259
builder.store_var(in2, convert(v2, builder.get_mlir_type(ty2)))
239260
cg(builder, outv, [in1, in2], ())
@@ -253,18 +274,16 @@ def _apply_masked_binary_op(
253274
*,
254275
inner_ty1: types.Type | None = None,
255276
inner_ty2: types.Type | None = None,
256-
ref_var: Var | None = None,
257277
) -> None:
258278
"""Apply ``op(v1, v2)`` to two scalar MLIR values, convert the result to
259279
the target Masked's value type, and pack it with the given validity bit.
260280
Numeric/boolean only at this layer.
261281
"""
262-
# datetime/timedelta add/sub: delegate to the unit-aware scalar
263-
# lowering when we know the operand inner types.
282+
# datetime/timedelta add/sub/comparisons: delegate to the unit-aware
283+
# scalar lowering when we know the operand inner types.
264284
if (
265285
inner_ty1 is not None
266286
and inner_ty2 is not None
267-
and ref_var is not None
268287
and _needs_datetimelike_delegate(op, inner_ty1, inner_ty2)
269288
):
270289
_apply_masked_datetimelike_binary(
@@ -277,7 +296,6 @@ def _apply_masked_binary_op(
277296
op,
278297
inner_ty1,
279298
inner_ty2,
280-
ref_var,
281299
)
282300
return
283301

@@ -320,7 +338,6 @@ def _lower(
320338
op,
321339
inner_ty1=ty1,
322340
inner_ty2=ty2,
323-
ref_var=args[0],
324341
)
325342

326343
return _lower
@@ -383,7 +400,6 @@ def _lower(
383400
op,
384401
inner_ty1=m_inner_ty,
385402
inner_ty2=s_inner_ty,
386-
ref_var=m_var,
387403
)
388404
else:
389405
_apply_masked_binary_op(
@@ -396,7 +412,6 @@ def _lower(
396412
op,
397413
inner_ty1=s_inner_ty,
398414
inner_ty2=m_inner_ty,
399-
ref_var=m_var,
400415
)
401416

402417
return _lower
@@ -424,14 +439,15 @@ def _make_temp_var(
424439
) -> Var:
425440
"""Create a fresh typed IR ``Var`` for staging an intermediate value.
426441
427-
Used to feed the masked payload into a registered scalar lowering (which
442+
Used to feed a masked payload into a registered scalar lowering (which
428443
operates on plain IR vars) and to receive its result. The name is derived
429-
from ``base_var`` plus ``name_suffix`` so typemap keys stay unique when one
430-
operand feeds several unary calls in a single expression.
444+
from ``base_var`` plus ``name_suffix``; callers pass a unique ``base_var``
445+
(e.g. the result ``target``) so typemap keys stay distinct when one operand
446+
feeds several op lowerings in a single expression.
431447
"""
432448
scope = getattr(base_var, "scope", None)
433449
loc = getattr(base_var, "loc", None)
434-
name = f"$masked_uop_{base_var.name}_{name_suffix}"
450+
name = f"$masked_tmp_{base_var.name}_{name_suffix}"
435451
temp = numba_ir.Var(scope=scope, name=name, loc=loc)
436452
builder.fndesc.typemap[temp.name] = numba_type
437453
return temp

‎python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_lowering.py‎

Lines changed: 38 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -817,7 +817,7 @@ def _bool(v):
817817

818818

819819
def test_masked_datetime_minus_datetime_is_timedelta():
820-
"""TODO: write docstring."""
820+
"""Masked ``datetime - datetime`` yields a Masked timedelta."""
821821

822822
@cuda.jit(
823823
types.void(
@@ -850,7 +850,7 @@ def k(out_v, out_valid, a, av, b, bv):
850850

851851

852852
def test_masked_datetime_plus_timedelta_is_datetime():
853-
"""TODO: write docstring."""
853+
"""Masked ``datetime + timedelta`` yields a Masked datetime."""
854854

855855
@cuda.jit(
856856
types.void(
@@ -870,7 +870,7 @@ def k(out_v, a, av, t, tv):
870870

871871

872872
def test_masked_timedelta_plus_timedelta_is_timedelta():
873-
"""TODO: write docstring."""
873+
"""Masked ``timedelta + timedelta`` yields a Masked timedelta."""
874874

875875
@cuda.jit(
876876
types.void(
@@ -889,12 +889,9 @@ def k(out_v, a, av, b, bv):
889889
assert int(out_v.get().view("int64")[0]) == 420
890890

891891

892-
@pytest.mark.parametrize(
893-
"op,ref",
894-
[(operator.lt, lambda a, b: a < b), (operator.gt, lambda a, b: a > b)],
895-
)
896-
def test_masked_datetime_comparison(op, ref):
897-
"""TODO: write docstring."""
892+
@pytest.mark.parametrize("op", [operator.lt, operator.gt])
893+
def test_masked_datetime_comparison(op):
894+
"""Masked datetime comparisons evaluate on the operand payloads."""
898895

899896
@cuda.jit(
900897
types.void(
@@ -910,11 +907,41 @@ def k(out, a, av, b, bv):
910907

911908
out = cp.zeros(1, dtype=np.bool_)
912909
_launch(k, out, _dt_in([400]), _bool(True), _dt_in([1000]), _bool(True))
913-
assert bool(out.get()[0]) == ref(400, 1000)
910+
assert bool(out.get()[0]) == op(400, 1000)
911+
912+
913+
def test_masked_datetime_comparison_mixed_units():
914+
"""Mixed-unit comparisons scale to a common unit rather than comparing
915+
raw payloads: ``1 s`` vs ``2 ns`` compares as ``1_000_000_000`` vs ``2``
916+
(True for ``>``), not the raw ``1 < 2``.
917+
"""
918+
dt_s = types.NPDatetime("s")
919+
dt_ns = types.NPDatetime("ns")
920+
921+
@cuda.jit(
922+
types.void(
923+
types.boolean[::1],
924+
dt_s[::1],
925+
types.boolean[::1],
926+
dt_ns[::1],
927+
types.boolean[::1],
928+
)
929+
)
930+
def k(out, a, av, b, bv):
931+
out[0] = (Masked(a[0], av[0]) > Masked(b[0], bv[0])).value
932+
933+
a = cp.asarray(np.array([1], dtype="int64")).view("datetime64[s]")
934+
b = cp.asarray(np.array([2], dtype="int64")).view("datetime64[ns]")
935+
out = cp.zeros(1, dtype=np.bool_)
936+
_launch(k, out, a, _bool(True), b, _bool(True))
937+
# 1 second (1e9 ns) > 2 ns is True; a raw i64 compare would give 1 > 2 = False.
938+
assert bool(out.get()[0]) is True
914939

915940

916941
def test_masked_datetime_arith_validity_propagates():
917-
"""TODO: write docstring."""
942+
"""Temporal arithmetic ANDs operand validity: an NA operand poisons the
943+
result.
944+
"""
918945

919946
@cuda.jit(
920947
types.void(

‎python/cudf/cudf/tests/private_objects/mlir_backend/test_masked_typing.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@ def test_masked_type_unsupported_value_becomes_poison():
5555

5656
@pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"])
5757
def test_masked_datetime_timedelta_not_poisoned(unit):
58-
"""TODO: write docstring."""
58+
"""datetime64/timedelta64 payloads are supported MaskedType value types
59+
(not poisoned) across the units cuDF columns use.
60+
"""
5961
dt = MaskedType(types.NPDatetime(unit))
6062
td = MaskedType(types.NPTimedelta(unit))
6163
assert dt.value_type == types.NPDatetime(unit)

0 commit comments

Comments
 (0)