Skip to content

Commit 27e1eac

Browse files
[MLIR] String affix methods: startswith / endswith
Add startswith/endswith over mlir_string, Masked(mlir_string) and string-literal arguments, returning boolean / Masked(bool) (validity carried from the operands). Pure-MLIR byte prefix/suffix match via a shared _bytes_equal helper that early-exits at the first differing byte; the affix length is clamped to 0 when it cannot fit so the scan never over-reads. Introduces the method-dispatch machinery (attribute typing + DeferredMethodCall) that the search/predicate ops will reuse.
1 parent cd7446a commit 27e1eac

3 files changed

Lines changed: 224 additions & 6 deletions

File tree

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,3 +153,59 @@ def _lower_compare(lhs: ir.Value, rhs: ir.Value) -> ir.Value:
153153
arith.select(rhs_longer, arith.constant(i32, -1), zero_i32),
154154
)
155155
return arith.select(found, byte_diff, len_cmp)
156+
157+
158+
def _bytes_equal(
159+
a_data: ir.Value, a_off: ir.Value, b_data: ir.Value, n: ir.Value
160+
) -> ir.Value:
161+
"""i1: whether ``a_data[a_off:a_off+n] == b_data[0:n]`` byte-for-byte.
162+
163+
Early-exits (``scf.while``) at the first differing byte; callers must pass
164+
``n == 0`` when the range would read out of bounds.
165+
"""
166+
i1 = ir.IntegerType.get_signless(1)
167+
zero_i64 = arith.constant(T.i64(), 0)
168+
loop = scf.WhileOp([T.i64(), i1], [zero_i64, arith.constant(i1, 1)])
169+
before = loop.before.blocks.append(T.i64(), i1)
170+
with ir.InsertionPoint(before):
171+
idx, eq = before.arguments
172+
in_bounds = arith.cmpi(arith.CmpIPredicate.ult, idx, n)
173+
scf.condition(arith.andi(in_bounds, eq), [idx, eq])
174+
after = loop.after.blocks.append(T.i64(), i1)
175+
with ir.InsertionPoint(after):
176+
idx, _eq = after.arguments
177+
av = _byte_at(a_data, arith.addi(a_off, idx))
178+
bv = _byte_at(b_data, idx)
179+
next_idx = arith.addi(idx, arith.constant(T.i64(), 1))
180+
scf.YieldOp([next_idx, arith.cmpi(arith.CmpIPredicate.eq, av, bv)])
181+
return loop.results[1]
182+
183+
184+
def _lower_startswith(str_view: ir.Value, prefix: ir.Value) -> ir.Value:
185+
"""``startswith(prefix)`` -> ``i1`` (byte prefix match)."""
186+
i1 = ir.IntegerType.get_signless(1)
187+
src_nb = llvm.zext(T.i64(), _view_nbytes(str_view))
188+
pfx_nb = llvm.zext(T.i64(), _view_nbytes(prefix))
189+
too_long = arith.cmpi(arith.CmpIPredicate.ugt, pfx_nb, src_nb)
190+
# Compare 0 bytes when the prefix cannot fit, so the scan never over-reads
191+
# the source; the result is false anyway.
192+
n = arith.select(too_long, arith.constant(T.i64(), 0), pfx_nb)
193+
eq = _bytes_equal(
194+
_view_data(str_view), arith.constant(T.i64(), 0), _view_data(prefix), n
195+
)
196+
return arith.andi(arith.xori(too_long, arith.constant(i1, 1)), eq)
197+
198+
199+
def _lower_endswith(str_view: ir.Value, suffix: ir.Value) -> ir.Value:
200+
"""``endswith(suffix)`` -> ``i1`` (byte suffix match)."""
201+
i1 = ir.IntegerType.get_signless(1)
202+
src_nb = llvm.zext(T.i64(), _view_nbytes(str_view))
203+
sfx_nb = llvm.zext(T.i64(), _view_nbytes(suffix))
204+
too_long = arith.cmpi(arith.CmpIPredicate.ugt, sfx_nb, src_nb)
205+
zero_i64 = arith.constant(T.i64(), 0)
206+
# Clamp both the start offset and the length when the suffix cannot fit so
207+
# the scan stays in bounds; the result is false anyway.
208+
n = arith.select(too_long, zero_i64, sfx_nb)
209+
offset = arith.select(too_long, zero_i64, arith.subi(src_nb, sfx_nb))
210+
eq = _bytes_equal(_view_data(str_view), offset, _view_data(suffix), n)
211+
return arith.andi(arith.xori(too_long, arith.constant(i1, 1)), eq)

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

