CUDA Tiled Matrix Multiplication
Problem Statement
Implement tiled matrix multiplication C = A @ B with A of shape (M, K) and B of shape (K, N), using shared-memory tiles of size TPB x TPB.
Background
Each block computes one TPB x TPB output tile. Loop over K in steps of TPB: cooperatively load a tile of A and a tile of B into shared memory, cuda.syncthreads(), accumulate the partial products, cuda.syncthreads() again before loading the next tiles. Shared memory lets every element be reused TPB times instead of re-read from DRAM.
Your Task
Implement matmul_kernel and run(M=64, N=64, K=64) comparing to A @ B.
How it is tested
Your solution must define a top-level function run(...) that allocates the inputs, copies them to the GPU, launches your @cuda.jit kernel, and returns a Python bool from np.allclose(gpu_result, reference). The grader prints run(...); the expected output is True.
Example
M = 64, N = 64, K = 64
True
- The input values
M = 64,N = 64, andK = 64define the dimensions of matricesAandBfor the matrix multiplicationC = A @ B. - The
matmul_kernelfunction is launched with these input matrices, utilizing shared-memory tiles to efficiently compute the product, with each block handling aTPB x TPBoutput tile. - The kernel accumulates partial products by cooperatively loading tiles of
AandBinto shared memory, allowing for reuse of elements and reducing DRAM accesses, ultimately producing the result matrixC. - The resulting matrix
Cfrom the GPU computation is compared to the reference result fromA @ Busingnp.allclose, yieldingTrueif the results match within a tolerance, indicating successful implementation of the CUDA tiled matrix multiplication.
Constraints
- Allocate two TPB x TPB shared tiles for A and B
- Loop over K in steps of TPB; cuda.syncthreads() before and after each inner product
- Bounds-check loads (pad with 0.0) and the final store