Login
📚 The Mojo Manual
Chapters ▾
1 Get Started2 Language Basics3 Value Ownership4 Value Lifecycle5 Metaprogramming6 Pointers7 GPU Programming8 TileTensor9 Layouts and LayoutTensor10 Python Interop
⇩ Download ▾
✎ Open in XYZ Type (editable)⇩ This section (Typst)

7.3 GPU block and warp operations and synchronization

When multiple GPU threads write to the same memory location without a defined order of execution, a race condition occurs. The final outcome of the computation becomes non-deterministic, depending on the scheduling and timing of execution of threads by the GPU hardware. Such bugs are notoriously difficult to debug because they may not appear consistently in every run.

To write correct and robust parallel programs, you need explicit mechanisms to coordinate the execution of threads and manage the visibility of their memory operations. These mechanisms are known as synchronization primitives. They are not merely performance optimizations; they are essential tools for correctness. Without them, threads operate in complete isolation, unable to safely share intermediate results, divide complex tasks, or perform the collective computations that are the hallmark of high-performance GPU algorithms.

A barrier is a fundamental synchronization primitive that creates a meeting point in the program where all participating threads must wait for each other. When a thread reaches a barrier, it pauses execution until every other thread in the group also arrives. This ensures that all threads proceed together past the barrier, maintaining consistent state and preventing race conditions when accessing shared data.

Mojo provides two complementary categories of GPU coordination tools. Synchronization primitives like barrier and syncwarp coordinate thread execution and enforce memory visibility, but they don't perform computation themselves—they're pure coordination mechanisms. In contrast, collective operations like reductions, broadcasts, and prefix sums combine synchronization with common computational patterns: they coordinate threads and compute a result. Use synchronization primitives when you need explicit control over when threads coordinate (such as managing access to shared memory between distinct phases of an algorithm), and use collective operations when you need to aggregate or distribute data across threads (such as computing a sum or maximum across a thread block or warp). Both types of tools are essential for writing correct and efficient GPU code, and understanding when to use each is key to building robust parallel algorithms.

This guide covers Mojo's low-level primitives for managing coordination at the thread block and warp levels. For foundational GPU architecture concepts and detailed explanations of the GPU execution model, see Intro to GPUs. For a discussion of basic kernel creation and device management, see GPU programming fundamentals. The tiled matrix multiplication example later in this chapter uses TileTensor and stack_allocation to allocate shared memory tiles. For an introduction to TileTensor itself—including how to create one in global, shared, or local memory—see Using TileTensor.

We'll explore Mojo's synchronization and collective communication primitives for coordinating parallel work on the GPU. Key topics include:

Block-level synchronization and operations

This section covers coordination mechanisms for all threads within a thread block:

These tools serve different but complementary purposes: barrier is a pure synchronization primitive for coordinating execution and memory visibility, while block reduction operations are collective computations that internally handle their own synchronization. You can use barrier to build custom coordination patterns, or use block reductions when you need both coordination and computation together.

The barrier primitive

The gpu.sync.barrier function is the primary mechanism for coordinating all threads within a single thread block. It creates a synchronization point in the kernel's execution flow that no thread can pass until every other thread in its block has also reached that point.

The barrier primitive does two things: it acts as both an execution barrier and a memory fence.

The most common use case for barrier is managing access to the fast, on-chip shared memory shared by all threads within a block. Here's how a typical algorithm works:

  1. Threads in a block cooperatively load a segment of data from the high-latency global memory into a shared memory array. Each thread is responsible for loading one or more elements.

  2. A call to barrier is made. This is essential to ensure that the entire data segment is fully loaded into shared memory before any thread attempts to use it.

  3. Threads perform computations, reading from and writing to the shared memory array. This phase leverages the low latency of shared memory to accelerate the algorithm.

  4. If the computation itself involves multiple stages of shared memory communication, another barrier call may be necessary to ensure the results of one stage are visible before the next begins.

  5. Finally, threads write their results from shared memory back to global memory.

The Mojo barrier function is functionally equivalent to the __syncthreads intrinsic in both NVIDIA CUDA and AMD HIP and threadgroup_barrier(mem_flags::mem_threadgroup) in Apple Metal, providing a portable syntax for this fundamental operation.

Block-level reduction operations

In addition to the basic barrier primitive, Mojo provides higher-level block-wide collective operations through the gpu.primitives.block module. These operations combine the coordination functionality of barrier with common computational patterns, offering both convenience and performance benefits.

The gpu.primitives.block module includes several reduction primitives:

