Skip to content

Commit 202cfd6

Browse files
Add strip/lstrip/rstrip to the MLIR UDF backend
Second batch of string-producing ops. s.strip(chars) scans from each requested end while the byte is in the strip set (empty set == ASCII whitespace, code <= 0x20, matching libcudf detail::strip), then allocates and copies the surviving byte range via NRT. - string_lowering_impl.py: _gep_byte_offset, _empty_view, _is_strip_char, _strip_core, and lower_strip/lstrip/rstrip. - strings_lowering.py: optional-arg typing (_make_strip_attr) and lowering (_make_lower_strip) over plain/literal/masked operands, reusing _view_and_valid/_store_maybe_masked; registered as getattr methods. Tests cover whitespace-mode strip/lstrip/rstrip and a character-set strip, verified device-side against the expected column then freed via NRT_decref. Verified under compute-sanitizer: memcheck clean, no NRT device-heap leaks.
1 parent 7b9de2b commit 202cfd6

3 files changed

Lines changed: 430 additions & 2 deletions

File tree

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

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,190 @@ def lower_concat(
199199
return build_mlir_string(mi, data, total)
200200

201201

202+
def _gep_byte_offset(base: ir.Value, offset: ir.Value) -> ir.Value:
203+
"""Pointer to ``base + offset`` bytes (GEP on an ``i8*`` buffer)."""
204+
return llvm.getelementptr(
205+
llvm.PointerType.get(),
206+
base,
207+
[offset],
208+
[GEP_DYNAMIC_INDEX],
209+
T.i8(),
210+
None,
211+
)
212+
213+
214+
def _empty_view() -> ir.Value:
215+
"""View with a null data pointer and ``nbytes == 0`` (whitespace mode)."""
216+
i32 = ir.IntegerType.get_signless(32)
217+
ptr = llvm.PointerType.get()
218+
zero = arith.constant(i32, 0)
219+
null = llvm.inttoptr(ptr, arith.constant(T.i64(), 0))
220+
view_ty = llvm.StructType.get_literal([ptr, i32, i32])
221+
val = llvm.UndefOp(view_ty)
222+
val = llvm.insertvalue(
223+
container=val, value=null, position=ir.DenseI64ArrayAttr.get([0])
224+
)
225+
val = llvm.insertvalue(
226+
container=val, value=zero, position=ir.DenseI64ArrayAttr.get([1])
227+
)
228+
return llvm.insertvalue(
229+
container=val, value=zero, position=ir.DenseI64ArrayAttr.get([2])
230+
)
231+
232+
233+
def _is_strip_char(
234+
byte_val: ir.Value, strip_data: ir.Value, strip_nb: ir.Value
235+
) -> ir.Value:
236+
"""i1: whether ``byte_val`` should be stripped.
237+
238+
Mirrors libcudf ``detail::strip``: with an empty strip set, strip bytes with
239+
code ``<= 0x20`` (ASCII whitespace); otherwise strip any byte present in the
240+
strip set.
241+
"""
242+
i1 = ir.IntegerType.get_signless(1)
243+
zero = arith.constant(T.i64(), 0)
244+
one = arith.constant(T.i64(), 1)
245+
strip_empty = arith.cmpi(arith.CmpIPredicate.eq, strip_nb, zero)
246+
is_ws = arith.cmpi(
247+
arith.CmpIPredicate.ule, byte_val, arith.constant(T.i8(), 0x20)
248+
)
249+
loop = scf.ForOp(zero, strip_nb, one, [arith.constant(i1, 0)])
250+
si = loop.body.arguments[0]
251+
found = loop.body.arguments[1]
252+
with ir.InsertionPoint(loop.body):
253+
sb = _byte_at(strip_data, si)
254+
eq = arith.cmpi(arith.CmpIPredicate.eq, byte_val, sb)
255+
scf.YieldOp([arith.ori(found, eq)])
256+
return arith.select(strip_empty, is_ws, loop.results[0])
257+
258+
259+
def _strip_core(
260+
gpu_module: ir.Module,
261+
str_view: ir.Value,
262+
strip_view: ir.Value,
263+
do_left: bool,
264+
do_right: bool,
265+
) -> ir.Value:
266+
"""Shared strip/lstrip/rstrip: trim the matching run from each requested end.
267+
268+
Scans from the left (and/or right) while the current byte is a strip char
269+
and no non-strip byte has yet been seen on that side, then allocates and
270+
copies the surviving ``[left_idx, right_idx)`` byte range.
271+
"""
272+
data = _view_data(str_view)
273+
nb = llvm.zext(T.i64(), _view_nbytes(str_view))
274+
strip_data = _view_data(strip_view)
275+
strip_nb = llvm.zext(T.i64(), _view_nbytes(strip_view))
276+
zero = arith.constant(T.i64(), 0)
277+
one = arith.constant(T.i64(), 1)
278+
279+
if do_left:
280+
loop_l = scf.ForOp(zero, nb, one, [zero])
281+
i = loop_l.body.arguments[0]
282+
left = loop_l.body.arguments[1]
283+
with ir.InsertionPoint(loop_l.body):
284+
should_strip = _is_strip_char(
285+
_byte_at(data, i), strip_data, strip_nb
286+
)
287+
still = arith.cmpi(arith.CmpIPredicate.eq, left, i)
288+
advance = arith.andi(should_strip, still)
289+
scf.YieldOp([arith.select(advance, arith.addi(left, one), left)])
290+
left_idx = loop_l.results[0]
291+
else:
292+
left_idx = zero
293+
294+
if do_right:
295+
loop_r = scf.ForOp(zero, nb, one, [nb])
296+
j = loop_r.body.arguments[0]
297+
right = loop_r.body.arguments[1]
298+
with ir.InsertionPoint(loop_r.body):
299+
ridx = arith.subi(arith.subi(nb, j), one)
300+
should_strip = _is_strip_char(
301+
_byte_at(data, ridx), strip_data, strip_nb
302+
)
303+
expected = arith.subi(nb, j)
304+
still = arith.cmpi(arith.CmpIPredicate.eq, right, expected)
305+
advance = arith.andi(should_strip, still)
306+
scf.YieldOp([arith.select(advance, ridx, right)])
307+
right_idx = loop_r.results[0]
308+
else:
309+
right_idx = nb
310+
311+
nonempty = arith.cmpi(arith.CmpIPredicate.ult, left_idx, right_idx)
312+
result_len = arith.select(nonempty, arith.subi(right_idx, left_idx), zero)
313+
mi, out_data = allocate_mlir_string(gpu_module, result_len)
314+
llvm.intr_memcpy(
315+
out_data, _gep_byte_offset(data, left_idx), result_len, False
316+
)
317+
return build_mlir_string(mi, out_data, result_len)
318+
319+
320+
def lower_strip(
321+
gpu_module: ir.Module, str_view: ir.Value, strip_view: ir.Value
322+
) -> ir.Value:
323+
"""``s.strip(chars)`` -> trimmed ``mlir_string`` (both ends).
324+
325+
Parameters
326+
----------
327+
gpu_module : ir.Module
328+
The GPU module the NRT functions are declared in.
329+
str_view : ir.Value
330+
View of the source string's bytes.
331+
strip_view : ir.Value
332+
View of the strip-character set (empty for whitespace mode).
333+
334+
Returns
335+
-------
336+
ir.Value
337+
The trimmed owned ``mlir_string``.
338+
"""
339+
return _strip_core(gpu_module, str_view, strip_view, True, True)
340+
341+
342+
def lower_lstrip(
343+
gpu_module: ir.Module, str_view: ir.Value, strip_view: ir.Value
344+
) -> ir.Value:
345+
"""``s.lstrip(chars)`` -> left-trimmed ``mlir_string``.
346+
347+
Parameters
348+
----------
349+
gpu_module : ir.Module
350+
The GPU module the NRT functions are declared in.
351+
str_view : ir.Value
352+
View of the source string's bytes.
353+
strip_view : ir.Value
354+
View of the strip-character set (empty for whitespace mode).
355+
356+
Returns
357+
-------
358+
ir.Value
359+
The left-trimmed owned ``mlir_string``.
360+
"""
361+
return _strip_core(gpu_module, str_view, strip_view, True, False)
362+
363+
364+
def lower_rstrip(
365+
gpu_module: ir.Module, str_view: ir.Value, strip_view: ir.Value
366+
) -> ir.Value:
367+
"""``s.rstrip(chars)`` -> right-trimmed ``mlir_string``.
368+
369+
Parameters
370+
----------
371+
gpu_module : ir.Module
372+
The GPU module the NRT functions are declared in.
373+
str_view : ir.Value
374+
View of the source string's bytes.
375+
strip_view : ir.Value
376+
View of the strip-character set (empty for whitespace mode).
377+
378+
Returns
379+
-------
380+
ir.Value
381+
The right-trimmed owned ``mlir_string``.
382+
"""
383+
return _strip_core(gpu_module, str_view, strip_view, False, True)
384+
385+
202386
def _is_begin_utf8_char(byte_val: ir.Value) -> ir.Value:
203387
"""i1: whether ``byte_val`` is a UTF-8 start byte (top 2 bits != ``10``)."""
204388
i32 = ir.IntegerType.get_signless(32)

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

Lines changed: 99 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44
55
Covers ``len`` (UTF-8 character count), the string comparison operators
66
(``==``/``!=``/``<``/``<=``/``>``/``>=``), the affix/search/predicate methods,
7-
and ``+`` (concatenation) --- the first string-*producing* op, which allocates a
8-
new buffer via NRT --- over ``mlir_string``, ``Masked(mlir_string)`` and
7+
``+`` (concatenation) and ``strip``/``lstrip``/``rstrip`` --- the
8+
string-*producing* ops, which allocate a new buffer via NRT --- over
9+
``mlir_string``, ``Masked(mlir_string)`` and
910
string-literal operands; masked operands yield a ``Masked`` result whose validity
1011
is the AND of the operand validity bits. All lowerings are pure MLIR (see
1112
:mod:`cudf.core.udf.mlir_backend.string_lowering_impl`). Registered with
@@ -95,6 +96,13 @@
9596
"count": _impl._lower_count,
9697
}
9798

