Mojo Gpu Fundamentals

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>>(args) | ctx.enqueue_function[kernel, 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] | LayoutTensor[...address_space=AddressSpace.SHARED].stack_allocation() | | threadIdx.x | thread_idx.x (returns UInt) | | blockIdx.x * blockDim.x + threadIdx.x | global_idx.x (convenience) | | __shfl_down_sync(mask, val, d) | warp.sum(val), warp.reduce[...]() | | atomicAdd(&ptr, val) | Atomic.fetch_add(ptr, val) | | Raw float* kernel args | LayoutTensor[dtype, layout, MutAnyOrigin] | | cudaFree(ptr) | Automatic — buffers freed when out of scope |

Imports

# 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                               # warp.sum, warp.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.os.atomic import Atomic                                  # atomics

# Layout system — NOT in std, separate package
from layout import Layout, LayoutTensor

Kernel definition

Kernels are plain functions — no decorator, no special return type. Parameters use MutAnyOrigin:

def my_kernel(
    input: LayoutTensor[DType.float32, layout, MutAnyOrigin],
    output: LayoutTensor[DType.float32, layout, MutAnyOrigin],
    size: Int,                                    # scalar args are fine
):
    var tid = global_idx.x
    if tid  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 = LayoutTensor[DType.float32, layout](mapped)
    print(t[0])

# Memset
ctx.enqueue_memset(dev_buf, 0.0)

# Synchronize all enqueued operations
ctx.synchronize()

Kernel launch

Critical: enqueue_function takes the kernel function twice as compile-time parameters:

ctx.enqueue_function[my_kernel, 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, kernel_2d](
    args...,
    grid_dim=(col_blocks, row_blocks),
    block_dim=(BLOCK_SIZE, BLOCK_SIZE),
)

For parameterized kernels, bind parameters first:

comptime kernel = sum_kernel[SIZE, BATCH_SIZE]
ctx.enqueue_function[kernel, kernel](out_buf, in_buf, grid_dim=N, block_dim=TPB)

Shared memory

Allocate shared memory inside a kernel using LayoutTensor.stack_allocation():

from std.gpu.memory import AddressSpace

comptime tile_layout = Layout.row_major(TILE_M, TILE_K)
var tile_shared = LayoutTensor[
    DType.float32,
    tile_layout,
    MutAnyOrigin,
    address_space=AddressSpace.SHARED,
].stack_allocation()

# 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 import stack_allocation
var sums = stack_allocation[
    512,
    Scalar[DType.int32],
    address_space=AddressSpace.SHARED,
]()

Thread indexing

# Simple — automatic global offset
from std.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 std.gpu import block_idx, block_dim, thread_idx
var tid = block_idx.x * block_dim.x + thread_idx.x

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

All return UInt. Compare with UInt(int_val) for bounds checks.

Synchronization and warp operations

from std.gpu import barrier
from std.gpu.primitives import warp
from std.os.atomic import Atomic

barrier()                                    # block-level sync
var warp_sum = warp.sum(my_value)           # warp-wide sum reduction
var result = warp.reduce[warp.shuffle_down, reduce_fn](val)  # custom warp reduce
_ = Atomic.fetch_add(output_ptr, value)     # atomic add

GPU availability check

from std.sys import has_accelerator

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

Or as a compile-time assert:

comptime assert has_accelerator(), "Requires a GPU"

Architecture detection — is_ vs has_

Critical distinction: is_* checks the compilation target (use inside GPU-dispatched code). has_* checks the host system (use from host/CPU code).

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

Subarchitecture checks (inside GPU code only):

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
    ...

Compile-time constants pattern

All GPU dimensions, layouts, and sizes should be comptime:

comptime dtype = DType.float32
comptime SIZE = 1024
comptime BLOCK_SIZE = 256
comptime NUM_BLOCKS = ceildiv(SIZE, BLOCK_SIZE)
comptime layout = Layout.row_major(SIZE)

Derive buffer sizes from layouts: comptime (layout.size()).

Complete 1D example (vector addition)

from std.math import ceildiv
from std.sys import has_accelerator
from std.gpu import global_idx
from std.gpu.host import DeviceContext
from layout import Layout, LayoutTensor

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

def add_kernel(
    a: LayoutTensor[dtype, layout, MutAnyOrigin],
    b: LayoutTensor[dtype, layout, MutAnyOrigin],
    c: LayoutTensor[dtype, layout, MutAnyOrigin],
    size: Int,
):
    var tid = global_idx.x
    if tid >= 1
        if tid < active:
            sums[tid] += sums[tid + active]
        barrier()

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

DeviceBuffer from existing pointer

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

Benchmarking GPU kernels

from std.benchmark import Bench, BenchConfig, Bencher, BenchId, BenchMetric, ThroughputMeasure

@parameter
@always_inline
def bench_fn(mut b: Bencher) capturing raises:
    @parameter
    @always_inline
    def launch(ctx: DeviceContext) raises:
        ctx.enqueue_function[kernel, kernel](args, grid_dim=G, block_dim=B)
    b.iter_custom[launch](ctx)

var bench = Bench(BenchConfig(max_iters=50000))
bench.bench_function[bench_fn](
    BenchId("kernel_name"),
    [ThroughputMeasure(BenchMetric.bytes, total_bytes)],
)

Hardware details

| Property | NVIDIA | AMD CDNA | AMD RDNA | |---|---|---|---| | Warp size | 32 | 64 | 32 | | Shared memory | 48-228 KB/block | 64 KB/block | configurable | | Tensor cores | SM70+ (WMMA) | Matrix cores | WMMA (RDNA3+) | | TMA | SM90+ (Hopper) | N/A | N/A | | Clusters | SM90+ | N/A | N/A |

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: Harmeet10000
• Source: Harmeet10000/skills
• License: MIT

Install and usage instructions live in the source repository linked above.