FusedMM: Unified Sparse Kernel Fusion

• FusedMM is a unified operation that fuses SDDMM and SpMM to eliminate intermediate storage and reduce memory traffic.
• It enables efficient graph embedding and GNN message passing through customizable, five-stage kernel decomposition and communication-eliding strategies.
• FusedMM achieves significant speedups by leveraging architecture-aware optimizations, including SIMD vectorization on CPUs, replication reuse in distributed setups, and explicit vectorization on GPUs.

FusedMM denotes the fusion of sampled dense-dense matrix multiplication and sparse-dense matrix multiplication into a single operation in which the SDDMM output is consumed by SpMM without materializing the intermediate sparse result. In the literature, the term appears both as a unified kernel for graph embedding and graph neural networks and as a broader distributed-memory and GPU execution pattern for back-to-back sparse ML kernels. Across these settings, the central objective is consistent: eliminate intermediate storage, reduce memory traffic or communication, and preserve the algebraic structure of message passing and sparse linear algebra workloads (Rahman et al., 2020, Bharadwaj et al., 2022, Zubair et al., 2023).

1. Algebraic definition and core formulations

The basic ingredients are SDDMM, SpMM, and their composition. In one formulation, given dense matrices $A \in \mathbb{R}^{m \times r}$, $B \in \mathbb{R}^{n \times r}$, and a sparse matrix $S \in \mathbb{R}^{m \times n}$, SDDMM is defined as

$\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$

where $*$ denotes element-wise multiplication and only the nonzeros of $S$ are updated. The corresponding sparse-dense products are

$$\textrm{SpMMA}(S, B) := S \cdot B, \qquad \textrm{SpMMB}(S, A) := S^T \cdot A.$$

FusedMM is then the back-to-back composition

$$\textrm{FusedMMA}(S, A, B) := \textrm{SpMMA}(\textrm{SDDMM}(A, B, S), B),$$

$$\textrm{FusedMMB}(S, A, B) := \textrm{SpMMB}(\textrm{SDDMM}(A, B, S), A).$$

For the $i$-th output row, one explicit form is

$B \in \mathbb{R}^{n \times r}$0

This formulation makes the dependence on the sampled dot products and subsequent aggregation explicit (Bharadwaj et al., 2022).

A closely related GPU-oriented expression writes SDDMM, SpMM, and FusedMM as

$B \in \mathbb{R}^{n \times r}$1

$B \in \mathbb{R}^{n \times r}$2

Here $B \in \mathbb{R}^{n \times r}$3 is element-wise multiplication and $B \in \mathbb{R}^{n \times r}$4 encodes the sparsity pattern. This presentation emphasizes that FusedMM is the composition of SDDMM with SpMM applied sequentially without materializing intermediate results (Zubair et al., 2023).

In graph learning, the same idea appears as a fused message-passing update. Given sparse adjacency $B \in \mathbb{R}^{n \times r}$5, node features $B \in \mathbb{R}^{n \times r}$6 and $B \in \mathbb{R}^{n \times r}$7, and output $B \in \mathbb{R}^{n \times r}$8, the update for node $B \in \mathbb{R}^{n \times r}$9 is

$S \in \mathbb{R}^{m \times n}$0

This is the graph-centric interpretation of the same algebraic pattern: generate edge-wise messages and aggregate them immediately, in one pass, rather than as two separately materialized kernels (Rahman et al., 2020).

2. Role in graph embedding and graph neural networks

FusedMM was introduced as a unified SDDMM-SpMM kernel for graph embedding and graph neural networks, motivated by the message passing paradigm. In that paradigm, message generation computes a message for each edge, for example through $S \in \mathbb{R}^{m \times n}$1, and message aggregation combines incoming messages at each vertex, for example through $S \in \mathbb{R}^{m \times n}$2. Frameworks such as DGL and PyTorch Geometric implement these phases using distinct sparse matrix operations, but the SDDMM output is nearly always followed by SpMM; the separation introduces extra memory for storing intermediate message tensors and additional memory reads and writes, especially for high-dimensional messages (Rahman et al., 2020).

To generalize beyond a fixed kernel template, FusedMM decomposes message passing into five parameterized stages supplied by user-defined functions.

Stage	Role	Representative operations
VOP	Elementwise operation on source/target features	elementwise multiplication, copying neighbor's feature
ROP	Reduction of a vector to a scalar	sum, dot product
SOP	Scalar transformation	sigmoid
MOP	Combination of message scalar/vector with target feature	scaling $S \in \mathbb{R}^{m \times n}$3
AOP	Aggregation across neighbors	sum, max

This decomposition is used to capture almost all computational patterns needed by popular graph embedding and GNN approaches. The cited examples include Graph Layout (FR), Node Embedding (Verse/Force2Vec), GCN, and GNN with MLP. The significance of the five-stage formulation is that it treats fusion not as a single hard-coded operator, but as a customizable kernel family in which the algebraic structure of message passing remains explicit (Rahman et al., 2020).