These operations automatically handle the necessary synchronization and shared memory management internally, making them both easier to use correctly and often more efficient than manually implementing equivalent functionality with barrier and shared memory operations.

Using block operations in practice

Block-level operations are commonly used in multi-stage algorithms where threads must coordinate through shared memory. A typical pattern involves:

  1. Load phase: Threads cooperatively load data into shared memory
  2. Synchronize: Use barrier to ensure all data is loaded
  3. Compute phase: Process data using shared memory

  4. Reduce phase: Use block reduction operations to aggregate results

This pattern appears in algorithms like tiled matrix multiplication, stencil operations, and parallel reductions, where the combination of shared memory and proper synchronization enables significant performance improvements over naive approaches.

Block synchronization example: tiled matrix multiplication

Matrix multiplication benefits from a technique called tiling, where we break large matrices into smaller tiles that fit in the GPU's fast shared memory. Instead of repeatedly reading from slow global memory, threads in a block cooperatively load a tile into shared memory once, then all threads can access it multiple times. This creates a classic producer-consumer pattern: threads work together to load data (producer phase), then all threads use that data to compute results (consumer phase). Without proper synchronization between these phases, the algorithm produces incorrect results. For a deeper understanding of the tiling strategy, see this section of our blog post on optimizing matrix multiplication on NVIDIA's Blackwell.

tiled_matmul.mojo
from std.math import ceildiv
from std.sys import exit, has_accelerator

# GPU programming imports from open source stdlib
from std.gpu.sync import barrier
from std.gpu.host import DeviceContext
from std.gpu import thread_idx, block_idx
from std.gpu.memory import AddressSpace

# TileTensor support from open source layout package
from layout import TileTensor, stack_allocation
from layout.tile_layout import row_major

# Data type selection: float32 provides good balance of precision and performance
comptime float_dtype = DType.float32

# Matrix dimensions: chosen to be small enough for easy understanding
# while still demonstrating tiling concepts effectively
comptime MATRIX_SIZE = 64  # 64x64 matrices
comptime MATRIX_M = MATRIX_SIZE  # Number of rows in matrices A and C
comptime MATRIX_N = MATRIX_SIZE  # Number of columns in matrices B and C
comptime MATRIX_K = MATRIX_SIZE  # Shared dimension (A cols = B rows)

# Tile dimensions: chosen to fit comfortably in GPU shared memory
# and demonstrate clear blocking behavior
comptime TILE_SIZE = 16  # 16x16 tiles balance memory usage and parallelism
comptime TILE_M = TILE_SIZE  # Tile height for matrix A and C
comptime TILE_N = TILE_SIZE  # Tile width for matrix B and C
comptime TILE_K = TILE_SIZE  # Tile depth for the K dimension

# Derived constants
comptime NUM_TILES_PER_SIDE = MATRIX_SIZE // TILE_SIZE  # Number of tiles per matrix side (4)
comptime THREADS_PER_TILE = TILE_SIZE * TILE_SIZE  # Threads needed per tile (256)
comptime TOTAL_TILES_TO_PROCESS = NUM_TILES_PER_SIDE  # Tiles to process in K dimension

# TileTensor provides type-safe multi-dimensional data access with automatic memory layout handling
# Layout definitions using example matrix dimensions
comptime matrix_a_layout = row_major[MATRIX_M, MATRIX_K]()  # A: M x K
comptime matrix_b_layout = row_major[MATRIX_K, MATRIX_N]()  # B: K x N
comptime matrix_c_layout = row_major[MATRIX_M, MATRIX_N]()  # C: M x N

# Layout definitions for tile access
comptime tile_a_layout = row_major[TILE_M, TILE_K]()
comptime tile_b_layout = row_major[TILE_K, TILE_N]()


