Skip to content

Commit 9d7bee6

Browse files
brandon-b-millerbrandon-b-miller
committed
Add review-udf-mlir-migration agent skill; drop in-development dev-doc section
Add a project agent skill (.agents/skills/review-udf-mlir-migration/) that teaches a reviewer's agent to review the cuDF UDF numba-cuda -> numba-cuda-mlir migration from first principles: SKILL.md plus FOUNDATIONS, ARCHITECTURE, HISTORY, and REVIEW_PLAYBOOK under references/ (matching the repo's .agents/skills convention, e.g. build-test-cudf-java). The in-development MLIR/NRT section previously appended to the UDF string memory-management developer guide is removed; that knowledge now lives in the skill, and the legacy developer guide is left unchanged since the legacy backend is still the one in place.
1 parent 6967131 commit 9d7bee6

6 files changed

Lines changed: 930 additions & 117 deletions

File tree

‎.agents/skills/review-udf-mlir-migration/SKILL.md‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
---
2+
name: review-udf-mlir-migration
3+
description: >-
4+
Reviews and explains cuDF UDF migration changes from numba-cuda to
5+
numba-cuda-mlir from first principles. Use for PRs or code touching
6+
cudf.core.udf, mlir_backend, MaskedType, NAType, mlir_string, string UDFs,
7+
NRT, row/groupby/rolling apply kernels, CooperativeArray, compiler extension
8+
typing/lowering/models, or the final MLIR UDF cutover.
9+
---
10+
11+
# Review the cuDF UDF MLIR migration
12+
13+
This migration preserves cuDF's Python UDF features while replacing the
14+
compiler substrate beneath them. It is not an import rename. cuDF is an unusually
15+
deep extension-API consumer: its nullable values, strings, group operations,
16+
kernel wrappers, and memory ownership are compiler extensions.
17+
18+
## Required review posture
19+
20+
Explain before judging. A cuDF reviewer should be able to understand:
21+
22+
1. the user-facing operation this code enables;
23+
2. the old execution path;
24+
3. the new execution path;
25+
4. why this PR is a separate layer;
26+
5. the correctness contracts at that layer;
27+
6. what the tests prove and do not prove.
28+
29+
Do not begin with MLIR syntax. Begin with a small UDF or dataframe operation and
30+
follow it down to the changed code.
31+
32+
Current code, branch ancestry, and the PR diff are authoritative. Historical
33+
branches and this skill explain intent; they do not override current source.
34+
35+
## Read the references
36+
37+
Always read:
38+
39+
- [FOUNDATIONS.md](references/FOUNDATIONS.md): first-principles mental model.
40+
- [REVIEW_PLAYBOOK.md](references/REVIEW_PLAYBOOK.md): investigation and report
41+
method.
42+
43+
Read as needed:
44+
45+
- [ARCHITECTURE.md](references/ARCHITECTURE.md): row, Masked/NA, strings/NRT, groupby,
46+
registration, ABI, and tests.
47+
- [HISTORY.md](references/HISTORY.md): why the work exists and why the PR sequence has its
48+
present shape.
49+
50+
## Start every review by locating the layer
51+
52+
Inspect the PR base and diff; do not infer scope from its title. Classify each
53+
change into one or more layers:
54+
55+
1. dependency and availability plumbing;
56+
2. extension type and physical data model;
57+
3. typing and type unification;
58+
4. lowering and casts;
59+
5. kernel ABI and host argument marshalling;
60+
6. ownership, NRT, and result cleanup;
61+
7. row/groupby/rolling kernel construction;
62+
8. production dispatch/cutover;
63+
9. removal of the legacy implementation.
64+
65+
Then identify the feature slice: Masked/NA, read-only string, allocating string,
66+
groupby/cooperative array, rolling, or generic compiler integration.
67+
68+
## Reconstruct the vertical path
69+
70+
Trace the changed feature through all applicable stages:
71+
72+
```text
73+
cuDF API call
74+
-> generated kernel / device UDF
75+
-> numba typing and extension registry
76+
-> extension data model
77+
-> MLIR lowering and linked device functions
78+
-> kernel argument/result ABI
79+
-> launch
80+
-> cuDF column + validity mask
81+
-> cleanup of temporary owned values
82+
```
83+
84+
When one stage changes, search every producer and consumer of that representation.
85+
A type declaration without a matching model, lowering, marshaller, or cleanup
86+
path is incomplete even if its unit test passes.
87+
88+
## Core review invariants
89+
90+
### Backend identity
91+
92+
- Types, signatures, registries, and models must come from one
93+
`numba_cuda_mlir`/vendored-Numba universe.
94+
- Equivalent-looking classes imported from stock Numba can fail dispatch by
95+
identity.
96+
- Registration must happen before compilation and against the registry the
97+
active compiler actually uses.
98+
99+
### Type and model agreement
100+
101+
- Logical type, MLIR representation, byte size, field order, alignment,
102+
argument handler, and output marshaller must agree.
103+
- Ask the active model manager for a type where possible; do not independently
104+
recreate an identified struct and accidentally create a distinct `.1` type.
105+
- An extension struct crossing a kernel boundary is normally pointer-marshalled,
106+
not represented as a memref/NumPy array element.
107+
108+
### Null semantics
109+
110+
- `Masked(value, valid)` carries a payload and validity bit.
111+
- Invalid payload contents are unspecified. Tests must not assert them.
112+
- Operator typing resolves the payload operation; lowering propagates validity
113+
according to the operation.
114+
- `cudf.NA`, branch unification, casts, and `pack_return` are part of the same
115+
semantic chain.
116+
117+
### Ownership
118+
119+
- Determine whether every value is borrowed or owned.
120+
- A borrowed string/group view must not be freed and must not outlive its source.
121+
- An owned value must expose its `MemInfo`, survive any escape into an output
122+
slot, and be decref'd exactly once per owned reference.
123+
- Balanced allocation/free statistics do not prove absence of use-after-free.
124+
Use behavioral checks and Compute Sanitizer where ownership changes.
125+
126+
### Cooperative execution
127+
128+
- Row apply is conceptually one thread per row.
129+
- Groupby apply is conceptually one block per group.
130+
- A block-shared allocation must have one allocator/owner, a pointer broadcast,
131+
and synchronization that prevents free-before-last-consumer.
132+
- Generic per-thread NRT reasoning does not by itself prove a cooperative
133+
operation safe.
134+
135+
### Test scope
136+
137+
- A typing test proves resolution, not representation or execution.
138+
- A tiny kernel proves a registration/lowering slice, not cuDF marshalling.
139+
- An integration test proves the path it reaches; verify it actually reaches the
140+
new backend.
141+
- New tests should fail on the PR base for the intended reason.
142+
- Test borrowed and owned cases, literal and runtime validity, nullable inputs,
143+
and unsupported types where applicable.
144+
145+
## Preserve the row/groupby migration strategy
146+
147+
The row/groupby feature PRs intentionally stay outside their production apply
148+
paths until one final cutover. Intermediate tests may compile focused kernels or
149+
monkeypatch the minimum pipeline seam. Do not request a partial production
150+
cutover merely to make a feature PR look end-to-end.
151+
152+
Rolling UDFs have a separate numba-cuda-mlir production path and must be
153+
reviewed as their own vertical slice.
154+
155+
Do not accept a final row/groupby cutover based only on focused extension tests.
156+
It must exercise the real cuDF API, marshalling, launch, output construction,
157+
cleanup, and legacy removal.
158+
159+
## Required review response
160+
161+
Lead with a plain-language verdict, then use this order:
162+
163+
1. **What users can do** — a minimal UDF/API example.
164+
2. **Where this PR sits** — feature and layer in the migration.
165+
3. **How it works** — old path versus new path, from type to launched kernel.
166+
4. **Why it matters** — prerequisite unlocked or legacy dependency removed.
167+
5. **Correctness contracts** — the specific invariants above that apply.
168+
6. **Findings** — concrete bugs or risks, ordered by severity, with paths/lines.
169+
7. **Verification** — what was run, base-versus-head behavior, and limitations.
170+
171+
Do not confuse unfamiliarity with a defect. Do not approve merely because the
172+
new code resembles the answer-key branch. Verify semantics against tests and
173+
the current compiler API.

