Naïve Matrix Multiplication

Overview

Implement a kernel that multiplies square matrices \(A\) and \(B\) and stores the result in \(\text{output}\). This is the most straightforward implementation where each thread computes one element of the output matrix.

Key concepts

This puzzle covers:

  • 2D thread organization for matrix operations
  • Global memory access patterns
  • Matrix indexing in row-major layout
  • Thread-to-output element mapping

The key insight is understanding how to map 2D thread indices to matrix elements and compute dot products in parallel.

Configuration

  • Matrix size: \(\text{SIZE} \times \text{SIZE} = 2 \times 2\)
  • Threads per block: \(\text{TPB} \times \text{TPB} = 3 \times 3\)
  • Grid dimensions: \(1 \times 1\)

Layout configuration:

  • Input A: row_major[SIZE, SIZE]()
  • Input B: row_major[SIZE, SIZE]()
  • Output: row_major[SIZE, SIZE]()

Code to complete

def naive_matmul[
    size: Int
](
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
):
    var row = block_dim.y * block_idx.y + thread_idx.y
    var col = block_dim.x * block_idx.x + thread_idx.x
    # FILL ME IN (roughly 6 lines)

View full file: problems/p16/p16.mojo

Tips
  1. Use the row and col the stub already computes from the thread indices
  2. Check if indices are within size
  3. Accumulate products in a local variable
  4. Write final sum to correct output position

Running the code

To test your solution, run the following command in your terminal:

pixi run p16 --naive

pixi run -e amd p16 --naive

pixi run -e apple p16 --naive

uv run poe p16 --naive

Your output will look like this if the puzzle isn’t solved yet:

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([4.0, 6.0, 12.0, 22.0])

Solution

def naive_matmul[
    size: Int
](
    output: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
    a: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
    b: TileTensor[mut=False, dtype, LayoutType, ImmutAnyOrigin],
):
    var row = block_dim.y * block_idx.y + thread_idx.y
    var col = block_dim.x * block_idx.x + thread_idx.x

    if row < size and col < size:
        var acc: output.ElementType = 0

        comptime for k in range(size):
            acc += a[row, k] * b[k, col]

        output[row, col] = acc

The naive matrix multiplication using TileTensor follows this basic approach:

Matrix layout (2×2 example)

Matrix A:          Matrix B:                   Output C:
[a[0,0] a[0,1]]    [b[0,0] b[0,1]]             [c[0,0] c[0,1]]
[a[1,0] a[1,1]]    [b[1,0] b[1,1]]             [c[1,0] c[1,1]]

Implementation details

  1. Thread mapping:

    var row = block_dim.y * block_idx.y + thread_idx.y
var col = block_dim.x * block_idx.x + thread_idx.x

  2. Memory access pattern:

    • Row-wise access: a[row, k] walks along one row of \(A\)
    • Column-wise access: b[k, col] walks down one column of \(B\), which in a row-major layout strides by SIZE on every step
    • Output writing: output[row, col]

  3. Computation flow:

    # Use var for mutable accumulator with tensor's element type
var acc: output.ElementType = 0

# comptime for compile-time loop unrolling
comptime for k in range(size):
    acc += a[row, k] * b[k, col]

Key language features

  1. Variable declaration:

    • Annotating the accumulator in var acc: output.ElementType = 0 ties its type to the output tensor’s element type, so the accumulation and the final store agree
    • Initialized to zero before accumulation

  2. Loop optimization:

    • comptime for unrolls the loop at compile time
    • Improves performance for small, known matrix sizes
    • Enables better instruction scheduling

Performance characteristics

  1. Memory access:

    • Each thread makes 2 x SIZE global memory reads
    • One global memory write per thread
    • No data reuse between threads

  2. Computational efficiency:

    • Simple implementation but suboptimal performance
    • Many redundant global memory accesses
    • No use of fast shared memory

  3. Limitations:

    • High global memory bandwidth usage
    • Poor data locality
    • Limited scalability for large matrices

This naive implementation serves as a baseline for understanding matrix multiplication on GPUs, highlighting the need for optimization in memory access patterns.