Scalable Block-Sparse Matrix Multiplication Using Template Task Graphs

Block-sparse matrix operations are a special case of general sparse algebra where the matrix is sparsely populated with dense blocks, e.g., in sparse tensor algebra for quantum chemistry. One of the challenges of implementing distributed matrix multiplication C = A × B in general is the management of communication flows since both input matrices A and B are readily available and must be distributed to the processes computing the relevant blocks of C. In this paper, we propose an addition to the Template Task Graph programming model that allows applications to constrain the execution of tasks using a flexible API. We show that such constraints can be used in a pure dataflow model to replace artificial control flow with a more structured approach. In the context of sparse matrix multiplication, we found that constraints allow us to limit the number of concurrent communications and thus avoid creating a bottleneck in the network.