Skip to content

Commit 4ad36b1

Browse files
[MLIR] len: use view-based int32 (size_type) impl
Switch len to the cutover's pure-MLIR view pattern: mlir_string is reduced to an internal view {ptr, i32 nbytes, i32 length} and the character count is returned as int32 (libcudf size_type), matching historical cudf str.len and the full cutover reference. Supersedes the earlier int64 result; the character count is bounded by size_type like the rest of the string API.
1 parent f97e06a commit 4ad36b1

2 files changed

Lines changed: 92 additions & 37 deletions

File tree

Lines changed: 77 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
3-
"""Pure-MLIR building blocks composed by the ``mlir_string`` lowerings."""
3+
"""Pure-MLIR building blocks composed by the ``mlir_string`` lowerings.
4+
5+
Helpers work on raw MLIR SSA values (no cuDF/NRT dependency). ``mlir_string``
6+
layout is ``{ptr meminfo, ptr data, i64 nbytes}``; ops run over an internal
7+
*view* value ``{ptr data, i32 nbytes, i32 length}`` built by
8+
:func:`_mlir_string_to_view`.
9+
"""
410

511
from __future__ import annotations
612

@@ -10,6 +16,59 @@
1016
from numba_cuda_mlir.lowering_utilities import GEP_DYNAMIC_INDEX
1117

1218

19+
def _byte_at(data: ir.Value, idx: ir.Value) -> ir.Value:
20+
"""Load ``data[idx]`` as ``i8`` via a dynamic GEP on the ``i8*`` buffer."""
21+
return llvm.load(
22+
T.i8(),
23+
llvm.getelementptr(
24+
llvm.PointerType.get(),
25+
data,
26+
[idx],
27+
[GEP_DYNAMIC_INDEX],
28+
T.i8(),
29+
None,
30+
),
31+
)
32+
33+
34+
# --- mlir_string / view field accessors ------------------------------------
35+
# view layout: {ptr data, i32 nbytes, i32 length}; only data + nbytes are used.
36+
def _view_data(view_val: ir.Value) -> ir.Value:
37+
return llvm.extractvalue(llvm.PointerType.get(), view_val, [0])
38+
39+
40+
def _view_nbytes(view_val: ir.Value) -> ir.Value:
41+
return llvm.extractvalue(ir.IntegerType.get_signless(32), view_val, [1])
42+
43+
44+
def _ms_extract_nbytes(ms_val: ir.Value) -> ir.Value:
45+
"""Extract ``nbytes`` (i64, field 2) from an ``mlir_string`` value."""
46+
return llvm.extractvalue(T.i64(), ms_val, [2])
47+
48+
49+
def _mlir_string_to_view(ms_val: ir.Value) -> ir.Value:
50+
"""Build an internal view ``{ptr data, i32 nbytes, i32 length}``.
51+
52+
``length`` is set equal to ``nbytes`` here; the character count is computed
53+
on demand by :func:`_lower_len`.
54+
"""
55+
i32 = ir.IntegerType.get_signless(32)
56+
ptr = llvm.PointerType.get()
57+
data = llvm.extractvalue(ptr, ms_val, [1])
58+
nbytes_i32 = arith.trunci(i32, _ms_extract_nbytes(ms_val))
59+
view_ty = llvm.StructType.get_literal([ptr, i32, i32])
60+
val = llvm.UndefOp(view_ty)
61+
val = llvm.insertvalue(
62+
container=val, value=data, position=ir.DenseI64ArrayAttr.get([0])
63+
)
64+
val = llvm.insertvalue(
65+
container=val, value=nbytes_i32, position=ir.DenseI64ArrayAttr.get([1])
66+
)
67+
return llvm.insertvalue(
68+
container=val, value=nbytes_i32, position=ir.DenseI64ArrayAttr.get([2])
69+
)
70+
71+
1372
def _is_begin_utf8_char(byte_val: ir.Value) -> ir.Value:
1473
"""i1: whether ``byte_val`` is a UTF-8 start byte (top 2 bits != ``10``)."""
1574
i32 = ir.IntegerType.get_signless(32)
@@ -19,35 +78,28 @@ def _is_begin_utf8_char(byte_val: ir.Value) -> ir.Value:
1978
)
2079

2180

