Tiled Matrix Multiplication

fn matmul_tiled[
    layout: Layout, size: Int
](
    out: LayoutTensor[mut=False, dtype, layout],
    a: LayoutTensor[mut=False, dtype, layout],
    b: LayoutTensor[mut=False, dtype, layout],
):
    local_i = thread_idx.x
    local_j = thread_idx.y
    global_i = block_idx.x * TPB + local_i
    global_j = block_idx.y * TPB + local_j

    a_shared = tb[dtype]().row_major[TPB, TPB]().shared().alloc()
    b_shared = tb[dtype]().row_major[TPB, TPB]().shared().alloc()

    var acc: out.element_type = 0

    # Iterate over tiles to compute matrix product
    @parameter
    for tile in range((size + TPB - 1) // TPB):
        # Reset shared memory tiles
        if local_i < TPB and local_j < TPB:
            a_shared[local_i, local_j] = 0
            b_shared[local_i, local_j] = 0

        barrier()

        # Load A tile - global row stays the same, col determined by tile
        if global_i < size and (tile * TPB + local_j) < size:
            a_shared[local_i, local_j] = a[
                global_i, tile * TPB + local_j
            ]

        # Load B tile - row determined by tile, global col stays the same
        if (tile * TPB + local_i) < size and global_j < size:
            b_shared[local_i, local_j] = b[
                tile * TPB + local_i, global_j
            ]

        barrier()

        # Matrix multiplication within the tile
        if global_i < size and global_j < size:

            @parameter
            for k in range(min(TPB, size - tile * TPB)):
                acc += a_shared[local_i, k] * b_shared[k, local_j]

        barrier()

    # Write out final result
    if global_i < size and global_j < size:
        out[global_i, global_j] = acc

The tiled implementation with LayoutTensor handles matrices efficiently by processing them in blocks. Here’s a comprehensive analysis:

Implementation Architecture

1. Thread and Block Organization:

   local_i = thread_idx.x
   local_j = thread_idx.y
   global_i = block_idx.x * TPB + local_i
   global_j = block_idx.y * TPB + local_j
   

   • Defines row-major layout for all tensors
   • Ensures consistent memory access patterns
   • Enables efficient 2D indexing

2. Shared Memory Setup:

   a_shared = tb[dtype]().row_major[TPB, TPB]().shared().alloc()
   b_shared = tb[dtype]().row_major[TPB, TPB]().shared().alloc()
   

   • Uses TensorBuilder for structured allocation
   • Maintains row-major layout for consistency
   • Enables efficient tile processing

Tile Processing Pipeline

1. Tile Iteration:

   @parameter
   for tile in range((size + TPB - 1) // TPB):
   

   • Compile-time unrolled loop
   • Handles matrix size not divisible by TPB
   • Processes matrix in TPB×TPB tiles

2. Shared Memory Reset:

   if local_i < TPB and local_j < TPB:
       a_shared[local_i, local_j] = 0
       b_shared[local_i, local_j] = 0
   

   • Clears previous tile data
   • Ensures clean state for new tile
   • Prevents data corruption

3. Tile Loading:

   # Load A tile - global row stays the same, col determined by tile
   if global_i < size and (tile * TPB + local_j) < size:
       a_shared[local_i, local_j] = a[global_i, tile * TPB + local_j]
   
   # Load B tile - row determined by tile, global col stays the same
   if (tile * TPB + local_i) < size and global_j < size:
       b_shared[local_i, local_j] = b[tile * TPB + local_i, global_j]
   

   • Handles boundary conditions
   • Uses LayoutTensor indexing
   • Maintains memory coalescing

4. Computation:

   @parameter
   for k in range(min(TPB, size - tile * TPB)):
       acc += a_shared[local_i, k] * b_shared[k, local_j]
   

   • Processes current tile
   • Uses shared memory for efficiency
   • Handles partial tiles correctly

Memory Access Optimization

1. Global Memory Pattern:

   A[global_i, tile * TPB + local_j] → Row-major access
   B[tile * TPB + local_i, global_j] → Row-major access
   

   • Maximizes memory coalescing:

     Coalesced Access (Good):          Non-Coalesced Access (Bad):
     Thread0: [M0][M1][M2][M3]         Thread0: [M0][ ][ ][ ]
     Thread1: [M4][M5][M6][M7]    vs   Thread1: [ ][M1][ ][ ]
     Thread2: [M8][M9][MA][MB]         Thread2: [ ][ ][M2][ ]
     Thread3: [MC][MD][ME][MF]         Thread3: [ ][ ][ ][M3]
     ↓                                 ↓
     1 memory transaction              4 memory transactions
     

     When threads access consecutive memory locations (left), the GPU can combine multiple reads into a single transaction. When threads access scattered locations (right), each access requires a separate transaction, reducing performance.

2. Shared Memory Usage:

   a_shared[local_i, k] → Row-wise access
   b_shared[k, local_j] → Row-wise access
   

   • Optimized for matrix multiplication
   • Reduces bank conflicts
   • Enables data reuse

Synchronization and Safety

1. Barrier Points:

   barrier()  # After shared memory reset
   barrier()  # After tile loading
   barrier()  # After computation
   

   • Ensures shared memory consistency
   • Prevents race conditions
   • Maintains thread cooperation

2. Boundary Handling:

   if global_i < size and global_j < size:
       out[global_i, global_j] = acc
   

   • Prevents out-of-bounds access
   • Handles matrix edges
   • Safe result writing

Key Optimizations

1. Layout Optimization:

   • Row-major layout for all tensors
   • Efficient 2D indexing
   • Proper alignment

2. Memory Access:

   • Coalesced global memory loads
   • Efficient shared memory usage
   • Minimal bank conflicts

3. Computation:

   • Register-based accumulation
   • Compile-time loop unrolling
   • Efficient tile processing

This implementation achieves high performance through:

• Efficient use of LayoutTensor for memory access
• Optimal tiling strategy
• Proper thread synchronization
• Careful boundary handling