Lines changed: 105 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22
# SPDX-License-Identifier: Apache-2.0
33
"""Typing and lowering for ``mlir_string`` operations.
44
5-
Covers ``len`` (UTF-8 character count) and the string comparison operators
6-
(``==``/``!=``/``<``/``<=``/``>``/``>=``) over ``mlir_string``,
7-
``Masked(mlir_string)`` and string-literal operands; masked operands yield a
8-
``Masked`` result whose validity is the AND of the operand validity bits. All
9-
lowerings are pure MLIR (see
5+
Covers ``len`` (UTF-8 character count), the string comparison operators
6+
(``==``/``!=``/``<``/``<=``/``>``/``>=``), and the affix methods
7+
``startswith``/``endswith`` over ``mlir_string``, ``Masked(mlir_string)`` and
8+
string-literal operands; masked operands yield a ``Masked`` result whose validity
9+
is the AND of the operand validity bits. All lowerings are pure MLIR (see
1010
:mod:`cudf.core.udf.mlir_backend.string_lowering_impl`). Registered with
1111
``numba_cuda_mlir`` at import time via :func:`_register`.
1212
"""
@@ -20,8 +20,10 @@
2020
from numba_cuda_mlir._mlir import ir
2121
from numba_cuda_mlir._mlir.dialects import arith, llvm
2222
from numba_cuda_mlir.extending import lowering_registry, typing_registry
23+
from numba_cuda_mlir.lowering_utilities import DeferredMethodCall
2324
from numba_cuda_mlir.numba_cuda.typing.templates import (
2425
AbstractTemplate,
26+
AttributeTemplate,
2527
Signature,
2628
)
2729
from numba_cuda_mlir.typing import signature as nb_signature
@@ -71,6 +73,12 @@
7173
(MLIRStringType, _masked_string),
7274
)
7375

76+
# affix method name -> its pure-MLIR impl (one string arg -> boolean).
77+
_BOOL_METHODS = {
78+
"startswith": _impl._lower_startswith,
79+
"endswith": _impl._lower_endswith,
80+
}
81+
7482

7583
def _is_plain_string(ty: types.Type) -> bool:
7684
"""Whether ``ty`` is an ``mlir_string`` or a compile-time string literal."""
@@ -161,6 +169,70 @@ def generic(self, args, kws):
161169
return _CmpOpTemplate
162170

163171

172+
def _make_method_attr(
173+
attrname: str, retty: types.Type, masked: bool
174+
) -> object:
175+
"""Build a ``resolve_<method>`` for a one-string-arg method.
176+
177+
Parameters
178+
----------
179+
attrname : str
180+
Method name (e.g. ``"startswith"``).
181+
retty : types.Type
182+
The unmasked return type.
183+
masked : bool
184+
Whether the receiver is ``Masked(mlir_string)`` (wraps the result).
185+
186+
Returns
187+
-------
188+
callable
189+
A ``resolve_`` function for the attribute template.
190+
"""
191+
result_ty = MaskedType(retty) if masked else retty
192+
recvr = _masked_string if masked else mlir_string
193+
194+
class _MethodTemplate(AbstractTemplate):
195+
key = f"MLIRString.{attrname}.{'m' if masked else 's'}"
196+
197+
def generic(self, args, kws):
198+
"""Resolve ``s.<method>(other)`` for a string ``other``."""
199+
if len(args) == 1 and not kws and _is_string_arg(args[0]):
200+
return nb_signature(result_ty, args[0], recvr=self.this)
201+
return None
202+
203+
def resolve(self, mod):
204+
return types.BoundFunction(_MethodTemplate, recvr)
205+
206+
return resolve
207+
208+
209+
@typing_registry.register_attr
210+
class MLIRStringAttrs(AttributeTemplate):
211+
"""Attribute typing for ``mlir_string`` methods."""
212+
213+
key = mlir_string
214+
215+
216+
@typing_registry.register_attr
217+
class MaskedMLIRStringAttrs(AttributeTemplate):
218+
"""Attribute typing for ``Masked(mlir_string)`` methods."""
219+
220+
key = _masked_string
221+
222+
223+
for _name in _BOOL_METHODS:
224+
setattr(
225+
MLIRStringAttrs,
226+
f"resolve_{_name}",
227+
_make_method_attr(_name, types.boolean, masked=False),
228+
)
229+
setattr(
230+
MaskedMLIRStringAttrs,
231+
f"resolve_{_name}",
232+
_make_method_attr(_name, types.boolean, masked=True),
233+
)
234+
235+
164236
def _lower_len(
165237
builder: MLIRLower, target: Var, args: list[Var], kwargs: list
166238
) -> None:
@@ -286,8 +358,29 @@ def _lower(builder, target, args, kwargs):
286358
return _lower
287359

288360

