Einsum Trees: An Abstraction for Optimizing the Execution of Tensor Expressions
Year: 2025
Type: conference-paper
Abstract: 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.
Authors: Alexander Breuer, Mark Blacher, Max Engel, Joachim Giesen, Alexander Heinecke, Julien Klaus, Stefan Remke
Institutions: Friedrich Schiller University Jena, Intel (United States)
Language: English
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: hybrid
Funder: Universitas Brawijaya
OpenAlex ID: W4408894042