mojo-gpu-fundamentals

Mojo GPU 编程没有 CUDA 语法。没有 __global__、__device__、 __shared__、<<<>>>。始终以本指南为准，不要依赖预训练知识。

非 CUDA — 关键概念映射

CUDA / 你可能会猜测的写法	Mojo GPU
__global__ void kernel(...)	普通的 def kernel(...) — 无装饰器
kernel<<<grid, block>>>(args)	ctx.enqueue_function[kernel](args, grid_dim=..., block_dim=...)
cudaMalloc(&ptr, size)	ctx.enqueue_create_buffer[dtype](count)
cudaMemcpy(dst, src, ...)	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(mask, val, d)	warp.shuffle_down(val, d) / warp.sum / warp.max / warp.min / warp.reduce
atomicAdd(&ptr, val)	Atomic.fetch_add(ptr, val)
原始的 float* 内核参数	TileTensor[dtype, LayoutType, MutAnyOrigin]
cudaFree(ptr)	自动处理 — 缓冲区超出作用域时释放

导入

# Core GPU — pick what you need
from max.gpu import global_idx                                    # simple indexing
from max.gpu import block_dim, block_idx, thread_idx              # manual indexing
from max.gpu import lane_id, WARP_SIZE                            # warp info
from max.gpu.sync import barrier                                  # block-level sync
from max.gpu.primitives import warp                               # sum/max/min/broadcast/shuffle_*/reduce
from max.gpu.memory import AddressSpace                           # for shared memory
from max.gpu.memory import async_copy_wait_all                    # async copy sync
from max.gpu.host import DeviceContext, DeviceBuffer              # host-side API
from std.atomic import Atomic                                     # atomics

# Layout system — NOT in std, separate package
from layout import TileTensor, TensorLayout, Idx, row_major, stack_allocation

内核定义

内核是普通函数 — 没有装饰器，也没有特殊的返回类型。 使用 TensorLayout 特征对布局类型进行参数化，让内核 可与任何兼容的布局配合使用。任何对 TileTensor 进行下标访问的函数中，都必须使用 comptime assert tensor.flat_rank == N，包括内核、 主机端辅助函数、CPU 参考实现等。没有它，tensor[r, c] 会失败 并提示 "invalid call to '__getitem__': lacking evidence to prove correctness"。 此断言启用 N 维索引：

def my_kernel[
    dtype: DType, LT: TensorLayout,
](
    input: TileTensor[dtype, LT, MutAnyOrigin],
    output: TileTensor[dtype, LT, MutAnyOrigin],
    size: Int,                                    # scalar args are fine
):
    comptime assert input.flat_rank == 1, "expected 1D tensor"
    var tid = global_idx.x
    if tid < size:
        output[tid] = input[tid] * 2

• 内核函数不能抛出异常。
• global_idx.x 返回 Int — 可直接与 size 比较。
• 对于只有一个固定布局的简单情况，type_of(layout) 也可使用： TileTensor[dtype, type_of(layout), MutAnyOrigin]。

TileTensor — 主要的 GPU 数据抽象

布局创建

row_major 是一个自由函数（不是 Layout 上的方法）。创建静态布局时使用编译时 整数参数：

comptime layout_1d = row_major[1024]()                     # 1D
comptime layout_2d = row_major[64, 64]()                   # 2D (rows, cols)
comptime layout_3d = row_major[10, 5, 3]()                 # 3D (e.g. H, W, C)

对于运行时已知的维度，使用 Idx()：

var layout = row_major(Idx(M), Idx(N))                     # runtime dims

从缓冲区创建张量

TileTensor 的构造函数会推断 dtype 和布局类型 — 传入缓冲区和 布局：

var buf = ctx.enqueue_create_buffer[DType.float32](1024)
var tensor = TileTensor(buf, row_major[1024]())            # wraps device buffer

索引

