9 Graphic Processing Unit (GPU)

In this section, we will understand how GPUs accelerate the execution of thousands of High-Performance Computing (HPC), scientific, financial, and machine learning applications. A deep understanding of low-level details, i.e., microarchitectural behavior, is primordial to optimize software to achieve peak performance. For example, a GPU instruction whose source registers access the same register bank may result in bank conflicts, since only one value from each register bank can be provided per clock cycle. This can lead to performance degradation because, due to microarchitectural constraints, the requested data cannot all be serviced in the same cycle. In this case, the behavior is related to register mapping; the compiler does what it can to help prevent such conflicts.

PIC

A GPU, short for Graphics Processing Unit, is a coprocessor (not stand-alone ) traditionally connected to the CPU via PCIe (chip-to-chip interconnect); the CPU gets the responsibility to launch the computation and transfer data from/to system to/from device memory. It is designed to maximize throughput and tolerate long latencies (e.g the amount of time required to service a memory request due to off-chip memory accesses) by quickly switching between ready hardware threads or warps, i.e fine-grained multi-threading, where there could be ten of thousands of scalar threads (grouped into warps) created and being executed by the GPU, because they are extremely lightweight it is OK, as context switching happens with zero cycle penalty, this also means large portion of silicon is invested on compute density and simple control logic. In contrast, CPUs hide latency using several sophisticated microarchitectural techniques, including instruction-level parallelism (ILP), out-of-order execution, speculative execution, branch prediction, prefetching, and large cache hierarchies.

A GPU is basically a set of independent multiprocessors or Streaming Multiprocessors (SMs) (tens to hundreds). All SMs are connected to the GPU’s global memory, also called device memory (accessed via 32, 64, 128 byte memory transaction). Each SM contains SIMD execution units/lanes for executing the same instruction on different data elements and supports hardware multithreading (SIMD hardware is not directly exposed to the user and in this case each thread can access arbitrary memory location which contrast with vanilla SIMD ISA). So, warps on a GPU are like threads on a CPU (when fine-grained multi-threading is used), while a thread on a GPU is like a lane on a SIMD CPU. There could be thousands of active threads running in each SM. Key differences between traditional SIMD and GPUs are that, in the latter, each thread within a warp can independently compute and access different memory locations, which is not generally exposed in the same way in vanilla SIMD programming. Additionally, individual threads within a warp can follow different execution paths through control-flow divergence while employing a SIMD datapath (even though this creates branch divergence and thus causes some lanes to be idle), as covered later.

PIC

In SIMD execution, one instruction is fetched and decoded and then issued for execution across multiple data operands. Similar to vector processors, each active SIMD lane executes the instruction for one thread (e.g it could be a load or store). The number of physical SIMD lanes does not necessarily have to equal the warp size; if there are fewer physical lanes than threads in a warp, execution of a warp instruction may require multiple cycles.

A SIMD execution lane is sometimes associated with what NVIDIA calls a CUDA core or Streaming Processor (SP). To sustain operations on many data elements simultaneously, the register file must provide high read and write bandwidth; GPU register files are therefore typically heavily banked.

Furthermore, each SM has several load and store units to enable issuing several memory transactions (from threads) per clock cycle.

In addition to SIMD execution, fine-grained hardware multithreading is employed by GPUs to tolerate long memory and operation latencies. Hardware thread switching is considerably faster than an OS-level thread context switch because the architectural state of many GPU threads is maintained simultaneously on-chip. In contrast, an OS context switch may require architectural state, such as registers, to be saved and restored through memory.

Maintaining many hardware thread contexts (e.g register context) simultaneously also requires a huge register file. For example, if a processor maintains 48 resident threads and each thread uses 16 registers, then at least \(48 \times 16\) thread-private register values must be maintained.

The code executed by GPU threads is called a kernel. GPU programming follows the Single Program, Multiple Data or SPMD model where many threads execute the same computing kernel on different data which is different from SIMD where ISA contains vector instructions. At the hardware level, threads are grouped into warps and warp instructions are executed in lock step (executing the same instruction and advancing at the same time) on a SIMD pipeline.

Threads belonging to a warp therefore normally execute instructions together. However, due to conditional branches, some threads may not participate in the execution of a particular control-flow path. In this case, only the lanes corresponding to participating threads are active (e.g via masking).

At each cycle, warp schedulers (there could be many warp schedulers within each SM) select a ready warp (mapping between the warp scheduler and the warp could be as simple as doing warpId modulo N), which has a current active program counter (PC), and issue an instruction into the execution pipeline (for the entire warp as a whole). If the instruction is a conditional branch for which different threads within the warp follow different control-flow paths, branch divergence occurs. Only the subset of threads belonging to the currently executed path remains active, while the other threads are temporarily inactive until a path in which they participate is executed or until the threads reconverge.

PIC

In a SIMD execution pipeline, threads belonging to the same warp cannot generally execute different instructions on the same SIMD execution resources in the same cycle. Therefore, to allow threads within a warp to follow distinct control-flow paths, the divergent paths can be executed serially using different active masks. For each path, only the SIMD lanes corresponding to participating threads are enabled, while the remaining lanes are masked off. Once the divergent paths have been executed, the threads can reconverge and resume collective execution.

