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

Alexander Breuer, Mark Blacher, Max Engel, Joachim Giesen, Alexander Heinecke, Julien Klaus and 1 other

Einsum is a concise notation that says which tensor to compute from several inputs but not how, so a backend must break the expression into a tree of pairwise tensor contractions and run each one — and each contraction runs fast only if the data is laid out well in memory, otherwise it pays for costly transposes; this paper (Friedrich Schiller University Jena and Intel — Alexander Breuer, Mark Blacher, Alexander Heinecke, and colleagues) introduces the Einsum Tree IR, an intermediate representation that optimizes the memory layout across the whole contraction tree at once rather than one contraction at a time, then lowers each contraction to fast matrix-multiply microkernels.

The argument in one glance

Problem

Einsum says what tensor to compute, but not how, and layout decides the cost

Einsum (Einstein summation) is a compact, declarative notation for tensor algebra: you write the indices of your input tensors and the indices of the output you want, and it specifies which output tensor results. A matrix multiply, for instance, is written by labeling the two inputs with shared and free index letters and naming the output indices. Einsum can contract any number of tensors at once, which makes it a favorite for everything from simulating quantum circuits to deep-learning models. But einsum only says what to compute, not how. A backend has to turn the declaration into actual work, and this happens in two stages. First, a contraction-path algorithm breaks the multi-tensor expression into a sequence of binary (two-tensor) contractions, forming a contraction tree whose leaves are the input tensors and whose internal nodes are intermediate results. Second, each binary contraction has to be executed efficiently. Here the crucial, easily-overlooked cost is the data layout: a binary tensor contraction runs fast when its tensors are laid out in memory the way the computation wants to read them, and slow when they are not, because a mismatch forces an expensive permutation (a transpose that physically reshuffles the data). So the same expression can run at very different speeds depending on how the tensors are arranged in memory at each step of the tree.
What people did

Find a good contraction path, and optimize the layout of each contraction locally

Prior work split the problem into two separately-studied pieces. One line of work computes good contraction paths — the order in which to combine tensors so the total arithmetic is small. A separate line optimizes the data layout for a single, local binary contraction — choosing how to arrange that one contraction's tensors so its matrix-multiply core runs fast.
Why it fell short

Optimizing each contraction in isolation forces permutations between them

Studying the two problems in isolation misses the interaction that actually governs performance. The layout that is best for one binary contraction in isolation is not necessarily the layout that the next contraction in the tree wants for its input — the same intermediate tensor is the output of one node and the input of another, and the two nodes may prefer different arrangements. When each contraction is optimized locally, the tree ends up needing a permutation between the nodes to convert the intermediate tensor from the layout one node produced to the layout the next node expects, and those permutations at the intermediate nodes are expensive, memory-bound reshuffles that can dominate the runtime. So local optimization of each contraction's layout, however good it is per contraction, leaves the overall execution paying for transposes that a coordinated, tree-wide layout choice could have avoided. What was missing was a way to optimize the data layout globally, across the whole contraction tree, so that intermediate tensors flow from one contraction to the next without a costly permutation in between.
This paper's idea

Optimize layout across the whole tree with an einsum-tree IR, then lower to GEMMs

The paper introduces the Einsum Tree IR, an intermediate representation for globally optimizing the data layout of a whole contraction tree. The IR makes layout a first-class, manipulable thing with three properties: the order of the index letters in a node's string encodes that tensor's memory layout (a general row-major order), swapping a node's left and right children yields a genuinely different tree, and permutation nodes — a node with a single child and single parent that transposes data — are explicit parts of the IR. On top of this it defines three transformations: swap the left and right children of a binary node, reorder the index string of an interior node, and insert a permutation node. A simple, fast heuristic then traverses the tree from the root downward, using these transformations to give every intermediate tensor a memory layout that lets it be consumed by the next contraction without a permutation, while preserving the layouts of the root (output) and leaf (input) tensors so the user's data is not forced to be transposed. Once the layouts are fixed, each binary contraction is lowered to nested loops over small matrix multiplies (GEMMs): a simple mask identifies which index dimensions become the GEMM dimensions, spatial loops are parallelized across cores with OpenMP, and the GEMM cores are highly-tuned Tensor Processing Primitives generated by LIBXSMM. The authors report that this einsum-tree approach with the tensor-primitive backend is competitive on single binary contractions and achieves the best performance on most complex contraction trees, evaluated on an Arm server (NVIDIA Grace), an x86 server (Intel Xeon), and an x86 desktop.

