closure-migration

闭包迁移

在删除参数化重载之前迁移调用方。优先使用接受值并支持统一闭包的 API。搭配 mojo-syntax（GPU 启动代码还需搭配 mojo-gpu-fundamentals）。

禁止事项

严格禁止——任何情况下都不得在嵌套闭包上添加 @__parameter / @parameter，以保留、引入或掩盖旧式闭包。 不得将其用作迁移桥接，不得为了满足仍在捕获的 API 而添加，不得为了“借用 imm”而添加，也不得将其放在精简的 *_value 包装器后面。此举一律禁止。

严格禁止——任何情况下都不得新增对捕获型参数化端点的调用（func: def(...) capturing [_] -> None 或不带 thin 的 capturing），也不得恢复已删除的捕获型重载，以掩盖类型检查或绑定失败。请查看被调用方的参数类型：thin 是函数指针，没有问题（add_function[vec_add]、 compile_function[kernel]()）。capturing 且不带 thin 才是旧式闭包。不得将统一闭包调用点重新转回 capturing[_]。

严格禁止——不得编写空捕获列表 {}。 如果嵌套 def 没有自由运行时捕获，则省略列表：def foo() -> Int:，而不是 def foo() {} -> Int:，也不是 raises {}。{} 并不表示“没有捕获”；它只是多余内容，评审者会指出这一点。{imm} 表示捕获外部所有状态，不能代替空列表。

不要	应改为
在嵌套 def 上使用 @__parameter / @parameter	使用统一的 def … {imm}: / {mut x, imm}: / 命名捕获
使用 @__parameter 函数体 + *_value 转发器	将函数体改为统一闭包；按值传递
使用 @__parameter“imm 借用”辅助函数	使用文件作用域函数 / 带 imm 参数的普通函数
删除 {imm} 代替原有的 @__copy_capture(x, y)	使用 {var x, var y, imm}（复制每个列出的名称）。只有在每个捕获都通过复制捕获时，才使用 {var} 全量捕获。{var^} / {var^ x} 表示移动，不是复制
因为 API 仍是 capturing[_] / comptime fn，就保留 @__parameter	迁移或扩展该 API，使其接受值；不要掩盖问题
将 MutUntrackedOrigin / MutAnyOrigin 指针包装进新的 DeviceBuffer，以便 memset 或启动“不产生别名”	将原始 DeviceBuffer 作为 imm 参数传入（enqueue_memset 接受 imm）
将 comptime if 分支缓冲区 / TileTensor 提到函数作用域（在其他分支中使用大小为 1 的占位符）	将它们留在分支中；在这些局部变量旁定义统一闭包
在计时的 call_fn 内重建 a_shape / layouts	在辅助函数体中构建一次；使用 {imm} 捕获
新增 api[fn](args)，而被调用方参数是 capturing 且不带 thin	保留按值传递的统一闭包。thin 函数指针形式的 api[fn] 可以使用
空 {}（def f() {} -> T / raises {}）	省略捕获列表。只有在确实要捕获外部状态时才使用 {imm}

如果被调用方仍然只接受 comptime capturing[_] 函数参数，且此次更改不会删除该重载，则在类型要求时，使用不带 @__parameter 的嵌套 def … capturing -> T，或迁移该 API。对于尚未重写的现有捕获型 api[fn] 调用点，予以保留；不要新增此类调用点。thin 函数指针参数不属于这种情况。切勿在调用方上使用 @__parameter。

@__copy_capture 会将值复制到闭包存储中。请保留这一语义；不要将其改为 {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}:

装饰器 / 意图	捕获方式
@__copy_capture(x, y)	{var x, var y, …}（复制）
每个捕获都通过复制捕获	可以使用 {var} 全量捕获
移动到存储中（消耗）	{var^ x} / {var^}——不是复制捕获的对应写法
旧列表中的 comptime 参数	省略它（value X is a parameter and does not need a capture convention）
将 {var …} 列表放在同一个 def 上，也就是原先带有 @__copy_capture 的那个 def。已经使用 {imm} 的嵌套子项会借用这些副本。如果将 kernel_launch 移出 bench_fn，{var …} 列表也要连同名称一起移到 launch 上——不要只移出而在任一 def 上保留 {imm}。

{var x} 不会转换为 capturing thin 被调函数，例如 elementwise_compute_lambda_type；而且 capturing def 上不会解析捕获列表。在该 API 迁出 capturing 之前，请在主机端从 imm residual_buf 参数构造 LayoutTensor，并在 def … capturing 上保留 @__copy_capture(residual_lt)（不要加 @__parameter）。绝不要在 GPU epilogue 内调用 DeviceBuffer 方法。