One mechanism for handling conditional branches is to introduce a divergence stack for each warp. Each stack entry contains three pieces of information, an execution program counter (PC), an active-vector bit mask, where each bit indicates whether the corresponding thread is active in the execution path starting at that PC, and a reconvergence PC, which indicates where the currently active threads will eventually rejoin the previously diverged threads.

PIC

Figure 3: source: [13]

The divergence stack initially contains a single entry corresponding to the first instruction of the kernel, together with a full active mask and a reconvergence point corresponding conceptually to the end of the kernel. If there is no divergence, the hardware does not need to actively use the reconvergence stack for ordinary instruction execution.

Otherwise, when a conditional branch causes divergence, with the help of compiler-provided reconvergence information, the stack is updated to represent the divergent execution paths. Entries corresponding to the different paths are pushed onto the stack, each containing its execution PC, an active mask indicating which threads participate in that path, and the associated reconvergence PC. Execution then proceeds along the path represented by the top-of-stack entry.

As warp instructions execute, the current warp PC advances normally. The warp’s next PC can be compared with the reconvergence PC associated with the currently executing path. When they match, the current divergent path has reached its reconvergence boundary. Its stack state is then popped, the active mask is updated according to the new top-of-stack entry, and execution proceeds from the PC stored in that entry (tos or top-of-stack entry ).

Consequently, after the first path reaches the reconvergence point, another pending path may become the top of the stack:

\[ PC \leftarrow PC_{\mathrm {TOS}}, \qquad Mask \leftarrow Mask_{\mathrm {TOS}} \]

After all divergent paths have reached the common reconvergence point, the corresponding divergent states have been removed and execution resumes at the reconvergence PC with the appropriate reconverged active mask.

The stack structure naturally supports nested control-flow divergence, when another divergent branch is encountered while executing an already divergent path, additional execution states can be pushed onto the same stack and processed before returning to the outer divergence.

During divergence, some SIMD lanes remain inactive or underutilized while each divergent path is being executed, which decreases SIMD utilization and results in a loss of parallelism; this is more pronounced for control-flow intensive applications. Ideally, we want all independent threads to be able to make progress (otherwise we degrade performance), this is why Independent Thread Scheduling (ITS) was introduced to overcome the serialization issue.

When threads within a warp access consecutive memory addresses, their memory accesses can be coalesced; that is, several loads or stores targeting the same memory transaction segment or cache line can be combined into a smaller number of memory transactions, thereby reducing memory traffic to the lower levels of the memory hierarchy. On the other hand, if the memory locations accessed by the threads have a large stride, performance may decrease due to the increased number of memory transactions generated, as multiple cache lines may need to be fetched. Thus, one must be careful about the memory access patterns employed.

PIC

Threads are grouped together into warps, which form the basic scheduling and instruction-issue units. A single warp instruction is fetched and decoded for the warp and executes collectively across its active lanes. Each SM contains a large register file (shared by all threads in the SM) in addition to on-chip shared memory (implemented as SRAM), which is fast to access but software-managed. Shared memory is accessible by threads belonging to the same thread block and remains allocated to that block while it is resident on the SM.

Shared memory is typically organized into banks (each bank has its own row/column decoder and sense amplifiers and can serve x-byte loads and stores from and to the addresses mapped to it). Only thread can be connected to the same memory bank at the same time, this means that if multiple threads access different addresses that map to the same memory bank, the resulting accesses must be serialized, resulting in a bank conflict (worst possible case is N-way bank conflict). In contrast, if multiple threads access the same shared-memory location, the value can be broadcast to the requesting threads without causing a bank conflict. For example, assume that shared memory contains 10 banks and that each bank can service one 8-byte word per access. If the 10 threads of a warp access 10 addresses that map to 10 different banks, there is no bank conflict, and all 10 accesses can be serviced simultaneously in the same cycle. On the other hand, if multiple threads access different addresses mapped to the same bank, those accesses must be serviced in multiple cycles, reducing the effective shared-memory bandwidth. A common way to resolve bank conflict when accessin a shared memory, e.g a 2d tensor, is to pad a column so that each thread memory location land in a different bank. For example, consider a warp of 32 threads and a 2D tensor stored in shared memory containing 32-bit floating-point elements. Suppose that, for each iteration \(i\), each thread accesses a different word location defined by

\[ \texttt {threadIdx.x} \times 32 + i \]

or if you will,

\[ \texttt {tensor[threadIdx.x * 32 + i]} \]

The compiler accounts for the size of each element when computing the actual memory address. Assuming a base address of zero, the byte address is therefore

\[ \text {Address} = (\texttt {threadIdx.x}\times 32+i) \times \operatorname {sizeof}(\texttt {float}) = (\texttt {threadIdx.x}\times 32+i)\times 4 \]

since memory is byte-addressable and a float occupies 4 bytes.

Assuming 32 shared-memory banks with a bank width of 4 bytes, the bank ID can be computed as

\[ \text {bankID} = \left ( \frac {\text {Address}}{4} \right ) \bmod 32 = (\text {Address} \gg 2)\bmod 32 \]

