用 Mojo 编写 GPU 内核时如何避免套用 CUDA…

在 Mojo 中编写面向 NVIDIA/AMD/Apple GPU 的加速代码时，用 Mojo 原生的 kernel 定义、TileTensor、DeviceContext 与布局系统替代 CUDA 习惯写法。

适用场景

解决思路

1. 不要使用任何 CUDA 语法：没有 `__global__`、`__device__`、`__shared__`、`<<<>>>`；kernel 就是普通 `def` 函数，无装饰器。
2. 按对照表替换概念：`__global__ void kernel(...)` → 普通 `def kernel(...)`；`kernel<<<grid, block>>>(args)` → `ctx.enqueue_function[kernel](args, grid_dim=..., block_dim=...)`；`cudaMalloc` → `ctx.enqueue_create_buffer[dtype](count)`；`cudaMemcpy` → `ctx.enqueue_copy(dst_buf, src_buf)` 或 `ctx.enqueue_copy(dst_buf=..., src_buf=...)`；`cudaDeviceSynchronize()` → `ctx.synchronize()`；`__syncthreads()` → 从 `max.gpu` 或 `max.gpu.sync` 导入的 `barrier()`；`__shared__ float s[N]` → `stack_allocation[dtype, address_space=AddressSpace.SHARED](layout)`；`threadIdx.x` → `thread_idx.x`；`blockIdx.x * blockDim.x + threadIdx.x` → `global_idx.x`（返回 `Int`）；`__shfl_down_sync` → `warp.shuffle_down(val, d)` / `warp.sum` / `warp.max` / `warp.min` / `warp.reduce`；`atomicAdd` → `Atomic.fetch_add(ptr, val)`；裸 `float*` 参数 → `TileTensor[dtype, LayoutType, MutAnyOrigin]`；`cudaFree` → 自动，buffer 超出作用域即释放。
3. 按需导入：核心 GPU 从 `max.gpu` 导入 `global_idx`、`block_dim`、`block_idx`、`thread_idx`、`lane_id`、`WARP_SIZE`；从 `max.gpu.sync` 导入 `barrier`；从 `max.gpu.primitives` 导入 `warp`；从 `max.gpu.memory` 导入 `AddressSpace`、`async_copy_wait_all`；从 `max.gpu.host` 导入 `DeviceContext`、`DeviceBuffer`；从 `std.atomic` 导入 `Atomic`；布局系统从独立的 `layout` 包导入 `TileTensor`、`TensorLayout`、`Idx`、`row_major`、`stack_allocation`。
4. 定义 kernel 时用 `TensorLayout` trait 参数化布局类型，并在任何对 `TileTensor` 做下标的函数（kernel、host 辅助函数、CPU 参考实现等）中强制写# `comptime assert tensor.flat_rank == N`，否则 `tensor[r, c]` 会报 `"invalid call to '__getitem__': lacking evidence to prove correctness"`；kernel 函数不能 raise；`global_idx.x` 返回 `Int`，可直接与 `size` 比较；单一固定布局的简单场景也可用 `type_of(layout)`。
5. 创建布局：`row_major` 是自由函数（不是 `Layout` 的方法），静态维度用编译期整数参数 `row_major[1024]()`、`row_major[64, 64]()`、`row_major[10, 5, 3]()`；运行时维度用 `row_major(Idx(M), Idx(N))`。
6. 从 buffer 创建张量：`TileTensor` 构造函数会推断 dtype 和布局类型，传入 buffer 与 layout，例如 `var tensor = TileTensor(buf, row_major[1024]())`。
7. 索引：`tensor[tid]`（1D）、`tensor[row, col]`（2D）、`tensor[row, col, channel]`（3D）；用 `tensor.dim[0]()` 查询维度大小（编译期索引），算术中需用 `Int(tensor.dim[1]())` 包裹。
8. 派生张量（`.tile(...)`、`.vectorize(...)`、`.distribute(...)`）会产生新布局，其 rank 不继承父张量的 assert，索引前必须重新 assert，例如 `var vec = tensor.vectorize[1, 4]()` 后写 `comptime assert vec.flat_rank == 2`。
9. Tiling：在 kernel 内用 `var tile = tensor.tile[block_size, block_size](block_idx.y, block_idx.x)` 提取子块，再用 `tile[thread_idx.y, thread_idx.x]` 访问。
10. Vectorize 与 distribute：`comptime thread_layout = row_major[WARP_SIZE // simd_width, simd_width]()`，然后 `var fragment = tensor.vectorize[1, simd_width]().distribute[thread_layout=thread_layout](lane_id())`；用 `fragment.copy_from_async(source_fragment)` 做异步拷贝，`fragment.copy_from(source_fragment)` 做同步拷贝。
11. 类型转换：`var val = tensor[row, col].cast[DType.float32]()`。
12. 不同布局的元素类型不统一时用 `rebind`：`tensor[idx]` 返回 `SIMD[dtype, layout_expr]`，不同布局的张量即使都是标量（宽度 1）也无法统一，导致 `__iadd__`/算术报错；应写 `var k_val = rebind[Scalar[dtype]](conv_kernel[k])` 和 `var s_val = rebind[Scalar[dtype]](s_data[idx])` 后再相加。`rebind` 是内置函数无需导入；当表达式中所有张量共享同一布局时不需要。读写单个元素做标量算术或传给辅助函数时也用 `rebind`：`var val = rebind[Scalar[dtype]](tensor[idx])`、`tensor[idx] = rebind[tensor.ElementType](computed_scalar)`；`tensor.ElementType` 是 `SIMD[dtype, element_size]`，基本布局下 `element_size=1`。
13. 内存管理：`var ctx = DeviceContext()`；用 `ctx.enqueue_create_buffer[DType.float32](1024)` 分配设备 buffer，用 `ctx.enqueue_create_host_buffer[DType.float32](1024)` 分配 host buffer；用 `dev_buf.enqueue_fill(0.0)` 直接初始化设备 buffer；用 `ctx.enqueue_copy(dst_buf=dev_buf, src_buf=host_buf)` 或位置形式 `ctx.enqueue_copy(dev_buf, host_buf)` 拷贝；用 `with dev_buf.map_to_host() as mapped:` 映射设备 buffer 到 host（上下文管理器自动同步）；用 `ctx.enqueue_memset(dev_buf, 0.0)` 做 memset；用 `ctx.synchronize()` 同步所有已入队操作。
14. 内核启动：`ctx.enqueue_function[my_kernel](input_tensor, output_tensor, size, grid_dim=num_blocks, block_dim=block_size)`；2D 网格/块用元组 `grid_dim=(col_blocks, row_blocks)`、`block_dim=(BLOCK_SIZE, BLOCK_SIZE)`。
15. 若 kernel 带任何 comptime 参数，必须先绑定再启动：错误写法 `ctx.enqueue_function[vector_add](a, b, c, N, grid_dim=G, block_dim=B)` 会产生大量 "no matching method"/"DevicePassable" 模板错误；正确写法是先 `comptime kernel = vector_add[type_of(layout)]`，再 `ctx.enqueue_function[kernel](a, b, c, N, grid_dim=G, block_dim=B)`。单态 kernel（签名直接用 `type_of(layout)`、无 `[LT: TensorLayout]` 等）可按名直接传入，无需绑定。
16. 共享内存：在 kernel 内用 `layout` 包的 `stack_allocation`，返回指定地址空间的 `TileTensor`：`var tile_shared = stack_allocation[DType.float32, address_space=AddressSpace.SHARED](row_major[TILE_M, TILE_K]())`；可链式 `.fill()` 零初始化（返回张量），如 `var regs = stack_allocation[DType.float32](row_major[TM, TN]()).fill(0)`；从全局加载到共享后必须 `barrier()` 再读取；也可用 `std.memory` 的裸指针版本 `stack_allocation[512, Scalar[DType.int32], address_space=AddressSpace.SHARED]()`。
17. 线程索引：简单场景用 `from max.gpu import global_idx`，`var tid = global_idx.x`（1D）、`var row = global_idx.y`、`var col = global_idx.x`；需要分别拿到 block/thread 时用 `from max.gpu import block_idx, block_dim, thread_idx`，写 `var tid = block_idx.x * block_dim.x + thread_idx.x`；warp 信息用 `from max.gpu import lane_id, WARP_SIZE`，`var my_lane = lane_id()`（0..WARP_SIZE-1）。这些返回值都是 `Int`，边界检查无需转换。
18. 同步与 warp 操作：`barrier()` 做块级同步；`_ = Atomic.fetch_add(output_ptr, value)` 做原子加；归约 `warp.sum(val)`、`warp.max(val)`、`warp.min(val)` 结果广播到所有 lane；`warp.broadcast(val)` 把 lane 0 的值广播到所有 lane；`warp.reduce[warp.shuffle_down, reduce_fn](val)` 做自定义归约（会广播）；shuffle 是逐 lane 移位/交换而非广播：`warp.shuffle_down(val, offset)`、`warp.shuffle_xor(val, offset)`（offset 与 lane id 异或，蝶形）、`warp.shuffle_xor(mask, val, offset)`（显式成员掩码，很少用）；`offset` 类型是 `UInt32`，传普通 `Int` 会类型不匹配，应写 `var off: UInt32 = 16` 再 `warp.shuffle_xor(val, off)`。
19. GPU 可用性检查：`from std.sys import has_accelerator`，在 `main` 中写 `comptime if not has_accelerator(): print("No GPU found")` 否则走 GPU 代码；或写编译期断言 `comptime assert has_accelerator(), "Requires a GPU"`（必须位于函数体内）。
20. 架构检测区分 `is_` 与 `has_`：`is_*` 检查编译目标（在 GPU 分派代码内使用），`has_*` 检查宿主系统（在 host/CPU 代码中使用）；从 `std.sys.info` 导入 `is_gpu, is_nvidia_gpu, is_amd_gpu, is_apple_gpu` 与 `has_nvidia_gpu_accelerator, has_amd_gpu_accelerator, has_apple_gpu_accelerator`，从 `std.sys` 导入 `has_accelerator`；host 侧用 `comptime if not has_accelerator():` 决定是否运行 GPU 代码，kernel 或 GPU 编译代码内用 `comptime if is_nvidia_gpu(): ... elif is_amd_gpu(): ...` 按架构分派；子架构检查（仅 GPU 代码内）用 `from std.sys.info import _is_sm_9x_or_newer, _is_sm_100x_or_newer` 和 `comptime if is_nvidia_gpu["sm_90"]():`。
21. 所有 GPU 维度、布局、大小都用 `comptime` 常量：`comptime dtype = DType.float32`、`comptime SIZE = 1024`、`comptime BLOCK_SIZE = 256`、`comptime NUM_BLOCKS = ceildiv(SIZE, BLOCK_SIZE)`、`comptime layout = row_major[SIZE]()`。
22. SIMD 加载：从裸指针做向量化加载用 `var val = ptr.unsafe_load[width=8](idx)`（`.load()` 已废弃），用 `var sum = val.reduce_add()` 做标量归约；TileTensor 向量化访问用 `var vec_tensor = tensor.vectorize[1, 4]()` 把元素分组为 `SIMD[4]`。
23. 归约模式：指针索引用 `ptr[unsafe_offset=i]`（裸 `ptr[i]` 已废弃），使用 `MutPointer` 而非已废弃的 `UnsafePointer`；块大小必须是 comptime（`comptime TPB = 512`）因为 `comptime for` 需要静态边界；在共享内存中做树形归约（`var active = TPB`，`comptime for _ in range(log2_floor(TPB)):` 内 `active >>= 1`，`if tid < active: sums[unsafe_offset=tid] += sums[unsafe_offset=tid + active]`，每轮后 `barrier()`），最后 `if tid < WARP_SIZE:` 做 `var v = warp.sum(sums[unsafe_offset=tid][0])`，`if tid == 0: _ = Atomic.fetch_add(output, v)`。
24. 从已有指针创建 DeviceBuffer：`var buf = DeviceBuffer[dtype](ctx, raw_ptr, count, owning=False)`（非拥有）。
25. GPU kernel 基准测试：从 `std.benchmark` 导入 `Bench, BenchConfig, Bencher, BenchId, BenchMetric, ThroughputMeasure`，从 `max.benchmark` 导入 `bencher_iter_custom`（GPU 形式是自由函数）；定义 `@inline(.always) def bench_fn(mut b: Bencher) raises capturing[_]:`，内部定义 `@inline(.always) def launch(ctx: DeviceContext) raises {imm}:` 调用 `ctx.enqueue_function[kernel](args, grid_dim=G, block_dim=B)`，创建 `var ctx = DeviceContext()`，用取值形式 `bencher_iter_custom(b, launch, ctx)`（不是 `b.iter_custom(...)`）；用 `var bench = Bench(BenchConfig(max_iters=50000))` 和 `bench.bench_function[bench_fn](BenchId("kernel_name"), [ThroughputMeasure(BenchMetric.bytes, total_bytes)])`；`Bencher.iter_custom` 不接受 `DeviceContext`，该形式在 `max.benchmark` 中；优先用 `bencher_iter_custom(b, launch, ctx)` 配统一闭包和显式捕获列表（`{imm}`、`{var}` 或具名捕获）；不要在这些 launch 闭包上用 `@__parameter` / `@parameter`。
26. 参考硬件细节表：NVIDIA warp 大小 32、共享内存 48-228 KB/block、SM70+ 有 tensor cores（WMMA）、SM90+ 有 TMA 和 Clusters；AMD CDNA warp 大小 64、共享内存 64 KB/block、有 Matrix cores、无 TMA/Clusters；AMD RDNA warp 大小 32、共享内存可配置、RDNA3+ 有 WMMA、无 TMA/Clusters。

适用边界 · 注意事项

• 目标不是 GPU 或其他加速器、只是写普通 Mojo 代码时（此时应改用 mojo-syntax 技能）。
• 没有可用的 GPU 硬件且代码路径依赖 `has_accelerator()` 为真时。
• 需要 CUDA 专有特性（如 TMA、Clusters）但目标架构不支持时（TMA/Clusters 仅 NVIDIA SM90+，AMD 均不支持）。