For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo package

gpu

GPU programming primitives.

These low level constructs allow you to write code that runs on the GPU with traditional programming style--partitioning work across threads that are mapped onto 1-, 2-, or 3-dimensional blocks. The thread blocks can subsequently be grouped into a grid of thread blocks.

A kernel is a function that runs on the GPU in parallel across many threads. Currently, the DeviceContext struct provides the interface for compiling and launching GPU kernels inside MAX custom operations.

The gpu.host package includes APIs to manage interaction between the host (that is, the CPU) and device (that is, the GPU or accelerator).

The gpu package exports aliases you can use to access information about the grid and the current thread, including block dimensions, block index in the grid, and thread index. Import these directly from gpu:

from max.gpu import block_dim, block_idx, thread_idx, global_idx

For an example of launching a GPU kernel from a MAX custom operation, see the vector addition example in the MAX repo.

comptime values

block_dim

comptime block_dim = _BlockDim()

Contains the dimensions of the block as x, y, and z values.

For example: block_dim.y.

block_id_in_cluster

comptime block_id_in_cluster = _ClusterBlockIdx()

Contains the block id of the threadblock within a cluster, as x, y, and z values.

block_idx

comptime block_idx = _BlockIdx()

Contains the block index in the grid, as x, y, and z values.

cluster_dim

comptime cluster_dim = _ClusterDim()

Contains the dimensions of the cluster, as x, y, and z values.

cluster_idx

comptime cluster_idx = _ClusterIdx()

Contains the cluster index in the grid, as x, y, and z values.

global_idx

comptime global_idx = _GlobalIdx()

Contains the global offset of the kernel launch, as x, y, and z values.

grid_dim

comptime grid_dim = _GridDim()

Provides accessors for getting the x, y, and z dimensions of a grid.

MAX_THREADS_PER_BLOCK_METADATA

comptime MAX_THREADS_PER_BLOCK_METADATA = _resolve_max_threads_per_block_metadata()

This is metadata tag that is used in conjunction with __llvm_metadata to give a hint to the compiler about the max threads per block that's used.

thread_idx

comptime thread_idx = _ThreadIdx()

Contains the thread index in the block, as x, y, and z values.

WARP_SIZE

comptime WARP_SIZE = _resolve_warp_size()

The number of threads that execute in lockstep within a warp on the GPU.

This constant represents the hardware warp size, which is the number of threads that execute instructions synchronously as a unit. The value is architecture-dependent:

• 32 threads per warp on NVIDIA GPUs
• 32 threads per warp on AMD RDNA GPUs
• 64 threads per warp on AMD CDNA GPUs
• 0 if no GPU is detected

The warp size is a fundamental parameter that affects:

• Thread scheduling and execution
• Memory access coalescing
• Synchronization primitives
• Overall performance optimization

Packages

• ​compute: GPU compute operations package - MMA and tensor core operations.
• ​host: Implements the gpu host package.
• ​memory: GPU memory operations package.
• ​primitives: GPU primitives package - warp, block, and cluster operations.
• ​sync: GPU synchronization primitives package.

Modules

• ​globals: This module provides GPU-specific global constants and configuration values.
• ​intrinsics: Provides low-level GPU intrinsic operations and memory access primitives.

Functions

• ​lane_id: Returns the lane ID of the current thread within its warp.
• ​sm_id: Returns the Streaming Multiprocessor (SM) ID of the current thread.
• ​warp_id: Returns the warp ID of the current thread within its block. The warp ID is a unique identifier for each warp within a block, ranging from 0 to BLOCK_SIZE/WARP_SIZE-1. This ID is commonly used for warp-level programming and synchronization within a block.