The right shift by two removes the two least-significant bits corresponding to the byte offset within a 4-byte word, while the next five bits select one of the \(32=2^5\) banks.

For \(i=0\), the bank mapping becomes

\[ \text {bankID} = (\texttt {threadIdx.x}\times 32)\bmod 32 = 0 \]

Thus, during the first iteration, all 32 threads access different memory locations that map to the same bank, resulting in a 32-way bank conflict. More generally, for any fixed \(i\),

\[ \text {bankID} = (\texttt {threadIdx.x}\times 32+i)\bmod 32 = i\bmod 32 \]

so all threads still map to the same bank.

On the other hand, if we add one dummy padding column, the row stride becomes 33 words instead of 32. The mapping then becomes

\[ \text {bankID} = (\texttt {threadIdx.x}\times 33+i)\bmod 32. \]

For \(i=0\),

\[ \text {bankID} = (\texttt {threadIdx.x}\times 33)\bmod 32 = \texttt {threadIdx.x}\bmod 32 \]

Therefore, thread 0 maps to bank 0, thread 1 to bank 1, and so on, up to thread 31 mapping to bank 31, thereby eliminating the bank conflict.

Each thread has private registers. A thread also has a private local memory address space backed by GPU device memory or DRAM (meaning high latency), for example when register spilling occurs, and is therefore not physically local on-chip memory (unless there is a cache). Similarly, threads within a block share a shared-memory address space.

GPUs also provide specialized matrix-multiplication units, or Tensor Cores, that perform matrix multiply-and-accumulate (MMA) operations of the form

\[ D = A \times B + C \]

where \(A\) and \(B\) are input matrix fragments with shapes \(m \times k\) and \(k \times n\), respectively, \(C\) is an \(m \times n\) accumulator matrix fragment, and \(D\) is the resulting \(m \times n\) matrix fragment. These specialized units can execute many multiply-and-accumulate operations in parallel and are used to accelerate general matrix multiplications (GEMMs), and consequently deep-learning workloads such as AI model training and inference.

Thus, a basic GPU pipeline consists of a scalar front end (instruction fetch and decode), SIMD backend/execution, and writeback (long latency operations are tolerated by executing/scheduling many warps on the same core). Because an SM maintains many resident warps, it maintains the execution state associated with these warps, including their program counters. Each warp has a corresponding slot in a warp scheduler. At every cycle, each warp scheduler selects a ready warp for instruction issue, that is, select the next instruction form the list of available warps (e.g using round-robin policy); each warp has a dedicated PC. When one warp encounters a dependency or another stalling event, all the threads in the warp would be stalled and the scheduler select another ready warp, thereby hiding latency (similar to fine grained multi-threading architecture). Hardware structures like scoreboards are used to track operand readiness (e.g only instructions with no read-after-write/write-after-write hazards are considered for scheduling meaning instructions from a given warp can be scheduled for execution even if past instructions have not finished yet) and outstanding dependencies and help determine whether instructions from a warp are ready for issue.

Once a warp instruction has been selected and decoded, its active threads obtain their source operands from the SM register file. Logically, each thread has its own architectural registers, while physically these registers reside in a large shared, banked register file on the SM. That is, the threads in the warp access the register file using their associated warp ID and the specified register IDs, and the retrieved register values are then supplied in parallel to the corresponding streaming processors or SIMD lanes. The instruction then enters the execution stage, where the active threads are mapped onto SIMD execution lanes. Depending on the instruction, execution may involve arithmetic units, load/store units or tensor units. The resulting values are eventually written back into the corresponding thread registers and global memory which can be then copied back to the system memory.

A GPU program contains host code, which executes on the CPU, and device code, or kernels, which execute on the GPU (the kernel code is loaded into the device memory by the GPU driver). A kernel code can only access GPU memory, can not define static variables and must have void return type. Additionally, each active GPU context/program has a dedicated page table (similar to CPU). The host code is compiled for the CPU architecture, whereas GPU device code is compiled for the GPU architecture. In CUDA, for example, device code may first be represented using PTX, an intermediate virtual instruction-set representation or ISA (PTX virtual ISA was introduced to provide backward compatibility), and subsequently translated into the native machine ISA (the actual ISA; partially documented) corresponding to a particular GPU architecture.

A kernel is executed by many threads. The number and organization of these threads are determined by the grid configuration specified when launching the kernel. A grid can be multidimensional, and each grid element corresponds to a thread block. A block itself contains multiple threads and can also be organized in one, two, or three dimensions. Threads are assigned to SMs in block granularity. In CUDA notation, a kernel can be launched as

\[ \texttt {kernel<<<dimGrid, dimBlock>>>(args)} \]

Thread blocks are dispatched to SMs. A block remains associated with an SM during its execution and is not migrated between SMs. Multiple blocks can simultaneously reside on the same SM. Threads within each block are further grouped into warps for scheduling and execution. Only a subset of all resident warps is issued in any particular cycle (depending on the size of warp schedulers). When all blocks belonging to the grid have completed, the kernel execution is complete.

Although kernel code is written from the perspective of an individual scalar thread, many thread instances execute the same kernel. Each thread has identifiers, such as

