Generating Two-Level, GPU-Aware Mappings for Distributed Tensor Computations

We introduce a two-level scheme to generate GPUaware MPI/NCCL code for distributed tensor computations. Our generator takes the specification of a linearized Directed Acyclic Graph (DAG) of tensor operators and produces a global mapping solution that considers MPI communication (inter and intranode) and the local computation. The core of our generator is a new bit-vector representation that compactly models mappings as well as communication directions along the grid. We incorporate the 2-level mapping decisions into a non-linear formulation which is optimized in an iterative fashion with the Z3 SMT solver. The new mapper supports both NVIDIA NCCL, MVAPICH-gdr, allowing for better portability. We demonstrate the efficiency of our mapping generator on a set of matrix- and tensor- DAGs, on two multi-GPU clusters with NVLink or PCIe intra-node interconnect, and compare against the COSMA library and CTF framework, achieving speedups ranging from 2.6× (over COSMA) to 18× (over CTF).