Parametric Mappings for Distributed-Memory Tensor Computations

Abstract

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 ×.