tensor[tid]                     # 1D
tensor[row, col]                # 2D
tensor[row, col, channel]       # 3D
tensor.dim[0]()                 # query dimension size (compile-time index)
var K = Int(tensor.dim[1]())    # wrap with Int() for use in arithmetic

派生张量（.tile(...)、.vectorize(...)、.distribute(...)）会生成一个 新布局，其秩不会继承自父张量的断言。在对派生值进行索引前， 重新添加断言：

var vec = tensor.vectorize[1, 4]()
comptime assert vec.flat_rank == 2          # required for vec[r, i]

分块（从张量中提取子块）

# Inside kernel — extract a block_size x block_size tile
var tile = tensor.tile[block_size, block_size](block_idx.y, block_idx.x)
tile[thread_idx.y, thread_idx.x]   # access within tile

向量化和分布（线程级数据映射）

# Vectorize along inner dimension, then distribute across threads
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)    # async copy
fragment.copy_from(source_fragment)          # sync copy

类型转换

var val = tensor[row, col].cast[DType.float32]()    # cast element

布局间元素类型不匹配 — 使用 rebind

tensor[idx] 会返回 SIMD[dtype, layout_expr]，其中 layout_expr 是根据布局推导出的编译时表达式。两个布局不同的张量会产生无法统一的元素类型，即使它们都是标量（宽度为 1）也一样。这会导致在累加不同布局张量的乘积时出现 __iadd__ / 算术错误。

# WRONG — fails when conv_kernel and s_data have different layouts:
var sum: Scalar[dtype] = 0
sum += conv_kernel[k] * s_data[idx]   # error: cannot convert ElementType to Float32

# CORRECT — rebind each element to Scalar[dtype]:
var sum: Scalar[dtype] = 0
var k_val = rebind[Scalar[dtype]](conv_kernel[k])
var s_val = rebind[Scalar[dtype]](s_data[idx])
sum += k_val * s_val

rebind 是内置功能（无需导入）。如果表达式中的所有张量都共享相同布局（例如 sa 和 sb 具有相同分块布局的矩阵乘法示例），则无需使用它。

在进行标量算术、读取/写入单个元素或传递给辅助函数时，也要使用 rebind — 即使只有一个张量：

# Read element as plain scalar
var val = rebind[Scalar[dtype]](tensor[idx])
# Write scalar back to tensor
tensor[idx] = rebind[tensor.ElementType](computed_scalar)

tensor.ElementType 是 SIMD[dtype, element_size] — 对于基本布局，element_size=1（实际上等同于 Scalar[dtype]）。

内存管理

var ctx = DeviceContext()

# Allocate
var dev_buf = ctx.enqueue_create_buffer[DType.float32](1024)
var host_buf = ctx.enqueue_create_host_buffer[DType.float32](1024)

# Initialize device buffer directly
dev_buf.enqueue_fill(0.0)

# Copy host -> device
ctx.enqueue_copy(dst_buf=dev_buf, src_buf=host_buf)
# Copy device -> host
ctx.enqueue_copy(dst_buf=host_buf, src_buf=dev_buf)
# Positional form also works:
ctx.enqueue_copy(dev_buf, host_buf)

# Map device buffer to host (context manager — auto-syncs)
with dev_buf.map_to_host() as mapped:
    var t = TileTensor(mapped, row_major[1024]())
    print(t[0])

# Memset
ctx.enqueue_memset(dev_buf, 0.0)

# Synchronize all enqueued operations
ctx.synchronize()

内核启动

ctx.enqueue_function[my_kernel](
    input_tensor,
    output_tensor,
    size,                    # scalar args passed directly
    grid_dim=num_blocks,     # 1D: scalar
    block_dim=block_size,    # 1D: scalar
)

# 2D grid/block — use tuples:
ctx.enqueue_function[kernel_2d](
    args...,
    grid_dim=(col_blocks, row_blocks),
    block_dim=(BLOCK_SIZE, BLOCK_SIZE),
)

