Multi-Block Coordination Basics

fn cluster_coordination_basics[
    in_layout: Layout, out_layout: Layout, tpb: Int
](
    output: LayoutTensor[mut=True, dtype, out_layout],
    input: LayoutTensor[mut=False, dtype, in_layout],
    size: Int,
):
    """Real cluster coordination using SM90+ cluster APIs."""
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x

    # DIAGNOSTIC: Check what's happening with cluster ranks
    my_block_rank = Int(block_rank_in_cluster())
    block_id = Int(block_idx.x)

    shared_data = tb[dtype]().row_major[tpb]().shared().alloc()

    # FIX: Use block_idx.x for data distribution instead of cluster rank
    # Each block should process different portions of the data
    var data_scale = Float32(
        block_id + 1
    )  # Use block_idx instead of cluster rank

    # Phase 1: Each block processes its portion
    if global_i < size:
        shared_data[local_i] = input[global_i] * data_scale
    else:
        shared_data[local_i] = 0.0

    barrier()

    # Phase 2: Use cluster_arrive() for inter-block coordination
    cluster_arrive()  # Signal this block has completed processing

    # Block-level aggregation (only thread 0)
    if local_i == 0:
        var block_sum: Float32 = 0.0
        for i in range(tpb):
            block_sum += shared_data[i][0]
        # FIX: Store result at block_idx position (guaranteed unique per block)
        output[block_id] = block_sum

    # Wait for all blocks in cluster to complete
    cluster_wait()

The cluster coordination solution demonstrates the fundamental multi-block synchronization pattern using a carefully orchestrated two-phase approach:

Phase 1: Independent block processing

Thread and block identification:

global_i = block_dim.x * block_idx.x + thread_idx.x  # Global thread index
local_i = thread_idx.x                               # Local thread index within block
my_block_rank = Int(block_rank_in_cluster())         # Cluster rank (0-3)
block_id = Int(block_idx.x)                          # Block index for reliable addressing

Shared memory allocation and data processing:

• Each block allocates its own shared memory workspace: tb[dtype]().row_major[tpb]().shared().alloc()
• Scaling strategy: data_scale = Float32(block_id + 1) ensures each block processes data differently
  • Block 0: multiplies by 1.0, Block 1: by 2.0, Block 2: by 3.0, Block 3: by 4.0
• Bounds checking: if global_i < size: prevents out-of-bounds memory access
• Data processing: shared_data[local_i] = input[global_i] * data_scale scales input data per block

Intra-block synchronization:

• barrier() ensures all threads within each block complete data loading before proceeding
• This prevents race conditions between data loading and subsequent cluster coordination

Phase 2: Cluster coordination

Inter-block signaling:

• cluster_arrive() signals that this block has completed its local processing phase
• This is a non-blocking operation that registers completion with the cluster hardware

Local aggregation (Thread 0 only):

if local_i == 0:
    var block_sum: Float32 = 0.0
    for i in range(tpb):
        block_sum += shared_data[i][0]  # Sum all elements in shared memory
    output[block_id] = block_sum        # Store result at unique block position

• Only thread 0 performs the sum to avoid race conditions
• Results stored at output[block_id] ensures each block writes to unique location

Final synchronization:

• cluster_wait() blocks until ALL blocks in the cluster have completed their work
• This ensures deterministic completion order across the entire cluster

Key technical insights

Why use block_id instead of my_block_rank?

• block_idx.x provides reliable grid-launch indexing (0, 1, 2, 3)
• block_rank_in_cluster() may behave differently depending on cluster configuration
• Using block_id guarantees each block gets unique data portions and output positions

Memory access pattern:

• Global memory: Each thread reads input[global_i] exactly once
• Shared memory: Used for intra-block communication and aggregation
• Output memory: Each block writes to output[block_id] exactly once

Synchronization hierarchy:

1. barrier(): Synchronizes threads within each block (intra-block)
2. cluster_arrive(): Signals completion to other blocks (inter-block, non-blocking)
3. cluster_wait(): Waits for all blocks to complete (inter-block, blocking)

Performance characteristics:

• Compute complexity: O(TPB) per block for local sum, O(1) for cluster coordination
• Memory bandwidth: Each input element read once, minimal inter-block communication
• Scalability: Pattern scales to larger cluster sizes with minimal overhead