GPU Matrix Multiplication: Tiling and Shared Memory

Introduction

Welcome to the first lesson of our new module, "Parallelizing Machine Learning Workloads." In the previous module, we concluded by examining how to tune GPU kernel launch parameters to manage occupancy and hide latency. We established that the ultimate goal is performance, which often means finding clever ways to overcome memory bottlenecks.

This lesson directly builds on that foundation by tackling one of the most fundamental and performance-critical operations in all of deep learning: matrix multiplication. Today's learning outcome is to explain how matrix multiplication is parallelized on GPUs using tiling and shared memory.

You'll learn why a straightforward parallel approach is inefficient and discover the elegant and powerful technique of tiling (also called blocking). This method uses the GPU's fast, on-chip shared memory as a programmer-managed cache to dramatically reduce data movement from the much slower main GPU memory. Understanding this algorithm is crucial, as it forms the computational heart of libraries like cuBLAS, which are called by frameworks like PyTorch and TensorFlow to accelerate everything from fully-connected layers to convolutions.

1. The Naive Approach and its Memory Bottleneck

Let's start by considering the most direct way to parallelize a matrix multiplication $C=A\times B$. The definition of an element $C_{ij}$ in the output matrix is the dot product of the i-th row of A and the j-th column of B:

$$C_{ij}=\sum _{k=0}^{N-1}{A}_{ik}\cdot B_{kj}$$

A natural parallel mapping is to assign one GPU thread to compute each element $C_{ij}$. While this creates a large amount of parallel work, it hides a severe performance problem related to memory access.

To compute its single output value, each thread must:

1. Read the entire i-th row of matrix A from global memory.
2. Read the entire j-th column of matrix B from global memory.

The issue becomes clear when we consider multiple threads working together. Threads calculating elements in the same row of C will all redundantly read the exact same row from matrix A. Similarly, threads computing elements in the same column of C will all redundantly read the exact same column of B. As we discussed, global memory is the largest but slowest memory on the GPU; this redundant traffic saturates memory bandwidth and forces the powerful arithmetic units to wait, negating much of the benefit of parallelism.

The following video introduces this problem and sets the stage for the solution.

2. The Solution: Tiling with Shared Memory

To solve this memory bottleneck, we can employ a classic optimization technique known as tiling or blocking. Given your background in computer science, this may be conceptually similar to cache-blocking optimizations for CPU algorithms, but here we will explicitly manage the data movement ourselves.

The core idea is to use the GPU's shared memory. As a refresher, shared memory is a small, fast, on-chip memory space that is private to a thread block. Its latency is orders of magnitude lower than global memory.

The tiling strategy works as follows:

1. We partition the large matrices in global memory into smaller, fixed-size sub-matrices called tiles.
2. We assign one thread block to be responsible for computing one tile of the output matrix C.
3. To do this, the thread block cooperatively loads the necessary tiles from input matrices A and B into shared memory.
4. Once the data is in shared memory, all threads in the block can perform their calculations by accessing this fast memory, avoiding the global memory bottleneck.

This animated GIF provides an excellent high-level visualization of the process.

This strategy is so fundamental that it's the standard approach used in production-level GPU libraries.

3. The Tiled Algorithm in Detail

Let's break down the execution flow for a single thread block, which is responsible for computing one tile of the output matrix C. Assume we are using square tiles of size TILE_WIDTH x TILE_WIDTH. The thread block will also be configured with TILE_WIDTH x TILE_WIDTH threads.

The computation proceeds in a loop that steps through the tiles of A and B.

Outer Loop (iterating through tiles): For each pair of tiles needed to compute the final output tile...

1. Cooperative Load from Global to Shared Memory: Each thread in the block calculates an index and loads one element from a tile of matrix A and one element from a tile of matrix B into two corresponding shared memory arrays. This step is a collaborative effort where all threads in the block work together to stage the required data.

2. Synchronize: A call to __syncthreads() is issued. This is a barrier that forces all threads in the block to wait at this point. This is absolutely critical to ensure that the entire shared memory tile is loaded before any thread attempts to use it for computation.

