Work

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

Year: 2025
Type: article
Abstract: 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.
Institution: Rice University
Language: English
Open Access status: green
OpenAlex ID: W4413825368
FWCI: 0.705
Cites: 27
Cited by: 1
Related to: 10