Why this problem is worth solving

Einsum sits underneath a huge range of tensor computations, so making its execution fast pays off broadly, and layout is where much of the performance hides.

Einsum is a lingua franca for tensor algebra. Deep-learning frameworks use it, quantum-circuit simulators use it, and the tensor-network methods that appear across physics and machine learning are naturally expressed as einsum contractions. Because it is declarative — it states the result, not the procedure — the same einsum expression can be executed in many ways with wildly different performance, and the person writing the expression usually neither knows nor wants to know how it is carried out. That makes the backend that turns einsum into work a high-leverage place to optimize: a better backend speeds up everything expressed in einsum without any change to the user’s code. Within that backend, the arithmetic (how many multiply-adds) is set by the contraction path, but a large share of the actual runtime is spent moving data into the right shape for each contraction. If those data movements can be minimized, the same computation runs substantially faster on the same hardware. So optimizing einsum execution, and specifically its data layout, is a broadly applicable win.

The subtlety is that the two halves of the backend interact, so optimizing them separately leaves performance on the table.

A contraction tree is a chain of dependencies: each intermediate tensor is produced by one binary contraction and consumed by another. The efficiency of a single binary contraction depends on how its tensors are laid out, and it is tempting to pick, for each contraction on its own, the layout that makes that contraction fastest. But those local choices collide at the boundaries between contractions. The layout that one node produces for an intermediate tensor may not be the layout the next node wants to read, and reconciling the two requires physically permuting the tensor — a transpose that touches all the data and is limited by memory bandwidth, not arithmetic. Under local optimization these permutations pile up at the internal nodes of the tree, and they can cost more than the contractions themselves. The realization is that layout must be optimized globally, treating the whole tree as one object: if you choose the intermediate layouts so that each producer emits exactly what its consumer wants, the permutations disappear. The difficulty is representing the tree so that layout is something you can reason about and transform coherently across nodes, and doing the optimization cheaply enough to be practical. That is what the Einsum Tree IR is for.

What people did before, and where it broke

Compute a good contraction path. Choose the order of binary contractions to minimize total arithmetic. Where it broke. This settles how much multiply-add work there is but says nothing about data layout, so the tree can still be dominated by permutation costs.
Optimize the layout of one binary contraction at a time. For each contraction in isolation, pick the layout that makes its matrix-multiply core fastest. Where it broke. Locally optimal layouts clash at node boundaries, forcing permutations between contractions that can dominate the runtime.
Materialize transposes as needed between contractions. Insert whatever data reshuffles are required to connect locally-optimized contractions. Where it broke. Those transposes are memory-bound and expensive, and they accumulate at the tree’s intermediate nodes.
Rely on general tensor or deep-learning backends. Use existing einsum executors (for example framework backends or auto-schedulers). Where it broke. Without globally coordinating layout across the tree, they leave the permutation overhead in place, which the paper’s comparisons target.
The gap. No approach globally optimized the data layout across an entire contraction tree so that intermediate tensors flow between contractions without costly permutations.

How Einsum Trees works

