High Performance Single Precision Matrix Multiplication on NVIDIA T4 GPU with CUDA

This repository contains my attempt to get write a high performance single precision matrix multiplication (SGEMM) kernel on NVIDIA T4 GPU using CUDA. The NVIDIA T4 GPU belongs to the Turing architecture so it does have access to recent features like async general memory loads for examples. Moreover, I choose to work with single precision, so I will not make use of tensor cores in this repository, as T4 only offers float16 tensorcores. Therefore, this repository is a study of "classical" (pre-Ampere) techniques for optimizing matrix multiplication on NVIDIA GPUs.

Results

I achieved 82% and 96% of cublas's throughput on $N=M=K=2048$ and $N=M=K=4096$ respectively.

Here are detailed results for $N=M=K=2048$, the case I optimized for:

Kernel	ms (med)	GFLOP/s
naive	37.05	463.7
tiled_smem	18.28	939.6
tiled_smem_coarsened	9.05	1897.6
tiled_regblk	6.15	2791.7
tiled_regblk2	4.20	4090.9
tiled_warp	3.74	4591.7
tiled_warp_vec_gmem	3.00	5734.4
tiled_warp_vec_smem	2.90	5928.2
cublas	2.40	7163.1

Here's a comparison of the different kernels with cuBlas:

We notice that cuBlas is much more efficient for smaller matrices, as it implements a specialized kernel for small matrices, using the split-k technique. For larger matrices, the problem becomes compute bound, and our kernel is able to almost reach the performance of cuBlas. At this point, tensor cores would be needed to get a higher throughput.

Kernels

Baseline Kernel

This is the most basic implementation of matrix multiplication (kernels/naive.cu). Each thread computes one element of the output matrix C.

Shared Memory Tiling

The next step is to implement tiling using shared memory (kernels/tiled_smem.cu). Each block computes a TILE x TILE tile of the output matrix C. The input matrices A and B are loaded into shared memory in tiles, and each thread computes one element of the output tile by using the shared memory tiles of A and B.

Thread Coarsening

Now each thread is reponsible for a RTILE x RTILE tile of the output matrix C. This kernel is implemented in kernels/tiled_smem_coarsened.cu. There are two benefits to this:

• There's less threads launched, so less scheduling overhead and more amorization of the fixed launch cost of the kernel.
• There's more reuse of the data loaded from shared memory.

Note that here, we using TILE=64 and RTILE=2 and each thread manages 4 output elements that are TILE/RTILE=32-strided in both directions. Making reads and writes coalesced.

Register Tiling

Now, we make explicit the data reuse in registers. Each thread loads its RTILE x RTILE tile of A and B into registers, and then computes the output tile using these register tiles. This increases instruction level parallelism. This kernel is implemented in kernels/tiled_regblk.cu.

• kernels/tiled_regblk2.cu enables this more, know we can have TILE_M x TILE_K and TILE_K x TILE_N tiles of A and B in shared memory, and each thread can have RTILE_M x RTILE_K and RTILE_K x RTILE_N tiles of A and B in registers. This allows more flexibility and more choice when doing auto-tuning. This was fruitful as we can see this kernel was more efficient than the previous one.

Warp Tiling

We make each warp responsible for a WARP_M x WARP_N tile of the output matrix C. This enables us to control more the shared memory access patterns. Here for example, each warp will have RTILE_M x RTILE_N tiles of contiguous 8x4 blocks. This has as effect to increase shared memory broudcasting (threads using same shared memory value at the same time) and therefore increase effective bandwidth. However, we also increase the number of bank conflicts in shared memory, so it's a trade-off. This is confirmed by the profiler:

• Mem Pipes Busy went from 78.8% to 63.1%
• Notifies us about 10% bank conflict.

Here's an illustration of what I think is happening:

This kernel is implemented in kernels/tiled_warp.cu.

Vectorization

Profiling shows we are stalling on memory loads, so we implement vectorized loads from GMEM and SMEM in kernels kernels/tiled_warp_vec_gmem and kernels/tiled_warp_vec_smem.