\[ \texttt {threadIdx}, \qquad \texttt {blockIdx}, \]

that can be combined with the block dimensions to construct a unique global thread index. In one dimension, for example,

\[ i = \texttt {blockIdx.x} \times \texttt {blockDim.x} + \texttt {threadIdx.x} \]

This index can then be mapped to a data element. One can map the elements of a tensor or array onto the threads of the grid. Because consecutive thread IDs are normally grouped into the same warp, arranging the mapping such that consecutive threads access consecutive memory locations enables memory coalescing as said ealier. The active threads of a warp can then execute the same instruction across SIMD lanes while the memory subsystem services their accesses using a small number of efficient memory transactions.

Let us consider first the 2D case. Each block in a grid can be uniquely identified using

\[ (\texttt {blockIdx.x},\texttt {blockIdx.y}) \in [0,\texttt {gridDim.x}-1] \times [0,\texttt {gridDim.y}-1] \]

Once a block is identified, a thread within the block can similarly be uniquely identified using

\[ (\texttt {threadIdx.x},\texttt {threadIdx.y}) \in [0,\texttt {blockDim.x}-1] \times [0,\texttt {blockDim.y}-1] \]

Given the coordinates

\[ (\texttt {blockIdx.x},\texttt {blockIdx.y}, \texttt {threadIdx.x},\texttt {threadIdx.y}) \]

we can construct a unique linear thread index in exactly the same way as when linearizing a multidimensional tensor. Since

\[ \text {blockSize} = \texttt {blockDim.x}\times \texttt {blockDim.y} \]

the linear index is

\[ \begin {aligned} \text {index} ={}& \texttt {blockIdx.y} \times \texttt {gridDim.x} \times \text {blockSize} \\ &+ \texttt {blockIdx.x} \times \text {blockSize} \\ &+ \texttt {threadIdx.y} \times \texttt {blockDim.x} + \texttt {threadIdx.x} \end {aligned} \]

Or if you will,

\[ \text {index} = \left ( \texttt {blockIdx.y}\texttt {gridDim.x} + \texttt {blockIdx.x} \right ) \left ( \texttt {blockDim.x}\texttt {blockDim.y} \right ) +\texttt {threadIdx.y}\texttt {blockDim.x} + \texttt {threadIdx.x} \]

More generally, we can borrow the row-major multidimensional tensor linearization formula. The important point is that x is the fastest-varying dimension, followed by y, and then z. The global linear index can therefore be decomposed as

\[ \text {index} = \text {linearized block index} \times \text {number of threads per block} + \text {linearized thread index within the block} \]

For example, in 3D, a block is identified by

\[ (\texttt {blockIdx.x},\texttt {blockIdx.y},\texttt {blockIdx.z}) \]

and a thread within that block by

\[ (\texttt {threadIdx.x},\texttt {threadIdx.y},\texttt {threadIdx.z}) \]

The linearized block index is

\[ \text {blockIndex} = \texttt {blockIdx.z} \times \texttt {gridDim.y} \times \texttt {gridDim.x} + \texttt {blockIdx.y}\times \texttt {gridDim.x} + \texttt {blockIdx.x} \]

while the linearized thread index within the block is

\[ \text {threadIndex} = \texttt {threadIdx.z} \times \texttt {blockDim.y}\times \texttt {blockDim.x} + \texttt {threadIdx.y}\times \texttt {blockDim.x} + \texttt {threadIdx.x} \]

Since the number of threads in each block is

\[ \text {blockSize} = \texttt {blockDim.x} \texttt {blockDim.y} \texttt {blockDim.z} \]

the complete global linear thread index is

\[ \begin {aligned} \text {index} ={}& \left ( \texttt {blockIdx.z} \texttt {gridDim.y} \texttt {gridDim.x} + \texttt {blockIdx.y} \texttt {gridDim.x} + \texttt {blockIdx.x} \right ) \\ &\times \left ( \texttt {blockDim.x} \texttt {blockDim.y} \texttt {blockDim.z} \right ) \\ &+ \texttt {threadIdx.z} \texttt {blockDim.y} \texttt {blockDim.x} + \texttt {threadIdx.y} \texttt {blockDim.x} + \texttt {threadIdx.x} \end {aligned} \]

Thus, we can uniquely identify every thread for arbitrary grid and block dimensions using the same principle employed to linearize multidimensional arrays. Blocks and threads can be viewed as elements of nested multidimensional iteration spaces, we first linearize the block coordinates, multiply by the number of threads per block (mimicing size of elements), and then add the linearized coordinates of the thread within that block (where size of each element is 1).

9.0.1 GEMM CUDA Implementation

When implementing a CUDA kernel, we first need to define what the grid looks like. Basically, the grid configuration is chosen based on the size of the tensors we are operating on. Each coordinate of the grid is mapped to a tensor location, which defines the traversal during computation (the grid configuration may lead to different program flow). This is where one should be very careful, as it may have a tremendous impact on performance, especially if one wants the observed performance of the kernel to be close to the theoretical peak performance. To illustrate this, we will look at GEMM,

\[ C += AB \]

assuming that \(C\) is initially zero. We start step by step until we reach a good kernel.

PIC