22-
def _lower_len(ms_val: ir.Value) -> ir.Value:
23-
"""``len(mlir_string) -> i64`` character count (pure MLIR).
81+
def _lower_len(str_view: ir.Value) -> ir.Value:
82+
"""``len`` -> ``i32`` character count (number of UTF-8 start bytes).
2483
25-
Counts UTF-8 start bytes over the ``nbytes``-long data buffer, so the result
26-
is the number of characters (code points), matching cuDF ``str.len``. The
27-
count is accumulated directly in ``i64`` to match the declared result type.
28-
The ``mlir_string`` layout is ``{ptr meminfo, ptr data, i64 nbytes}`` --
29-
fields 1 (``data``) and 2 (``nbytes``) are read below.
84+
``int32`` matches libcudf's ``size_type`` used for string lengths.
3085
"""
31-
data = llvm.extractvalue(llvm.PointerType.get(), ms_val, [1])
32-
nb = llvm.extractvalue(T.i64(), ms_val, [2])
33-
zero_i64 = arith.constant(T.i64(), 0)
34-
one_i64 = arith.constant(T.i64(), 1)
35-
36-
loop = scf.ForOp(zero_i64, nb, one_i64, [zero_i64])
86+
i32 = ir.IntegerType.get_signless(32)
87+
data = _view_data(str_view)
88+
nb = llvm.zext(T.i64(), _view_nbytes(str_view))
89+
zero_i32 = arith.constant(i32, 0)
90+
loop = scf.ForOp(
91+
arith.constant(T.i64(), 0),
92+
nb,
93+
arith.constant(T.i64(), 1),
94+
[zero_i32],
95+
)
3796
with ir.InsertionPoint(loop.body):
3897
idx = loop.induction_variable
3998
acc = loop.inner_iter_args[0]
40-
# data[idx] -- GEP a dynamic byte offset into the i8* buffer
41-
byte_ptr = llvm.getelementptr(
42-
llvm.PointerType.get(),
43-
data,
44-
[idx],
45-
[GEP_DYNAMIC_INDEX],
46-
T.i8(),
47-
None,
99+
inc = arith.select(
100+
_is_begin_utf8_char(_byte_at(data, idx)),
101+
arith.constant(i32, 1),
102+
zero_i32,
48103
)
49-
b = llvm.load(T.i8(), byte_ptr)
50-
inc = arith.select(_is_begin_utf8_char(b), one_i64, zero_i64)
51104
scf.yield_([arith.addi(acc, inc)])
52-
53105
return loop.results[0]

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,16 @@
3434
from numba_cuda_mlir.mlir_lowering import MLIRLower
3535
from numba_cuda_mlir.numba_cuda.core.ir import Var
3636

37+
# libcudf size_type; the width of a string length result.
38+
size_type = types.int32
39+
3740

3841
class LenMLIRStringTemplate(AbstractTemplate):
3942
"""``len`` over strings.
4043
41-
``len(mlir_string)`` -> ``int64`` (UTF-8 character count), and
42-
``len(Masked(mlir_string))`` -> ``Masked(int64)`` (validity carried from the
43-
operand).
44+
``len(mlir_string)`` -> ``int32`` (UTF-8 character count), and
45+
``len(Masked(mlir_string))`` -> ``Masked(int32)`` (validity carried from the
46+
operand). ``int32`` matches libcudf's ``size_type`` for string lengths.
4447
"""
4548

4649
key = len
@@ -60,34 +63,34 @@ def generic(
6063
Returns
6164
-------
6265
Signature or None
63-
``int64`` for a bare ``mlir_string``, ``Masked(int64)`` for a
66+
``int32`` for a bare ``mlir_string``, ``Masked(int32)`` for a
6467
``Masked(mlir_string)``, else ``None``.
6568
"""
6669
if len(args) != 1 or kws:
6770
return None
6871
arg = args[0]
6972
if isinstance(arg, MLIRStringType):
70-
return nb_signature(types.int64, mlir_string)
73+
return nb_signature(size_type, mlir_string)
7174
if isinstance(arg, MaskedType) and isinstance(
7275
arg.value_type, MLIRStringType
7376
):
74-
return nb_signature(MaskedType(types.int64), arg)
77+
return nb_signature(MaskedType(size_type), arg)
7578
return None
7679

7780

7881
def _lower_len(
7982
builder: MLIRLower, target: Var, args: list[Var], kwargs: list
8083
) -> None:
81-
"""``len(mlir_string)``: count UTF-8 characters, returned as ``int64``."""
82-
ms_val = builder.load_var(args[0])
83-
builder.store_var(target, _impl._lower_len(ms_val))
84+
"""``len(mlir_string)``: count UTF-8 characters, returned as ``int32``."""
85+
view = _impl._mlir_string_to_view(builder.load_var(args[0]))
86+
builder.store_var(target, _impl._lower_len(view))
8487

8588

8689
def _lower_masked_len(
8790
builder: MLIRLower, target: Var, args: list[Var], kwargs: list
8891
) -> None:
8992
"""``len(Masked(mlir_string))``: character count packed with the operand's
90-
validity bit (``Masked(int64)``).
93+
validity bit (``Masked(int32)``).
9194
9295
The payload is scanned unconditionally; this is safe because null rows carry
9396
``nbytes == 0`` (the marshaller leaves a null ``data`` pointer with zero
@@ -97,9 +100,9 @@ def _lower_masked_len(
97100
m = builder.load_var(args[0])
98101
st = llvm.StructType(m.type)
99102
ms_val, m_valid = _extract_masked_value_valid(m, st.body[0], st.body[1])
100-
count_i64 = _impl._lower_len(ms_val)
103+
count = _impl._lower_len(_impl._mlir_string_to_view(ms_val))
101104
target_type = builder.get_numba_type(target.name)
102-
packed = _pack_masked(builder, target_type, count_i64, m_valid)
105+
packed = _pack_masked(builder, target_type, count, m_valid)
103106
builder.store_var(target, packed)
104107

105108

0 commit comments

Comments
 (0)