Tiled Matrix Multiplication

The tiled matrix multiplication implementation demonstrates efficient handling of large matrices \((8 \times 8)\) using small tiles \((3 \times 3)\). Here’s how it works:

1. Thread indexing setup

   • Global position calculation:

     tiled_row = block_idx.y * TPB + thread_idx.y
     tiled_col = block_idx.x * TPB + thread_idx.x
     

   • Local position in tile:

     local_row = thread_idx.y
     local_col = thread_idx.x
     

2. Shared memory allocation

   Input matrices (8×8):
   A = [0  1  2  3  4  5  6  7 ]    B = [0  2  4  6  8  10 12 14]
       [8  9  10 11 12 13 14 15]        [16 18 20 22 24 26 28 30]
       [16 17 18 19 20 21 22 23]        [32 34 36 38 40 42 44 46]
       [24 25 26 27 28 29 30 31]        [48 50 52 54 56 58 60 62]
       [32 33 34 35 36 37 38 39]        [64 66 68 70 72 74 76 78]
       [40 41 42 43 44 45 46 47]        [80 82 84 86 88 90 92 94]
       [48 49 50 51 52 53 54 55]        [96 98 100 102 104 106 108 110]
       [56 57 58 59 60 61 62 63]        [112 114 116 118 120 122 124 126]
   
   Shared memory per block (3×3):
   a_shared[TPB, TPB]  b_shared[TPB, TPB]
   

3. Tile processing loop

   Number of tiles = (8 + 3 - 1) // 3 = 3 tiles
   
   For each tile:
   1. Reset shared memory
   2. Load tile from A and B
   3. Compute partial products
   4. Accumulate in register
   

4. Memory loading pattern

   • Loading A tile:

     if tiled_row < size and (tile * TPB + local_col) < size:
         a_shared[local_row, local_col] = a[tiled_row, tile * TPB + local_col]
     

   • Loading B tile:

     if (tile * TPB + local_row) < size and tiled_col < size:
         b_shared[local_row, local_col] = b[tile * TPB + local_row, tiled_col]
     

5. Computation within tile

   For k in range(min(TPB, size - tile * TPB)):
       acc += a_shared[local_row, k] * b_shared[k, local_col]
   

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

6. Synchronization points

   barrier() after:
   1. Shared memory reset
   2. Tile loading
   3. Tile computation
   

Key performance features:

• Processes 8×8 matrix using 3×3 tiles
• Uses shared memory for fast tile access
• Minimizes global memory transactions
• Handles matrix boundaries correctly
• Maintains coalesced memory access

2. Boundary handling:

   if row < size and col < size:
       out[row, col] = 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

2. Memory access:

   • Coalesced global memory loads
   • Efficient shared memory usage

3. Computation:

   • Register-based accumulation i.e. var acc: out.element_type = 0
   • Compile-time loop unrolling via @parameter

This implementation achieves high performance through:

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

The idiomatic tiled matrix multiplication leverages Mojo’s LayoutTensor API and asynchronous memory operations for a cleaner implementation:

1. LayoutTensor tile API

   out_tile = out.tile[TPB, TPB](block_idx.y, block_idx.x)
   a_tile = a.tile[TPB, TPB](block_idx.y, idx)
   b_tile = b.tile[TPB, TPB](idx, block_idx.x)
   

   This directly expresses “get the tile at position (block_idx.y, block_idx.x)” without manual coordinate calculation. See the documentation for more details.

2. Asynchronous memory operations

   copy_dram_to_sram_async[thread_layout=load_a_layout](a_shared, a_tile)
   copy_dram_to_sram_async[thread_layout=load_b_layout](b_shared, b_tile)
   async_copy_wait_all()
   

   These operations:

   • Launch asynchronous memory transfers that may overlap with computation via copy_dram_to_sram_async
   • Use specialized thread layouts for optimal memory access patterns
   • Eliminate the need for manual memory initialization

3. Specialized compile-time load layouts

   alias load_a_layout = Layout.row_major(1, TPB)
   alias load_b_layout = Layout.row_major(TPB, 1)
   

   These layouts optimize how threads cooperate during memory transfers:

   • load_a_layout: Each thread loads a slice of a row (coalesced access)
   • load_b_layout: Each thread loads a slice of a column (transposed access)

4. Efficient thread synchronization

   // Wait for async operations to complete
   async_copy_wait_all()
   // Ensure all threads can see the shared memory contents
   barrier()
   

   The barriers ensure proper synchronization:

   • After memory transfers complete
   • After computation for each tile

5. Proper boundary handling

   if block_idx.y * TPB + local_row < size and block_idx.x * TPB + local_col < size:
       out_tile[local_row, local_col] = acc
   

   This critical check prevents out-of-bounds writes for blocks at the matrix boundaries.

6. Tile processing loop

   for idx in range((size + TPB - 1) // TPB):
      // Process one tile
   

   Uses ceiling division to handle matrices whose dimensions aren’t perfect multiples of the tile size.

Performance considerations

The idiomatic implementation maintains the performance benefits of tiling while providing cleaner abstractions:

1. Memory locality: Exploits spatial and temporal locality through tiling
2. Coalesced access: Specialized load layouts ensure coalesced memory access patterns
3. Compute-memory overlap: Potential overlap through asynchronous memory operations
4. Shared memory efficiency: No redundant initialization of shared memory
5. Register pressure: Uses accumulation registers for optimal compute throughput

This implementation shows how high-level abstractions can express complex GPU algorithms without sacrificing performance. It’s a prime example of Mojo’s philosophy: combining high-level expressiveness with low-level performance control.

Key differences from manual tiling

The idiomatic approach is not just cleaner but also potentially more performant due to the use of specialized memory layouts and asynchronous operations.