To illustrate why memory pattern is critical, we start with a small experiment. Suppose that the tensor \(C\) is small enough to fit within one block, that is, the grid dimension is \(1\times 1\), and the block dimensions match those of \(C\). Now suppose we decide that the row index \(i\) is mapped to \(\texttt {threadIdx.x}\) and the column index \(j\) to \(\texttt {threadIdx.y}\). That is,

\[ i=\texttt {threadIdx.x}, \qquad j=\texttt {threadIdx.y} \]

The corresponding loop nest is

\[ \begin {aligned} &\text {for all } \texttt {threadIdx.y}=0:1:\texttt {blockDim.y}-1\\ &\qquad \text {for all } \texttt {threadIdx.x}=0:W:\texttt {blockDim.x}-W\\ &\qquad \qquad i=\texttt {threadIdx.x},\qquad j=\texttt {threadIdx.y}\\ &\qquad \qquad \text {for all } k=0:1:K-1\\ &\qquad \qquad \qquad c_{ij} \mathrel {+}=a_{ik}b_{kj} \end {aligned} \]

Here, \(W\) is the warp size, and

\[ \texttt {threadIdx.x}=0:W:\texttt {blockDim.x}-W \]

basically advances by \(W\), since execution happens in lockstep within a warp. Thus, the inner computation updates one subcolumn of \(C\), in \(\mathbb {R}^{W}\), by multiplying a row panel of \(A\), in \(\mathbb {R}^{W\times K}\), by a column of \(B\), in \(\mathbb {R}^{K}\):

\[ C_{i:i+W,j} = A_{i:i+W,:}\,B_{:,j} \]

PIC

While the computation is correct, the access pattern is terrible. We start with \(C\). Each \(c_{ij}\) is reused across the \(k\)-loop, so it can be stored in a register for each thread. However, because \(C\) is traversed column by column, consecutive threads of a warp access elements belonging to different rows of the same column. Assuming row-major storage, these accesses are strided rather than contiguous, resulting in poorly coalesced memory transactions when \(C\) is loaded or stored.

Now consider \(B\). Each thread traverses the same column of \(B\). At a fixed \(k\), all threads in the warp require the same value \(b_{kj}\), so this value can potentially be broadcast. For \(A\), consecutive threads within the warp access elements belonging to different rows of \(A\), so the accesses are not coalesced across the warp. Thus, the access pattern for \(A\) has good spatial locality along the \(k\)-dimension for each thread, although but does not provide warp-level memory coalescing.

What we need to fix is the traversal by instead mapping the column index to \(\texttt {threadIdx.x}\) and the row index to \(\texttt {threadIdx.y}\). The loop nest becomes

\[ \begin {aligned} &\text {for all } \texttt {threadIdx.y}=0:1:\texttt {blockDim.y}-1\\ &\qquad \text {for all } \texttt {threadIdx.x}=0:W:\texttt {blockDim.x}-W\\ &\qquad \qquad j=\texttt {threadIdx.x},\qquad i=\texttt {threadIdx.y}\\ &\qquad \qquad \text {for all } k=0:1:K-1\\ &\qquad \qquad \qquad c_{ij} \mathrel {+}=a_{ik}b_{kj} \end {aligned} \]

This corresponds to

\[ C_{i,j:j+W} = A_{i,:}\,B_{:,j:j+W} \]

Now each warp generates much less memory traffic. At each \(k\)-iteration, the warp accesses the subrow \(B_{k,j:j+W}\), so consecutive threads access consecutive elements of \(B\), allowing the memory requests to be coalesced. The same happens for \(C\); each thread keeps its \(c_{ij}\) accumulator in a register throughout the \(k\)-loop, and when the loop finishes, consecutive threads store consecutive elements \( C_{i,j:j+W}\) allowing the stores to be coalesced.

For \(A\), all threads in the warp use the same element \(a_{ik}\) at a given \(k\), so the value can be shared across the warp through the memory system’s broadcast behavior.

PIC

As the results demonstrate, the memory-access pattern has a significant impact on performance, even though the underlying arithmetic operation remains unchanged.

Table 1: Effect of memory coalescing for a \(32\times 32\times 512\) matrix multiplication.
Kernel Time (ms) GFLOP/s Max Error
Global memory + non-coalesced mapping 0.4378 2.40 \(0.000\times 10^{0}\)
Global memory + coalesced mapping 0.0404 25.92 \(0.000\times 10^{0}\)

GPU: NVIDIA A100-SXM4-40GB, compute capability 8.0, warp size 32.

Usually, we create many blocks and end up with the following loop nest (block tiling):

\[ \begin {aligned} &\text {for all } \texttt {blockIdx.y}=0:1:\texttt {gridDim.y}-1\\ &\qquad \text {for all } \texttt {blockIdx.x}=0:1:\texttt {gridDim.x}-1\\ &\qquad \qquad \text {for all } \texttt {threadIdx.y}=0:1:\texttt {blockDim.y}-1\\ &\qquad \qquad \qquad \text {for all } \texttt {threadIdx.x}=0:W:\texttt {blockDim.x}-W\\ &\qquad \qquad \qquad \qquad j=\texttt {blockIdx.x}\,\texttt {blockDim.x}+\texttt {threadIdx.x},\\ &\qquad \qquad \qquad \qquad i=\texttt {blockIdx.y}\,\texttt {blockDim.y}+\texttt {threadIdx.y},\\ &\qquad \qquad \qquad \qquad \text {for all } k=0:1:K-1\\ &\qquad \qquad \qquad \qquad \qquad c_{ij}\mathrel {+}=a_{ik}b_{kj}. \end {aligned} \]

