The Evolution of Performance: Matrix Multiplication on Tensor Cores#

Early GPU Programming#

In the early 2000s, a discipline known as GPGPU (General-Purpose Computing on Graphics Processing Units) started to gain traction. Researchers and practitioners began leveraging the shaders and limited programmability of graphics pipelines for scientific computing tasks. While these techniques were sometimes cumbersome (using graphics APIs like OpenGL or DirectX for computation), they laid the groundwork for dedicated GPU programming frameworks such as NVIDIA’s CUDA (Compute Unified Device Architecture) and the OpenCL standard.

CUDA and Beyond#

With the introduction of CUDA in 2007, developers gained more direct control over GPU hardware. CUDA provided a C-like language extension, making it more intuitive to write parallel kernels targeting the GPU. This opened new frontiers for HPC, with GPUs accelerating not only graphics workloads but also large-scale matrix multiplications, deep learning training, and more.

The Key Idea#

Traditional GPU cores excel at Single-Precision (FP32) or Double-Precision (FP64) computations, but deep learning and many HPC applications can often work with lower precisions (e.g., FP16, BF16) without significantly impacting accuracy. In certain domains, using half-precision or mixed-precision arithmetic can dramatically speed up calculations and reduce memory bandwidth requirements. Tensor Cores capitalize on this by performing multiple multiply-and-accumulate operations in a specialized circuit designed for matrix math.

The Tensor Operation#

A single Tensor Core is capable of performing a fused multiply-add operation over small matrix tiles, typically 4×4 or 8×8, in a single clock cycle (depending on the hardware generation). This specialized design gives Tensor Cores their massive throughput advantage, especially for matrix-heavy tasks like deep learning training and inference.

Hardware Requirements#

• Tensor Cores were first introduced in NVIDIA’s Volta GPU architecture (e.g., Tesla V100).
• Subsequent architectures such as Turing (e.g., RTX 20 series, Quadro RTX), Ampere (e.g., A100, RTX 30 series), and Hopper continue to expand and refine Tensor Core capabilities.

If you have a GPU from the Volta series or later, you can leverage Tensor Cores. Make sure you have the appropriate driver and CUDA toolkit installed.

Software Libraries#

Many software libraries can help you integrate Tensor Core operations without writing low-level intrinsics:

• cuBLAS: NVIDIA’s CUDA Basic Linear Algebra Subroutine library provides convenience functions for matrix multiplication (SGEMM, DGEMM, HGEMM, etc.) that can automatically utilize Tensor Cores for FP16 operations.
• cuBLASLt: This newer library extends cuBLAS with more flexible APIs, making it easier to control matrix layouts and precision settings for Tensor Cores.
• CUTLASS: A collection of templated C++ compute kernels for GEMM (General Matrix Multiply), convolution, and more, designed for learning and customizing advanced GPU operations, including Tensor Cores.
• cuDNN: NVIDIA’s library for deep neural networks also uses Tensor Cores behind the scenes for convolution and other operations.

You can often exploit Tensor Cores by simply ensuring your data is in half precision (FP16) or another supported format, then calling the relevant library function.

Example 1: A Simple cuBLAS Matmul#

Here is a brief example of how you might call cuBLAS to perform half-precision matrix multiplication (HGEMM). Note that actual error-checking and memory management details are typically more extensive.

1
#include <iostream>
2
#include <cuda_fp16.h>
3
#include <cublas_v2.h>
4

5
int main() {
6
    // Dimensions
7
    const int M = 512, K = 512, N = 512;
8
    // Allocate host memory (using half precision)
9
    half *h_A, *h_B, *h_C;
10
    h_A = (half*)malloc(M*K*sizeof(half));
11
    h_B = (half*)malloc(K*N*sizeof(half));
12
    h_C = (half*)malloc(M*N*sizeof(half));
13

14
    // Initialize host memory with some values
15
    for(int i = 0; i < M*K; i++) {
16
        h_A[i] = __float2half(static_cast<float>(i % 5));
17
    }
18
    for(int i = 0; i < K*N; i++) {
19
        h_B[i] = __float2half(static_cast<float>(i % 3));
20
    }
21

22
    // Device pointers
23
    half *d_A, *d_B, *d_C;
24
    cudaMalloc((void**)&d_A, M*K*sizeof(half));
25
    cudaMalloc((void**)&d_B, K*N*sizeof(half));
26
    cudaMalloc((void**)&d_C, M*N*sizeof(half));
27

28
    // Copy from host to device
29
    cudaMemcpy(d_A, h_A, M*K*sizeof(half), cudaMemcpyHostToDevice);
30
    cudaMemcpy(d_B, h_B, K*N*sizeof(half), cudaMemcpyHostToDevice);
31

32
    // Create cuBLAS handle
33
    cublasHandle_t handle;
34
    cublasCreate(&handle);
35

36
    // Set math mode to use Tensor Cores
37
    cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);