‎.agents/skills/review-udf-mlir-migration/references/ARCHITECTURE.md‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
# Architecture map
2+
3+
Paths and symbols evolve. Search current source before relying on a name below.
4+
5+
## Extension registration
6+
7+
The MLIR backend's recurring shape is:
8+
9+
- a frontend type in a `*_typing.py` module;
10+
- typing templates registered with `typing_registry`;
11+
- a physical model registered with `register_model`;
12+
- operation and cast implementations registered with `lowering_registry` and
13+
`lower_cast`;
14+
- import-time registration, either through decorators/direct calls or through
15+
a module `_register()` helper.
16+
17+
Use types and templates from `numba_cuda_mlir` and its vendored Numba namespace.
18+
Do not mix stock Numba classes into signatures or registrations.
19+
20+
## Masked and NA
21+
22+
Primary modules:
23+
24+
- `python/cudf/cudf/core/udf/api.py`
25+
- `python/cudf/cudf/core/udf/mlir_backend/masked_typing.py`
26+
- `python/cudf/cudf/core/udf/mlir_backend/masked_lowering.py`
27+
28+
Logical model:
29+
30+
```text
31+
MaskedType(T) -> { T value, i1 valid }
32+
NAType -> payload-free sentinel
33+
```
34+
35+
Review together:
36+
37+
- constructor typing and lowering;
38+
- `.value`/`.valid` attribute typing and extraction;
39+
- `NAType` and `MaskedType.unify`;
40+
- scalar/NA/Masked casts;
41+
- binary and unary result typing;
42+
- validity propagation;
43+
- datetime/timedelta unit-aware delegation;
44+
- truth semantics;
45+
- membership;
46+
- `pack_return`.
47+
48+
Important details:
49+
50+
- `Masked op Masked`: validity is normally `lhs.valid AND rhs.valid`.
51+
- `Masked op scalar`: preserve Masked validity.
52+
- an invalid payload is poison/undefined and must not be inspected by tests;
53+
- comparison returns a Masked boolean when nullability remains;
54+
- truth testing returns a plain boolean according to the API's invalid-value
55+
semantics;
56+
- datetime-like operations must reuse unit-aware scalar lowering rather than
57+
blindly applying an integer operation to raw payloads.
58+
59+
## Strings
60+
61+
Primary modules:
62+
63+
- `mlir_backend/strings_typing.py`
64+
- `mlir_backend/strings_lowering.py`
65+
- `mlir_backend/string_lowering_impl.py`
66+
- host-side apply/marshalling utilities found by searching for
67+
`ManagedStrArrayWrapper`, `MLIRStringArgHandler`, or `mlir_string`.
68+
69+
Current incremental design:
70+
71+
```text
72+
mlir_string -> { ptr meminfo, ptr data, i64 nbytes }
73+
```
74+
75+
`MLIRStringType` is not a NumPy-internal scalar and is pointer-marshalled to
76+
kernels. `ManagedStrArrayWrapper` preserves the intended extension type instead
77+
of allowing CUDA Array Interface inference to reduce it to a byte array.
78+
79+
Review strings by operation class:
80+
81+
### Borrowed/read-only
82+
83+
Examples: `len`, comparisons, `contains`, `find`, `rfind`, `count`,
84+
`startswith`, `endswith`, and character predicates.
85+
86+
- input `meminfo` may be null;
87+
- no allocation should be introduced;
88+
- literals should become views without unnecessary owning-string allocation;
89+
- nullable variants must preserve/AND validity;
90+
- `len` is a character count, not blindly `nbytes`; test multi-byte UTF-8.
91+
92+
### Allocating/string-producing
93+
94+
Examples: concatenation, case conversion, replacement, and strip variants.
95+
96+
- allocate variable-sized storage;
97+
- construct a non-null owner;
98+
- preserve returned ownership across output assignment;
99+
- release non-escaping temporaries;
100+
- materialize or transfer final results into a cuDF column;
101+
- release temporary result owners afterward.
102+
103+
### Legacy production representation
104+
105+
The first port retained `StringView`, `UDFString`, and `ManagedUDFString` plus
106+
C++ shims. These types remain active in the legacy row/groupby backend while
107+
the incremental MLIR design replaces them with `mlir_string`. The project
108+
deliberately stopped making the new-backend PR sequence teach this three-type
109+
design only to delete it later. Do not mistake "legacy in the MLIR design" for
110+
"already removed from production."
111+
112+
## Row apply
113+
114+
Find the current row pipeline from:
115+
116+
- `core/indexed_frame.py`;
117+
- `core/udf/row_function.py` and/or `mlir_backend/row_function.py`;
118+
- `scalar_function.py`, `templates.py`, and `udf_kernel_base.py`.
119+
120+
Conceptual flow:
121+
122+
1. infer column and UDF types;
123+
2. build a row record/schema;
124+
3. marshal data and mask arguments;
125+
4. generate a kernel that loads one row per logical thread;
126+
5. create Masked/string extension values;
127+
6. inline or call the user function;
128+
7. normalize the return;
129+
8. scatter payload and validity;
130+
9. build the output column and clean managed temporaries.
131+
132+
Check both type-inference compilation and execution compilation. Transitional
133+
branches sometimes used different compilers for these phases.
134+
135+
## Groupby apply and cooperative arrays
136+
137+
Find current code by searching:
138+
139+
- `groupby_function`, `groupby_typing`, and `groupby_lowering`;
140+
- `CooperativeArray`;
141+
- grouped apply templates and group slot globals.
142+
143+
Conceptual flow:
144+
145+
1. grouped rows are contiguous and described by offsets;
146+
2. one CUDA block handles one group;
147+
3. `group["x"]` creates a borrowed cooperative view;
148+
4. elementwise operations may allocate a managed block-shared result;
149+
5. reductions use block-wide shuffle/shared-memory/atomic algorithms;
150+
6. one thread writes a block-uniform result.
151+
152+
Review:
153+
154+
- allocation by one thread;
155+
- shared pointer and owner broadcast;
156+
- barriers before consumers and before lifetime end;
157+
- empty-group semantics;
158+
- shared-memory sizing, especially multi-stage reductions such as variance/std;
159+
- dtype promotion and reduction accumulator type;
160+
- ownership of intermediate arrays;
161+
- whether compiler releases include the required cooperative/linking support.
162+
163+
Groupby NRT and arena allocation have historically been later than basic
164+
cooperative reductions. Confirm landed state rather than assuming every
165+
prototype feature is production.
166+
167+
## Rolling apply
168+
169+
Rolling UDFs have a production numba-cuda-mlir path in `rolling_utils.py`.
170+
Treat it as its own vertical slice. Review:
171+
172+
- return-type inference and execution compilation both use the intended backend;
173+
- window start/stop bounds and indexing;
174+
- `min_periods` and output-validity construction;
175+
- empty-input behavior;
176+
- kernel cache keys and reuse;
177+
- dtype/return-type conversion and output allocation.
178+
- preservation of the current public rejection of nullable inputs with
179+
`NotImplementedError`, unless the PR deliberately changes and tests that API.
180+
181+
Do not infer that row/groupby support proves rolling support, or vice versa.
182+
183+
## Compilation and linking
184+
185+
Search current numba-cuda-mlir for the APIs actually used by the PR:
186+
187+
- extension typing/lowering/model registries;
188+
- compiler entry points and requested output, such as
189+
`compiler.compile(..., output="ltoir")`;
190+
- library/link inputs and post-load callbacks;
191+
- NRT function emission or object linking;
192+
- runtime setup globals;
193+
- argument handler hooks.
194+
195+
Review symbol ownership:
196+
197+
- every required device symbol is defined;
198+
- no runtime is linked twice;
199+
- destructors remain callable across module boundaries;
200+
- post-load setup targets the loaded module/library;
201+
- architecture/toolkit support matches cuDF's supported matrix.
202+
203+
## Test layers
204+
205+
Typical focused tests live under:
206+
207+
`python/cudf/cudf/tests/private_objects/mlir_backend/`
208+
209+
Interpret them by layer:
210+
211+
- `test_*_typing.py`: signatures, attributes, unification, rejection;
212+
- `test_*_lowering.py`: direct tiny kernels and runtime values;
213+
- marshalling helpers: ABI and host/device boundary;
214+
- dataframe/series/groupby apply tests: production integration;
215+
- NRT stats tests: allocation/refcount accounting;
216+
- Compute Sanitizer: illegal access/use-after-free evidence.
217+
218+
For focused tests, verify the intended extension registration is loaded. A
219+
generic template can accidentally steal an operation while leaving the test
220+
green.

0 commit comments

Comments
 (0)