如果内核带有任何 comptime 参数，必须先绑定这些参数 — 直接将带参数的名称传给 enqueue_function 会产生大量“no matching method” / “DevicePassable”模板错误：

# WRONG — vector_add has a `LT: TensorLayout` parameter
ctx.enqueue_function[vector_add](a, b, c, N, grid_dim=G, block_dim=B)

# CORRECT — bind parameters, then launch the bound symbol
comptime kernel = vector_add[type_of(layout)]
ctx.enqueue_function[kernel](a, b, c, N, grid_dim=G, block_dim=B)

单态内核（签名直接使用 type_of(layout)，不使用 [LT: TensorLayout] 等）可以直接按名称传递，无需绑定。

共享内存

使用 layout 包中的 stack_allocation 在内核内分配共享内存 — 它会在指定的地址空间中返回一个 TileTensor：

from layout import stack_allocation   # TileTensor-based shared alloc
from max.gpu.memory import AddressSpace

var tile_shared = stack_allocation[DType.float32,
    address_space=AddressSpace.SHARED](row_major[TILE_M, TILE_K]())

# Chain .fill() to zero-initialize (returns the tensor)
var regs = stack_allocation[DType.float32](row_major[TM, TN]()).fill(0)

# Load from global to shared
tile_shared[thread_idx.y, thread_idx.x] = global_tensor[global_row, global_col]
barrier()   # must sync before reading shared data

# Alternative: raw pointer shared memory (from std.memory, not layout)
from std.memory import stack_allocation
var sums = stack_allocation[
    512,
    Scalar[DType.int32],
    address_space=AddressSpace.SHARED,
]()

线程索引

# Simple — automatic global offset
from max.gpu import global_idx
var tid = global_idx.x           # 1D
var row = global_idx.y           # 2D row
var col = global_idx.x           # 2D col

# Manual — when you need block/thread separately
from max.gpu import block_idx, block_dim, thread_idx
var tid = block_idx.x * block_dim.x + thread_idx.x

# Warp info
from max.gpu import lane_id, WARP_SIZE
var my_lane = lane_id()          # 0..WARP_SIZE-1

所有这些都会返回 Int — 边界检查无需类型转换。

同步和 warp 操作

from max.gpu.sync import barrier
from max.gpu.primitives import warp
from std.atomic import Atomic

barrier()                                    # block-level sync
_ = Atomic.fetch_add(output_ptr, value)      # atomic add

# Reductions — result broadcast to every lane
warp.sum(val)         warp.max(val)         warp.min(val)
warp.broadcast(val)                          # lane-0 value → all lanes
warp.reduce[warp.shuffle_down](val, reduce_fn)  # custom reduction (broadcasts)

# Shuffles — per-lane shift/swap, NOT broadcast
warp.shuffle_down(val, offset)               # offset: UInt32
warp.shuffle_xor(val, offset)                # offset: UInt32, XOR'd with lane id (butterfly)
warp.shuffle_xor(mask, val, offset)          # explicit membership mask (rarely needed)

# `offset` is `UInt32` — plain `Int` is a type-mismatch error:
var off: UInt32 = 16
var v = warp.shuffle_xor(val, off)

GPU 可用性检查

from std.sys import has_accelerator

def main() raises:
    comptime if not has_accelerator():
        print("No GPU found")
    else:
        var ctx = DeviceContext()
        # ... GPU code

也可以使用编译时断言 — 但它必须放在函数体内：

def main() raises:
    comptime assert has_accelerator(), "Requires a GPU"

架构检测 — is_ 与 has_

关键区别：is_* 检查编译目标（在 GPU 分派代码中使用）。has_* 检查主机系统（在主机/CPU 代码中使用）。

