closure-migration

Imported from zharia/Semantic-Computational-Runtime (.agents/skills/closure-migration/SKILL.md). Install upstream with npx skills add zharia/Semantic-Computational-Runtime --skill closure-migration. Copyright stays with the author.

Closure migration

Migrate callers before deleting parametric overloads. Prefer value-taking APIs with unified closures. Pair with mojo-syntax (and mojo-gpu-fundamentals for GPU launch code).

Forbidden

Hard ban — do not under any circumstance add @__parameter / @parameter to a nested closure to persist, introduce, or paper over a legacy closure. Not as a migration bridge, not to satisfy a still-capturing API, not to “borrow imm”, not behind a thin *_value wrapper. That is forbidden.

Hard ban — do not under any circumstance add a new caller of a capturing parametric endpoint (func: def(...) capturing [_] -> None or capturing without thin), or restore a deleted capturing overload, to paper over a typecheck or bind failure. Look at the callee parameter type: thin is a function pointer and is fine (add_function[vec_add], compile_function[kernel]()). capturing and not thin is the legacy closure. Do not route a unified-closure site back through capturing[_].

Hard ban — do not write an empty capture list {}. If the nested def has no free runtime captures, omit the list: def foo() -> Int: not def foo() {} -> Int: and not raises {}. {} is not a spelling of “no captures”; it is noise that reviewers will flag. {imm} is capture-all of outer state, not a stand-in for an empty list.

If a callee still only accepts a comptime capturing[_] function parameter and that overload is not being deleted in this change, use a nested def … capturing -> T without @__parameter when that is what the type requires, or migrate that API. Leave an existing capturing api[fn] site you have not rewritten; do not add new capturing ones. thin function-pointer parameters are not that case. Never put @__parameter on the caller.

@__copy_capture copies into closure storage. Transfer it; do not drop it for {imm}:

# WRONG — copy-capture became an imm borrow
@__copy_capture(gamma, epsilon, input_fn)
def bench_fn(mut b: Bencher) raises:
    ...
# became
def bench_fn(mut b: Bencher) raises {imm}:

# CORRECT — each listed name is a named copy; trailing imm for the rest
def bench_fn(mut b: Bencher) raises {var gamma, var epsilon, var input_fn, imm}:

Put the {var …} list on the same def that had @__copy_capture. Nested children that already used {imm} then borrow those copies. If you hoist kernel_launch out of bench_fn, the {var …} list moves with the names onto the launch — do not hoist and leave {imm} on either def.

{var x} does not convert to a capturing thin callee such as elementwise_compute_lambda_type, and a capture list is not parsed on a capturing def. Until that API is migrated off capturing, form the LayoutTensor on the host from an imm residual_buf parameter and keep @__copy_capture(residual_lt) on def … capturing (no @__parameter). Never call DeviceBuffer methods inside the GPU epilogue.

Target shapes

@__parameter nested defs type as capturing thin and do not convert to a unified FuncType. Call rewrite alone is not enough — change how the closure is declared (and never re-add @__parameter).

Checklist

1. Inventory: rg 'api_name\[' --glob '*.mojo' and rg '@__parameter|@parameter' --glob '*.mojo'
2. Rewrite calls: api[fn](a, b) → api(a, fn, b)
3. On every nested closure in scope: drop @__parameter / @parameter / @__copy_capture; add a capture list. Each @__copy_capture(x) name must become {var x} (or {var} capture-all if that was every capture). Do not replace the decorator with {imm}. Skip comptime params. Self-check the pre-image: every runtime name in @__copy_capture(...) appears as var name in the new list.
4. If a callee still needs a comptime capturing param and that overload is not being deleted → migrate that API first, or leave an existing capturing api[fn] site you have not rewritten. Do not add new capturing api[fn] callers or keep @__parameter on the caller. thin function-pointer parameters (add_function[vec_add], compile_function[kernel]()) are not that case.
5. Delete parametric overloads only after callers typecheck
6. Update skills/docs that still teach the legacy path
7. Typecheck: mojo build --emit llvm <file> -o /tmp/x.ll (filters Metal noise)
8. Offload bind: if the change passes a kernel into DeviceFunction / add_function / compile_info, run kgen -elaborate (below). Do not bt-b200 a compile-time bind assert.
9. Self-check: rg '@__parameter|@parameter' --glob '<touched>.mojo' → zero on nested closures you own
10. NFC self-check: allocation sites, comptime vs Coord(IndexList) layouts, and timed call_fn bodies match the pre-migration code except capture lists and api[fn](…) → api(…, fn, …). No new .as_unsafe_any_origin() / unsafe_origin_cast / DeviceBuffer(…, some.ptr, owning=False) used to dodge aliasing

Verify compile_offload bind without a GPU

compile_offload_closure func must be fully bound is an elaborator assert in Mojo/lib/Elaborator/IREvaluatorContext.cpp (evaluateCompileOffloadClosureAttr). It fires while folding #kgen.compile_offload_closure on CompiledFunctionInfo.populate — the comptime field DeviceFunction and compile_info instantiate. Offload codegen has not started. A CI SIGSEGV during compiling mojo object with that message is a host bind failure, not a GPU or PTX failure.

Do not bt-b200 / bt-mi355 it. Elaboration is host-only.

source ./utils/start-modular.sh
# Matching compiler: bazel-run kgen, not a stale PATH kgen vs std.mojoc
./bazelw run //Mojo/tools/kgen -- -elaborate path/to/file.mojo -o /dev/null

