Batch-Reduce GEMM: Methods and Applications

• Batch-reduce GEMM is a matrix multiplication primitive that fuses multiple GEMM operations with reduction to produce a consolidated output.
• It exploits uniform memory layouts and vectorized microkernels to efficiently tackle deep learning, tensor contractions, and scientific computing challenges.
• Auto-tuning and architecture-specific optimizations enable BRGEMM implementations to achieve up to 95% of peak hardware throughput on modern CPUs and GPUs.

Batch-reduce GEMM (BRGEMM) refers to a class of matrix multiplication primitives that generalize standard batched GEMM to support efficient summation (reduction) over a sequence of matrix pairs, with extensive application in deep learning, scientific computing, and tensor contractions. These kernels efficiently fuse the accumulation of multiple independent GEMM operations with reduction into a shared result buffer, exploiting modern hardware vectorization, memory hierarchy, and software auto-tuning to deliver both high throughput and broad applicability.

1. Mathematical Formulation and Core Kernel Interface

The canonical batch-reduce GEMM computes: $$C \leftarrow \beta C + \alpha \sum_{i=0}^{N-1} A_i B_i,$$ where each $A_i \in \mathbb{R}^{m \times k}$, $B_i \in \mathbb{R}^{k \times n}$, and $C \in \mathbb{R}^{m \times n}$. Here, $N$ is the batch count; $A$ and $B$ may be provided as explicit pointer arrays or as base pointers with uniform strides. In the special case where the reduction is omitted (i.e., each $C_i$ is independent), BRGEMM degrades to standard strided-batched GEMM.

A typical BRGEMM API admits the signature:

batchreduce_gemm(m, n, k, α, A₀, str_A, B₀, str_B, β, C₀, str_C, N);

where $A_i = A_0 + i \cdot str_A$, $B_i = B_0 + i \cdot str_B$, and $C_0$ is updated in-place (usually $str_C = 0$ unless a batch of $C_i$ is used). The core distinction is that all the $A_i B_i$ products are reduced into the same $C$, enabling maximal accumulator reuse and register blocking (Georganas et al., 2019).

Strided-batched GEMM, as standardized in cuBLAS and oneMKL, generalizes to: $$C^{(p)} \leftarrow \alpha\,\operatorname{op}(A^{(p)})\,\operatorname{op}(B^{(p)}) + \beta\,C^{(p)}, \;\;\; p=0,\ldots,P-1,$$ with each operand at fixed strides in memory (Shi et al., 2016, Jhurani et al., 2013).

2. Memory Layout, Uniform-Striding, and Kernel Organization

Efficient BRGEMM critically depends on uniform memory layout. All $A_i$, $B_i$ (and potentially $C_i$) must be placed in contiguous buffers at compile-time or known, fixed strides. This enables implementations to map the batch dimension directly to vector lanes or thread blocks. Three data organizations recur across systems:

• Strided: All batch matrices are laid out consecutively, enabling $A_i = A_0 + i \cdot stride_A$ pointer arithmetic (Boukaram et al., 2021, Shi et al., 2016).
• Indexed: An array of pointers (less optimal on GPUs).
• Constant: A or B reused for the entire batch (for example, outer products with a shared factor).

For example, in the H2OPUS-TLR factorization, adaptive-rank tiles are first sorted and bucketed into groups of matching dimensions. Within each bucket, $m$, $n$, $k$ are fixed, and all tiles are packed contiguously into 1-D arrays with explicit stride formulas:

• $A_{buffer}$ size = $B \cdot (m \cdot k)$, with stride $m \cdot k$,
• $B_{buffer}$ size = $B \cdot (k \cdot n)$, with stride $k \cdot n$,
• $C_{buffer}$ size = $B \cdot (m \cdot n)$, with stride $m \cdot n$.

The resulting buffers enable the launch of a single uniform-stride batched GEMM kernel invocation (Boukaram et al., 2021).

Batched GEMM implementations such as TGEMM_multi_uniform (NVIDIA GPUs) or cublasSgemmStridedBatched (cuBLAS) follow this scheme, with the pointer offsets and memory layout fully described by structure-of-arrays concepts (Jhurani et al., 2013, Shi et al., 2016).

3. Kernel Implementation, Vectorization, and Microarchitecture Awareness

BRGEMM leverages highly-tuned microkernels, exploiting wide vector units by vectorizing across the batch. On architectures like RISC-V RVV or x86 AVX-512, the principal vectorization strategy is to unroll the entire $m \times n \times k$ computation, exposing the maximum batch dimension to the auto-vectorizer (Banchelli et al., 10 Jan 2025). For instance:

• Inner $M, N, K$ loops are fully unrolled; the only remaining dynamic loop is the batch index $e \in [0, E)$.
• For constant $A$ (across batch), all its elements are in scalar registers, while $B_i$, $C_i$ are accessed with vector-gather (for indexed), strided loads (for strided), or broadcast operations.
• The compiler emits vector FMA instructions such as RVV's vfmacc.vf or AVX-512's _mm512_fmadd_pd, enabling a single FMA to update scores of output tiles in parallel (Banchelli et al., 10 Jan 2025).

GPU kernels achieve similar effects: per-tile thread blocks loop over the batch dimension, staging $A$ and $B$ panels in shared memory and unrolling the $k$ loop. On Volta/Turing GPUs, tensor core WMMA APIs are directly employed for $k$ an integer multiple of 8/16 (Boukaram et al., 2021).

A typical CUDA GEMM call for uniform-stride batches:

cublasDgemmStridedBatched( handle, transA, transB, m, n, k, &alpha, A_base, lda, stride_A, B_base, ldb, stride_B, &beta, C_base, ldc, stride_C, batchCount );

with all matrices residing in pitch-aligned slabs (Shi et al., 2016, Boukaram et al., 2021).