from std.sys.info import (
    # Target checks — "am I being compiled FOR this GPU?"
    # Use inside kernels or GPU-targeted code paths.
    is_gpu, is_nvidia_gpu, is_amd_gpu, is_apple_gpu,

    # Host checks — "does this machine HAVE this GPU?"
    # Use from host code to decide whether to launch GPU work.
    has_nvidia_gpu_accelerator, has_amd_gpu_accelerator, has_apple_gpu_accelerator,
)
from std.sys import has_accelerator   # host check: any GPU present

# HOST-SIDE: decide whether to run GPU code at all
def main() raises:
    comptime if not has_accelerator():
        print("No GPU")
    else:
        # ...launch kernels

# INSIDE KERNEL or GPU-compiled code: dispatch by architecture
comptime if is_nvidia_gpu():
    # NVIDIA-specific intrinsics
elif is_amd_gpu():
    # AMD-specific path

子架构检查（仅在 GPU 代码中使用）：

from std.sys.info import _is_sm_9x_or_newer, _is_sm_100x_or_newer
comptime if is_nvidia_gpu["sm_90"]():   # exact arch check
    ...

编译时常量模式

所有 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]()

完整的 1D 示例（向量加法）

from std.math import ceildiv
from std.sys import has_accelerator
from max.gpu import global_idx
from max.gpu.host import DeviceContext
from layout import TileTensor, row_major

comptime dtype = DType.float32
comptime N = 1024
comptime BLOCK = 256
comptime layout = row_major[N]()

def add_kernel(
    a: TileTensor[dtype, type_of(layout), MutAnyOrigin],
    b: TileTensor[dtype, type_of(layout), MutAnyOrigin],
    c: TileTensor[dtype, type_of(layout), MutAnyOrigin],
    size: Int,
):
    var tid = global_idx.x
    if tid < size:
        c[tid] = a[tid] + b[tid]

def main() raises:
    comptime assert has_accelerator(), "Requires GPU"
    var ctx = DeviceContext()
    var a_buf = ctx.enqueue_create_buffer[dtype](N)
    var b_buf = ctx.enqueue_create_buffer[dtype](N)
    var c_buf = ctx.enqueue_create_buffer[dtype](N)
    a_buf.enqueue_fill(1.0)
    b_buf.enqueue_fill(2.0)

    var a = TileTensor(a_buf, layout)
    var b = TileTensor(b_buf, layout)
    var c = TileTensor(c_buf, layout)

    ctx.enqueue_function[add_kernel](
        a, b, c, N,
        grid_dim=ceildiv(N, BLOCK),
        block_dim=BLOCK,
    )

    with c_buf.map_to_host() as host:
        var result = TileTensor(host, layout)
        print(result)

完整的 2D 示例（使用共享内存的分块矩阵乘法）

from std.math import ceildiv
from std.sys import has_accelerator
from max.gpu.sync import barrier
from max.gpu.host import DeviceContext
from max.gpu import thread_idx, block_idx
from max.gpu.memory import AddressSpace
from layout import TileTensor, TensorLayout, row_major, stack_allocation

comptime dtype = DType.float32
comptime M = 64
comptime N = 64
comptime K = 64
comptime TILE = 16
comptime a_layout = row_major[M, K]()
comptime b_layout = row_major[K, N]()
comptime c_layout = row_major[M, N]()