Optimizing layout across a contraction tree, then lowering to matrix multiplies

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Give the computation an intermediate representation built for optimization. The Einsum Tree IR exists so that a high-level, declarative tensor expression can be reshaped and improved by a compiler, turning layout decisions into explicit, transformable structure rather than hidden runtime choices. Introducing a purpose-built intermediate representation that makes the important optimizations of a tensor program expressible and automatic is the same undertaking the work on an abstraction for automatic tensorized program optimization (TensorIR) covered elsewhere on this site pursues, letting the representation itself carry what the optimizer needs to manipulate. The Einsum Tree IR encodes memory layout in the index strings and makes permutations first-class, so a tree of contractions becomes something a heuristic can optimize globally.

Make performance a question of data layout, and eliminate the needless transposes. The heart of the work is that the runtime is governed by how tensors are arranged in memory, and that many of the transposes a naive execution performs are avoidable by choosing layouts coherently. Recognizing that memory layout is the lever and reworking it to remove the transformations that would otherwise be forced is the same principle the work on layout transformation elimination and adaptation (SmartMem) covered elsewhere on this site is built on, treating layout as a first-class optimization target and cutting the conversions that waste bandwidth. Einsum Trees does this across a whole contraction tree, so intermediate tensors pass between contractions in the layout the consumer wants, with no permutation in between.

Lower to tuned microkernels so each piece runs near peak. Once the tree’s layouts are settled, the paper does not hand-write inner loops; it lowers each contraction to small, highly-tuned matrix-multiply kernels generated by a specialized backend. Compiling a high-level tensor operation down to carefully generated, hardware-tuned kernels so each core computation reaches near-peak throughput is the same idea the work on high-performance library generation for tensor kernels (Heron) covered elsewhere on this site develops, trusting a kernel generator to produce the fast inner loops. Einsum Trees uses LIBXSMM Tensor Processing Primitives as those GEMM cores and parallelizes the surrounding spatial loops across cores.

Following a contraction tree through the IR

Optimizing and executing one contraction tree

A trace, drawn from the paper's abstract and the authors' poster, of how the Einsum Tree IR optimizes and lowers a tree. Exact algorithms and results should be read in the paper.

Step 1 of 5

Start from an einsum expression and its contraction path

A multi-tensor einsum expression has already been broken into a contraction path: a tree of binary contractions whose leaves are the (already reordered) input tensors and whose internal nodes are intermediate results.

The path fixes how much arithmetic there is; layout is still open.

Why it works

Global layout optimization removes the permutations local optimization forces. Picking each contraction’s layout in isolation clashes at node boundaries → but choosing intermediate layouts across the whole tree so each producer emits what its consumer wants → eliminates the transposes at intermediate nodes.
Encoding layout in the index string makes it manipulable. Layout is usually an implicit runtime property → but putting the memory order into the order of the index letters → lets the optimizer reason about and change layouts directly in the IR.
Explicit permutation nodes make transpose cost visible. A hidden transpose is easy to ignore → but representing permutations as real nodes → lets the heuristic insert them only where genuinely needed and avoid them everywhere else.
Preserving root and leaf layouts avoids boundary costs. Reworking every tensor’s layout could force transposing the user’s inputs and output → but keeping the leaf and root layouts fixed → means the optimization is free of hidden costs at the tree’s boundaries.
Lowering to tuned GEMM primitives runs each contraction near peak. Naive inner loops underuse the hardware → but mapping each contraction onto LIBXSMM tensor primitives with parallelized spatial loops → makes the core matrix multiplies fast on both Arm and x86.

What it costs (trade-offs and scope)

The layout optimization is a fast heuristic, not a guaranteed optimum. The authors emphasize that the tree-optimization heuristic is remarkably simple and quick, which is a virtue, but it means the chosen layouts are good rather than provably best; a harder-to-compute global optimum could in principle do better on some trees.
It assumes the contraction path is given. Einsum Trees optimizes layout for a provided contraction tree (with input tensors assumed already reordered to fit the tree); finding the contraction path itself is a separate, complementary problem the approach builds on rather than solves.
It targets CPUs with matrix-multiply microkernel backends. The lowering uses LIBXSMM Tensor Processing Primitives (and the paper also evaluates a TBLIS backend), and the evaluation is on Arm and x86 CPUs; other hardware such as GPUs would need a different lowering and kernel backend.
The advantage is largest on complex trees. On single binary contractions the approach is reported as competitive, while its clearest wins are on complex contraction trees with layout-sensitive intermediates, so the benefit depends on the structure of the expression.
This explainer is drawn from limited source material. The full paper text was not accessible for this write-up; the description here is based on the paper’s abstract and the authors’ conference poster, so the precise algorithms, the heuristic’s details, and the exact quantitative results should be confirmed in the paper.