4. Performance Characteristics and Hardware Efficiency

BRGEMM, when implemented with optimal kernel organization and memory layout, can sustain a large fraction of peak hardware throughput:

• On NVIDIA V100, H2OPUS-TLR achieves 3–4 TFLOP/s in double precision for typical 3D covariance tile updates (tile size $m=512$, $k \approx 32$, $n=512$, $batchCount \approx 512$), with the overall Cholesky at $\sim$1.2 TFLOP/s and 80–90% of the runtime in strided-batched GEMMs (Boukaram et al., 2021).
• On AVX-512, BRGEMM microkernels reach 80–95% of peak FMA rates for both deep learning and scientific codes, outperforming pointer-array-based batched GEMM by up to 30–60% (Georganas et al., 2019, Jhurani et al., 2013).
• On RISC-V RVV and Intel AVX-512, explicit vectorization across the batch yields speedups of 3.5×–32.6× (FPGA prototype) and up to 16.2× on Sapphire Rapids for extremely small ($2 \times 2 \times 2$) DGEMMs versus cblas_dgemm loops—the precise gain depends on matrix and batch sizes, memory hierarchy, and register spilling (Banchelli et al., 10 Jan 2025).

Compared to explicit data transformations plus monolithic GEMM, or pointer-array batched GEMM, BRGEMM/strided-batched GEMM eliminates costly pointer-fetch, transpose, and launch overhead (15–20%), improving effective throughput in batched workloads (Boukaram et al., 2021, Shi et al., 2016).

The following table summarizes representative speedups:

Matrix Size	Platform	Batch-Reduce GEMM Speedup over Baseline
$2\times2\times2$	FPGA RVV0.7	32.6× (vs. cblas_dgemm loop)
$4\times4\times4$	AVX-512	9.1×
$16\times16$	NVIDIA K20c	1.61× (vs. cuBLAS batched)

5. Applications in Deep Learning, Tensor Contraction, and Scientific Computing

BRGEMM acts as a unifying primitive for a range of workloads:

• Deep Learning: All major primitives—LSTMs, CNNs, and MLPs—can be expressed as series of BRGEMM calls. Time-step and gate multiplication in RNNs, convolution sums in CNNs, and blocked multiplications in MLPs reduce to BRGEMMs with suitable blocking and tiling. This enables replacing the traditional “combinatorial explosion” of hand-tuned kernels with a single, robust, and auto-tunable primitive (Georganas et al., 2019).
• Tensor Contractions: Higher-order tensor contractions, such as in Tucker/HOOI decompositions, are mapped to sequences of strided-batched GEMMs. By eliminating explicit transposes/data movement and expressing batch iteration in the batch-reduce loop, up to 10× speedup on K40c and 2–3× on multicore CPUs (vs. classic BLAS-based approaches) is observed (Shi et al., 2016).
• Finite Element and DG Methods: In codes like SeisSol, which make heavy use of small fixed-size DGEMMs, batching thousands of small GEMMs and vectorizing across the batch exposes much higher instruction-level parallelism—even on architectures otherwise unable to saturate wide vector units—realizing substantial speedups and full code portability (Banchelli et al., 10 Jan 2025).

6. Auto-tuning, Microkernel Generation, and Architectural Recommendations

Exposing optimal performance with batch-reduce GEMM requires judicious kernel and memory parameterization:

• Tile/block sizes ($m_b$, $n_b$) must fit register files and caches. For GPUs, $m=512$–$1024$ on V100 maximizes L2/SMEM reuse (Boukaram et al., 2021).
• Adaptive-rank batching: Sorting and bucketing tiles by rank (dimension) ensures all matrices in a batch are handled by a single uniform kernel, avoiding costly kernel branching (Boukaram et al., 2021).
• Precision and tensor core utilization: Setting $k$ (or batch) as a multiple of 8/16 ensures access to mixed-precision WMMA/Tensor Core units, with FP16 packing double the memory bandwidth (Boukaram et al., 2021).
• Auto-generated microkernels: JIT-generated assembly/reg-kernels for specific $m_b, n_b$ (e.g. LIBXSMM) can deliver up to 95% of theoretical throughput, with register blocking and prefetch strategies auto-tuned for architecture (Georganas et al., 2019).
• Full code portability: With well-defined uniform-stride strategies, a single source can vectorize to both RVV and AVX-512 out of the box, requiring only batch size and stride parameters to be compile-time known (Banchelli et al., 10 Jan 2025).

7. Extensions, Limitations, and Integration Strategies

BRGEMM’s main limitation is the assumption of uniform strides and dense matrix tiles. Exceptional cases with irregular sparsity or scatter/gather–heavy access require either explicit pointer arrays or further kernel specialization. Extended-operand modes (e.g., swapping stride direction or batches on inner dimension) are needed for certain tensor contractions, with a 2–3× runtime penalty versus plain SB-GEMM but still outperforming explicit copy+GEMM for small sizes (Shi et al., 2016).

For highly irregular adaptive workloads (e.g. variable-rank tiles in TLR Cholesky), the state-of-the-art is to dynamically regroup tiles into uniform-stride “buckets” as the computation evolves (Boukaram et al., 2021).

Integration into software stacks is accomplished by refactoring explicit tile- and contraction-loop nests to call BRGEMM or SB-GEMM at the innermost batched level, flattening indices wherever possible to maximize uniformity and enable hardware/BLAS autotuning (Shi et al., 2016, Georganas et al., 2019).

A plausible implication is that as more modern hardware exposes longer and more flexible vector units, and as deep learning frameworks evolve towards unified tensor IR, BRGEMM primitives will further abstract away layer/operation specialization, enabling consistent, tunable performance across CPUs, GPUs, and emerging vector architectures.