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"""
2020from numba_cuda_mlir ._mlir import ir
2121from numba_cuda_mlir ._mlir .dialects import arith , llvm
2222from numba_cuda_mlir .extending import lowering_registry , typing_registry
23+ from numba_cuda_mlir .lowering_utilities import DeferredMethodCall
2324from numba_cuda_mlir .numba_cuda .typing .templates import (
2425 AbstractTemplate ,
26+ AttributeTemplate ,
2527 Signature ,
2628)
2729from numba_cuda_mlir .typing import signature as nb_signature
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
7583def _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+
164236def _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+
289382def _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 ()
0 commit comments