99+
# strip method name -> its pure-MLIR impl (optional chars arg -> mlir_string).
100+
_STRIP_METHODS = {
101+
"strip": _impl.lower_strip,
102+
"lstrip": _impl.lower_lstrip,
103+
"rstrip": _impl.lower_rstrip,
104+
}
105+
98106
# character-class predicate method name -> its pure-MLIR impl (no arg -> bool).
99107
_IS_METHODS = {
100108
"isalpha": _impl._lower_isalpha,
@@ -279,6 +287,51 @@ def resolve(self, mod):
279287
return resolve
280288

281289

290+
def _make_strip_attr(attrname: str, masked: bool) -> object:
291+
"""Build a ``resolve_<strip>`` for an optional-arg method returning a string.
292+
293+
Parameters
294+
----------
295+
attrname : str
296+
Method name (``"strip"``/``"lstrip"``/``"rstrip"``).
297+
masked : bool
298+
Whether the receiver is ``Masked(mlir_string)`` (wraps the result).
299+
300+
Returns
301+
-------
302+
callable
303+
A ``resolve_`` function for the attribute template.
304+
"""
305+
recvr = _masked_string if masked else mlir_string
306+
307+
class _StripTemplate(AbstractTemplate):
308+
key = f"MLIRString.{attrname}.{'m' if masked else 's'}"
309+
310+
def generic(self, args, kws):
311+
"""Resolve ``s.strip()`` / ``s.strip(chars)``.
312+
313+
Zero args (whitespace mode) or one string arg (chars set); the
314+
result is ``Masked(mlir_string)`` when the receiver or the argument
315+
is masked, else ``mlir_string``.
316+
"""
317+
if kws or len(args) > 1:
318+
return None
319+
arg_masked = len(args) == 1 and _is_masked_string(args[0])
320+
result_ty = (
321+
MaskedType(mlir_string)
322+
if masked or arg_masked
323+
else mlir_string
324+
)
325+
if len(args) == 0:
326+
return nb_signature(result_ty, recvr=self.this)
327+
return nb_signature(result_ty, args[0], recvr=self.this)
328+
329+
def resolve(self, mod):
330+
return types.BoundFunction(_StripTemplate, recvr)
331+
332+
return resolve
333+
334+
282335
@typing_registry.register_attr
283336
class MLIRStringAttrs(AttributeTemplate):
284337
"""Attribute typing for ``mlir_string`` methods."""
@@ -360,6 +413,17 @@ def resolve_valid(self, typ: MaskedType) -> types.Type:
360413
f"resolve_{_name}",
361414
_make_is_attr(_name, masked=True),
362415
)
416+
for _name in _STRIP_METHODS:
417+
setattr(
418+
MLIRStringAttrs,
419+
f"resolve_{_name}",
420+
_make_strip_attr(_name, masked=False),
421+
)
422+
setattr(
423+
MaskedMLIRStringAttrs,
424+
f"resolve_{_name}",
425+
_make_strip_attr(_name, masked=True),
426+
)
363427

