Einsum Trees: An Abstraction for Optimizing the Execution of Tensor Expressions

Einsum is a declarative language for tensor expressions that specifies an output tensor in terms of several input tensors. However, it does not specify how to compute the output tensor from the input tensors. A typical computational backend for the einsum language comprises two parts: First, a contraction path algorithm that breaks down an einsum expression into a sequence of binary tensor contractions. Second, the execution of the binary contractions. For efficient binary contractions, the data layout of the tensors must be optimized. So far, the computation of contraction paths and the optimization of the data layout for single, that is, local, binary tensor contractions have been studied in isolation. For optimizing the overall execution times of einsum expressions, we introduce Einsum Tree IR, an intermediate representation for globally optimizing the data layout for a given contraction path. We illustrate the effectiveness of the approach on a state-of-the-art Arm server processor, an x86 server processor, and an x86 desktop system.