Tiled Matrix Multiplication

def matmul_tiled[
    size: Int
](
    output: TileTensor[mut=True, dtype, LayoutTiledType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutTiledType, ImmutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutTiledType, ImmutAnyOrigin],
):
    var local_row = thread_idx.y
    var local_col = thread_idx.x
    var tiled_row = block_idx.y * TPB + local_row
    var tiled_col = block_idx.x * TPB + local_col

    var a_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB, TPB]())
    var b_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB, TPB]())

    var acc: output.ElementType = 0

    # Iterate over tiles to compute matrix product
    comptime for tile in range((size + TPB - 1) // TPB):
        # Load A tile - global row stays the same, col determined by tile
        if tiled_row < size and (tile * TPB + local_col) < size:
            a_shared[local_row, local_col] = a[
                tiled_row, tile * TPB + local_col
            ]

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

        barrier()

        # Matrix multiplication within the tile
        if tiled_row < size and tiled_col < size:
            comptime for k in range(min(Int(TPB), Int(size - tile * TPB))):
                acc += a_shared[local_row, k] * b_shared[k, local_col]

        barrier()

    # Write out final result
    if tiled_row < size and tiled_col < size:
        output[tiled_row, tiled_col] = acc

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

1. Shared memory allocation

   Input matrices (9×9) - Perfect fit for (3×3) tiling:
   A = [0  1  2  3  4  5  6  7  8 ]    B = [0  2  4  6  8  10 12 14 16]
       [9  10 11 12 13 14 15 16 17]        [18 20 22 24 26 28 30 32 34]
       [18 19 20 21 22 23 24 25 26]        [36 38 40 42 44 46 48 50 52]
       [27 28 29 30 31 32 33 34 35]        [54 56 58 60 62 64 66 68 70]
       [36 37 38 39 40 41 42 43 44]        [72 74 76 78 80 82 84 86 88]
       [45 46 47 48 49 50 51 52 53]        [90 92 94 96 98 100 102 104 106]
       [54 55 56 57 58 59 60 61 62]        [108 110 112 114 116 118 120 122 124]
       [63 64 65 66 67 68 69 70 71]        [126 128 130 132 134 136 138 140 142]
       [72 73 74 75 76 77 78 79 80]        [144 146 148 150 152 154 156 158 160]
   
   Shared memory per block (3×3):
   a_shared[TPB, TPB]  b_shared[TPB, TPB]
   

2. Tile processing loop

   Number of tiles = 9 // 3 = 3 tiles (perfect division!)
   
   For each tile:
   1. Load tile from A and B
   2. Compute partial products
   3. Accumulate in register
   

3. Memory loading pattern

   • With perfect \((9 \times 9)\) tiling, bounds check is technically unnecessary but included for defensive programming and consistency with other matrix sizes.

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

4. Computation within tile

   if tiled_row < size and tiled_col < size:
       comptime for k in range(min(Int(TPB), Int(size - tile * TPB))):
           acc += a_shared[local_row, k] * b_shared[k, local_col]
   

   • Avoids shared memory bank conflicts:

     Bank Conflict Free (Good):        Bank Conflicts (Bad):
     Thread0: a_shared[0,k] b_shared[k,0]  Thread0: a_shared[k,0] b_shared[0,k]
     Thread1: a_shared[0,k] b_shared[k,1]  Thread1: a_shared[k,0] b_shared[1,k]
     Thread2: a_shared[0,k] b_shared[k,2]  Thread2: a_shared[k,0] b_shared[2,k]
     ↓                                     ↓
     Parallel access to different banks    Serialized access to same bank of b_shared
     (or broadcast for a_shared)           if shared memory was column-major
     

     Shared memory bank conflicts explained:

     • Left (Good): b_shared[k,threadIdx.x] accesses different banks, a_shared[0,k] broadcasts to all threads
     • Right (Bad): If b_shared were column-major, threads would access same bank simultaneously
     • Key insight: This is about shared memory access patterns, not global memory coalescing
     • Bank structure: Shared memory has 32 banks; conflicts occur when multiple threads access different addresses in the same bank simultaneously

5. Synchronization points

   barrier() after:
   1. Tile loading
   2. Tile computation
   

6. Result writing:

   if tiled_row < size and tiled_col < size:
      output[tiled_row, tiled_col] = acc
   

   • Defensive bounds checking included for other matrix sizes and tiling strategies
   • Direct assignment to output matrix
   • All threads write valid results

Key performance features:

• Processes \((9 \times 9)\) matrix using \((3 \times 3)\) tiles (perfect fit!)
• Uses shared memory for fast tile access
• Minimizes global memory transactions with coalesced memory access
• Optimized shared memory layout and access pattern to avoid shared memory bank conflicts

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: output.ElementType = 0
   • Compile-time loop unrolling via comptime for

This implementation achieves high performance through:

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

from max.gpu.memory import async_copy_wait_all
from layout.layout_tensor import copy_dram_to_sram_async
from layout import Layout as IntTupleLayout

comptime NUM_THREADS = TPB * TPB
comptime BLOCK_DIM_COUNT = 2


def matmul_idiomatic_tiled[
    size: Int
](
    output: TileTensor[mut=True, dtype, LayoutTiledType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutTiledType, ImmutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutTiledType, ImmutAnyOrigin],
):
    var local_row = thread_idx.y
    var local_col = thread_idx.x
    var tiled_row = block_idx.y * TPB + local_row
    var tiled_col = block_idx.x * TPB + local_col

    # Get the tile of the output matrix that this thread block is responsible for
    var out_tile = output.tile[TPB, TPB](block_idx.y, block_idx.x)
    var a_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB, TPB]())
    var b_shared = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB, TPB]())

    var acc: output.ElementType = 0

    comptime load_a_layout = IntTupleLayout.row_major(
        1, TPB
    )  # Coalesced loading
    comptime load_b_layout = IntTupleLayout.row_major(
        1, TPB
    )  # Coalesced loading
    # Note: Both matrices stored in same orientation for correct matrix multiplication
    # Transposed loading would be useful if B were pre-transposed in global memory

    comptime for idx in range(
        size // TPB
    ):  # Perfect division: 9 // 3 = 3 tiles
        # Get tiles from A and B matrices
        var a_tile = a.tile[TPB, TPB](block_idx.y, Int(idx))
        var b_tile = b.tile[TPB, TPB](Int(idx), block_idx.x)

        # Asynchronously copy tiles to shared memory with consistent orientation
        copy_dram_to_sram_async[
            thread_layout=load_a_layout,
            num_threads=NUM_THREADS,
            block_dim_count=BLOCK_DIM_COUNT,
        ](a_shared.to_layout_tensor(), a_tile.to_layout_tensor())
        copy_dram_to_sram_async[
            thread_layout=load_b_layout,
            num_threads=NUM_THREADS,
            block_dim_count=BLOCK_DIM_COUNT,
        ](b_shared.to_layout_tensor(), b_tile.to_layout_tensor())

        # Wait for all async copies to complete
        async_copy_wait_all()
        barrier()

        # Compute partial matrix multiplication for this tile
        comptime for k in range(TPB):
            acc += a_shared[local_row, k] * b_shared[k, local_col]

        barrier()

    # Write final result to output tile
    if tiled_row < size and tiled_col < size:
        out_tile[local_row, local_col] = acc

