Handling Non-Divisible Dimensions in Matrix Tiling

The Problem of Ragged Edges

When a matrix dimension is not perfectly divisible by the tile size, the grid of tiles laid over the matrix will have partial tiles, or "ragged edges," along its borders.

Threads in a block are assigned to load and compute a full tile's worth of data. If a block is processing a ragged edge, some of its threads will be assigned to coordinates that lie outside the actual matrix dimensions. Without any checks, these threads will attempt to read from or write to invalid global memory addresses. This can lead to segmentation faults, unpredictable behavior, or silent data corruption.

The following video provides a concise explanation of this problem.

The Solution: Guarding Memory Accesses

The solution is to "guard" every memory access that could potentially fall outside the matrix bounds. This means adding conditional checks before loading from the input matrices and before storing to the output matrix.

There are three critical points where we need to add these checks:

1. Loading Matrix A Tile: Before a thread loads an element of matrix A into shared memory, it must verify that the calculated source row and column indices are within the M x K bounds of A.
2. Loading Matrix B Tile: Similarly, before loading an element of B, a thread must check that its source indices are within the K x N bounds of B.
3. Storing Result C Element: After the block has computed its portion of the result, each thread must verify that its target destination indices are within the M x N bounds of the output matrix C before writing its final value.

What should a thread do if its assigned element is out of bounds?

• For loads (A and B): The thread should write a 0.0f into its assigned spot in the shared memory tile. This "zero-padding" ensures that these "ghost" elements do not corrupt the dot-product accumulation, as any value multiplied by zero is zero.
• For stores (C): The thread should simply do nothing. It has computed a value for a non-existent element of C, so that value is discarded.

This strategy is explained clearly in the continuation of the previous video.

Implementing Boundary Checks in the CUDA Kernel

Let's now translate this strategy into code. We'll be modifying the tiled GEMM kernel to support arbitrary dimensions for an operation $C_{M\times N}=A_{M\times K}\times B_{K\times N}$. In our code, we'll use N1 for M, N2 for K, and N3 for N.

The article "Memory Coalescing and Tiled Matrix Multiplication" provides a complete, generalized kernel. We will focus on the section that introduces these boundary conditions.

Let's break down the key modifications from that reading:

1. Adjusting the Phase Loop

Our loop that iterates through tiles along the common dimension K (which is N2 in the code) must run enough times to process any final partial tile. We achieve this by using the ceil function.

// Before, for a matrix of size N x N:
for (int phase = 0; phase < N / TILE_WIDTH; phase++)

// After, for a general M x K matrix multiplication (N2 = K):
for (int phase = 0; phase < ceil((float)N2 / TILE_WIDTH); phase++)

This ensures that if N2 is, for example, 33 and TILE_WIDTH is 32, the loop runs for ceil(33/32) = 2 phases, covering both the full tile and the remaining sliver.

2. Guarded Loads into Shared Memory

Inside the phase loop, we replace the direct memory loads with guarded if-else statements.

• Loading sh_A: We check if the thread's global row index i is less than N1 (dimension M) and if its global column index phase*TILE_WIDTH+tx is less than N2 (dimension K).

  int i = TILE_WIDTH * by + ty; // Global row index for A and C
  int tx = threadIdx.x;
  
  if ((i < N1) && ((phase * TILE_WIDTH + tx) < N2)) {
      sh_A[ty][tx] = A[i * N2 + phase * TILE_WIDTH + tx];
  } else {
      sh_A[ty][tx] = 0.0f; // Pad with zero
  }
  

• Loading sh_B: We check if the thread's global row index phase*TILE_WIDTH + ty is less than N2 (dimension K) and if its global column index j is less than N3 (dimension N).

  int j = TILE_WIDTH * bx + tx; // Global column index for B and C
  int ty = threadIdx.y;
  
  if (((phase * TILE_WIDTH + ty) < N2) && (j < N3)) {
      sh_B[ty][tx] = B[(phase * TILE_WIDTH + ty) * N3 + j];
  } else {
      sh_B[ty][tx] = 0.0f; // Pad with zero
  }
  

3. Guarded Store into Global Memory

Finally, after all phases are complete and the value is accumulated, we guard the write to matrix C.

// After the phase loop
if ((i < N1) && (j < N3)) {
    C[i * N3 + j] = value;
}

This prevents threads corresponding to locations outside the M x N output matrix from writing their results.

This pattern of boundary checking is fundamental to writing robust GPU kernels. You can see it applied in nearly every production-level implementation. For example, the advanced kernels in "CUDA Matrix Multiplication Optimization" by Lei Mao, while more complex, all incorporate this same principle of guarding loads and stores. In the gemm_v02 kernel from that article, the checks if (A_row_idx < m && A_col_idx < k) inside the loading function and if (C_row_idx < m && C_col_idx < n) before the final write serve the exact same purpose.

Conclusion

In this lesson, you have taken the final and most critical step in creating a truly general-purpose GEMM kernel. By moving beyond the simplifying assumption of perfectly divisible matrix dimensions, you have learned to write robust code that can handle real-world inputs without error.

Key Takeaways:

• The Problem: When matrix dimensions are not multiples of the tile size, some threads in blocks at the "ragged edges" will be mapped to out-of-bounds memory locations.
• The Solution: Add conditional "guards" to every memory access that could cross a matrix boundary.
• Guarded Loads: When loading from input matrices A and B, check if the source address is valid. If it's not, load 0.0f into shared memory to "pad" the tile without corrupting the computation.
• Guarded Stores: Before writing the final result to the output matrix C, check if the destination address is within the matrix's valid dimensions. If not, simply discard the result.
• Loop Boundary: The main loop over the K dimension must use ceil() to ensure it processes the final partial tile.

We now have a functionally correct and general matrix multiplication kernel. The major optimization steps are complete. The next, and final, step in our GEMM journey is to rigorously evaluate its performance.

In the next lesson, we will benchmark each version of our kernel—from the naive implementation to this generalized, tiled version—against NVIDIA's highly-optimized cuBLAS library. This will allow us to quantify the impact of each optimization and analyze any remaining performance gaps, providing deep insight into what separates a good kernel from a world-class one.