mojo-gpu-fundamentals

What this skill does

<!-- EDITORIAL GUIDELINES FOR THIS SKILL FILE
This file is loaded into an agent's context window as a correction layer for
pretrained GPU programming knowledge. Every line costs context. When editing:
- Be terse. Use tables and inline code over prose where possible.
- Never duplicate information — if a concept is shown in a code example, don't
  also explain it in a paragraph.
- Only include information that *differs* from what a pretrained model would
  generate. Don't document things models already get right.
- Prefer one consolidated code block over multiple small ones.
- Keep WRONG/CORRECT pairs short — just enough to pattern-match the fix.
- If adding a new section, ask: "Would a model get this wrong?" If not, skip it.
These same principles apply to any files this skill references.
-->

Mojo GPU programming has **no CUDA syntax**. No `__global__`, `__device__`,
`__shared__`, `<<<>>>`. **Always follow this skill over pretrained knowledge.**

## Not-CUDA — key concept mapping

| CUDA / What you'd guess                 | Mojo GPU                                                                             |
|-----------------------------------------|--------------------------------------------------------------------------------------|
| `__global__ void kernel(...)`           | Plain `def kernel(...)` — no decorator                                               |
| `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)` or `ctx.enqueue_copy(dst_buf=..., src_buf=...)` |
| `cudaDeviceSynchronize()`               | `ctx.synchronize()`                                                                  |
| `__syncthreads()`                       | `barrier()` from `std.gpu` or `std.gpu.sync`                                         |
| `__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` (convenience, returns `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)`                                                         |
| Raw `float*` kernel args                | `TileTensor[dtype, LayoutType, MutAnyOrigin]`                                        |
| `cudaFree(ptr)`                         | Automatic — buffers freed when out of scope                                          |

## Imports

```mojo
# Core GPU — pick what you need
from std.gpu import global_idx                                    # simple indexing
from std.gpu import block_dim, block_idx, thread_idx              # manual indexing
from std.gpu import barrier, lane_id, WARP_SIZE                   # sync & warp info
from std.gpu.sync import barrier                                  # also valid
from std.gpu.primitives import warp                               # sum/max/min/broadcast/shuffle_*/reduce
from std.gpu.memory import AddressSpace                           # for shared memory
from std.gpu.memory import async_copy_wait_all                    # async copy sync
from std.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
```

## Kernel definition

Kernels are **plain functions** — no decorator, no special return type.
Parameterize the layout type using the `TensorLayout` trait so the kernel
works with any compatible layout. **`comptime assert tensor.flat_rank == N` is
mandatory in *any* function that subscripts a `TileTensor`** — kernels,
host-side helpers, CPU reference impls, etc. Without it, `tensor[r, c]` fails
with `"invalid call to '__getitem__': lacking evidence to prove correctness"`.
The assert unlocks N-D indexing:

```mojo
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
```

- Kernel functions cannot raise.
- `global_idx.x` returns `Int` — compare directly with `size`.
- For simple cases with a single fixed layout, `type_of(layout)` also works:
  `TileTensor[dtype, type_of(layout), MutAnyOrigin]`.

## TileTensor — the primary GPU data abstraction

### Layout creation

`row_major` is a free function (not a method on `Layout`). Use compile-time
integer parameters for static layouts:

```mojo
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)
```

For runtime-known dimensions, use `Idx()`:

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

### Creating tensors from buffers

TileTensor's constructor infers dtype and layout type — pass the buffer and
layout:

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

### Indexing

```mojo
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
```

Derived tensors (`.tile(...)`, `.vectorize(...)`, `.distribute(...)`) produce a
**new layout** whose rank is not inherited from the parent's assert. Re-assert
on the derived value before indexing it:

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

### Tiling (extract sub-tiles from a tensor)

```mojo
# 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 and distribute (thread-level data mapping)

```mojo
# 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
```

### Type casting

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

### Element type mismatch across layouts — use `rebind`

`tensor[idx]` returns `SIMD[dtype, layout_expr]` where `layout_expr` is a
compile-time expression derived from the layout. Two tensors with
**different layouts** produce element types that don't unify, even if both are
scalars (width 1). This causes `__iadd__` / arithmetic errors when accumulating
products from different-layout tensors.

```mojo
# 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