EinDecomp: Decomposition of Declaratively-Specified Machine Learning and Numerical Computations for Parallel Execution

We consider the problem of automatic parallelism in high-performance, tensor-based systems. Our focus is on intra-operator parallelism for inference tasks on a single GPU server or CPU cluster, where each operator is automatically broken op so that it runs on multiple devices. We assert that tensor-based systems should offer a programming abstraction based on an extended Einstein summation notation , which is a fully declarative, mathematical specification for tensor computations. We show that any computation specified in the Einstein summation notation can be re-written into an equivalent tensor-relational computation that facilitates intra-operator parallelism, and this re-write generalizes existing notations of tensor parallelism such as "data parallel" and "model parallel." We consider the algorithmic problem of optimally computing a tensor-relational decomposition of a graph of operations specified in our extended Einstein summation notation.