Uniform-Stride Batched GEMM
- Uniform-stride batched GEMM is a method that computes many small matrix multiplications using a contiguous memory layout and fixed strides.
- It reduces overhead by replacing pointer-to-pointer arrays with a single base pointer and stride, optimizing kernel launches and memory access.
- This approach is particularly effective for scientific simulations, deep learning primitives, and high-throughput tensor contractions with small matrix dimensions.
A uniform-stride batched GEMM (General Matrix-Matrix Multiplication) computes a batch of independent small matrix multiplications using a memory layout and interface that exploits a regular (constant) stride between matrices. This approach delivers superior performance, reduces overhead compared to pointer-to-pointer batching, and is particularly effective for workloads with many small GEMMs (dimensions typically under 16 or 32), as encountered in modern scientific simulation, deep learning primitives, and high-throughput tensor contractions on both NVIDIA GPUs and vector CPUs (Jhurani et al., 2013, Shi et al., 2016, Banchelli et al., 10 Jan 2025).
1. Interface Specification and Semantics
Uniform-stride batched GEMM exposes an API where each batch of the matrix operands (A, B, C) is arranged in contiguous flat buffers with fixed strides. Instead of arrays of pointers to matrix slices, a single base pointer and stride for each operand suffices.
For example, on a CUDA-enabled GPU (cuBLAS), the prototypical interface is:
1 2 3 4 5 6 7 8 9 10 |
cublasStatus_t cublasSgemmStridedBatched(
cublasHandle_t handle,
cublasOperation_t transA, cublasOperation_t transB,
int m, int n, int k,
const float *alpha,
const float *A0, int lda, long long strideA,
const float *B0, int ldb, long long strideB,
const float *beta,
float *C0, int ldc, long long strideC,
int batchCount ); |
strideA, strideB, and strideC denote the element offsets between consecutive A, B, C matrices in memory (Shi et al., 2016). TGEMM_multi_uniform, as realized by Jhurani & Mullowney (Jhurani et al., 2013), utilizes a similar interface where, for each matrix , the base pointer and stride manage selection: , etc. This enables the kernel to process hundreds of thousands of matrices in a single, tightly-packed launch, reducing kernel-launch overhead and improving coalesced memory access.
The semantics are:
$$
Cp \leftarrow \alpha \ \mathrm{op}(A