目标形式

旧形式	首选形式
api[fn](a, b)	api(a, fn, b)
Comptime 参数 fn: def(...) raises capturing[_] -> None	FuncType: def(...) raises -> None + runtime ref func: FuncType（或 func: FuncType）
嵌套的 @__parameter / @parameter def	统一的 def ...(…) raises {imm}: / {mut buf, imm}: / 命名捕获

嵌套的 @__parameter def 会被定型为 capturing thin，不会转换为统一的 FuncType。仅重写调用还不够——必须更改闭包的声明方式（且绝不要重新添加 @__parameter）。

检查清单

1. 清点：rg 'api_name\[' --glob '*.mojo' 和 rg '@__parameter|@parameter' --glob '*.mojo'
2. 重写调用：api[fn](a, b) → api(a, fn, b)
3. 对范围内的每个嵌套闭包：移除 @__parameter / @parameter / @__copy_capture；添加捕获列表。每个 @__copy_capture(x) 名称都必须变为 {var x}（或者在捕获所有变量时使用 {var}）。不要将该装饰器替换为 {imm}。跳过 comptime 参数。检查修改前代码：@__copy_capture(...) 中的每个 runtime 名称，都应在新列表中以 var name 的形式出现。
4. 如果被调函数仍需要 comptime capturing 参数，且不会删除该重载 → 先迁移该 API，或保留尚未重写的现有 capturing api[fn] 调用点。不要新增 capturing api[fn] 调用方，也不要在调用方保留 @__parameter。thin 函数指针参数（add_function[vec_add]、compile_function[kernel]()）不属于这种情况。
5. 只有在调用方通过类型检查后，才删除参数化重载
6. 更新仍在讲解旧路径的 skills/docs
7. 类型检查：mojo build --emit llvm <file> -o /tmp/x.ll（过滤 Metal 噪声）
8. 卸载绑定：如果更改将 kernel 传入 DeviceFunction / add_function / compile_info，请运行 kgen -elaborate（见下文）。不要对编译期绑定断言运行 bt-b200。
9. 自检：rg '@__parameter|@parameter' --glob '<touched>.mojo' → 你负责的嵌套闭包中应为零
10. NFC 自检：分配点、comptime 与 Coord(IndexList) 布局，以及计时的 call_fn 主体，都应与迁移前的代码一致，只有捕获列表和 api[fn](…) → api(…, fn, …) 除外。不得新增用于规避别名问题的 .as_unsafe_any_origin() / unsafe_origin_cast / DeviceBuffer(…, some.ptr, owning=False)

在没有 GPU 的情况下验证 compile_offload 绑定

compile_offload_closure func must be fully bound 是 Mojo/lib/Elaborator/IREvaluatorContext.cpp（evaluateCompileOffloadClosureAttr）中的 elaborator 断言。 在对 CompiledFunctionInfo.populate 执行折叠时触发，comptime 字段 DeviceFunction 和 compile_info 会在此实例化。Offload 代码生成尚未开始。CI 在 compiling mojo object 期间出现带有此消息的 SIGSEGV，表示主机绑定失败，而非 GPU 或 PTX 故障。

不要在 bt-b200 / bt-mi355 上调试。Elaboration 仅在主机上进行。

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

项目内的先例：KGEN/test/mojo-integration/compile_offload/internal/compile_different_emissionoptions_same_gpu.mojo 中的 kgen -elaborate。

如果 PATH 中的 kgen 报错 Mojo precompiled file is incompatible with the current version of the Mojo compiler，说明该二进制文件相对于 .derived 软件包已过时。请使用 bazel-run（或 ./bazelw run //:install），不要使用 GPU 机器。

在 Metal/CPU 主机上执行 NVIDIA offload：在 compile_info / DeviceFunction 上固定 target=（A100.target()、B200 等），或使用 mojo build --target-accelerator=nvidia:sm_100a（参见 compile_offload_gpu_cross_compilation.mojo）。kgen 没有 --target-accelerator；该标志属于 mojo build。

对崩溃的文件进行 elaboration，或使用一个精简文件实例化相同的 DeviceFunction[F.__call__, …] / CompiledFunctionInfo 类型。已绑定的 非泛型内核（Kernels.vec_add）不能替代失败的特化。

