COSTA: Communication-Optimal Shuffle and Transpose Algorithm with\n Process Relabeling

Communication-avoiding algorithms for Linear Algebra have become increasingly\npopular, in particular for distributed memory architectures. In practice, these\nalgorithms assume that the data is already distributed in a specific way, thus\nmaking data reshuffling a key to use them. For performance reasons, a\nstraightforward all-to-all exchange must be avoided.\n Here, we show that process relabeling (i.e. permuting processes in the final\nlayout) can be used to obtain communication optimality for data reshuffling,\nand that it can be efficiently found by solving a Linear Assignment Problem\n(Maximum Weight Bipartite Perfect Matching). Based on this, we have developed a\nCommunication-Optimal Shuffle and Transpose Algorithm (COSTA): this\nhighly-optimised algorithm implements $A=\\alpha\\cdot \\operatorname{op}(B) +\n\\beta \\cdot A,\\ \\operatorname{op} \\in \\{\\operatorname{transpose},\n\\operatorname{conjugate-transpose}, \\operatorname{identity}\\}$ on distributed\nsystems, where $A, B$ are matrices with potentially different (distributed)\nlayouts and $\\alpha, \\beta$ are scalars. COSTA can take advantage of the\ncommunication-optimal process relabeling even for heterogeneous network\ntopologies, where latency and bandwidth differ among nodes. The implementation\nnot only outperforms the best available ScaLAPACK redistribute and transpose\nroutines multiple times, but is also able to deal with more general matrix\nlayouts, in particular it is not limited to block-cyclic layouts. Finally, we\nuse COSTA to integrate a communication-optimal matrix multiplication algorithm\ninto the CP2K quantum chemistry simulation package. This way, we show that\nCOSTA can be used to unlock the full potential of recent Linear Algebra\nalgorithms in applications by facilitating interoperability between algorithms\nwith a wide range of data layouts, in addition to bringing significant\nredistribution speedups.\n

Paper

Similar papers

© 2026 NYSGPT2525 LLC