Space Filling Curves is All You Need: Communication-Avoiding Matrix Multiplication Made Simple

TL;DR

This paper addresses the data-movement bottleneck and tuning burden in GEMM on modern CPUs by introducing a Space Filling Curve based, communication-avoiding GEMM (SFC-CA GEMM). The core idea is to partition GEMM computation with generalized Hilbert curves to preserve locality, and to replicate inputs to minimize on-critical-path data movement across memory hierarchies. The approach is implemented as a compact (~30 LOC) framework built on BRGEMM tensor primitives, with a performance model and extensive tests across multiple CPU platforms, achieving up to about 2x speedups over vendor libraries and near-roofline performance for a wide range of GEMM shapes. The work also provides a rigorous analysis of communication costs and demonstrates broad portability, with potential extensions to GPUs and higher-dimensional contractions via tensor compilers.

Abstract

General Matrix Multiplication (GEMM) is the cornerstone of Deep Learning and HPC workloads; accordingly, academia and industry have heavily optimized this kernel. Modern platforms with matrix multiplication accelerators exhibit high FLOP/Byte machine balance, which makes implementing optimal matrix multiplication challenging. On modern CPU platforms with matrix engines, state-of-the-art vendor libraries tune input tensor layouts, parallelization schemes, and cache blocking to minimize data movement across the memory hierarchy and maximize throughput. However, the best settings for these parameters depend strongly on the target platform (number of cores, memory hierarchy, cache sizes) and on the shapes of the matrices, making exhaustive tuning infeasible; in practice this leads to performance "glass jaws". In this work we revisit space filling curves (SFC) to alleviate the problem of this cumbersome tuning. SFC convert multi-dimensional coordinates (e.g. 2D) into a single dimension (1D), keeping nearby points in the high-dimensional space close in the 1D order. We partition the Matrix Multiplication computation space using recent advancements in generalized SFC (Generalized Hilbert Curves), and we obtain platform-oblivious and shape-oblivious matrix-multiplication schemes that exhibit inherently high degree of data locality. Furthermore, we extend the SFC-based work partitioning to implement Communication-Avoiding (CA) algorithms that replicate the input tensors and provably minimize communication/data-movement on the critical path. The integration of CA-algorithms is seamless and yields compact code (~30 LOC), yet it achieves state-of-the-art results on multiple CPU platforms, outperforming vendor libraries by up to 2x(geometric-mean speedup) for a range of GEMM shapes.

Figures (10)

• Figure 1: Multi-core GEMM performance (Bfloat16) of a vendor-optimized library (oneDNN) and this work (SFC-CA) on a platform with Matrix Multiplication Accelerator (64-core Intel Xeon Emerald Rapids). The x-axis shows the $M\times N\times K$ GEMM configurations. OneDNN illustrates performance "glass jaws" and is far-off the roofline.
• Figure 2: SFC-based traversal of 16$\times$16 grid. The SFC used is the generalized 2D Hilbert curve. The numbers in the grid boxes indicate the order (i.e. timestamp) of the traversal, also color-coded according to the heatmap bar on the right.
• Figure 3: SFC-based partitioning of $C$ matrix using 64 cores. The $C$ dimensions are 4096$\times$4096, and by using blocks of size 32$\times$32 we obtain a grid of 128$\times$128 $C$ blocks. Left: 2D $C$ decomposition, the SFC yields a 2D core decomposition (grid of 8$\times$8 cores). Each core is implicitly assigned a $C$ tile consisting of 16$\times$16 $C$ blocks (each having size 32$\times$32). Each core-local grid of 16$\times$16 $C$ blocks/$C$-tasks is traversed by the core using the SFC ordering illustrated in Figure \ref{['fig:hilbert']}. Middle: 2.5D decomposition by replicating $C$ by a factor of 2 (i.e. 2 layers of $C$), within each layer the SFC yields a 2D core decomposition with rectangular $C$ tiles. Within each layer of $C$ the work is partitioned among 32 cores. Right: 3D decomposition of $C$ by replicating it by a factor of 4. Within each layer the SFC yields a 2D core decomposition (4$\times$4 core decomposition within each plane), for a logical 4$\times$4$\times$4 3D grid of cores.
• Figure 4: SFC-based Communication-Avoiding GEMM
• Figure 5: SFC-based partitioning of $C$ matrices with different $M$:$N$ aspect ratios using 64 cores. The SFC-based partitioning yields 2D core decompositions with aspect ratios matching the corresponding $M$:$N$ ratio.