361+
def _make_lower_binary(impl_fn: object) -> object:
362+
"""Build the ``lower_getattr`` for a one-string-arg method (startswith/endswith).
363+
364+
The returned getattr defers to a method lowering that handles plain, literal
365+
and masked operands uniformly and packs a ``Masked`` result when required.
366+
"""
367+
368+
def _impl_lower(builder, target, args, kwargs):
369+
valid_ty = builder.get_mlir_type(types.boolean)
370+
true_val = arith.constant(valid_ty, 1)
371+
recv, vr = _view_and_valid(builder, args[0], true_val)
372+
other, vo = _view_and_valid(builder, args[1], true_val)
373+
result = impl_fn(recv, other)
374+
_store_maybe_masked(builder, target, result, arith.andi(vr, vo))
375+
376+
def _getattr(context, builder, target, value, attr=None):
377+
builder.store_var(target, DeferredMethodCall(value, _impl_lower))
378+
379+
return _getattr
380+
381+
289382
def _register() -> None:
290-
"""Register ``len`` + comparison typing and lowering with ``numba_cuda_mlir``."""
383+
"""Register ``len`` + comparison + affix typing/lowering with numba_cuda_mlir."""
291384
typing_registry.register_global(len)(LenMLIRStringTemplate)
292385
lowering_registry.lower(len, mlir_string)(_lower_len)
293386
lowering_registry.lower(len, MaskedType)(_lower_masked_len)
@@ -298,5 +391,11 @@ def _register() -> None:
298391
for lhs_ty, rhs_ty in _STR_COMBOS:
299392
lowering_registry.lower(cmp_op, lhs_ty, rhs_ty)(impl)
300393

394+
# startswith / endswith methods on mlir_string and Masked(mlir_string)
395+
for name, impl_fn in _BOOL_METHODS.items():
396+
getattr_fn = _make_lower_binary(impl_fn)
397+
lowering_registry.lower_getattr(mlir_string, name)(getattr_fn)
398+
lowering_registry.lower_getattr(_masked_string, name)(getattr_fn)
399+
301400

302401
_register()

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,3 +367,66 @@ def k(o, ov, sa, sav, sb, sbv):
367367
for i in range(n):
368368
if got_valid[i]:
369369
assert bool(out.get()[i]) is (a[i] == b[i])
370+
371+
372+
def test_string_affix():
373+
"""``startswith`` / ``endswith`` against a literal, over an array."""
374+
strings = ["abc", "abcd", "xabc", "ab", "", "cab", "h\u00e9llo"]
375+
arr, _k = _make_mlir_string_array(strings)
376+
n = len(strings)
377+
starts = cp.zeros(n, dtype=np.bool_)
378+
ends = cp.zeros(n, dtype=np.bool_)
379+
380+
@cuda.jit(
381+
types.void(
382+
types.boolean[::1],
383+
types.boolean[::1],
384+
types.CPointer(mlir_string),
385+
),
386+
extensions=[mlir_string_arg_handler],
387+
)
388+
def k(so, eo, s):
389+
i = cuda.grid(1)
390+
if i < so.size:
391+
so[i] = s[i].startswith("ab")
392+
eo[i] = s[i].endswith("bc")
393+
394+
with MLIRNumbaCudaConfig():
395+
k[1, n](starts, ends, arr)
396+
cuda.synchronize()
397+
assert starts.get().tolist() == [s.startswith("ab") for s in strings]
398+
assert ends.get().tolist() == [s.endswith("bc") for s in strings]
399+
400+
401+
def test_masked_string_affix_propagates_validity():
402+
"""``Masked(str).startswith(literal)`` -> ``Masked(bool)``; validity carried."""
403+
strings = ["abc", "abd", "xyz"]
404+
valid = [True, False, True]
405+
arr, _k = _make_mlir_string_array(strings)
406+
n = len(strings)
407+
out = cp.zeros(n, dtype=np.bool_)
408+
out_valid = cp.zeros(n, dtype=np.bool_)
409+
410+
@cuda.jit(
411+
types.void(
412+
types.boolean[::1],
413+
types.boolean[::1],
414+
types.CPointer(mlir_string),
415+
types.boolean[::1],
416+
),
417+
extensions=[mlir_string_arg_handler],
418+
)
419+
def k(o, ov, s, sv):
420+
i = cuda.grid(1)
421+
if i < o.size:
422+
r = Masked(s[i], sv[i]).startswith("ab")
423+
o[i] = r.value
424+
ov[i] = r.valid
425+
426+
with MLIRNumbaCudaConfig():
427+
k[1, n](out, out_valid, arr, cp.array(valid, dtype=np.bool_))
428+
cuda.synchronize()
429+
assert out_valid.get().tolist() == valid
430+
for i in range(n):
431+
if valid[i]:
432+
assert bool(out.get()[i]) is strings[i].startswith("ab")

0 commit comments

Comments
 (0)