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
511from __future__ import annotations
612
1016from 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+
1372def _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 ]
0 commit comments