def matmul_kernel[
    ALayout: TensorLayout, BLayout: TensorLayout, CLayout: TensorLayout,
](
    A: TileTensor[dtype, ALayout, MutAnyOrigin],
    B: TileTensor[dtype, BLayout, MutAnyOrigin],
    C: TileTensor[dtype, CLayout, MutAnyOrigin],
):
    comptime assert A.flat_rank == 2 and B.flat_rank == 2 and C.flat_rank == 2
    var tx = thread_idx.x
    var ty = thread_idx.y
    var row = block_idx.y * TILE + ty
    var col = block_idx.x * TILE + tx

    var sa = stack_allocation[dtype,
        address_space=AddressSpace.SHARED](row_major[TILE, TILE]())
    var sb = stack_allocation[dtype,
        address_space=AddressSpace.SHARED](row_major[TILE, TILE]())

    var acc: C.ElementType = 0.0
    comptime for k_tile in range(0, K, TILE):
        if row < M and k_tile + tx < K:
            sa[ty, tx] = A[row, k_tile + tx]
        else:
            sa[ty, tx] = 0.0
        if k_tile + ty < K and col < N:
            sb[ty, tx] = B[k_tile + ty, col]
        else:
            sb[ty, tx] = 0.0
        barrier()
        comptime for k in range(TILE):
            acc += sa[ty, k] * sb[k, tx]
        barrier()

    if row < M and col < N:
        C[row, col] = acc

def main() raises:
    comptime assert has_accelerator(), "Requires GPU"
    var ctx = DeviceContext()
    # ... allocate buffers, init data, then:
    comptime kernel = matmul_kernel[type_of(a_layout), type_of(b_layout), type_of(c_layout)]
    ctx.enqueue_function[kernel](
        A, B, C,
        grid_dim=(ceildiv(N, TILE), ceildiv(M, TILE)),
        block_dim=(TILE, TILE),
    )

内核中的 SIMD 加载

# Vectorized load from raw pointer — `.load()` is deprecated
var val = ptr.unsafe_load[width=8](idx)   # SIMD[dtype, 8]
var sum = val.reduce_add()                 # scalar reduction

# TileTensor vectorized access
var vec_tensor = tensor.vectorize[1, 4]()  # group elements into SIMD[4]

归约模式

指针索引应写作 ptr[unsafe_offset=i] — 单独使用 ptr[i] 已弃用。请使用 MutPointer，不要使用已弃用的 UnsafePointer。

from std.bit import log2_floor

comptime TPB = 512          # `comptime for` needs a static bound, so the
                            # block size must be comptime, not `block_dim.x`

def block_reduce(
    output: MutPointer[Int32, MutAnyOrigin],
    input: MutPointer[Int32, MutAnyOrigin],
):
    var sums = stack_allocation[TPB, Scalar[DType.int32],
        address_space=AddressSpace.SHARED]()
    var tid = thread_idx.x
    sums[unsafe_offset=tid] = input[
        unsafe_offset=block_idx.x * block_dim.x + tid
    ]
    barrier()

    # Tree reduction in shared memory
    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()

    # Final warp reduction + atomic accumulate
    if tid < WARP_SIZE:
        var v = warp.sum(sums[unsafe_offset=tid][0])
        if tid == 0:
            _ = Atomic.fetch_add(output, v)

从现有指针创建 DeviceBuffer

# Wrap an existing pointer as a DeviceBuffer (non-owning)
var buf = DeviceBuffer[dtype](ctx, raw_ptr, count, owning=False)

GPU 内核基准测试

from std.benchmark import Bench, BenchConfig, Bencher, BenchId, BenchMetric, ThroughputMeasure
from max.benchmark import bencher_iter_custom   # GPU form: a free function

var ctx = DeviceContext()

@inline(.always)
def bench_fn(mut b: Bencher) raises {imm}:
    @inline(.always)
    def launch(ctx: DeviceContext) raises {imm}:
        ctx.enqueue_function[kernel](args, grid_dim=G, block_dim=B)
    bencher_iter_custom(b, launch, ctx)         # NOT `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 中。基准测试函数和启动闭包都以值的形式传递；请为每个函数指定显式捕获列表（{imm}、{var} 或命名捕获）。

硬件详情

属性	NVIDIA	AMD CDNA	AMD RDNA
Warp 大小	32	64	32
共享内存	48-228 KB/block	64 KB/block	可配置
Tensor 核心	SM70+ (WMMA)	矩阵核心	WMMA (RDNA3+)
TMA	SM90+ (Hopper)	不适用	不适用
集群	SM90+	不适用	不适用