def tiled_matmul_kernel(
    matrix_a: TileTensor[float_dtype, type_of(matrix_a_layout), MutAnyOrigin],
    matrix_b: TileTensor[float_dtype, type_of(matrix_b_layout), MutAnyOrigin],
    matrix_c: TileTensor[float_dtype, type_of(matrix_c_layout), MutAnyOrigin],
):
    # Thread and block indices
    var thread_x = thread_idx.x
    var thread_y = thread_idx.y
    var block_x = block_idx.x
    var block_y = block_idx.y

    # Global matrix coordinates
    var global_row = block_y * TILE_M + thread_y
    var global_col = block_x * TILE_N + thread_x

    # Tile starting positions
    var tile_row_start = block_y * TILE_M
    var tile_col_start = block_x * TILE_N

    # Allocate shared memory tiles for fast on-chip access
    var tile_a_shared = stack_allocation[
        float_dtype, address_space=AddressSpace.SHARED
    ](tile_a_layout)

    var tile_b_shared = stack_allocation[
        float_dtype, address_space=AddressSpace.SHARED
    ](tile_b_layout)

    # Initialize accumulator and start tiling loop
    var accumulator: matrix_c.ElementType = 0.0

    # Iterate through tiles along K dimension
    # Use comptime for to unroll the loop at compile time
    comptime for k_tile in range(0, MATRIX_K, TILE_K):
        # Cooperative tile loading
        # Calculate global coordinates for tile loading
        var a_global_row = tile_row_start + thread_y
        var a_global_col = k_tile + thread_x
        var b_global_row = k_tile + thread_y
        var b_global_col = tile_col_start + thread_x

        # Bounds checking
        var load_a_valid = (a_global_row < MATRIX_M) and (
            a_global_col < MATRIX_K
        )
        var load_b_valid = (b_global_row < MATRIX_K) and (
            b_global_col < MATRIX_N
        )

        # Load tiles into shared memory with bounds checking
        if load_a_valid:
            tile_a_shared[thread_y, thread_x] = matrix_a[
                a_global_row, a_global_col
            ]
        else:
            tile_a_shared[thread_y, thread_x] = 0.0

        if load_b_valid:
            tile_b_shared[thread_y, thread_x] = matrix_b[
                b_global_row, b_global_col
            ]
        else:
            tile_b_shared[thread_y, thread_x] = 0.0

        # Ensure all threads finish loading tiles before any thread starts computing
        barrier()

        # Compute dot product using shared memory tiles
        comptime for k in range(TILE_K):
            var a_element = tile_a_shared[thread_y, k]
            var b_element = tile_b_shared[k, thread_x]
            accumulator += a_element * b_element

        # Ensure all threads finish computing before any thread loads next tiles
        barrier()

    # Write final result to global memory with bounds checking
    if (global_row < MATRIX_M) and (global_col < MATRIX_N):
        matrix_c[global_row, global_col] = accumulator

This tiled algorithm leverages the GPU's memory hierarchy for better performance. Shared memory is an on-chip cache that's much faster than global memory, but it's limited in size—a typical block might have only 48KB available. We break the computation into stages: threads cooperatively load small tiles from global memory into this fast shared memory, perform computations on those tiles, then repeat for the next set of tiles. Each thread loads one element per tile, creating coalesced memory accesses that maximize bandwidth. Once a tile sits in shared memory, all threads in the block can access it repeatedly without triggering expensive global memory reads.

The first barrier call appears immediately after the cooperative tile loading phase. This synchronization point is critical: it ensures that all threads in the block finish writing their elements to shared memory before any thread begins reading from it to compute results. Without this barrier, you'd have a classic read-before-write race condition. Fast threads could race ahead and start reading from shared memory locations that slow threads haven't populated yet, leading to incorrect results from uninitialized data. Even worse, the bug would be non-deterministic—sometimes the code would work (if threads happened to execute in a favorable order), and sometimes it would fail, making debugging extremely difficult. The barrier eliminates this unpredictability by establishing a clear happens-before relationship: all writes complete before any reads begin.

The second barrier call appears at the end of the computation phase, right before the loop continues to load the next set of tiles. This barrier solves the opposite problem: it prevents write-during-read races. Without it, fast threads could finish their computations and start loading new tile data into shared memory while slow threads are still reading the old data for their calculations. This would corrupt the shared memory with partially overwritten values, again producing incorrect results. The pattern is symmetric: the first barrier protects readers from seeing incomplete writes, while the second protects readers from concurrent overwrites. Together, these two barriers implement a safe producer-consumer cycle: load → barrier → compute → barrier → repeat. Both barriers are absolutely essential—removing either one breaks the algorithm's correctness.

Warp-level operations

While the thread block is the scope for shared memory communication, the warp is the fundamental unit of execution scheduling. Because threads within a warp are executed simultaneously by the hardware, communication between them happens much faster than communication that requires coordination across different warps. Mojo provides a suite of primitives for these high-speed, intra-warp operations, which are essential for many performance-critical optimization patterns.

This section covers high-speed coordination and data exchange within a single warp:

Warp-level synchronization

The gpu.sync.syncwarp function provides a more granular synchronization barrier that operates only on the threads within a single warp.