文件作用域的 @no_inline 包装器可以绑定剩余的内核 comptime 参数；当 func 是该包装器时（compile_function[kernel]() / add_function[kernel]），可修复 DeviceFunction[func, declared_arg_types]。 它不会绑定 DeviceFunction[F.__call__, TypeList.of[T0,…]()]。 即使运行时值已经是完全特化的薄内核，F.__call__ 仍然是 _PtrWrapper::__call__[AnyType, …]。对于薄内核，请使用薄端点（add_function[kernel](*args) / compile_function[kernel]()）。DeviceGraphBuilder.add_function 仅支持薄内核；捕获内核请使用 enqueue_function / recording_context()。不要用 capturing[_] / @__parameter 来掩盖捕获调用点的问题。

捕获选择

默认情况／症状	选择
只读外部状态	{imm}（捕获全部）——如果主体调用 offset_ptr 或构造可变 TileTensor，则不要使用（见下方冻结说明）
修改某些外部状态；同时读取 Int／其他可通过寄存器传递的值	{mut buf, imm}——不要使用捕获全部的 {mut}
修改多个外部名称	{mut a, mut b, imm}
替换 @__copy_capture(x, y)	{var x, var y, imm}（复制）。只有在每个捕获都应为副本时才使用 {var}。{var^} 表示移动——这不是对应写法。绝不要使用 {imm}
仅指定名称的精度	{mut count}、{imm buf, imm shape}——{imm} 捕获主体使用的名称，而不是作用域内的所有局部变量
没有自由运行时捕获	省略捕获列表。绝不要使用 {}
Could not infer capture convention	添加 {imm}／{mut name, imm}／命名列表
捕获项出现 expression must be mutable in assignment	该名称需要 mut（或者如果视图数组会在每次迭代时重新构建，则在闭包内局部声明临时视图数组）
register passible value … can not be captured by 'mut'	捕获全部的 {mut} 捕获了 Int（等）——使用 {mut buf, imm}
.unsafe_ptr()／TileTensor 出现 .mut … is 'False' but … is 'True'	缓冲区以 {imm} 捕获——使用 {mut out_buf, imm}；不要用 unsafe_mut_cast 充数
offset_ptr／类似调用出现 'lit.call' op callee expected call argument #0	{imm} 冻结了用作 self 的缓冲区——使用 {mut cb_a, mut cb_b, …, imm}
bench_func 出现 cannot bind an RValue to a reference	kernel_launch 在 bench_func 内部嵌套，并使用了 @__copy_capture——将 kernel_launch 提到外层作用域，并为每个以副本捕获的名称添加 {var x, var y, …}。不要将其提到外层后再写 {imm}
raises {…} 处出现 expected ':' in function definition	@__parameter def 上的捕获列表——移除 @__parameter，保留列表
aliasing values passed immutably…mutably／说明中提到 origin_of(buf)	闭包结构体会持有指向同一来源的字段，其中一个为 mut，另一个为 imm——见下方 别名
Int（或其他可通过寄存器传递的）局部变量出现 aliasing … origin_of(n)	@__copy_capture(n) 变成了对 var n 的 imm 引用。使用 {var n}（复制）。除非这样更简单，否则不要删除 n 并改用另一个 Int
出现 aliasing values passed mutably to 'x' … and passed mutably to 'y'	对同一来源进行了两次 mut 捕获（视图和底层缓冲区，或一个值以及也捕获该值的嵌套辅助函数）——见下方 别名
bencher_iter_custom 值调用出现 cannot bind an RValue to a reference	接受值的 bencher_iter_custom 声明为 ref func: FuncType，因而拒绝临时／提升后的闭包值——将 func: FuncType 改为按值传递
bench_function[fn] 出现 'lit.call' op invalid symbol use／origin<false> 与 origin<true>	@__parameter 捕获轻量型 bench_fn 捕获了统一的 {imm} 子项——移除 @__parameter；bench_function(fn, id, …) 按值传递
_dispatch_ag_norm 中的 two_launch／capturing 参数	如果外层 API 仍接受 comptime capturing 函数参数，则写成 def two_launch() raises capturing:，不要使用 @__parameter 或捕获列表
cannot capture … not copyable／不是参数引用	如果可行则使用 {imm}；否则使用 {var}／命名的 var x；绝不要使用 @__parameter
包装器指针已是 MutUntrackedOrigin 的 Memset／launch 别名	将原始 DeviceBuffer 作为 imm 参数传递。不要包装未跟踪指针