While the inner loop now has a good memory access pattern, within a given block, i.e., for fixed \(\texttt {blockIdx.y}\) and \(\texttt {blockIdx.x}\), the same panel of \(B\) is reused across different values of \(\texttt {threadIdx.y}\). Without explicit caching, the same elements of \(B\) may therefore need to be fetched repeatedly through the memory hierarchy.

Similarly, elements of \(A\) are reused across threads spanning different columns of the output tile ( i.e when we advance\(\texttt {threadIdx.x}\)).

Fortunately, we can exploit shared memory, which is private to a block and persists throughout the block’s lifetime, by first loading the required panels of \(A\) and \(B\) into shared memory. Their reuse then benefits from explicit on-chip locality. For a large \(K\), however, these panels may be too large to fit entirely in shared memory. We therefore tile the \(K\)-dimension and operate on smaller panels or blocks of \(A\) and \(B\):

\[ \begin {aligned} &\text {for all } \texttt {blockIdx.y}=0:1:\texttt {gridDim.y}-1\\ &\qquad \text {for all } \texttt {blockIdx.x}=0:1:\texttt {gridDim.x}-1\\ &\qquad \qquad \text {for all } k'=0:k_c:K-k_c\\ &\qquad \qquad \qquad \text {load } A_r \text { and } B_r \text { into shared memory}\\ &\qquad \qquad \qquad \text {for all } \texttt {threadIdx.y}=0:1:\texttt {blockDim.y}-1\\ &\qquad \qquad \qquad \qquad \text {for all } \texttt {threadIdx.x}=0:W:\texttt {blockDim.x}-W\\ &\qquad \qquad \qquad \qquad \qquad j=\texttt {blockIdx.x}\,\texttt {blockDim.x}+\texttt {threadIdx.x},\\ &\qquad \qquad \qquad \qquad \qquad i=\texttt {blockIdx.y}\,\texttt {blockDim.y}+\texttt {threadIdx.y},\\ &\qquad \qquad \qquad \qquad \qquad \text {for all } k=k':1:k'+k_c-1\\ &\qquad \qquad \qquad \qquad \qquad \qquad C_{ij}\mathrel {+}=A^{(r)}_{ik}B^{(r)}_{kj} \end {aligned} \]

Thus, instead of operating over the entire \(K\)-dimension at once, each block processes \(K\) in tiles of size \(k_c\), with the corresponding panels of \(A\) and \(B\) staged in shared memory.

As the experimental results demonstrate, explicitly exploiting data reuse through shared memory improves the performance of the kernel.

Table 2: Effect of shared-memory tiling on GEMM performance.
Kernel Time (ms) GFLOP/s Max Error
Block tiling + global memory 6.0887 2821.60 \(0.000\times 10^{0}\)
Block tiling + shared memory 4.1651 4124.70 \(0.000\times 10^{0}\)

Now we will look at warp and register tiling. Basically, each block is partitioned so that each warp handles a tile of size

\[ M_W \times N_W \]

Each block base address is given by

\[ C_{\rm blockBase} = \mathrm {blockIdx}.y\,\mathrm {blockDim}.y\, \mathrm {gridDim}.x\,\mathrm {blockDim}.x + \mathrm {blockIdx}.x\,\mathrm {blockDim}.x \]

It is not difficult to see that the base address of the warp tile is

\[ C_{\rm tileBase} = C_{\rm blockBase} + \mathrm {warpRow}\,M_W\, \mathrm {gridDim}.x\,\mathrm {blockDim}.x + \mathrm {warpCol}\,N_W \]

where

\[ \mathrm {warpId} = \left \lfloor \frac {\mathrm {threadId}}{W} \right \rfloor , \]
\[ \mathrm {warpRow} = \left \lfloor \frac {\mathrm {warpId}} {\mathrm {blockDim}.x/N_W} \right \rfloor , \]

and

\[ \mathrm {warpCol} = \mathrm {warpId} \bmod \left ( \frac {\mathrm {blockDim}.x}{N_W} \right ) \]

PIC

Once the base address \(C_{\rm tileBase}\) of the selected warp tile has been computed, the warp tile of size

\[ M_W\times N_W \]

is further divided into register microtiles of size

\[ M_R\times N_R \]

Each thread of the warp is responsible for one microtile of the output matrix \(C\). Therefore,

\[ W= \frac {M_W}{M_R} \frac {N_W}{N_R} \]

when the warp tile is exactly partitioned into one \(M_R\times N_R\) microtile per thread. The number of thread microtiles along the column direction of the warp tile is

\[ N_{\rm micro,col} = \frac {N_W}{N_R} \]

Using the lane index inside the warp (\(\texttt {tid} \operatorname {mod} W\)), the microtile row and column are

\[ \mathrm {microRow} = \left \lfloor \frac {\mathrm {lane}}{N_W/N_R} \right \rfloor \]

and

\[ \mathrm {microCol} = \mathrm {lane} \bmod \left ( \frac {N_W}{N_R} \right ) \]

Hence, the row and column offsets of the thread microtile inside the warp tile are

\[ m_{\rm micro} = \mathrm {microRow}\,M_R, \qquad n_{\rm micro} = \mathrm {microCol}\,N_R \]

Since

\[ N = \mathrm {gridDim}.x\,\mathrm {blockDim}.x \]

the base address of the \(M_R\times N_R\) output microtile assigned to the thread is

\[ C_{\rm microBase} = C_{\rm tileBase} + \mathrm {microRow}\,M_R\,N + \mathrm {microCol}\,N_R \]

That is,

\[ C_{\rm microBase} = C_{\rm tileBase} + \left \lfloor \frac {\mathrm {lane}}{N_W/N_R} \right \rfloor M_R \left ( \mathrm {gridDim}.x\,\mathrm {blockDim}.x \right ) + \left ( \mathrm {lane} \bmod \frac {N_W}{N_R} \right )N_R \]

For local coordinates

\[ (r,c)\in [0,M_R-1]\times [0,N_R-1] \]

the corresponding output address is therefore

\[ C_{\rm addr}^{(R)}(r,c) = C_{\rm microBase} + rN+c \]

Once the base address \(C_{\rm microBase}\) has been determined, the thread keeps the corresponding \(M_R\times N_R\) output microtile resident in its register file during the reduction over \(K\). Denote this register-resident output tile by

\[ C_R\in \mathbb {R}^{M_R\times N_R} \]

Thus, each thread owns \( M_RN_R\) register accumulators,

\[ C_R= \begin {bmatrix} c_{00} & c_{01} & \cdots & c_{0,N_R-1}\\ c_{10} & c_{11} & \cdots & c_{1,N_R-1}\\ \vdots & \vdots & \ddots & \vdots \\ c_{M_R-1,0} & c_{M_R-1,1} & \cdots & c_{M_R-1,N_R-1} \end {bmatrix} \]

Assuming that \(C\) is initially zero, these registers are initialized as

\[ C_R[r,c]=0, \qquad 0\leq r<M_R,\quad 0\leq c<N_R \]

They remain resident in the register file of the thread throughout the \(K\)-loop.

The block tiles of \(A\) and \(B\) are first staged in shared memory. During the inner \(K\)-loop, each thread streams from shared memory the fragments required to update its \(M_R\times N_R\) register tile.

For a fixed \(k\), the thread loads an \(M_R\)-element fragment of \(A\),

\[ A_R= \begin {bmatrix} a_0\\ a_1\\ \vdots \\ a_{M_R-1} \end {bmatrix} \in \mathbb {R}^{M_R} \]

and an \(N_R\)-element fragment of \(B\),

\[ B_R= \begin {bmatrix} b_0 & b_1 & \cdots & b_{N_R-1} \end {bmatrix} \in \mathbb {R}^{N_R} \]

The \(M_R\) values of \(A_R\) and the \(N_R\) values of \(B_R\) are loaded into registers and used to update the register-resident output tile through the outer product

\[ C_R \mathrel {+}= A_R B_R \]

That is,

\[ C_R[r,c] \mathrel {+}= A_R[r]B_R[c], \qquad 0\leq r<M_R,\quad 0\leq c<N_R \]

or if you will,

\[ \begin {bmatrix} c_{00} & \cdots & c_{0,N_R-1}\\ c_{10} & \cdots & c_{1,N_R-1}\\ \vdots & \ddots & \vdots \\ c_{M_R-1,0} & \cdots & c_{M_R-1,N_R-1} \end {bmatrix} \mathrel {+}= \begin {bmatrix} a_0\\ a_1\\ \vdots \\ a_{M_R-1} \end {bmatrix} \begin {bmatrix} b_0 & b_1 & \cdots & b_{N_R-1} \end {bmatrix} \]

Therefore, at each iteration of the \(k\)-loop, \(M_R+N_R\) values are streamed from shared memory into registers and are used to perform \(M_RN_R\) FMAs on the register-resident output tile.

After the reduction has completed, the \(M_RN_R\) register accumulators are written back to their corresponding locations in \(C\).

PIC

\[ \begin {aligned} &\text {for all } \texttt {blockIdx.y}=0:1:\texttt {gridDim.y}-1\\ &\qquad \text {for all } \texttt {blockIdx.x}=0:1:\texttt {gridDim.x}-1\\[1mm] &\qquad \qquad \text {for all } k'=0:k_c:K-k_c\\ &\qquad \qquad \qquad \text {load } A_{\rm block} \text { and } B_{\rm block} \text { into shared memory}\\[1mm] &\qquad \qquad \qquad \text {for all } \mathrm {warpId} =0:1:\mathrm {warpsPerBlock}-1\\ &\qquad \qquad \qquad \qquad \mathrm {warpRowTile} = \left \lfloor \frac {\mathrm {warpId}}{N_{\rm warp,col}} \right \rfloor \\ &\qquad \qquad \qquad \qquad \mathrm {warpColTile} = \mathrm {warpId}\bmod N_{\rm warp,col}\\[1mm] &\qquad \qquad \qquad \qquad C_{\rm tileBase} = C_{\rm blockBase} + \mathrm {warpRowTile}\,M_W\,N + \mathrm {warpColTile}\,N_W\\[1mm] &\qquad \qquad \qquad \qquad \text {for all } \mathrm {lane}=0:1:W-1\\ &\qquad \qquad \qquad \qquad \qquad \mathrm {microRow} = \left \lfloor \frac {\mathrm {lane}}{N_W/N_R} \right \rfloor \\ &\qquad \qquad \qquad \qquad \qquad \mathrm {microCol} = \mathrm {lane}\bmod (N_W/N_R)\\[1mm] &\qquad \qquad \qquad \qquad \qquad C_{\rm microBase} = C_{\rm tileBase} + \mathrm {microRow}\,M_R\,N + \mathrm {microCol}\,N_R\\[1mm] &\qquad \qquad \qquad \qquad \qquad \text {keep } C_R\in \mathbb {R}^{M_R\times N_R} \text { resident in registers}\\[1mm] &\qquad \qquad \qquad \qquad \qquad \text {for all } k=k':1:k'+k_c-1\\ &\qquad \qquad \qquad \qquad \qquad \qquad \text {load } A_R\in \mathbb {R}^{M_R} \text { from shared memory into registers}\\ &\qquad \qquad \qquad \qquad \qquad \qquad \text {load } B_R\in \mathbb {R}^{N_R} \text { from shared memory into registers}\\[1mm] &\qquad \qquad \qquad \qquad \qquad \qquad \text {for all } r=0:1:M_R-1\\ &\qquad \qquad \qquad \qquad \qquad \qquad \qquad \text {for all } c=0:1:N_R-1\\ &\qquad \qquad \qquad \qquad \qquad \qquad \qquad \qquad C_R[r,c] \mathrel {+}= A_R[r]B_R[c] \end {aligned} \]

Now we will introduce Tensor Cores. As discussed at the beginning of this section, each Streaming Multiprocessor, or SM (the fundamental processing unit of the GPU), in addition to SIMD hardware, hosts Tensor Cores, specialized computation units used to accelerate matrix multiplication. The operands for a Tensor Core operation are four matrices \(D\), \(A\), \(B\), and \(C\), and the computation performed (Warp Matrix Multiply and Accumulate, or WMMA) is a multiply-add operation:

\[ D = AB + C \]

where \(A\) and \(B\) are input matrices with shapes \(m \times k\) and \(k \times n\), respectively; \(C\) is the accumulator, and \(D\) is the result (which can overwrite \(C\) for MMA).

If the Tensor Core performs a \(4 \times 4 \times 4\) matrix multiplication, then \(m=n=k=4\). Each element in the resulting \(N \times N\) matrix \(C\) (assuming \(C=D\)) is computed via FMA operations as

\[ c_{i,j} = \sum _{p=1}^{N} A_{i,p}B_{p,j} = A_{i,1}B_{1,j} + A_{i,2}B_{2,j} +\cdots + A_{i,N}B_{N,j} \]

The Tensor Core performs these multiply-accumulate operations using dedicated hardware. Each operand is declared using a fragment, for example:

wmma::fragment a_frag;

wmma::fragment b_frag;

wmma::fragment acc_frag;

wmma::fragment c_frag;

To initialize the accumulator, we can use

wmma::fill_fragment(acc_frag, 0.0f);

Each warp is responsible for one output tile. Kernel code perform the matrix multiplication using warp-level primitives. This means that all threads within the warp participate in computing the result, and one matrix MMA operation is performed collectively by the warp per WMMA instruction.

The first step is to load data from memory into the fragments:

wmma::load_matrix_sync(a_frag, a + aRow + aCol * lda, lda);
wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);

Then, the warp performs the matrix multiply-accumulate operation:

wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);

Finally, the primitive wmma::store_matrix_sync is used to write the result from the accumulator fragment back to memory:

wmma::store_matrix_sync(c + cRow + cCol * ldc,
                        acc_frag,
                        ldc,
                        wmma::mem_col_major);

The performance achieved using the different optimization mechanisms discussed in this section is summarized below.

Table 3: Performance comparison of the GEMM implementations for \(M=N=K=2048\).
Kernel Time (ms) GFLOP/s Max Error
Block tiling + global memory 6.0887 2821.60 \(0.000\times 10^{0}\)
Block tiling + shared memory 4.1651 4124.70 \(0.000\times 10^{0}\)
Block + shared + warp tiling 6.3380 2710.59 \(0.000\times 10^{0}\)
Block + shared + warp + register tiling 1.6128 10652.20 \(0.000\times 10^{0}\)
Block + shared + warp + Tensor Cores 1.5160 11332.13 \(6.166\times 10^{-5}\)
cuBLAS SGEMM FP32 1.0788 15925.22 \(0.000\times 10^{0}\)

NVIDIA A100-SXM4-40GB, compute capability 8.0, warp size 32.