一句话：提出 Einsum Tree IR，在给定收缩路径下全局优化张量数据布局，提升 einsum 表达式整体执行性能。
Einsum 是一种声明式张量表达式语言，只规定输出与输入张量的关系，不规定计算方式。典型后端先用收缩路径算法将表达式分解为二元张量收缩序列，再执行这些收缩；而高效执行需要优化张量数据布局。已有工作分别研究收缩路径生成和单个局部二元收缩的数据布局优化，缺乏对整体执行时间的全局布局优化。本文提出 Einsum Tree IR，一种在给定收缩路径下全局优化数据布局的中间表示，以降低 einsum 表达式的整体执行开销。作者在 Arm 服务器、x86 服务器和 x86 桌面系统上评估了该方法的有效性。
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.
@inproceedings{BreuerBEGHKR25,
title = {Einsum Trees: An Abstraction for Optimizing the Execution of Tensor Expressions},
author = {Alexander Breuer and Mark Blacher and Max Engel and Joachim Giesen and Alexander Heinecke and Julien Klaus and Stefan Remke},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3676641.3716254}
}