3. Compute from Shared Memory: Each thread now performs a series of multiplications and additions, accumulating a partial result in its own private register. Crucially, all data for this step is read from the fast shared memory arrays.

4. Synchronize Again: Before the next iteration of the outer loop (which will overwrite the shared memory with new data), another __syncthreads() is necessary to ensure all threads have finished their computations with the current data.

After the loop finishes, each thread contains a final, accumulated value in its private register. The thread then writes this single value to its corresponding position in the output matrix C in global memory.

The following video provides a detailed walkthrough of this algorithm, complete with code examples.

4. Analysis: Quantifying the Improvement

The benefit of tiling can be quantified by analyzing the reduction in global memory accesses. Let's consider the computation of a single $B\times B$ output tile, where $B$ is the tile dimension.

• Naive Approach: To compute $B^2$ output elements, each requiring a dot product of $N$ elements, the threads for this tile would naively perform approximately $B^2\times 2N$ reads from global memory.

• Tiled Approach: The thread block makes $N/B$ iterations. In each iteration, it reads a $B\times B$ tile from A and a $B\times B$ tile from B.

  • Number of global memory reads per iteration: $2\times B^2$
  • Total global memory reads: $(N/B)\times (2\times B^2)=2NB$

The tiled approach reduces global memory reads from $O(N{B}^2)$ to $O(NB)$, a reduction by a factor of B. For a typical tile size of 16 or 32, this represents a 16x to 32x reduction in traffic to the slowest part of the memory system, resulting in a dramatic performance improvement.

This blog post provides a very clear summary of this common GPU programming pattern and the associated analysis.

5. Tiling in Practice: Performance and Libraries

While the code can be complex, the performance gains are undeniable. The key is that we are trading a small amount of extra on-chip work (indexing into shared memory) for a massive reduction in slow off-chip memory traffic.

The following video shows a practical demonstration of this speedup using NVIDIA's profiler, and frames the shared memory concept in a way that should resonate with your CS background.

This tiled matrix multiplication algorithm is the basis for the highly optimized GEMM implementations found in libraries like NVIDIA's cuBLAS. When you call torch.matmul in PyTorch, you are invoking a kernel that uses this principle, but with many additional years of architecture-specific tuning.

These professional libraries also have to manage the trade-offs we discussed in our last lesson. For example:

• Tile Size Selection: Larger tiles offer more data reuse (higher arithmetic intensity), but result in fewer, larger thread blocks. This might lead to higher efficiency per block but lower overall parallelism (occupancy), which can hurt performance on smaller matrices. cuBLAS uses internal heuristics to select an optimal tile size for a given problem.
• Quantization Effects: Performance can degrade if matrix dimensions are not perfectly divisible by the tile size. The threads in blocks at the matrix edges will perform wasted work, an effect known as tile quantization.

Conclusion

In this lesson, you have dissected one of the most important algorithms in high-performance computing. By moving from a memory-bound naive approach to a compute-centric tiled strategy, we unlock the massive arithmetic potential of the GPU.

Key Takeaways:

• A naive parallel matrix multiplication is severely bottlenecked by redundant global memory accesses.
• Tiling is a blocking strategy that uses fast on-chip shared memory as a programmer-managed cache to improve data locality.
• The algorithm requires cooperative loading of data tiles by a thread block, followed by synchronization (__syncthreads()) before and after computation on the shared data.
• This technique dramatically reduces global memory traffic, trading it for much faster access to shared memory.
• Tiled GEMM is the fundamental building block for performance in deep learning, forming the core of libraries like cuBLAS.

Preview of the Next Lesson:

Now that you understand how to efficiently implement matrix multiplication on a GPU, we will explore how this powerful primitive is leveraged to accelerate other core deep learning operations. In the next lesson, we will investigate a classic technique called im2col that answers the question: How are convolution operations commonly mapped to matrix multiplication for GPU efficiency?