Skip to content

Commit 8315bc9

Browse files
Add in-development MLIR/NRT backend section to the UDF string memory guide
The shim-based backend documented here is still the one in place, so leave it unchanged. Append a clearly-marked 'in development' section that uses the existing legacy explanation as a lens and describes only the deltas for the pure-MLIR backend: the flat mlir_string {meminfo, data, nbytes} representation (no cudf::udf_string), creation via allocate_mlir_string / NRT_MemInfo_new_varsize instead of C++ shims, teardown via free_mlir_string_array / _mlir_string_array_to_column, plus borrowed-vs-owned, enabling NRT, and nullable (Masked) strings. The NRT model and the str1 + str2 lifecycle are unchanged.
1 parent 9033943 commit 8315bc9

1 file changed

Lines changed: 121 additions & 0 deletions

File tree

‎docs/cudf/source/cudf/developer_guide/udf_memory_management.md‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -400,3 +400,124 @@ the effect of freeing the entire `mi_str_allocation`.
400400
- **Stack**: Original `ManagedUDFString` becomes invalid/out-of-scope
401401
- **Reference Count**: N/A (object destroyed)
402402
- **cuDF** A `cudf::column` of string type containing the result of the UDF
403+
404+
405+
## The MLIR backend (in development): what changes
406+
407+
```{note}
408+
Everything above describes the **current, shipping** string UDF backend: C++
409+
shim functions plus an owning ``cudf::udf_string`` class. A reimplementation is
410+
in progress that lowers string operations **entirely in MLIR** and removes the
411+
C++ shim layer. This section does **not** supersede the material above --- it
412+
uses it as a lens. The NRT model itself is unchanged; only the representation of
413+
a device string and the mechanism that creates and frees it change. When the
414+
MLIR backend becomes the default, this section and the sections above will be
415+
reconciled into one.
416+
```
417+
418+
### What stays exactly the same
419+
420+
Re-read ["Numba memory management and the Numba Runtime (NRT)"](#numba-memory-management-and-the-numba-runtime-nrt), the ``MemInfo`` struct, ``NRT_decref``, the three [NRT Requirements](#nrt-requirements), and the five-phase
421+
``str1 + str2`` [lifecycle](#string-lifecycle-details). **All of it carries over
422+
verbatim.** NRT is still the memory manager; a reference-counted type still
423+
declares `has_nrt_meminfo` and exposes `get_nrt_meminfo`; operators that produce
424+
a string still initialize a `MemInfo` and place it where numba expects; the
425+
destructor still runs when `refct` hits zero. If you understand the model above,
426+
you understand the MLIR backend --- what follows is only the delta.
427+
428+
### What changes (1): the string representation
429+
430+
The owning C++ container goes away. Instead of ``ManagedUDFString`` wrapping a
431+
``cudf::udf_string``, the MLIR backend uses a flat value called ``mlir_string``:
432+
433+
```text
434+
legacy: ManagedUDFString { void* meminfo, cudf::udf_string } # owns a C++ class
435+
MLIR: mlir_string { ptr meminfo, ptr data, i64 nbytes } # raw UTF-8 bytes, 24 bytes
436+
```
437+
438+
There is no ``cudf::udf_string`` and no ``string_view`` class --- ``data`` is just
439+
a pointer to UTF-8 bytes and ``nbytes`` its length. NRT requirement (2) is
440+
satisfied even more directly than before: the ``MemInfo`` *is* field 0 of the
441+
value, so `get_nrt_meminfo` is a single `extractvalue`:
442+
443+
```python
444+
class MLIRStringModel(PrimitiveModel):
445+
_fields = ("meminfo", "data", "nbytes")
446+
def has_nrt_meminfo(self): return True
447+
def get_nrt_meminfo(self, value):
448+
return llvm.extractvalue(llvm.PointerType.get(), value, [0])
449+
```
450+
451+
One concept the flat layout makes explicit --- **borrowed vs. owned**:
452+
453+
- **Borrowed** (``meminfo == null``): ``data`` points into an existing libcudf
454+
column buffer. Every *input* string is borrowed; `NRT_decref` on a null
455+
`MemInfo` is a no-op, so there is nothing to free.
456+
- **Owned** (``meminfo != null``): ``data`` was allocated by NRT; this is every
457+
string *produced* by an operation.
458+
459+
### What changes (2): string creation (Phase 2)
460+
461+
This is the only phase of the lifecycle that changes. The legacy Phase 2 calls a
462+
C++ shim (``concat_shim``) that constructs a ``cudf::udf_string`` via
463+
placement-new and wraps it with ``make_meminfo_for_new_udf_string``. The MLIR
464+
backend does it all in MLIR through a single allocation helper:
465+
466+
```python
467+
def allocate_mlir_string(gpu_module, nbytes_i64):
468+
# NRT_MemInfo_new_varsize installs _nrt_varsize_dtor (frees the buffer at
469+
# refct 0, or a no-op under the arena); at least 1 byte to avoid malloc(0).
470+
alloc_size = arith.maxui(nbytes_i64, _const_i64(1))
471+
mi = _call_nrt(gpu_module, "NRT_MemInfo_new_varsize", [i64], [ptr], [alloc_size])
472+
data = _call_nrt(gpu_module, "NRT_MemInfo_data_fast", [ptr], [ptr], [mi])
473+
return mi, data
474+
```
475+
476+
`concat` then becomes allocate-once, copy, assemble --- no shim, no C++ class:
477+
478+
```python
479+
def lower_concat(gpu_module, lhs_view, rhs_view):
480+
total = arith.addi(lhs_nbytes, rhs_nbytes)
481+
mi, data = allocate_mlir_string(gpu_module, total) # owning MemInfo in field 0
482+
llvm.intr_memcpy(data, lhs_data, lhs_nbytes, False)
483+
llvm.intr_memcpy(_gep_byte_offset(data, lhs_nbytes), rhs_data, rhs_nbytes, False)
484+
return build_mlir_string(mi, data, total)
485+
```
486+
487+
The ``MemInfo`` returned by ``NRT_MemInfo_new_varsize`` is exactly the one that
488+
goes in field 0, satisfying NRT requirement (3). Everything the legacy Phase 2
489+
said about reference count 1 on creation, and Phases 3-4 about incref on
490+
assignment/return, is unchanged.
491+
492+
### What changes (3): teardown and column conversion (Phase 5)
493+
494+
The shapes match the legacy Phase 5; only the names and the per-element struct
495+
change:
496+
497+
| Legacy | MLIR backend |
498+
|---|---|
499+
| output array of ``ManagedUDFString`` | output ``rmm.DeviceBuffer`` of ``n`` × 24-byte ``mlir_string`` |
500+
| ``free_managed_udf_string_array`` (final `NRT_decref`) | ``free_mlir_string_array`` (final `NRT_decref` per entry) |
501+
| ``column_from_managed_udf_string_array`` | ``_mlir_string_array_to_column`` |
502+
503+
For an **owned** result the final `NRT_decref` drops `refct` to zero and runs
504+
`_nrt_varsize_dtor` (a no-op under the arena); for a **borrowed** result
505+
(`meminfo == null`, e.g. a UDF that returns its input unchanged) it is a no-op.
506+
`_mlir_string_array_to_column` copies each ``{data, nbytes}`` into a new libcudf
507+
strings column, so the column does not depend on any UDF-owned allocation after
508+
the freeing kernel runs.
509+
510+
### New concept: enabling NRT
511+
512+
Reference counting is only woven in when a UDF actually uses a heap-managed type.
513+
``CaptureNRTUsage`` detects this during a trial compile and ``nrt_enabled()``
514+
toggles ``numba_config.CUDA_ENABLE_NRT`` for the real compile (and the freeing
515+
kernel), so purely numeric UDFs pay none of this cost.
516+
517+
### New concept: nullable strings (``Masked``)
518+
519+
A nullable element is typed ``Masked(mlir_string)``. The refcount rules extend
520+
with no special-casing: casting into/out of ``Masked`` and reading ``.value``
521+
each insert an `incref` at the ownership-transfer point exactly as for the
522+
unmasked type, and an **invalid** element has a null ``MemInfo`` --- so "null" and
523+
"nothing allocated" coincide and every incref/decref on it is a no-op.

0 commit comments

Comments
 (0)