This function handles thread divergence. On some GPU architectures, threads within a warp can follow different execution paths due to conditional branching. syncwarp forces the specified threads in the warp to reconverge at a single point before proceeding.

The syncwarp function takes an optional mask argument. This is a 32-bit or 64-bit integer (depending on the warp size of the architecture) that acts as a bitmask. The ith bit of the mask corresponds to the thread at lane i within the warp. If a bit is set to 1, the corresponding thread participates in the synchronization; if it is 0, it does not. The default value of -1 (all bits set to 1) synchronizes all threads in the warp.

Understanding syncwarp requires knowing its platform-dependent behavior, which Mojo's portable API abstracts away:

This difference highlights a key benefit of Mojo. You write code against a single, portable API. The compiler is responsible for generating the correct, architecture-specific code. Therefore, if an algorithm relies on syncwarp for correctness on NVIDIA hardware, it still behaves as expected on other vendors' hardware.

Warp-level data exchange

Shuffle operations are the cornerstone of high-performance warp-level algorithms. These primitives enable threads within a warp to exchange data directly through registers, making them essential for implementing efficient parallel patterns like reductions, stencil computations, and sliding window operations.

Unlike shared memory communication that requires explicit synchronization and memory transactions, shuffle operations use the warp's simultaneous execution to achieve near-zero latency data exchange. This makes them ideal for:

Warp shuffle operations

The gpu.primitives.warp module provides five shuffle primitives, each optimized for specific data movement patterns:

All of these primitives other than broadcast take an optional mask argument that serves a dual purpose:

  1. Thread participation: The mask specifies which threads participate in the shuffle operation. It is a 32-bit or 64-bit integer (depending on the warp size) where the ith bit corresponds to lane i. If a bit is set to 1, that thread participates; if 0, it does not.

  2. Implicit synchronization: The mask also provides automatic synchronization for all participating threads. All threads whose bits are set in the mask will be synchronized before the shuffle completes, ensuring correct data exchange even after divergent control flow.

The default value of -1 (all bits set to 1) includes all threads in the warp.

These five primitives form the foundation for complex warp-level algorithms and serve as building blocks for higher-level collective operations.

Choosing the right shuffle primitive

While each shuffle primitive can technically perform any data exchange pattern, certain operations naturally fit specific use cases. Understanding these patterns helps you write more efficient and readable code.

Here are some specific patterns where each primitive excels.

shuffle_idx and broadcast work well for:

shuffle_up and shuffle_down are perfect for:

shuffle_xor excels at:

When optimizing performance, prefer shuffle operations over shared memory for register-sized data, and remember that shuffle_xor typically has the best instruction scheduling characteristics for reduction patterns.

Warp-level reduction operations

The gpu.primitives.warp module also provides higher-level functions for performing common reduction computations across all threads in a warp. These functions take advantage of hardware-specific intrinsics where possible, and fall back to shuffle-based reduction on other architectures:

You'll find these primitives most useful when computing aggregates across threads that are already working closely together. Use sum for computing totals, averages, or accumulating values across small data segments. The max and min functions work well for finding extremes in your data or implementing voting mechanisms where threads need consensus. prefix_sum is particularly valuable for scan operations—computing running totals or building cumulative results as you process data. It's essential for algorithms that need to track "how much have we processed so far?" at each step. These operations are significantly faster than a block-level reduction that uses shared memory and barrier calls.

Using warp operations in practice

Warp operations excel in algorithms that require frequent, fine-grained data exchange between nearby threads. Common patterns include:

These operations are particularly valuable in algorithms where the overhead of block-level synchronization would be prohibitive, such as in inner loops of compute-intensive kernels or when processing data that naturally aligns with warp boundaries.

Advanced synchronization mechanisms

Beyond the fundamental barrier and syncwarp primitives, Mojo provides additional synchronization mechanisms for specialized use cases and architecture-specific optimizations. These advanced primitives enable fine-grained control over memory ordering, asynchronous operations, and instruction scheduling. However, most of these mechanisms are available only on specific architectures. Consult the Mojo API reference documentation for the latest information on availability.

Mechanisms currently available only on NVIDIA GPUs:

Mechanisms currently available only on AMD GPUs:

Best practices and common pitfalls

Now that we've covered the core synchronization primitives, let's focus on using them correctly to write reliable and portable GPU code. Understanding common pitfalls is essential for avoiding bugs that are difficult to reproduce and debug.

This section provides guidance for writing correct, portable, and efficient GPU code:

Writing correct synchronized code

Correctness should always be your first priority. The following issues are common sources of bugs in parallel programs.

