Tiled Matrix Multiplication with Shared Memory

The Memory Hierarchy and Our Optimization Strategy

Before diving into the algorithm, let's quickly refresh our understanding of the GPU memory hierarchy. The performance difference between memory types is the entire motivation for tiling.

The Tiling Algorithm: Work Smarter, Not Harder

The central idea of tiling is to break a large problem into smaller, manageable chunks that can fit into a fast, local cache. Instead of each thread in a block fetching the entire row of matrix A and column of B it needs for its dot product, the threads will work together as a team.

The team (a thread block) will:

1. Load a small sub-matrix, or tile, of A from global memory into shared memory.
2. Load a corresponding tile of B from global memory into shared memory.
3. Perform a partial matrix multiplication using only the data in fast shared memory.
4. Repeat this process, accumulating the partial results, until the full dot product is complete.

This "cooperative fetching" strategy is a pure algorithmic optimization. The following resources explain the logic behind it.

For a more visual, step-by-step breakdown of how the tiles are loaded and used to compute partial results, the following article is an excellent supplement.

This animated diagram beautifully visualizes the core loop of the tiled algorithm. For each phase, tiles from A and B are loaded into shared memory, where threads use them to compute a partial result for the output tile in C.

Implementing the Tiled Kernel in CUDA

Now we'll translate this algorithm into a CUDA kernel. The core of the implementation is the load -> synchronize -> compute pattern executed in a loop.

Here is what the overall structure of our tiled kernel will look like. We will break it down piece by piece.

template<int TILE_SIZE>
__global__ void tiled_matmul_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
    // 1. Declare shared memory tiles
    __shared__ float shared_A[TILE_SIZE][TILE_SIZE];
    __shared__ float shared_B[TILE_SIZE][TILE_SIZE];

    // 2. Identify thread's position (global and local)
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    float sum = 0.0f;

    // 3. Loop over tiles in the K dimension
    for (int tile = 0; tile < (K + TILE_SIZE - 1) / TILE_SIZE; tile++) {
        // 4. Load a tile of A and B from global to shared memory
        // ... loading logic ...
        
        // 5. Synchronize to ensure all data is loaded
        __syncthreads();

        // 6. Compute partial dot product from shared memory
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += shared_A[ty][k] * shared_B[k][tx];
        }

        // 7. Synchronize again before loading the next tile
        __syncthreads();
    }

    // 8. Write final result to global memory
    if (row < M && col < N) {
        C[row * N + col] = sum;
    }
}

Let's dissect the most critical parts of this implementation. The following video and article provide detailed walkthroughs of the code.

Step 1: Cooperative Loading

This is the most complex part of the kernel. Each thread in the block is responsible for fetching one element of A and one element of B from global memory and placing them into the correct position in the shared_A and shared_B tiles.

The indexing to load the correct global memory element can be tricky. The key is to combine the tile loop index with the thread's local index (tx, ty).

• For shared_A, a thread at (ty, tx) within the block loads the element A[row][tile * TILE_SIZE + tx]. All threads in the same row of the block (ty is constant) load elements from the same row of matrix A.
• For shared_B, the same thread loads element B[tile * TILE_SIZE + ty][col]. All threads in the same column of the block (tx is constant) load elements from the same column of matrix B.

The following resource provides a complete, well-explained implementation.

Step 2: Synchronization with __syncthreads()

As the resources emphasize, __syncthreads() is not optional. It acts as a barrier, forcing all threads in a block to wait at that line of code until every single thread in the block has arrived.

• The first __syncthreads() is essential to prevent a race condition. Without it, a fast thread might start computing from shared_A before a slower thread has finished writing its value into shared_A. The result would be incorrect. It ensures the "workbench is fully stocked before anyone starts cooking."
• The second __syncthreads() is also important. It ensures that all threads have finished their computations with the current tile before any thread races ahead to the next for loop iteration and starts overwriting the shared memory tiles with new data.

The Performance Payoff (A Preview)

By implementing this tiled approach, we dramatically reduce the number of slow global memory reads. Instead of each thread performing K reads from A and K reads from B (all from global memory), each thread now performs (K / TILE_SIZE) reads from A and (K / TILE_SIZE) reads from B. The rest of the K multiplications happen using ultra-fast shared memory.

Conclusion

In this lesson, you have learned and implemented one of the most important optimization techniques in CUDA programming. By replacing the naive, brute-force memory access pattern with a cooperative tiling strategy, you have taken a massive leap from a functionally correct kernel to a genuinely high-performance one.

Key Takeaways:

• Tiling is an algorithmic technique that improves performance by increasing data locality, breaking a large problem into smaller chunks that fit into fast cache memory.
• Shared memory is the fast, on-chip GPU memory that enables efficient tiling. It is accessible to all threads within a block.
• The fundamental pattern for tiled algorithms in CUDA is a loop that repeatedly executes a load-synchronize-compute sequence.
• __syncthreads() is essential for correctness. It synchronizes all threads in a block, preventing race conditions where threads might access data before it's ready or overwrite data that's still in use.
• This approach drastically reduces traffic to slow global memory, shifting the workload from waiting for memory to performing actual computation.

You have now built the "hero" kernel. In our next lesson, we will put it to the test. You will learn how to use the NVIDIA Nsight Compute profiler to scientifically measure the performance improvement, analyze the reduction in global memory traffic, and see just how much faster our new kernel is.