如果闭包还会引用可通过寄存器传递的外部值（Int、索引、长度），不要使用捕获全部的 {mut}。将显式的 mut 名称与末尾的默认 imm 混用（{mut tt_in, imm}）。对于未列出的捕获项，最多只能有一个裸约定（imm／mut／var）作为默认值。

{imm} 会冻结捕获的缓冲区。以下情况使用它是错误的：

• 写入输出（TileTensor / .unsafe_ptr()，此时需要 mut=True），或
• 调用内部执行 self._buf.unsafe_ptr()，然后调用 unsafe_mut_cast[True]() 的方法——通常是 CacheBustingBuffer.offset_ptr。 方法的 self 可能声明为 self（imm）；捕获的字段仍然不匹配（lit.call 参数 #0）。对此类缓冲区都使用可变捕获（{mut cb_a, mut cb_b, mut cb_c, mut cb_a_scales, mut cb_b_scales, imm}）。

不要用额外的 unsafe_mut_cast 来“修复”此问题（方法里已经有该操作），也不要恢复 @__parameter。只有当添加 mut 会与另一个捕获的同源字段形成别名对时，才跳过 mut（见下文）。

别名

当多个捕获值或嵌套值引用同一个源，且其中一个是可变的，闭包结构体就会包含别名字段。两个捕获只有在均为只读时，才允许指向相同内存。

对可通过寄存器传递的局部变量（Int 索引、长度、行数）使用 @__copy_capture(n) 会变为 {var n}——复制到闭包存储中。对 var n = … 使用 {imm}，引用的是这个可变局部变量；一旦统一的同级闭包嵌入 origin_of(n)，就会产生别名。请使用 {var n}。不要删除 n，也不要仅为这个整数包装函数。

请按以下优先顺序处理：

1. 按只读方式捕获（imm），只要函数体没有实际修改该源。删除 {mut …}；传递 imm 参数。enqueue_memset 接受 imm DeviceBuffer。只读取 config.rank_units 的辅助函数不应以 mut 方式捕获 config。当外层 var buf 保持可变，而嵌套的统一闭包会捕获 buf 的视图时，也优先采用此方式：调用一个普通函数（文件作用域或其他非捕获函数），将 buf: DeviceBuffer[…] 作为 imm 参数传入，并在该函数中构造启动操作。
2. 拆解实例，使源标记更细粒度。如果编译器因捕获整个对象（config、结构体、缓冲区 List）而将两处使用视为重叠，但修改实际上不会触及另一处使用的内存，就提取所需字段（例如，将 rank_units / rank_unit_start 提前放入 Array[Int, ngpus]），并捕获或传递这些字段，而不是父对象。
3. 将可变源作为参数传递（在参数列表中使用 mut buf: …），使其不再是捕获字段。只提升必须为 mut 且仍与另一处使用存在别名的捕获——只读捕获可以保留为捕获。如果值传递 API 的 FuncType 已固定，请使用精简的 {mut buf, imm} 适配器，只将 buf 转发给全 imm 函数体。在可行时，优先扩展 API。

不要通过擦除源来“修复”别名问题。这包括在原始缓冲区上使用 .as_unsafe_any_origin()、unsafe_origin_cast[MutUntrackedOrigin] / MutAnyOrigin，以及将已取消跟踪的指针（EPLocalSyncCounters.ptr、offset_ptr 结果）包装进新的 DeviceBuffer(ctx, ptr, …, owning=False) 中，仅为执行 memset 或启动操作。应改为将原始缓冲区作为参数传递。 不要使用 @__parameter“修复”任何问题。

安全规则

• 按值参数名称定位批量修改，不要按所有同名定义定位
• 不要批量替换捕获列表（这会破坏 {mut count} 等内容）
• with … as name 中的 name[i] = 是局部变量，不是外层的 mut 捕获
• 迁移代码中的嵌套闭包不得使用 @__parameter / @parameter
• 只有为确保源唯一性时，才将捕获提升为参数（或在普通函数中使用 imm DeviceBuffer 参数）
• 保持 NFC：不要改变分配生命周期、comptime 与动态布局的区别，也不要改变计时路径上的主机工作。comptime if 分支中的局部变量（缓冲区、主机副本、TileTensor、row_major[M, N]() 布局）必须保留在该分支中——在函数作用域分配大小为 1 的内存会造成语义变化。只构造一次的布局（a_shape）应留在 call_fn 之外

更多详情

完整的分步流程、错误目录和验证方法： process.md。