The idiomatic tiled matrix multiplication leverages Mojo’s TileTensor API and asynchronous memory operations for a beautifully clean implementation.

🔑 Key Point: This implementation performs standard matrix multiplication A × B using coalesced loading for both matrices.

What this implementation does:

• Matrix operation: Standard \(A \times B\) multiplication (not \(A \times B^T\))
• Loading pattern: Both matrices use Layout.row_major(1, TPB) for coalesced access
• Computation: acc += a_shared[local_row, k] * b_shared[k, local_col]
• Data layout: No transposition during loading - both matrices loaded in same orientation

What this implementation does NOT do:

• Does NOT perform \(A \times B^T\) multiplication
• Does NOT use transposed loading patterns
• Does NOT transpose data during copy operations

With the \((9 \times 9)\) matrix size, we get perfect tiling that eliminates all boundary checks:

1. TileTensor tile API

   var out_tile = output.tile[TPB, TPB](block_idx.y, block_idx.x)
   var a_tile = a.tile[TPB, TPB](block_idx.y, idx)
   var 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,
      num_threads = NUM_THREADS,
      block_dim_count = BLOCK_DIM_COUNT
   ](a_shared.to_layout_tensor(), a_tile.to_layout_tensor())
   copy_dram_to_sram_async[
      thread_layout = load_b_layout,
      num_threads = NUM_THREADS,
      block_dim_count = BLOCK_DIM_COUNT
   ](b_shared.to_layout_tensor(), b_tile.to_layout_tensor())
   async_copy_wait_all()
   

   These operations:

   • Issue the copy through the GPU’s asynchronous copy path (cp.async on NVIDIA), which moves global memory straight into shared memory without staging through registers, via copy_dram_to_sram_async
   • Use specialized thread layouts for optimal memory access patterns
   • Eliminate the need for manual memory initialization
   • Important:
     • Standard GPU loads are already asynchronous; these provide better resource utilization and register bypass
     • copy_dram_to_sram_async assumes a 1D thread block (block_dim.y == block_dim.z == 1) holding exactly thread_layout.size() threads. Neither assumption holds here, so both defaults are overridden:
       • block_dim_count: the dimensionality of the thread block (2 for the 2D thread block THREADS_PER_BLOCK_TILED = (TPB, TPB)), so the copy derives each thread’s index as thread_idx.y * block_dim.x + thread_idx.x instead of thread_idx.x
       • num_threads: the number of threads in the thread block (TPB * TPB == 9). The copy itself is carried out by the thread_layout.size() == TPB threads the layout describes; telling the function how large the block really is lets it turn the remaining threads off instead of letting them run off the end of the layout

