💻 computer science

Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

This paper presents the first formal study of layout selection in machine learning compilers by formulating it as a combinatorial optimization problem, proving its computational hardness, and proposing both optimal algorithms for bounded treewidth graphs and a weighted MaxSAT encoding for general instances to demonstrate that simple heuristics can degrade performance by up to 5x compared to optimal solutions.

Original authors: Clemens Eisenhofer, Yuwen Jia, Daniel Kroening, Sergey Pupyrev

Published 2026-08-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Clemens Eisenhofer, Yuwen Jia, Daniel Kroening, Sergey Pupyrev

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence relies on massive mathematical models that process vast amounts of data to recognize speech, translate languages, or generate images. To run these models quickly, engineers use specialized computer chips designed specifically for this heavy lifting. However, these chips do not simply execute instructions; they must also move data around efficiently. The speed at which a model runs often depends less on the raw power of the chip and more on how the data is arranged in its memory. Imagine a library where books are stored on shelves. If a reader needs to find a specific set of books, the time it takes depends entirely on whether those books are scattered across different aisles or neatly grouped together on a single shelf. In the world of computer chips, this arrangement is called a "layout." When a computer program performs a calculation, it expects data in one specific arrangement, but the previous step in the program might have left the data in a different one. If the two do not match, the computer must stop and rearrange the data before it can proceed, a process that wastes time and energy.

For years, the software that prepares these models for the chips has relied on a collection of rough guesses and rules of thumb to decide how to arrange this data. These rules work well enough for simple tasks, but as models have grown more complex, the guesses have started to fail, leading to significant slowdowns. A team of researchers from the Technical University of Vienna and Amazon set out to change this approach. Instead of relying on intuition, they treated the problem of arranging data as a formal mathematical puzzle. They built a precise model that calculates the exact cost of every possible arrangement, including the time it takes to move data between different formats. By doing this, they could determine the single best way to organize the data for any given model, rather than hoping a set of rules would get close enough.

The researchers discovered that finding this perfect arrangement is an incredibly difficult task. In the language of computer science, the problem is so complex that no computer can solve it quickly for every possible situation, especially as the models grow larger. They proved that even for a simplified version of the problem involving only basic matrix calculations, the number of possibilities is so vast that a standard computer would struggle to find the answer in a reasonable time. This finding ruled out the idea that a single, fast, universal algorithm could solve the problem for all future models. However, the team also found a way forward. They showed that while the problem is hard in general, it becomes manageable when the structure of the model resembles a tree with limited branching. For these specific structures, which are common in many real-world applications, they designed a method that finds the perfect solution quickly. For more complex structures that do not fit this pattern, they developed a way to translate the problem into a format that existing powerful solvers can handle, allowing them to find the best possible arrangement even when a perfect mathematical shortcut does not exist.

To test their ideas, the researchers implemented their new method inside a real-world compiler used for Amazon's Trainium chips, which are designed to run artificial intelligence models. They compared their new approach against the standard methods currently used in the industry, which rely on the older rules of thumb. The results were striking. On some complex models, particularly those used for image recognition, the old rules of thumb caused the models to run up to five times slower than necessary. This happened because the simple rules failed to see the big picture; they would arrange data perfectly for one step but create a mess for the next, forcing the computer to waste time rearranging data constantly. The new method, by looking at the entire sequence of steps at once, avoided these costly rearrangements and kept the data flowing smoothly.

However, the study also revealed a crucial limitation. While the new method could always find the mathematically best arrangement according to its own calculations, this did not always translate to the fastest speed on the actual hardware. In some cases, the new method produced a result that was theoretically perfect but performed worse than the older, simpler rules. The researchers traced this discrepancy to the cost model itself. The software used to predict how long a task would take was not perfectly accurate; it underestimated the time required for certain types of data movement. Because the new method was so good at finding the lowest cost according to its own flawed predictions, it sometimes chose an arrangement that looked cheap on paper but was actually expensive in reality. This finding suggests that the biggest hurdle for future improvements is not better search algorithms, but better ways to predict how long tasks will actually take.

The work provides a clear path forward for the field. It demonstrates that treating layout selection as a formal optimization problem is a viable and powerful strategy, capable of delivering massive speedups where simple rules fail. It also clarifies that the ultimate limit of performance is not the ability to find the best solution, but the accuracy of the predictions used to guide that search. For models with regular, predictable structures, the new solver-based approach is already a superior choice. For more chaotic and complex models, the focus must shift to refining the cost models so that the mathematical optimum aligns with the physical reality of the chip. By separating the problem of finding the best solution from the problem of predicting the cost, the researchers have given compiler developers a new tool to measure their progress and a clear target for where to focus their efforts next.

Technical Summary: Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

1. Problem Statement

Modern machine learning (ML) compilers must select memory layouts for tensors to minimize execution costs while adhering to hardware constraints. A tensor layout defines the physical representation of a multi-dimensional array (e.g., dimension ordering like NHWC vs. NCHW, or partition dimensions for systolic arrays). This selection is inherently global: an operator may execute fastest under one layout, while its consumers prefer another. Aligning these preferences often requires explicit layout conversions (transposes, copies, or reformatting), which introduce significant overhead.

Current compilers rely on ad-hoc heuristics—such as local greedy passes or rule-based propagation—to resolve these conflicts. However, these approaches lack a formal theoretical basis, often leading to suboptimal performance. The paper identifies a critical gap: the absence of a unified mathematical model for layout selection that accounts for the trade-off between operator execution costs and the cost of inserting layout conversions.