How it compares

Executing einsum tensor expressions

Property	Good path + local layout	General einsum backends	Einsum Tree IR (this work)
Contraction path (arithmetic)
Data layout across the tree
Permutations at intermediate nodes
Inner contraction execution
Reported performance

A qualitative comparison, based on the paper's abstract and the authors' poster, of ways to execute einsum expressions. By optimizing the memory layout globally across a contraction tree with the Einsum Tree IR and lowering each contraction to tuned matrix-multiply primitives, the approach avoids the permutations that local layout optimization forces, and the authors report it is competitive on single binary contractions and best on most complex contraction trees across Arm and x86 CPUs. Specific per-benchmark figures should be read from the paper.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s abstract and the authors’ conference poster, not its full text, which was not accessible for this write-up. The system and claims described are the authors’: einsum is a declarative language for tensor expressions that does not specify how to compute the output, so a backend comprises a contraction-path algorithm (breaking an expression into a sequence of binary tensor contractions) and the execution of those contractions, for which data layout must be optimized. The paper observes that contraction-path computation and the data-layout optimization of single, local binary contractions had been studied in isolation, and introduces the Einsum Tree IR, an intermediate representation for globally optimizing the data layout for a given contraction path. Per the authors’ poster, the IR encodes each tensor’s memory layout in the order of its index string, treats swapping a node’s children as a distinct tree, and makes permutation nodes explicit; three transformations (swap children, reorder index string, insert permutation node) and a simple, fast root-to-leaf heuristic optimize intermediate layouts while preserving root and leaf layouts; and each contraction is lowered to nested loops over small GEMMs using Tensor Processing Primitives from LIBXSMM with OpenMP parallelization of spatial loops. The authors report the approach is competitive on binary contractions and best on most complex contraction trees, evaluated on an Arm server (NVIDIA Grace), an x86 server (Intel Xeon), and an x86 desktop, against ATen, opt_einsum (PyTorch backend), TVM-Ansor, and TBLIS. No specific quantitative results are asserted here; exact algorithms and numbers should be read from the paper. This page is marked as a draft for that reason.

Background

Why moving data, not multiplying it, often decides tensor performance

It helps to see Einsum Trees within a broad truth of modern numerical computing: for many tensor and matrix computations, the bottleneck is not the arithmetic but the data movement around it. Processors can multiply and add far faster than they can shuttle operands into the right place, so a computation that is arithmetically identical can run at very different speeds depending on how its data is laid out and how often it must be reshuffled. This is why so much of high-performance tensor work is really about layout, blocking, and avoiding transposes, and why declarative interfaces like einsum — which free the user from specifying execution — put so much responsibility on the backend to get those data decisions right. Tensor networks and deep-learning models compound the challenge because they chain many contractions together, so a layout choice good for one step can be bad for the next, and the reshuffles between steps accumulate. The productive response, seen across tensor compilers, is to represent the whole computation in a form where these data decisions are explicit and globally optimizable, and then to lower the pieces onto hardware-tuned kernels. The broad effort to orchestrate the kernels of a tensor program so the whole thing runs efficiently rather than each piece in isolation is the same terrain the work on optimal kernel orchestration for tensor programs (Korch) covered elsewhere on this site occupies, coordinating across the computation rather than locally. Einsum Trees contributes a layout-centric intermediate representation for contraction trees specifically, arguing that optimizing the memory layout globally, and only then generating fast kernels, is what unlocks einsum performance.

