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
Background Knowledge
Introduction to CUDA and Matrix Multiplication
CUDA is a parallel computing platform and programming model developed by NVIDIA. It allows developers to harness the power of GPU (Graphics Processing Unit) to perform computationally intensive tasks. One such task is matrix multiplication, which is a fundamental operation in linear algebra. Matrix multiplication involves multiplying two matrices A and B to produce a resulting matrix C. The number of columns in A must be equal to the number of rows in B.
Shared Memory and Tiling
In the context of CUDA, shared memory refers to a small amount of memory that is shared among threads within a block. Tiling is a technique used to divide a large matrix into smaller sub-matrices, called tiles, to improve memory access patterns and reduce memory traffic. By using shared memory tiles, we can cooperatively load a tile of A and a tile of B into shared memory, allowing every element to be reused TPB times instead of re-reading from DRAM (Dynamic Random Access Memory).
CUDA Thread Hierarchy
In CUDA, the thread hierarchy consists of a grid of blocks, with each block containing a group of threads. The threads within a block can cooperate with each other using shared memory and synchronization primitives like cuda.syncthreads(). This allows for efficient parallelization of tasks like matrix multiplication. The TPB (threads per block) is a crucial parameter that determines the number of threads in a block, which in turn affects the performance of the kernel.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.