Implementing Tiled Computation

Welcome back to Shared Memory Optimization. We are now at Lesson 3, which means we already have two key tools in hand: we know how to place data in shared memory, and we know how to make threads wait safely with __syncthreads(). In this lesson, we will combine those ideas into a full tiled matrix multiplication.

Our goal is practical: we will build a kernel that processes a matrix in multiple tiles, handles edge cases safely, and still produces the correct answer when the matrix width is not divisible by the tile width. By the end, we will have a complete pattern that we can reuse in many CUDA programs.

A matrix multiplication entry is computed as:

If we do this directly from global memory, we read many values repeatedly. Tiling fixes that by loading small chunks of A and B into shared memory, then reusing them many times before moving on.

The tricky part appears at the edges. With TILE_WIDTH = 16 and a matrix width = 34, we need 3 tile steps, not 2:

1. Columns or rows 0 to 15
2. Columns or rows 16 to 31
3. Columns or rows 32 to 33, plus unused positions

That last step is only partially full, so we must guard those loads carefully. The common solution is simple: if a thread would read outside the matrix, it loads 0.0f instead.

Now we begin the kernel. Each block will compute one TILE_WIDTH x TILE_WIDTH region of the output matrix C. To do that, the block needs one shared tile from A, one shared tile from B, and a private accumulator in each thread.

/**
 * Optimized tiled matrix multiplication kernel.
 * Uses shared memory and guarded edge loads.
 */
__global__ void tiledMatrixMultiply(const float* A, const float* B, float* C, int width) {
    __shared__ float s_A[TILE_WIDTH][TILE_WIDTH];
    __shared__ float s_B[TILE_WIDTH][TILE_WIDTH];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row = blockIdx.y * TILE_WIDTH + ty;
    int col = blockIdx.x * TILE_WIDTH + tx;
    float value = 0.0f;

There are two coordinate systems here, and both matter. tx and ty tell us where a thread lives inside its block; row and col tell us which output element that thread is responsible for in the full matrix. The variable value will collect the dot product for C[row, col] over several tile iterations.

The next step is deciding how many tile pairs we must process. Because the width may not divide evenly by TILE_WIDTH, we use ceiling division. Then, in each loop iteration, every thread loads one value from A and one from B into shared memory.

    int numTiles = (width + TILE_WIDTH - 1) / TILE_WIDTH;

    for (int m = 0; m < numTiles; ++m) {
        int aCol = m * TILE_WIDTH + tx;
        int bRow = m * TILE_WIDTH + ty;

        s_A[ty][tx] = (row < width && aCol < width) ? A[row * width + aCol] : 0.0f;
        s_B[ty][tx] = (bRow < width && col < width) ? B[bRow * width + col] : 0.0f;

        __syncthreads();

This is the heart of the boundary-safe design. For width = 34, numTiles becomes 3. On the last iteration, some threads ask for valid matrix elements, while others would step outside the matrix. The ternary expressions prevent bad reads by writing 0.0f into those unused shared memory positions. After that, __syncthreads() ensures the whole tile is ready before any thread starts using it.

Once both shared tiles are fully loaded, each thread multiplies across the current tile and adds the result into its running sum. Then we need a second synchronization point before the next iteration starts.

        for (int k = 0; k < TILE_WIDTH; ++k)
            value += s_A[ty][k] * s_B[k][tx];

        __syncthreads();
    }

The inner loop always runs for TILE_WIDTH steps, even on the last partial tile. That works because any missing values were replaced with 0.0f, so they contribute nothing to the sum. The second __syncthreads() is just as important as the first one: it prevents one thread from overwriting s_A or s_B for the next tile while another thread is still reading the current tile.

After all tile pairs have been processed, each thread has its final dot product in value. Even then, we still need one last bounds check because some threads in the edge blocks correspond to positions outside the real matrix.

    if (row < width && col < width)
        C[row * width + col] = value;
}

