Cache Blocking of Distributed-Memory Parallel Matrix Power Kernels

TL;DR

The paper addresses the memory-bandwidth bottleneck of repeated SpMV operations by introducing a distributed cache-blocked Matrix Power Kernel that preserves cross-process data dependencies. It extends the level-based RACE approach to MPI, organizing local vertices by distance from the halo into a bulk structure and interleaving cache-blocked computation with halo communication in a three-phase workflow. The Distributed Level-Blocked MPK (DLB-MPK) achieves substantial speedups over traditional distributed MPK (TRAD), with node-wide gains around 1.6–2.7x on three modern CPUs and up to 4x in weak scaling for Chebyshev time propagation in an Anderson model, while maintaining compatibility with existing libraries. Overall, the method provides scalable, hardware-aware cache blocking for MPK on distributed systems, enabling faster sparse linear algebra kernels in quantum physics and related domains.

Abstract

Sparse matrix-vector products (SpMVs) are a bottleneck in many scientific codes. Due to the heavy strain on the main memory interface from loading the sparse matrix and the possibly irregular memory access pattern, SpMV typically exhibits low arithmetic intensity. Repeating these products multiple times with the same matrix is required in many algorithms. This so-called matrix power kernel (MPK) provides an opportunity for data reuse since the same matrix data is loaded from main memory multiple times, an opportunity that has only recently been exploited successfully with the Recursive Algebraic Coloring Engine (RACE). Using RACE, one considers a graph based formulation of the SpMV and employs s level-based implementation of SpMV for reuse of relevant matrix data. However, the underlying data dependencies have restricted the use of this concept to shared memory parallelization and thus to single compute nodes. Enabling cache blocking for distributed-memory parallelization of MPK is challenging due to the need for explicit communication and synchronization of data in neighboring levels. In this work, we propose and implement a flexible method that interleaves the cache-blocking capabilities of RACE with an MPI communication scheme that fulfills all data dependencies among processes. Compared to a "traditional" distributed memory parallel MPK, our new Distributed Level-Blocked MPK yields substantial speed-ups on modern Intel and AMD architectures across a wide range of sparse matrices from various scientific applications. Finally, we address a modern quantum physics problem to demonstrate the applicability of our method, achieving a speed-up of up to 4x on 832 cores of an Intel Sapphire Rapids cluster.

Figures (4)

• Figure 1: Graph (a) and sparsity pattern (b) of the matrix associated with a modified 5-point stencil. Graph (c) shows the permuted graph and (d) the sparsity pattern of the matrix after applying Breadth First Search (BFS) reordering. The vertices (rows) of the graph (matrix) that belong to a level are represented with the same color.
• Figure 2: $Lp$ diagram with 10 levels ($L(0),\ldots,L(9)$) and a maximum power of $p_m=5$. Level colors are the same as in Figure \ref{['fig:stencil_graph_permuted']}. Each node in the $Lp$ diagram is numbered according to the execution order. For $p=4$ and level $L(6)$, the explicit dependencies to levels at $p=3$ are indicated with red arrows. The nodes highlighted in orange fulfill $i+p=6$ ("diagonal").
• Figure 3: The global matrix $A$ from Figure \ref{['fig:stencil_matrix']} and some RHS vector $x$ are partitioned in a row-wise manner over two MPI processes in (a). The gray boxed-out regions show, on each MPI process, which elements are "local" (inside the gray region). The edge corresponding to the remote data dependency, i.e., the edge crossing the MPI boundary, is highlighted in blue in (b). The rows at global indices $8$ and $12$ are highlighted as examples of rows which contain remote data dependencies for the SpMV. In order to fulfill these data dependency, another MPI process must supply the appropriate "halo elements." Shown in (c) is the process of data exchange on the x-vector for our two example rows, where incoming halo elements are received into an appropriately resized buffer.
• Figure :