Skip to content

Commit 30b0e26

Browse files
[MLIR] Introduce MLIRString scalar string type
Add the cuDF-side MLIRString extension type (mlir_string): an NRT-managed owned string modeled as {ptr meminfo, ptr data, i64 nbytes}, plus the ManagedStrArrayWrapper host wrapper, its typeof registration, and the MLIRStringArgHandler kernel-launch marshaller. Type/marshalling surface only -- no string ops yet (len arrives next). Independent of the numeric MaskedType stack.
1 parent cba0893 commit 30b0e26

2 files changed

Lines changed: 189 additions & 0 deletions

File tree

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
from __future__ import annotations
4+
5+
import numpy as np
6+
from numba_cuda_mlir import types
7+
from numba_cuda_mlir.numba_cuda.extending import typeof_impl
8+
9+
10+
class MLIRStringType(types.Type):
11+
"""NRT-managed owned string: ``{ptr meminfo, ptr data, i64 nbytes}``.
12+
13+
``data`` points at a raw UTF-8 byte buffer and ``nbytes`` is its byte
14+
length. ``meminfo`` is the NRT ``MemInfo`` pointer used for lifetime
15+
management of *owned* (allocated) strings; read-only strings that borrow
16+
an existing buffer (e.g. a column element consumed by ``len``) carry a
17+
null ``meminfo`` and are never freed.
18+
"""
19+
20+
# ABI size: 8 (meminfo) + 8 (data) + 8 (nbytes) = 24 bytes. Must match the
21+
# device struct layout the marshalling helpers and string lowerings use.
22+
_extensionty_size = 24
23+
np_dtype: np.dtype[np.object_] = np.dtype("object")
24+
25+
def __init__(self) -> None:
26+
super().__init__(name="mlir_string")
27+
28+
@property
29+
def is_internal(self) -> bool:
30+
# ``mlir_string`` is an extension type that must be passed to kernels
31+
# as a ``CPointer`` (not an Array/memref), so it is never "internal".
32+
return False
33+
34+
@property
35+
def return_as(self) -> MLIRStringType:
36+
return self
37+
38+
39+
mlir_string = MLIRStringType()
40+
41+
42+
class ManagedStrArrayWrapper:
43+
"""Host-side wrapper whose ``typeof`` is ``CPointer(mlir_string)``.
44+
45+
Wraps a device buffer holding an array of ``mlir_string`` structs so that
46+
numba-cuda-mlir's dispatch types it as a pointer to ``mlir_string`` rather
47+
than inferring a raw ``uint8`` array from ``__cuda_array_interface__``.
48+
"""
49+
50+
def __init__(self, buffer) -> None:
51+
self._buffer = buffer
52+
53+
@property
54+
def ptr(self):
55+
return self._buffer.ptr
56+
57+
58+
@typeof_impl.register(ManagedStrArrayWrapper)
59+
def _typeof_mlir_str_array_wrapper(val, c) -> types.CPointer:
60+
return types.CPointer(mlir_string)
61+
62+
63+
class MLIRStringArgHandler:
64+
"""Kernel-launch extension that marshals ``mlir_string`` pointers.
65+
66+
A ``CPointer(mlir_string)`` argument is passed to the kernel as the raw
67+
device address (``uint64``). The tuple form carries an accompanying
68+
validity mask for nullable string columns.
69+
"""
70+
71+
def prepare_args(self, ty: types.Type, val, **kwargs) -> tuple:
72+
if isinstance(ty, types.CPointer) and isinstance(
73+
ty.dtype, MLIRStringType
74+
):
75+
return types.uint64, val.ptr
76+
if isinstance(ty, types.Tuple) and len(ty) >= 1:
77+
first_ty = ty[0]
78+
if isinstance(first_ty, types.CPointer) and isinstance(
79+
first_ty.dtype, MLIRStringType
80+
):
81+
ptr_val = val[0].ptr
82+
if len(ty) == 2:
83+
return ty, (ptr_val, val[1])
84+
return ty, (ptr_val,)
85+
return ty, val
86+
87+
88+
mlir_string_arg_handler = MLIRStringArgHandler()
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Python-surface tests for the MLIR backend strings typing module.
4+
5+
These exercise the host-side surface added by
6+
``cudf.core.udf.mlir_backend.strings_typing`` without compiling a kernel:
7+
8+
* ``MLIRStringType`` identity / ABI size / dtype / repr / return_as.
9+
* ``typeof()`` of the wrapper used to marshal ``mlir_string`` arrays.
10+
* ``MLIRStringArgHandler.prepare_args`` for the bare ``CPointer`` case and
11+
the ``Tuple(CPointer, mask)`` case used for nullable string columns.
12+
13+
Kernel-level (lowering) tests live in ``test_strings_lowering.py``.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
import numpy as np
19+
from numba_cuda_mlir import types
20+
from numba_cuda_mlir.numba_cuda.typing.typeof import typeof
21+
22+
from cudf.core.udf.mlir_backend.strings_typing import (
23+
ManagedStrArrayWrapper,
24+
MLIRStringArgHandler,
25+
MLIRStringType,
26+
mlir_string,
27+
mlir_string_arg_handler,
28+
)
29+
30+
31+
class _FakeBuffer:
32+
"""Minimal stand-in for a Buffer-like object exposing ``ptr``."""
33+
34+
def __init__(self, ptr):
35+
self.ptr = ptr
36+
37+
38+
def test_extensionty_size_constant():
39+
# 8 (meminfo) + 8 (data) + 8 (nbytes)
40+
assert MLIRStringType._extensionty_size == 24
41+
42+
43+
def test_mlir_string_repr_and_return_as():
44+
assert repr(mlir_string) == "mlir_string"
45+
assert mlir_string.return_as is mlir_string
46+
47+
48+
def test_mlir_string_np_dtype_object():
49+
assert MLIRStringType.np_dtype == np.dtype("O")
50+
51+
52+
def test_mlir_string_is_not_internal():
53+
# Passed to kernels as a CPointer, never as an Array/memref.
54+
assert mlir_string.is_internal is False
55+
56+
57+
def test_typeof_managed_str_array_wrapper():
58+
"""``typeof(ManagedStrArrayWrapper(...))`` -> ``CPointer(mlir_string)``."""
59+
wrapper = ManagedStrArrayWrapper(_FakeBuffer(0xDEAD_BEEF))
60+
ty = typeof(wrapper)
61+
assert isinstance(ty, types.CPointer)
62+
assert ty.dtype is mlir_string
63+
64+
65+
def test_arg_handler_cpointer_mlir_string():
66+
"""Bare ``CPointer(mlir_string)`` is rewritten to ``(uint64, ptr)``."""
67+
handler = MLIRStringArgHandler()
68+
wrapper = ManagedStrArrayWrapper(_FakeBuffer(0xCAFE_F00D))
69+
ty = types.CPointer(mlir_string)
70+
new_ty, new_val = handler.prepare_args(ty, wrapper)
71+
assert new_ty is types.uint64
72+
assert new_val == 0xCAFE_F00D
73+
74+
75+
def test_arg_handler_tuple_with_mask():
76+
"""``Tuple(CPointer(mlir_string), mask)`` -> tuple with the ptr extracted.
77+
78+
This is the shape produced for nullable string columns: the string-view
79+
pointer alongside its validity mask array.
80+
"""
81+
handler = MLIRStringArgHandler()
82+
wrapper = ManagedStrArrayWrapper(_FakeBuffer(0xAA))
83+
mask_value = "<opaque-mask-array>"
84+
tuple_ty = types.Tuple((types.CPointer(mlir_string), types.boolean[::1]))
85+
new_ty, new_val = handler.prepare_args(tuple_ty, (wrapper, mask_value))
86+
assert new_ty is tuple_ty
87+
assert new_val == (0xAA, mask_value)
88+
89+
90+
def test_arg_handler_passthrough_for_non_string_types():
91+
"""Types unrelated to ``mlir_string`` pass through unchanged."""
92+
handler = MLIRStringArgHandler()
93+
plain_ty = types.float64[::1]
94+
sentinel = object()
95+
new_ty, new_val = handler.prepare_args(plain_ty, sentinel)
96+
assert new_ty is plain_ty
97+
assert new_val is sentinel
98+
99+
100+
def test_module_exposes_arg_handler_singleton():
101+
assert isinstance(mlir_string_arg_handler, MLIRStringArgHandler)

0 commit comments

Comments
 (0)