2. Methodology and Formalization

2.1 Mathematical Formulation

The authors formulate layout selection as a combinatorial optimization problem over dataflow graphs ( $G = (V, E)$ ).

Inputs: A directed acyclic graph where vertices represent operators and edges represent tensors. Each tensor $t$ has a finite set of admissible layouts $L(t)$ .
Costs:
- Operator Cost ( $C^{op}_o$ ): A function mapping input and output layout combinations to execution costs (or $\infty$ for infeasible configurations).
- Transpose Cost ( $C^{tr}_t$ ): A function mapping a tensor's producer layout to a consumer's required layout, representing the cost of conversion.
Objective: Find a layout assignment $\ell(t, o)$ for every operator $o$ and tensor $t$ that minimizes the sum of operator execution costs and the costs of all necessary conversions.

The formulation unifies existing strategies, including XLA's layout assignment (dimension ordering), systolic array partition dimension selection (e.g., AWS Trainium), and mobile GPU optimizations.

2.2 Theoretical Analysis

The paper establishes the computational complexity of the problem:

NP-Hardness: Layout selection is proven to be NP-hard, even when restricted to programs containing only matrix multiplications over 2D tensors with a constant number of inputs and two layouts per tensor.
Inapproximability: Unless P=NP, no polynomial-time approximation algorithm exists with a finite approximation ratio for the general case.
Fixed-Parameter Tractability: The problem admits an optimal polynomial-time algorithm for dataflow graphs with bounded treewidth. The authors provide an explicit dynamic programming (DP) algorithm for this case, which is linear in the number of operators for fixed treewidth.
General Solving: For general instances, the problem is encoded as a weighted MaxSAT instance, allowing off-the-shelf solvers to find optimal solutions.

2.3 Implementation

The authors implemented the formalization in a production compiler for the AWS Trainium AI accelerator. They evaluated four strategies:

Local: Minimizes operator cost only, ignoring conversions.
Greedy: Iteratively refines a topological assignment to minimize the total objective (operator + conversion costs).
Rule-Based: The compiler's existing strategy using domain-specific rules and propagation.
Solver-Based: Uses the exact Treewidth DP and MaxSAT encodings to compute provably optimal solutions under the cost model.

3. Key Results

3.1 Practical Impact of Layout Selection

Performance Gains: Simple heuristics can degrade execution time by up to 5× compared to optimized strategies on complex workloads (e.g., ResNet-50, Vision Perceiver).
Workload Sensitivity: For regular dense transformer architectures (e.g., BERT, ALBERT), even the simple "Local" baseline performs nearly as well as optimal strategies because layout preferences are consistent. However, for complex architectures with heterogeneous operators (convolutions, cross-attention, skip connections), global optimization is critical.
Mechanism of Degradation: Poor layout selection leads to a significant increase in inserted transpose operations. On systolic arrays like Trainium, these conversions involve physical data movement across on-chip partitions that cannot be fused into consuming operators, causing substantial DMA overhead.

3.2 Viability of Exact Solvers

Treewidth: Empirical analysis shows that 92% of real-world ML dataflow graphs have a treewidth $\le 4$ . The Treewidth-based DP algorithm solves these instances in seconds. However, it fails on the remaining 8% (treewidth > 4) due to memory constraints.
MaxSAT: The MaxSAT encoding is viable for most models, solving instances within seconds. For large models where the solver times out, the best solution found within the first few seconds is often close to optimal, as the solver spends most time proving optimality rather than finding better assignments.

3.3 The Role of Cost Model Fidelity

A critical finding is that provably optimal solutions under the cost model do not always yield the best hardware performance.

Cost Model Error: The solver minimizes the formal objective exactly. When the solver underperforms the heuristic (e.g., on ResNet-50), the gap is attributed entirely to inaccuracies in the cost model, specifically the underestimation of DMA transfer costs resulting from tiling decisions made after layout selection.
Heuristic Advantage: Heuristics like the Rule-Based strategy sometimes outperform the solver because they incorporate domain-specific knowledge (e.g., preferred partition dimensions) that the analytical cost model fails to capture. They effectively "average out" model errors by staying close to conventional layouts, whereas the solver aggressively exploits the model's underestimations.

4. Significance and Claims

The paper claims to be the first formal study of layout selection in ML compilers. Its primary contributions are:

Unified Formalization: It provides a single mathematical objective that subsumes various existing "ad-hoc" layout optimizations, allowing for rigorous comparison of design choices.
Complexity Bounds: It proves NP-hardness and inapproximability, while identifying bounded treewidth as a tractable regime with a concrete dynamic programming algorithm.
Empirical Validation: Through integration into a production compiler, it demonstrates that while exact solvers are viable for many workloads, the ultimate performance ceiling is determined by the fidelity of the cost model, not just the search algorithm.
Practical Guidance: The authors argue that compiler developers should formalize layout selection as an optimization problem even when using heuristics. This provides a reference point for debugging and measuring optimality gaps. Furthermore, they suggest that once a heuristic's optimality gap is smaller than the cost model's error, further investment should shift from search algorithms to improving cost model fidelity (e.g., via calibration or learned models).

The paper concludes that while exact solvers can match or beat rule-based strategies on regular workloads, the "gap" on complex workloads highlights the need for more accurate cost models rather than more sophisticated search algorithms.