Parametric Mappings for Distributed-Memory Tensor Computations
Tensor computations are an important class of operations widely used in domains such as computational chemistry, machine learning, and various types of physical simulations that demand distributed-memory clusters. Recent work has shown that generating efficient mappings for multi-operator Directed Acyclic Graphs of distributed-memory tensor computations is possible by leveraging non-linear formulations underpinned by Satisfiability Modulo Theories (SMT) solvers. However, this approach is sensitive to the problem size, grid shape, and count of Processing Elements (PEs) given. We introduce a new method for automatically generating parametric mappings agnostic of the input problem sizes for the class of DAGs of tensor computations. The approach relies on three building blocks: the insight that in parametric DAGs not all parameters are equally important, an offline algorithm to generate a small bounded number of mapping classes – forming a family of mappings – and a sampling space that enables the generation of sufficiently large and diverse family of mapping classes. We show the effectiveness of our parametric mappings on a cluster of multi-GPUs, comparing against NVIDIA cuTensorNet, the Cyclops Tensor Framework (CTF) and PyTorch. Our mappings achieve speedups up to 22.1 × and 2.35 × over NVIDIA cuTensorNet, on matrix-based and tensor-based DAGs, respectively. When benchmarked against PyTorch, we observe speedups of 3.41 ×.