If you remember only three things

1. Einsum is a declarative notation that says which tensor to compute from several inputs but not how, so a backend must break the expression into a contraction tree of binary tensor contractions and execute each one. A binary contraction runs fast only when its tensors are laid out in memory the way the computation wants; otherwise it pays for an expensive permutation (a transpose). Prior work optimized the contraction path (arithmetic) and the layout of single, local contractions separately, so optimizing each contraction in isolation forces costly permutations at the intermediate nodes where one contraction feeds the next.
2. The paper introduces the Einsum Tree IR to optimize data layout globally across the whole tree. The IR encodes each tensor's memory layout in the order of its index letters, treats swapping a node's children as a distinct tree, and makes permutation nodes explicit; three transformations (swap children, reorder index string, insert permutation) plus a simple, fast heuristic that traverses the tree from the root choose intermediate layouts so each contraction produces exactly what the next one wants — eliminating intermediate permutations — while preserving the input and output layouts. Each contraction is then lowered to nested loops over small GEMMs, using tuned Tensor Processing Primitives from LIBXSMM with parallelized spatial loops.
3. The authors report the einsum-tree approach with the tensor-primitive backend is competitive on single binary contractions and achieves the best performance on most complex contraction trees, evaluated on an Arm server (NVIDIA Grace), an x86 server (Intel Xeon), and an x86 desktop against ATen, opt_einsum, TVM-Ansor, and TBLIS. The trade-offs: the layout optimization is a fast heuristic rather than a guaranteed optimum, it assumes the contraction path is given, it targets CPUs with matrix-multiply microkernel backends, and its clearest wins are on complex trees. This explainer is a draft based on the abstract and the authors' poster; exact algorithms and numbers are in the paper.

Check yourself

Check your understanding

1. Einsum is declarative. What does a backend still have to decide, and why does data layout matter?
2. Why is optimizing the data layout of each binary contraction in isolation not enough?
3. How does the Einsum Tree IR let a compiler reason about and change data layout?
4. Predict: after the Einsum Tree IR fixes the layouts, how does the paper actually run each binary contraction, and why does that choice help?

Sources

The paper is “Einsum Trees: An Abstraction for Optimizing the Execution of Tensor Expressions” by Alexander Breuer, Mark Blacher, Max Engel, Joachim Giesen, Alexander Heinecke, Julien Klaus, and Stefan Remke (Friedrich Schiller University Jena and Intel Corporation), published at ASPLOS 2025, DOI 10.1145/3676641.3716254 (pages 275 to 292).

This explainer is written from the paper’s abstract and the authors’ conference poster (hosted by the Jena scalable-computing group), because the full paper text was not accessible for this write-up (the ACM version was not retrievable and no author preprint was found). Einsum is a declarative language for tensor expressions that does not specify how to compute the output; a backend comprises a contraction-path algorithm that breaks an einsum expression into a sequence of binary tensor contractions and the execution of those contractions, for which data layout must be optimized. Because contraction-path computation and single, local binary-contraction layout optimization had been studied in isolation, the paper introduces the Einsum Tree IR, an intermediate representation for globally optimizing the data layout for a given contraction path. Per the authors’ poster, the IR encodes memory layout in the index-string order, makes swapping a node’s children a distinct tree, and includes explicit permutation nodes; three transformations (swap children, reorder index string, insert permutation) and a simple fast root-to-leaf heuristic optimize intermediate layouts while preserving root and leaf layouts; and each contraction is lowered to nested loops over small GEMMs using LIBXSMM Tensor Processing Primitives with OpenMP parallelization. The authors report the approach is competitive on binary contractions and best on most complex contraction trees, evaluated on an Arm server (NVIDIA Grace), an x86 server (Intel Xeon), and an x86 desktop, against ATen, opt_einsum, TVM-Ansor, and TBLIS. Numbers and mechanisms are the authors’; exact details and results should be confirmed in the paper. This page is marked draft because the full text was unavailable.