In-tree precedent: kgen -elaborate in KGEN/test/mojo-integration/compile_offload/internal/compile_different_emissionoptions_same_gpu.mojo.

If PATH kgen errors with Mojo precompiled file is incompatible with the current version of the Mojo compiler, the binary is stale relative to .derived packages. Use bazel-run (or ./bazelw run //:install), not a GPU box.

NVIDIA offload on a Metal/CPU host: pin target= on compile_info / DeviceFunction (A100.target(), B200, …), or mojo build --target-accelerator=nvidia:sm_100a (see compile_offload_gpu_cross_compilation.mojo). kgen has no --target-accelerator; that flag is mojo build.

Elaborate the crashing file, or a reduced file that instantiates the same DeviceFunction[F.__call__, …] / CompiledFunctionInfo type. A bound non-generic kernel (Kernels.vec_add) is not a substitute for the specialization that failed.

A file-scope @no_inline wrapper that binds leftover kernel comptime parameters fixes DeviceFunction[func, declared_arg_types] when func is that wrapper (compile_function[kernel]() / add_function[kernel]). It does not bind DeviceFunction[F.__call__, TypeList.of[T0,…]()]. F.__call__ is _PtrWrapper::__call__[AnyType, …] even when the runtime value is already a fully specialized thin kernel. For a thin kernel, use the thin endpoint (add_function[kernel](*args) / compile_function[kernel]()). DeviceGraphBuilder.add_function is thin-only; capturing kernels use enqueue_function / recording_context(). Do not paper a capturing site over with capturing[_] / @__parameter.

Capture choice

Do not use capture-all {mut} when the closure also mentions register-passable outer values (Int, indices, lengths). Mix an explicit mut name with a trailing default imm ({mut tt_in, imm}). At most one bare convention (imm / mut / var) may appear as the default for unlisted captures.

{imm} freezes captured buffers. That is wrong whenever the body:

• writes an output (TileTensor / .unsafe_ptr() where mut=True is required), or
• calls a method that internally does self._buf.unsafe_ptr() then unsafe_mut_cast[True]() — CacheBustingBuffer.offset_ptr is the usual case. The method may be declared self (imm); the captured field still fails to match (lit.call argument #0). Mut-capture every such buffer ({mut cb_a, mut cb_b, mut cb_c, mut cb_a_scales, mut cb_b_scales, imm}).

Do not “fix” this with extra unsafe_mut_cast (the method already has that) or by restoring @__parameter. Only skip mut when that would create an aliasing pair with another captured field of the same origin (see below).

Aliasing

When several captured / nested values reference the same origin and one of them is mutable, the closure struct would contain aliasing fields. Two captures that alias the same memory are permitted only when both are read-only.

@__copy_capture(n) on a register-passable local (Int index, length, row count) becomes {var n} — a copy into closure storage. {imm} of var n = … is a reference to that mutable local and aliases once a unified sibling embeds origin_of(n). Use {var n}. Do not delete n or wrap the function just for that integer.

Prefer, in order:

1. Capture as read (imm) whenever the body does not actually mutate that origin. Drop {mut …}; pass an imm parameter. enqueue_memset takes an imm DeviceBuffer. A helper that only reads config.rank_units must not mut-capture config. Prefer this also when an outer var buf stays mutable while a nested unified closure would capture a view of buf: call a normal function (file-scope or otherwise non-capturing) that takes buf: DeviceBuffer[…] as an imm parameter, and form the launch there.
2. Disassemble the instance so origin tags are finer-grained. If the compiler treats two uses as overlapping because they capture a whole object (a config, a struct, a List of buffers) but the mutation does not actually touch the other use's memory, pull out the fields you need (e.g. hoist rank_units / rank_unit_start into Array[Int, ngpus]) and capture or pass those instead of the parent.
3. Pass the mutable origin as an argument (mut buf: … in the parameter list) so it is not a captured field. Lift only the capture that must be mut and still aliases another use — read-only captures can stay captures. When the value-taking API's FuncType is fixed, use a thin {mut buf, imm} adapter that only forwards buf into the all-imm body. Prefer widening the API when practical.

Do not “fix” aliasing by erasing the origin. That includes .as_unsafe_any_origin(), unsafe_origin_cast[MutUntrackedOrigin] / MutAnyOrigin on the original buffer, and wrapping an already-untracked pointer (EPLocalSyncCounters.ptr, offset_ptr result) in a new DeviceBuffer(ctx, ptr, …, owning=False) just to memset or launch. Pass the original buffer as an argument instead. Do not “fix” anything with @__parameter.

Safety rules

• Key bulk edits off the value-argument name, not every def with that name
• Do not bulk-replace capture lists (destroys {mut count} etc.)
• name[i] = inside with … as name is a local, not an outer mut capture
• Zero @__parameter / @parameter on nested closures in migrated code
• Lift a capture to an argument only for origin exclusivity (or a normal function with an imm DeviceBuffer parameter)
• Stay NFC: do not change allocation lifetime, comptime vs dynamic layouts, or host work on the timed path. comptime if arm locals (buffers, host copies, TileTensors, row_major[M, N]() layouts) stay in that arm — a size-1 alloc at function scope is a semantic change. Layouts built once (a_shape) stay outside call_fn

More detail

Full step-by-step process, error catalog, and verification: process.md.