A plausible implication is that FusedMM serves as both a semantic abstraction and an implementation strategy: semantic, because it unifies message generation and aggregation under a single operator schema; implementation-oriented, because the same schema exposes opportunities for eliminating intermediate storage and specializing low-level execution.

3. Shared-memory kernel design and CPU optimization

The original FusedMM kernel work emphasizes low-level vectorized kernels, a suitable load balancing scheme, and an efficient utilization of the memory bandwidth. The core routine operates on rows of the adjacency matrix, processing all neighbors of a vertex in a tight loop. SIMD vectorization is used extensively: node features and accumulators are loaded into SIMD registers, reused across neighbor computations, and written back once after all neighbor processing is complete. Register blocking keeps repeatedly used vectors in registers and minimizes memory traffic, while non-temporal writes can be used for result storage (Rahman et al., 2020).

Load balancing is handled through 1D matrix partitioning. The adjacency matrix is partitioned across threads such that each thread processes a submatrix with roughly equal numbers of nonzeros. Because aggregation is per-row, this organization avoids synchronization and intermediate partial results. The overall implementation is explicitly architecture-aware: an automatic code generator produces specialized SIMD intrinsic implementations for AVX512 or AVX on x86 and ASIMD or NEON on ARM, and the register blocking factor and kernel structure are auto-tuned for each processor (Rahman et al., 2020).

The reported performance gains are substantial. FusedMM is described as an order of magnitude faster than its equivalent kernels in Deep Graph Library, and the optimized variant is up to 34× faster than DGL on Intel CPUs, with typical speedups of 5–20× across tasks and datasets. On ARM and AMD CPUs, speedups range from 4–19× over DGL. When specialized to SpMM mode, FusedMM matches or exceeds the performance of vendor-optimized MKL SpMM. Memory behavior is also central: for tasks requiring high-dimensional messages, DGL runs out of memory storing edge-level tensors, whereas FusedMM does not. Roofline analysis on Intel Skylake reports up to 66% of theoretical roofline performance, and end-to-end graph embedding can be sped up by up to 28× on different processors (Rahman et al., 2020).

These results situate FusedMM within the broader class of bandwidth-bound sparse ML kernels. The published CPU design treats fusion as a means of reducing data movement first and arithmetic optimization second, which is consistent with the performance profile of message-passing workloads.

4. Distributed-memory formulations and communication-eliding strategies

Distributed-memory FusedMM extends the same back-to-back SDDMM-plus-SpMM pattern to multi-node settings. A central theoretical result is that any distributed-memory, sparsity-agnostic algorithm for SpMM, including classic 1.5D and 2.5D variants, can be converted to perform SDDMM with identical communication costs and input/output data layouts. This equivalence allows the same distributed data distribution to be used for SDDMM, SpMMA, and SpMMB, paying only the cost of possibly replicating $S \in \mathbb{R}^{m \times n}$4 to store its transpose (Bharadwaj et al., 2022).

The distributed work introduces two communication-eliding strategies for FusedMM. The first is replication reuse: when a dense matrix is already replicated for SDDMM, the replication is reused for the subsequent SpMM. The second is local kernel fusion: the local SDDMM and SpMM kernels are combined into a single phase before communication, reducing intermediate storage and avoiding redundant communication. For the 1.5D dense shifting algorithm, the communication costs are given as

$S \in \mathbb{R}^{m \times n}$5

for sequential SDDMM then SpMMA without elision,

$S \in \mathbb{R}^{m \times n}$6

with replication reuse, and

$S \in \mathbb{R}^{m \times n}$7

with local kernel fusion. The respective optimal replication factors are $S \in \mathbb{R}^{m \times n}$8, $S \in \mathbb{R}^{m \times n}$9, and $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$0, and both eliding strategies yield up to $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$1 communication reduction (Bharadwaj et al., 2022).

The analysis is parameterized by

$\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$2

the normalized sparsity. For low $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$3, 1.5D sparse shifting with replication reuse is best, with communication cost

$\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$4

and optimal $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$5. As $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$6 increases, 1.5D dense shifting with elision or 2.5D algorithms become preferable. The 2.5D dense replicating algorithm with replication reuse has communication cost

$\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$7

with optimal $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$8 (Bharadwaj et al., 2022).

Empirically, on 256 nodes with 68 cores each of a Cray XC40 at LBNL, the communication-eliding 1.5D FusedMM algorithms save at least 30% of communication time relative to sequential distributed-memory SpMM and SDDMM, and communication-eliding 2.5D algorithms save 21%. On real-world matrices with hundreds of millions of edges, all variants achieve at least a 10× speedup over the SpMM algorithm in PETSc, and the communication-eliding techniques are up to 1.6 times faster than an unoptimized sequence of SDDMM and SpMM. The algorithms are embedded in collaborative filtering via alternating-least-squares and inference for attention-based graph neural networks, with strong and weak scaling up to 256 nodes and matrices with over a billion nonzeros (Bharadwaj et al., 2022).

5. Intel Max Series GPU realization with oneAPI ESIMD

A later line of work studies FusedMM together with SDDMM and SPMM on Intel Max Series GPUs using Intel oneAPI's Explicit SIMD SYCL extension. The premise is that ESIMD allows explicit vectorization by the programmer rather than relying on compiler vectorization. The sparse matrices are stored in CSR format using ia, ja, and avalue, and dense matrices use a row-major layout. Parallelism is mapped to rows of the output matrix, and for rows with many nonzeros those nonzeros are chunked so that each thread within an ESIMD work-group processes a vector of nonzeros (Zubair et al., 2023).

For each output row $\textrm{SDDMM}(A, B, S) := S * (A \cdot B^T),$9, the fused algorithm first computes the SDDMM dot products $*$0 at the columns indicated by the sparsity pattern, then immediately performs the SPMM phase by taking a weighted sum over selected rows of $*$1. The ESIMD implementation fuses these two passes within a single kernel for maximum performance. Each ESIMD work-group is assigned a row; threads within the work-group process chunks; simd types and block-load or block-store intrinsics are used to optimize L1, L3, and HBM usage; chunk size $*$2 is adjustable; and kernel launches are chosen so that the number of ESIMD work-items aligns with hardware, including multiples of 4096 for a single Xe stack (Zubair et al., 2023).

The key optimization strategies are explicit vectorization, loop unrolling, prefetching, occupancy tuning, and in-place fusion. The paper argues that ESIMD gives direct control over register usage, memory access patterns, and better management of thread divergence compared to CUDA or default SYCL, making it well-suited for irregular, memory-bound workloads such as sparsity. Operator fusion is treated as crucial because it avoids extra memory transfers by fusing SDDMM and SPMM phases and eliminating global memory writes in intermediate steps (Zubair et al., 2023).

The resulting performance is reported as close to the peak of the targeted Intel Data Center GPU. Against Intel oneMKL, SDDMM is on average 1.84× faster than dense GEMM and up to 3.4× for some configurations; SPMM is on average 1.37× faster than MKL’s sparse GEMM and up to 2.16× faster; and FusedMM is on average 2× as fast as the sum of MKL’s sparse and dense kernels, with best cases up to 3.25×. Peak FusedMM reaches 12,647 Gflops/s, corresponding to 48% of single-precision peak on Intel Max GPU. The roofline analysis reports arithmetic intensity up to 5.6 for large problems whose data fit in L3 after the first iteration, and SPMM achieves between 37–77% of the device’s theoretical peak for different arithmetic intensities. Against a recent CUDA implementation on NVIDIA V100, SDDMM is reported as up to 10× faster, and device utilization reaches up to 49% of theoretical single-precision peak versus approximately 27% on V100 (Zubair et al., 2023).

6. Performance interpretation, scope, and recurrent points of confusion

Across the cited works, FusedMM is not merely a shorthand for invoking SDDMM and SpMM in sequence. The published formulations consistently tie fusion to the removal of an intermediate sparse tensor or message tensor, and the performance gains depend on architecture-specific mechanisms: SIMD register blocking and 1D partitioning on CPUs, replication reuse and local kernel fusion in distributed memory, and explicit vectorization, prefetching, and occupancy tuning on Intel GPUs (Rahman et al., 2020, Bharadwaj et al., 2022, Zubair et al., 2023).

A second recurrent point is scope. In graph learning, FusedMM is a unified kernel that can capture almost all computational patterns needed by popular graph embedding and GNN approaches through user-defined functions. In distributed-memory sparse linear algebra, it is a communication-aware composition of SDDMM and SpMM that can be analyzed through 1.5D and 2.5D communication models. In GPU optimization, it is a fused sparse kernel whose effectiveness depends on avoiding intermediate global memory traffic and on matching the execution strategy to the targeted architecture. These are distinct levels of abstraction, but they are algebraically compatible because all are built on the same back-to-back SDDMM-plus-SpMM structure (Rahman et al., 2020, Bharadwaj et al., 2022, Zubair et al., 2023).

The aggregate evidence suggests a stable interpretation of FusedMM as a performance pattern for sparse machine learning. On CPUs it delivers up to 34× kernel-level speedups over DGL and up to 28× end-to-end speedups for graph embedding; in distributed memory it reduces communication time by at least 30% for 1.5D algorithms and attains at least 10× speedup over PETSc on large real-world matrices; and on Intel Max GPUs it reaches 12,647 Gflops/s and outperforms both oneMKL baselines and a recent CUDA implementation under the reported conditions (Rahman et al., 2020, Bharadwaj et al., 2022, Zubair et al., 2023).

Taken together, these results indicate that FusedMM is best understood not as a single implementation artifact, but as a unifying sparse-kernel concept linking graph message passing, distributed communication optimization, and hardware-aware kernel design.