364428

365429
class ConcatMLIRStringTemplate(AbstractTemplate):
@@ -609,6 +673,33 @@ def _getattr(context, builder, target, value, attr=None):
609673
return _getattr
610674

611675

676+
def _make_lower_strip(impl_fn: object) -> object:
677+
"""Build the ``lower_getattr`` for a strip method (optional chars arg).
678+
679+
With no argument the strip set is empty (whitespace mode); with one string
680+
argument the argument supplies the strip set. A ``Masked`` result is packed
681+
when the receiver or the argument is masked.
682+
"""
683+
684+
def _impl_lower(builder, target, args, kwargs):
685+
valid_ty = builder.get_mlir_type(types.boolean)
686+
true_val = arith.constant(valid_ty, 1)
687+
recv, vr = _view_and_valid(builder, args[0], true_val)
688+
if len(args) == 2:
689+
strip_view, vo = _view_and_valid(builder, args[1], true_val)
690+
valid = arith.andi(vr, vo)
691+
else:
692+
strip_view = _impl._empty_view()
693+
valid = vr
694+
result = impl_fn(builder.mlir_gpu_module, recv, strip_view)
695+
_store_maybe_masked(builder, target, result, valid)
696+
697+
def _getattr(context, builder, target, value, attr=None):
698+
builder.store_var(target, DeferredMethodCall(value, _impl_lower))
699+
700+
return _getattr
701+
702+
612703
def _lower_contains_op(
613704
builder: MLIRLower, target: Var, args: list[Var], kwargs: list
614705
) -> None:
@@ -760,5 +851,11 @@ def _register() -> None:
760851
lowering_registry.lower_getattr(mlir_string, name)(getattr_fn)
761852
lowering_registry.lower_getattr(_masked_string, name)(getattr_fn)
762853

854+
# strip/lstrip/rstrip (optional chars arg) -> (masked) mlir_string
855+
for name, impl_fn in _STRIP_METHODS.items():
856+
getattr_fn = _make_lower_strip(impl_fn)
857+
lowering_registry.lower_getattr(mlir_string, name)(getattr_fn)
858+
lowering_registry.lower_getattr(_masked_string, name)(getattr_fn)
859+
763860

764861
_register()

0 commit comments

Comments
 (0)