3. Optimized memory access layouts

   comptime load_a_layout = IntTupleLayout.row_major(1, TPB)  # Coalesced loading
   comptime load_b_layout = IntTupleLayout.row_major(1, TPB)  # Coalesced loading
   # Note: Both matrices use the same layout for standard A × B multiplication
   

   Thread layouts are IntTuple Layout values built with the function form Layout.row_major(1, TPB), which the solution imports as IntTupleLayout to keep them distinct from the row_major[...]() tensor layouts used for the tensors themselves.

   Memory Access Analysis for Current Implementation:

   Both matrices use Layout.row_major(1, TPB) for coalesced loading from global memory:

   • load_a_layout: Threads cooperate to load consecutive elements from matrix A rows
   • load_b_layout: Threads cooperate to load consecutive elements from matrix B rows
   • Key insight: Thread layout determines how threads cooperate during copy, not the final data layout

   Actual Computation Pattern (proves this is A × B):

   # This is the actual computation in the current implementation
   acc += a_shared[local_row, k] * b_shared[k, local_col]
   
   # This corresponds to: C[i,j] = Σ(A[i,k] * B[k,j])
   # Which is standard matrix multiplication A × B
   

   Why both matrices use the same coalesced loading pattern:

   Loading tiles from global memory:
   - Matrix A tile: threads load A[block_row, k], A[block_row, k+1], A[block_row, k+2]... (consecutive)
   - Matrix B tile: threads load B[k, block_col], B[k, block_col+1], B[k, block_col+2]... (consecutive)
   
   Both patterns are coalesced with Layout.row_major(1, TPB)
   

   Three separate memory concerns:

   1. Global-to-shared coalescing: Layout.row_major(1, TPB) ensures coalesced global memory access
   2. Shared memory computation: a_shared[local_row, k] * b_shared[k, local_col] avoids bank conflicts
   3. Matrix operation: The computation pattern determines this is A × B, not A × B^T

4. Perfect tiling eliminates boundary checks

   comptime for idx in range(size // TPB):  # Perfect division: 9 // 3 = 3
   

   With \((9 \times 9)\) matrices and \((3 \times 3)\) tiles, every tile is exactly full-sized. No boundary checking needed!

5. Clean tile processing with defensive bounds checking

   # Defensive bounds checking included even with perfect tiling
   if tiled_row < size and tiled_col < size:
       out_tile[local_row, local_col] = acc
   

   With perfect \((9 \times 9)\) tiling, this bounds check is technically unnecessary but included for defensive programming and consistency with other matrix sizes.

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: Tiles land in shared memory without a round trip through registers
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.

Educational: When would transposed loading be useful?

The current implementation does NOT use transposed loading. This section is purely educational to show what’s possible with the layout system.

Current implementation recap:

• Uses Layout.row_major(1, TPB) for both matrices
• Performs standard A × B multiplication
• No data transposition during copy

Educational scenarios where you WOULD use transposed loading:

While this puzzle uses standard coalesced loading for both matrices, the layout system’s flexibility enables powerful optimizations in other scenarios:

# Example: reading a pre-transposed B with one thread distribution and writing
# shared memory with another
# (This is NOT what the current implementation does)
comptime load_b_layout = IntTupleLayout.row_major(TPB, 1)   # Threads walk a column
comptime store_b_layout = IntTupleLayout.row_major(1, TPB)  # Threads walk a row
copy_dram_to_sram_async[
    src_thread_layout=load_b_layout,
    dst_thread_layout=store_b_layout,
    num_threads=NUM_THREADS,
    block_dim_count=BLOCK_DIM_COUNT,
](b_shared.to_layout_tensor(), b_tile.to_layout_tensor())

Note that the two thread layouts are IntTuple Layout values, the same form as load_a_layout above, and that copy_dram_to_sram_async takes LayoutTensor arguments. The src_thread_layout/dst_thread_layout pair must describe the same number of threads.

Use cases for transposed loading (not used in this puzzle):

1. Pre-transposed input matrices: When \(B\) is already stored transposed in global memory
2. Different algorithms: Computing \(A^T \times B\), \(A \times B^T\), or \(A^T \times B^T\)
3. Memory layout conversion: Converting between row-major and column-major layouts
4. Avoiding transpose operations: Loading data directly in the required orientation

Key distinction:

• Current implementation: Both matrices use Layout.row_major(1, TPB) for standard \(A \times B\) multiplication
• Transposed loading example: Would use different layouts to handle pre-transposed data or different matrix operations

This demonstrates Mojo’s philosophy: providing low-level control when needed while maintaining high-level abstractions for common cases.

---

Summary: Key takeaways

What the idiomatic tiled implementation actually does:

1. Matrix Operation: Standard A × B multiplication
2. Memory Loading: Both matrices use Layout.row_major(1, TPB) for coalesced access
3. Computation Pattern: acc += a_shared[local_row, k] * b_shared[k, local_col]
4. Data Layout: No transposition during loading

Why this is optimal:

• Coalesced global memory access: Layout.row_major(1, TPB) ensures efficient loading
• Bank conflict avoidance: Shared memory access pattern avoids conflicts
• Standard algorithm: Implements the most common matrix multiplication pattern