In this post, we will talk briefly about GPU and CUDA programming. Specifically, we will present the basics needed to understand the microarchitecture components involved during kernel computation. We will next discuss how to perform GEMM or General Matrix Multiplication using GPUs. Deep learning inference and training are computationally expensive.
In this post, we will talk briefly about GPU and CUDA programming. Specifically, we will present the basics needed to understand the microarchitecture components involved during kernel computation. We will next discuss how to perform GEMM or General Matrix Multiplication using GPUs. Deep learning inference and training are computationally expensive. To exploit locality, deep neural network layers (i.e., during the forward and backward passes for gradient propagation) are expressed in terms of GEMMs as much as possible (e.g., linear, convolutional, and Transformer layers). Thus, improving the performance of a GEMM kernel can directly improve the performance of deep learning models.
The goal is to understand how a performant kernel works (i.e how one can built it) and not to compete with cuBLAS or any other deeply optimized kernel.
For more details including the CUDA execution model refer to my notes https://books.deep-kondah.com (GPU section).
Each program, hosted in memory, consists of a set of instructions or machine code. The chief cpu starts with an initial program counter; it first fetches the instruction, followed by decoding to figure out what needs to be done. After that, the instruction is executed depending on the control unit output, which could be an arithmetic operation using the ALU or a memory operation (latency of memory transactions can take several hundreds of cycles). Finally, the last stage is write-back, where the register file could be updated with the result of the memory load or ALU operation.
Now, a naive single-cycle CPU processor is illustrated below, definitely not something we want to implement, as the clock cycle must be very long (to execute an entire instruction from start to finish) and the hardware units stay idle during the processing of the instruction. That is, the fetch or decode stage sit completely idle for the rest of the clock cycle which is a waste of resources.
Typically, a modern CPU uses pipelining, where the processing is decomposed into multiple stages, such as fetch (IF), decode (ID), execute (EX)., memory access (MEM) and write-back (WB); depending on the microarchitecture. Each stage requires one clock cycle and handles an instruction, and thus multiple instructions exist in the pipeline (they flow sequentially), in contrast to single-cycle machines. Of course, there are many challenges (pipeline hazards) we need to cope with due to the introduction of pipeline stages, such as data hazards and conditional branches.
We can think of performance in terms of the time it takes to complete the execution of a program:
\[ \text{Instruction Count}
\times
\text{CPI}
\times
\text{Clock Frequency}^{-1}
\]
Because memory, or DRAM, is relatively slow (e.g due to row-activation and precharge overheads), i.e., it can take many cycles to serve a request, and because of operand dependencies on previous instructions that have not yet completed, modern microprocessors feature large caches, make use of out-of-order execution, register renaming, branch prediction (and thus complex control logic) to reduce the number of cycles required to finish a program and thereby improve performance and reduce memory latency.
In scientific programs, there is typically a loop that operates on different elements of data; that is, the same instruction is executed several times. Instead of fetching that instruction again and again and decoding it, we can fetch and decode it once and execute it on the set of elements provided by the program, which here means updating the ISA to communicate this intent (to include vector instructions). Now, the program code has fewer dynamic instructions generated, and the branch overhead is decreased. This is vector hardware or vector processing, also known as SIMD.
There is also fine-grained multithreading, where several kernels or threads are launched, each with a different program counter, and the hardware juggles between them during execution. When one thread stalls for a memory request in the memory subsystem, another thread (with different program counter) can be picked by a scheduler using very fast(zero cycles penalty) hardware context switching(as the hardware context is hosted).
GPUs basically combine the two aforementioned concepts, that is, SIMD + fine-grained multithreading.
They have many multiprocessors, or SMs, each having warp schedulers and execution pipelines operating in a SIMD-like fashion. The critical difference between SIMD and what GPUs use (SIMT) is that each lane or thread can access a different memory location; that is, memory locations do not have to be consecutive as with vector loads or stores. Of course, having consecutive memory locations enables coalescing and generates less memory traffic.
A GPU kernel is written much like scalar code. Here is what the kernel looks like for matmul:
__global__ void matrixMul(float *A, float *B, float *C, int N)
{
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row < N && col < N) {
float sum = 0.0f;
for (int k = 0; k < N; k++) {
sum += A[row * N + k] * B[k * N + col];
}
C[row * N + col] = sum;
}
}The kernel is executed and invoked N times with N threads in the grid. The grid is basically a way to organize threads and define the program flow. Each thread is identified by coordinates, which can be 3-dimensional, and can use these coordinates to determine, for example, the memory locations that need to be handled by the thread in question. Specifically, threads are grouped into warps and are assigned to warp schedulers inside an SM. Each SM handles groups of warps belonging to blocks, and thus a grid is a multidimensional array of blocks (execution occurs in a hierarchical model).
The program and data flow vary depending on how you map the grid coordinates to memory locations used by the kernel. Inside a warp, threads are typically executed in lockstep, SIMD-like fashion. Therefore, if the threads in a warp access consecutive memory locations, fewer memory transactions are generated, and thus performance is better.
Additionally, threads can diverge during kernel execution, giving rise to warp divergence, which is handled by the GPU's control-flow mechanisms.
Each thread, and thus each warp, has access to registers. This is part of what enables fast switching between warps. Anything that spills from registers goes to local memory, which is mapped to global memory.
Each block has access to shared memory, which is on-chip, fast, and accessible during the lifetime of the block. If your scientific code enjoys high data reuse within the block (e.g matrix multiplication), you can use shared memory so that warps can fill their registers quickly instead of repeatedly hitting the GPU memory subsystem, thereby increasing performance.
GPUs also include specialized units such as Tensor Cores, which can be used to perform a fused matrix operation. Basically, threads in a warp collectively use their registers to hold operands, perform the operation, and finally store the result back to memory, i.e., the output matrix C.
Now, if we look at GEMM, a simple approach is to use block tiling by dividing the grid into blocks, with blocks scheduled onto SMs.
Next, we can improve performance using warp tiling and register tiling to increase the number of operations performed on loaded data by keeping a micro-tile resident in registers while streaming rows and columns of operand micro-tiles.
Shared memory is typically used to enable reuse, e.g., reuse of panels of A and B across warps.
Finally, Tensor Cores can be used at the warp level, where each warp works on an output micro-tile. The results of the experiments are listed below.
For more technical details, you can read my notes https://books.deep-kondah.com and also experiment with the code shared in https://github.com/mouadk/gemm-cuda.
You can also try to improve the kernel so that it first gets a fair comparison with cuBLAS and, why not, reaches similar performance if not better.
The goal of this post was to understand why GPUs are good at matrix multiplication (e.g by investing on compute and simpler control units), which is one of the cornerstones of forward and backward propagation in deep learning.
As planned, the next post will be about inference engines including vLLM.