Understanding and avoiding race conditions

As a quick reminder, a race condition occurs when multiple threads write to the same memory location without a defined order of execution, leading to a non-deterministic outcome. Here's a simple example where threads attempt to update a shared counter:

# INCORRECT: Race condition
shared_counter[0] += my_value  # Multiple threads modify same location

This leads to "lost updates" because the read-modify-write sequence isn't atomic. To prevent this, you must use synchronization primitives like barrier to coordinate access or use Atomic operations for simple updates. For example, you could use the Atomic.fetch_add method to atomically increment the counter:

# CORRECT: Atomic increment
_ = Atomic.fetch_add(shared_counter[0], my_value)

Avoiding deadlocks with barrier

A barrier must be encountered by all threads within a block to avoid a deadlock. Placing a barrier inside a conditional statement is a frequent source of bugs. If the condition causes some threads to execute the barrier while others skip it, the threads that reach the barrier will wait indefinitely for the others to arrive, causing the kernel to hang.

Therefore, barrier should be used in conditional code only if it's guaranteed that all threads in the block will evaluate the condition identically and follow the same execution path.

When to use syncwarp

The syncwarp primitive is needed when coordinating access to shared or global memory after divergent control flow within a warp. However, it is not needed before warp shuffle operations or warp reduction operations, as those operations provide their own implicit synchronization via the mask parameter.

Use syncwarp when:

Do not use syncwarp before:

Here's an example where syncwarp is needed (for shared memory coordination):

if thread_idx.x < 16:
    shared_data[thread_idx.x] = compute_something()
else:
    shared_data[thread_idx.x] = compute_something_else()

# syncwarp() needed here because threads diverged before writing to shared memory
syncwarp()
var result = shared_data[some_index]  # Now safe to read

And here's an example where syncwarp is not needed (shuffle operations):

if thread_idx.x < 16:
    value = compute_something()
else:
    value = compute_something_else()

# No syncwarp() needed - shuffle_down() synchronizes automatically via its mask
result = warp.shuffle_down(value, 1)

Handling shuffle boundary conditions

When using shuffle_up and shuffle_down, be mindful of edge cases. A thread will receive an undefined value if the source lane is out of bounds (for example, current_lane - offset < 0). When implementing patterns like sliding windows, you must add logic to handle these boundary conditions correctly.

Choosing the right synchronization level

The core principle for designing efficient GPU algorithms is to coordinate between warps with barrier and shared memory, and optimize within warps with gpu.primitives.warp primitives.

This hierarchical approach mirrors the GPU's architecture. Intra-warp communication is extremely fast, while cross-warp communication is more expensive.

Writing portable GPU code

Mojo is designed to write portable GPU code, but it's helpful to understand how.

First, Mojo's GPU operations have automatic fallback mechanisms. For example, a gpu.primitives.warp.max call will automatically use specialized redux instructions on the newest NVIDIA hardware but will fall back to a shuffle-based implementation that works on any other GPU. You get performance where available and correctness everywhere else.

Second, always avoid hardcoding hardware-specific values. The most common mistake is assuming a warp size of 32. Use the gpu.WARP_SIZE constant to ensure your code works correctly on all vendors' hardware.

Finally, for highly-tuned kernels, you can use comptime if blocks to write architecture-specific code paths while keeping a single source file.

from std.sys import is_amd_gpu, is_apple_gpu, is_nvidia_gpu

def adaptive_algorithm():
    comptime if is_nvidia_gpu():
        nvidia_optimized_path()
    elif is_amd_gpu():
        amd_optimized_path()
    elif is_apple_gpu():
        apple_optimized_path()
    else:
        # Conservative fallback for future hardware support
        portable_path()

Debugging synchronization issues

Synchronization bugs can be tricky. Here are some strategies to find them:

Conclusion and key takeaways

Summary of primitives and patterns

We've covered Mojo's low-level toolkit for managing concurrency and communication in GPU kernels. These primitives are the fundamental building blocks for writing correct and high-performance parallel algorithms.

The core mental model

The effective use of these primitives stems from a hierarchical approach to GPU algorithm design. Your key takeaway should be to coordinate between warps with barrier and shared memory, and optimize within warps with gpu.primitives.warp primitives.

This principle encourages you to structure algorithms to maximize intra-warp computation and communication, which is extremely fast, and to use the more costly block-level synchronization only when necessary to combine results or manage data dependencies between warps.

Next steps

To gain hands-on experience with the concepts in this guide, we encourage you to explore the following resources: