# Distributed-Memory Sparse Kernels for Machine Learning

arXiv: https://arxiv.org/abs/2203.07673 · HTML: https://mlsanity.com/paper/2203.07673/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 6 of its 25 references that are in the map.
- First posted on arXiv 15 Mar 2022.

## Builds on (most cited first, 6 of 6)

- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [FusedMM: A Unified SDDMM-SpMM Kernel for Graph Embedding and Graph Neural Networks](https://mlsanity.com/paper/2011.06391/) (2020)
- [Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication](https://mlsanity.com/paper/1908.09606/) (2019)
- [An Efficient Shared-memory Parallel Sinkhorn-Knopp Algorithm to Compute the Word Mover's Distance](https://mlsanity.com/paper/2005.06727/) (2020)
- [Combinatorial BLAS 2.0: Scaling combinatorial algorithms on distributed-memory systems](https://mlsanity.com/paper/2106.14402/) (2021)

## Built on by (most cited first, 3 of 3)

- [PyGim: An Efficient Graph Neural Network Library for Real Processing-In-Memory Architectures](https://mlsanity.com/paper/2402.16731/) (2024)
- [Performance Optimization of Deep Learning Sparse Matrix Kernels on Intel Max Series GPU](https://mlsanity.com/paper/2311.00368/) (2023)
- [The Ubiquitous Sparse Matrix-Matrix Products](https://mlsanity.com/paper/2508.04077/) (2025)

## Abstract (arXiv)

Sampled Dense Times Dense Matrix Multiplication (SDDMM) and Sparse Times Dense Matrix Multiplication (SpMM) appear in diverse settings, such as collaborative filtering, document clustering, and graph embedding. Frequently, the SDDMM output becomes the input sparse matrix for a subsequent SpMM operation. Existing work has focused on shared memory parallelization of these primitives. While there has been extensive analysis of communication-minimizing distributed 1.5D algorithms for SpMM, no such analysis exists for SDDMM or the back-to-back sequence of SDDMM and SpMM, termed FusedMM. We show that distributed memory 1.5D and 2.5D algorithms for SpMM can be converted to algorithms for SDDMM with identical communication costs and input / output data layouts. Further, we give two communication-eliding strategies to reduce costs further for FusedMM kernels: either reusing the replication of an input dense matrix for the SDDMM and SpMM in sequence, or fusing the local SDDMM and SpMM kernels. We benchmark FusedMM algorithms on Cori, a Cray XC40 at LBNL, using Erdos-Renyi random matrices and large real-world sparse matrices. On 256 nodes with 68 cores each, 1.5D FusedMM algorithms using either communication eliding approach can save at least 30% of time spent exclusively in communication compared to executing a distributed-memory SpMM and SDDMM kernel in sequence. On real-world matrices with hundreds of millions of edges, all of our algorithms exhibit at least a 10x speedup over the SpMM algorithm in PETSc. On these matrices, our communication-eliding techniques exhibit runtimes up to 1.6 times faster than an unoptimized sequence of SDDMM and SpMM. We embed and test the scaling of our algorithms in real-world applications, including collaborative filtering via alternating-least-squares and inference for attention-based graph neural networks.