This final guard protects the output write. For a 34 x 34 matrix with 16 x 16 blocks, the grid covers a 48 x 48 region in total. That means many edge threads are useful for cooperation during loading, but they do not own valid output elements. The if statement ensures only real matrix positions are written back to global memory.

With the kernel complete, we can set up a test in main(). The matrix width is deliberately chosen as 34, so we can prove that the tiled approach works even when the size does not line up neatly with the tile width.

int main() {
    const int width = 34; // intentionally not divisible by TILE_WIDTH
    const size_t bytes = width * width * sizeof(float);

    std::vector<float> h_A(width * width, 1.0f);
    std::vector<float> h_B(width * width, 2.0f);
    std::vector<float> h_C(width * width, 0.0f);

    float *d_A, *d_B, *d_C;
    CUDA_CHECK(cudaMalloc(&d_A, bytes));
    CUDA_CHECK(cudaMalloc(&d_B, bytes));
    CUDA_CHECK(cudaMalloc(&d_C, bytes));

    CUDA_CHECK(cudaMemcpy(d_A, h_A.data(), bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B, h_B.data(), bytes, cudaMemcpyHostToDevice));

The input values make verification easy. Every entry of A is 1.0f and every entry of B is 2.0f; therefore, each product inside a dot product is 2.0f. We also allocate device memory for all three matrices, then copy only A and B to the GPU, because C will be produced by the kernel.

Next, we choose the block and grid dimensions. The block is a single 16 x 16 tile, and the grid uses the same ceiling division logic as the kernel so that every output position is covered, including the edge region.

    dim3 block(TILE_WIDTH, TILE_WIDTH);
    dim3 grid((width + TILE_WIDTH - 1) / TILE_WIDTH,
              (width + TILE_WIDTH - 1) / TILE_WIDTH);

    tiledMatrixMultiply<<<grid, block>>>(d_A, d_B, d_C, width);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());

    CUDA_CHECK(cudaMemcpy(h_C.data(), d_C, bytes, cudaMemcpyDeviceToHost));

For width = 34, both grid dimensions become 3, so we launch a 3 x 3 set of blocks. That gives us full coverage of the matrix, plus extra edge threads that are safely handled by the guards inside the kernel. As in the previous lesson, cudaGetLastError() checks launch issues, and cudaDeviceSynchronize() waits for the kernel to finish before we copy results back.

Now we confirm that the result is correct. Since each output entry is the sum of 34 products, and each product is 1.0f * 2.0f = 2.0f, the expected value is 2.0f * width, which is 68.0f. After checking every value, the program prints the result, frees device memory, and returns a status code.

    const float expected = 2.0f * width;
    bool success = true;
    for (float value : h_C) {
        if (std::fabs(value - expected) > 1e-5f) {
            success = false;
            break;
        }
    }

    std::cout << "Boundary-safe tiled matrix multiplication: "
              << (success ? "SUCCESS" : "FAILURE") << std::endl;

    CUDA_CHECK(cudaFree(d_A));
    CUDA_CHECK(cudaFree(d_B));
    CUDA_CHECK(cudaFree(d_C));
    return success ? 0 : 1;
}

This final check is small, but it tells us a lot. If even one tile step, one synchronization point, or one bounds guard were wrong, the comparison would fail. When everything is correct, we get:

Boundary-safe tiled matrix multiplication: SUCCESS

That message confirms the full tiled computation is working, including the difficult edge tiles.

In this lesson, we turned the ideas from the first two lessons into a complete tiled matrix multiplication. We used shared memory for reuse, __syncthreads() for safe cooperation, ceiling division for the tile count, and guarded loads and stores for boundary safety. Most importantly, we verified the result on a matrix size that does not divide evenly by the tile width.

This is a very common CUDA pattern, and it is worth knowing well. In the upcoming practice section, we will reinforce it step by step, so we can write tiled kernels with much more confidence and independence.