Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication

Abstract

We propose COSMA: a parallel matrix-matrix multiplication algorithm that is near communication-optimal for all combinations of matrix dimensions, processor counts, and memory sizes. The key idea behind COSMA is to derive an optimal (up to a factor of 0.03\% for 10MB of fast memory) sequential schedule and then parallelize it, preserving I/O optimality. To achieve this, we use the red-blue pebble game to precisely model MMM dependencies and derive a constructive and tight sequential and parallel I/O lower bound proofs. Compared to 2D or 3D algorithms, which fix processor decomposition upfront and then map it to the matrix dimensions, it reduces communication volume by up to $\sqrt{3}$ times. COSMA outperforms the established ScaLAPACK, CARMA, and CTF algorithms in all scenarios up to 12.8x (2.2x on average), achieving up to 88\% of Piz Daint's peak performance. Our work does not require any hand tuning and is maintained as an open source implementation.

arXiv comment18 pages, 29 figures, short version submitted to the SC'19 conference

References in corpus (1)

• Convergence Analysis for Rectangular Matrix Completion Using Burer-Monteiro Factorization and Gradient Descent

Cited by in corpus (6)

• Fast Parallel Hypertree Decompositions in Logarithmic Recursion Depth
• Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
• TAMM: Tensor Algebra for Many-body Methods
• A High-Fidelity Flow Solver for Unstructured Meshes on Field-Programmable Gate Arrays
• Subspace recursive Fermi-operator expansion strategies for large-scale DFT eigenvalue problems on HPC architectures
• Communication Bounds for Convolutional Neural Networks