38

39
    // Perform matrix multiplication: C = alpha * A * B + beta * C
40
    float alpha = 1.0f;
41
    float beta  = 0.0f;
42
    cublasHgemm(handle,
43
                CUBLAS_OP_N, CUBLAS_OP_N,
44
                M, N, K,
45
                (const __half*)&alpha,
46
                (const __half*)d_A, M,
47
                (const __half*)d_B, K,
48
                (const __half*)&beta,
49
                d_C, M);
50

51
    // Copy result back to host
52
    cudaMemcpy(h_C, d_C, M*N*sizeof(half), cudaMemcpyDeviceToHost);
53

54
    // Cleanup
55
    cublasDestroy(handle);
56
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
57
    free(h_A); free(h_B); free(h_C);
58

59
    std::cout << "Matrix multiplication completed with Tensor Cores (HGEMM)." << std::endl;
60
    return 0;
61
}

Key Points in This Example:

• We used half-precision (__half) for our matrices.
• A call to cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH) indicates that we want to use Tensor Cores for this operation, if the GPU supports it and if the data types match.
• We used the cublasHgemm function, which expects half-precision inputs and outputs.

Depending on your hardware, you can tweak the math mode or switch to CUBLAS_TENSOR_OP_MATH_ALLOW_CONVERSION for additional flexibility in data type conversions.

Mixed-Precision Training and Computation#

One of the biggest drivers of Tensor Cores has been the explosive growth of deep learning. A technique called mixed-precision training involves using half-precision for most calculations while keeping select variables (like weight updates) in single-precision to maintain stability. NVIDIA’s Automatic Mixed Precision (AMP) feature in frameworks like PyTorch and TensorFlow helps automate this process.

Memory Layout#

Tensor Cores prefer data in specific layouts (for example, row-major vs. column-major, or custom layouts that match the 8×8 or 16×16 tile structures). Using libraries like cuBLAS or CUTLASS can abstract away some of these details. However, performance-sensitive kernels sometimes require reformatting data to align with Tensor Core tiling, maximizing occupancy and throughput.

Streamlining Kernel Launch#

For HPC codes that rely on many small kernels, kernel launch overhead can become a bottleneck. Employing large, batched GEMM calls or fusing multiple operations into a single kernel can help reduce overhead. Modern GPU architectures also support mechanisms like CUDA graphs, which can optimize repeated kernel launches to further reduce overhead.

Warp-Level Primitives#

For those looking to dive even deeper, NVIDIA provides warp-level primitives that allow direct invocation of Tensor Core operations. Programmers can use intrinsics like wmma.mma.sync to conduct matrix multiply-accumulate at warp-level granularity, controlling the formatting and precision of operands explicitly.

Custom Kernel Tuning with CUTLASS#

CUTLASS (CUDA Templates for Linear Algebra Subroutines) offers a highly configurable framework for composing GEMM kernels on NVIDIA GPUs. This library exposes a hierarchical tiling strategy that maps threads to sub-tiles of matrices, organizes data movement through shared memory, and ultimately orchestrates warp-level matrix multiply-accumulate. Advanced users can modify block sizes, tile sizes, and pipeline depths to squeeze out higher performance or support specialized workflows.

Emulating Higher Precision#

Although Tensor Cores handle half-precision inputs natively, some HPC applications demand the accuracy of at least single precision, if not double precision. One technique is to split computations across multiple half-precision operations, effectively emulating a higher precision. This approach can offer improvements over purely single-precision implementations, though it requires careful attention to rounding and accumulation errors.

Integration with MPI and Multi-GPU Systems#

Large-scale HPC clusters often harness multiple GPUs per node, with nodes connected over high-speed interconnects (e.g., InfiniBand). When using libraries like MPI (Message Passing Interface) for distributed-memory parallelism, scaling matrix multiplication across multiple GPUs can be a huge challenge. Libraries like NCCL (NVIDIA Collective Communications Library) handle multi-GPU collectives efficiently, and HPC frameworks can orchestrate distributed matrix multiplication with multi-GPU and Tensor-Core-accelerated kernels.

Multi-Node SGEMM and HPC Benchmarks#

Outside of deep learning, HPC teams may rely heavily on the HPCG (High Performance Conjugate Gradients) or HPL (High Performance LINPACK) benchmarks. Modern supercomputers often incorporate GPUs with Tensor Cores to push the boundaries of performance in these benchmarks. Achieving peak performance in HPL or HPCG involves careful data partitioning, overlap of communication and computation, and advanced multi-precision strategies.