Source-linked AI summary
Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers
Clemens Eisenhofer, Yuwen Jia, Daniel Kroening, Sergey Pupyrev
TL;DR
Layout selection is formalized as a global cost-minimization problem for ML compiler dataflow graphs. The paper combines exact algorithms and solver-based optimization, finding that solver methods work for transformers but cost-model error limits gains on vision and multimodal workloads.
Problem
ML compilers lack a formal basis for globally selecting tensor layouts while balancing operator execution costs against layout-conversion costs.
Method
The paper formulates layout selection over dataflow graphs, proves NP-hardness, provides a bounded-treewidth dynamic program, and encodes general instances as weighted MaxSAT.
Results
Solver-based methods match or sometimes outperform the production heuristic on transformer workloads but do not on vision and multimodal models.
Takeaways & Limitations
The remaining performance gap on vision and multimodal workloads lies in cost-model error rather than search quality, while accurate models make exact optimization useful.
Takeaways & Limitations
The evaluation uses one hardware target and compiler, and the formulation isolates layout selection from joint tiling and operator-fusion decisions.
Abstract
from arXiv · showhide
Modern machine learning compilers select tensor memory layouts to minimize execution cost under hardware constraints. Layout selection is global: an operator may be fastest under one layout while its consumers prefer another, and aligning these preferences requires explicit layout conversions that can hurt model performance. Despite its practical importance, layout selection lacks a formal basis, so current compilers rely on ad-hoc heuristics. This paper presents the first formal study of layout selection in machine learning compilers. We formulate the problem as combinatorial optimization over dataflow graphs, minimizing the sum of operator execution costs and the per-tensor cost of these conversions. Our theoretical analysis shows that optimal layout selection is computationally hard, even for programs containing only matrix multiplications over two-dimensional tensors. We design an optimal polynomial-time algorithm for dataflow graphs of bounded treewidth. For general instances, we give a weighted MaxSAT encoding that an off-the-shelf solver can optimize. The formulation unifies several existing layout optimization strategies, including XLA's layout assignment, partition dimension selection in systolic array compilers, and layout planning in mobile GPU optimizers. We implement the formalization in a production compiler for an AI accelerator and measure the execution time of the compiled models under greedy heuristics, the compiler's rule-based strategy, and an optimal solver. Simple heuristics degrade execution time by up to $5\times$ on some workloads. Where the compiler's cost model is accurate, the solver matches or beats the rule-based strategy. On workloads with complex data movement it falls behind, and since the solver minimizes the stated objective exactly, that gap isolates cost-model error from search quality, showing where compiler effort actually pays off.
1 Introduction
Layout selection is a global optimization problem because operator preferences interact through shared tensors and layout conversions. The paper formalizes this problem across compiler settings and motivates exact and solver-based approaches.
- Motivation: Data movement, including reshaping and layout conversion, contributes substantially to end-to-end speedups on specialized ML hardware.Compilers therefore make performance decisions that extend beyond arithmetic kernel selection.
- Layout Selection: A tensor layout specifies its physical representation, such as dimension ordering or mapping to hardware resources, and directly affects execution performance.Examples include row-major and column-major storage and formats such as NHWC and NCHW.
- Layout Selection: Shared tensors couple local operator preferences, so a layout that speeds one operator may require conversions or slower execution elsewhere in the graph.The Figure 1 computation illustrates this tradeoff between a cheap matmul, reduction cost, and an explicit output conversion.
- Scope and Unification: Layout selection appears across XLA dimension orderings, systolic-array partition dimensions, and mobile or GPU tensor formats.These settings can be represented under a common objective involving operator and data-transformation costs.
- Contributions: The paper formalizes layout selection as graph optimization, establishes hardness and tractable structure, and evaluates heuristic and solver-based strategies in a production compiler.The framework is intended to replace ad-hoc layout optimization with an explicit mathematical model.
2 A Mathematical Formulation of Layout Selection
The formulation assigns admissible layouts to tensor uses and minimizes operator costs plus conversions between producer and consumer layouts. It abstracts over hardware-specific layout meanings while representing both absorbable and explicit conversions.
- Tensors, operators, and dataflow graphs: An ML program is modeled as a directed acyclic dataflow graph whose vertices are operators and whose edges carry produced tensors.Each operator consumes an input tuple and produces output tensors.
- Tensor layouts: A tensor layout describes its hardware representation, with each tensor having a finite admissible layout set.The abstraction covers dimension permutations, memory placement, and partition dimensions.
- Operator costs and feasibility: Each operator has a local state consisting of layouts for all input and output tensors it touches.The state space is finite and bounded when operator arity and admissible layout counts are bounded.
- Operator costs and feasibility: Unsupported layout combinations receive infinite cost, while finite operator costs encode relative performance among feasible choices.This lets the model represent both hardware constraints and performance preferences.
- Layout conversion costs: A layout conversion is charged when a tensor’s producer layout differs from a consumer layout, and one conversion serves consumers requesting the same layout.Conversion costs abstract data movement, reformatting, and synchronization overhead.
- The layout selection problem: The optimization problem minimizes operator execution costs plus conversion costs over all feasible per-use layout assignments.Consumers agreeing with the producer incur no conversion cost.
- Instantiation: The formulation covers XLA dimension orderings and Trainium partition dimensions as instances of the same layout-selection model.It also provides a formal objective corresponding to several heuristic systems.
- Layout conversions versus loop transformations: On general-purpose hardware, a layout mismatch may be absorbed by loop transformations or transposed-operand interfaces, whereas Trainium generally requires physical data movement.The framework models absorbed mismatches as operator configurations and explicit mismatches as conversions.
3 Theoretical Considerations
Layout Selection is hard in general, including restricted matmul-only settings, but bounded-treewidth graphs admit exact dynamic programming and general instances admit weighted MaxSAT optimization. These methods encode operator configurations and shared-tensor conversion costs directly.
- Computational hardness: Layout Selection is NP-hard even when operators have at most three inputs.Unless P = NP, no polynomial-time approximation algorithm with a finite approximation ratio exists for the general problem.
- Inapproximability: The 3-SAT reduction makes zero-cost layout assignments equivalent to satisfying truth assignments.Each variable tensor uses layouts T or F, and each clause operator has zero cost exactly on satisfying input-layout triples.
- Inapproximability: It is NP-hard to distinguish optimum cost 0 from optimum cost at least 1, ruling out finite-ratio approximation guarantees.Any positive-cost solution has an unbounded ratio relative to an optimum of zero.
- Restricted hardness: NP-hardness persists when all non-source operators are two-dimensional matrix multiplications with two layouts, binary transpose costs, and 0-or-infinite operator costs.The reduction uses opposite-layout constraints corresponding to edges of an Odd Cycle Transversal instance.
- Restricted hardness: The hardness reduction represents graph vertices as two-layout tensors and edges as matmuls feasible only when their inputs use different layouts.An odd cycle forces at least one tensor to be read in both layouts, incurring a transpose cost.
- Treewidth-based exact algorithm: For bounded-treewidth dataflow graphs, dynamic programming solves Layout Selection optimally in time O(|V| · tw^2 · M^(tw+2) · 4M^(tw+1)).The algorithm is fixed-parameter tractable in treewidth and the number of operator configurations, with linear runtime in |V| when those parameters are fixed.
- Treewidth-based exact algorithm: The dynamic program stores frontier operator configurations and consumer-requested layouts, charging costs when operators leave the processed region.Introduce, forget, and join transitions record configurations, register requests, charge costs, and merge independently processed subtrees.
- MaxSAT optimization: A weighted MaxSAT encoding assigns soft-clause weights so that minimizing violated weight exactly matches the Layout Selection objective.An off-the-shelf solver therefore returns an optimal layout assignment for general instances.
4 Experimental Evaluation
The evaluation compares local, greedy, rule-based, and solver-based layout selection across diverse models, measuring objective quality, compilation practicality, hardware execution, and cost-model fidelity. Rule-based strategies often perform best on complex architectures, while exact methods are practical in many cases but depend on cost-model accuracy.
- Experimental setup: The evaluation spans diverse model families, including language, vision, diffusion, mixture-of-experts, multimodal, and audio architectures.
- Practical impact: Simple strategies lose about 20% of performance on average versus Rule-Based, reaching a 5× slowdown in the worst decile.
- Practical impact: 20% more transpose operations on average, rising to 2× at p90, explain much of the regression from simpler layout strategies.On Trainium, partition-dimension changes physically move elements across on-chip memory partitions and cannot be fused into consuming operators.
- Solver practicality: 92% of instances have treewidth at most 4, where dynamic programming solves the problem within seconds, but larger treewidth requires over 10 GB of memory.MaxSAT completes within seconds on most transformer configurations with compilation overhead usually under 5% of total compilation time.
- Algorithm comparison: Greedy solutions are typically within 1% of the Solver-Based optimum, while Rule-Based is within 3% on average and 5% at p90.These small objective differences can still produce significantly different hardware performance.
- Algorithm comparison: Rule-Based wins on most complex architectures, while Solver-Based can underperform heuristics despite optimizing the formal objective exactly.Examples include ResNet-50, Language Perceiver, and PaLiGemma-3B.
- Cost-model fidelity: The cost model correctly ranks approximately 87% of pairwise algorithm comparisons, but transpose estimates miss downstream DMA memory-access patterns.This mismatch explains why exact optimization can fail to translate into better hardware execution.
5 Threats to Validity
The evaluation’s validity is limited by approximate cost modeling, a single hardware/compiler setting, and formulation assumptions about optimization passes and execution.
- Internal validity: Approximate cost models can make solver-optimal layouts differ from the best wall-clock layouts.Downstream memory allocation and instruction scheduling are estimated when layout selection occurs.
- External validity: The evaluation uses one AWS Trainium target and one production compiler, so hardware-specific layout and transpose costs may not generalize.Other accelerators can expose different cost tradeoffs and may absorb layout mismatches differently.
- Construct validity: The formulation isolates layout selection, although joint decisions with fusion and tiling can change the optimal assignment.Fusion before layout assignment can eliminate intermediates and their transposes, producing a different problem instance.
- Construct validity: An additive sequential-cost objective omits overlapping computation and data movement, which would require a scheduling-aware formulation.
6 Related Work
Prior compiler systems demonstrate that layout matters but mainly use specialized heuristics or representations. This paper adds formal complexity results and tractable or solver-based optimization perspectives.
- Heuristic layout optimization: XLA assigns layouts by locally optimizing influential operators and propagating choices through the computational graph.
- Heuristic layout optimization: SmartMem reduces mobile-GPU layout transformations through operator sensitivity classification, restricted layout domains, and rule-based propagation.
- Formal optimization: Existing systems rely on hand-tuned heuristics without formally defining the layout-selection problem they approximate.
- Formal optimization: The paper formally proves NP-hardness under realistic restrictions and provides an explicit dynamic program for bounded-treewidth graphs.The approach builds on parameterized-complexity machinery while avoiding the large constants of generic constructions.
- Formal optimization: Constraint-based systems use ILP or constraint programming, while this paper explains their need through hardness results and identifies tractable cases.
- Layout representations: Linear and integer-set frameworks represent layouts expressively, whereas this work studies the computational problem of selecting them.
- Related compiler optimizations: Layout selection intersects with loop transformations and scheduling, with interchange sometimes substituting for physical transposition on general-purpose hardware.
- Specialized accelerators: The formulation is complementary to accelerator compilers that generate systolic arrays or expose explicit memory hierarchies.
7 Conclusions
The paper formalizes layout selection as optimization over dataflow graphs, proves its hardness, and provides algorithms for bounded-treewidth and general instances. In production evaluation, solver methods work best when the cost model is reliable, while remaining gaps reflect modeling limitations.
- Layout selection is formulated as cost minimization over dataflow graphs, with NP-hardness results, a bounded-treewidth dynamic program, and a MaxSAT encoding.
- Solver-based methods match and sometimes outperform the production heuristic on transformer workloads but not on vision and multimodal models.
- Because the solver is optimal under the objective, performance gaps indicate cost-model error rather than search-quality limitations.
- Future work includes joint layout-and-tiling models, operator parallelism, hybrid guaranteed solvers, and learned cost models from hardware profiling.
A Details on Trainium Graph Compiler
Trainium’s NeuronCore maps tensors through partitioned on-chip buffers and a systolic array, making layout choices central to hardware utilization and conversion costs.
- Figure 8 depicts tensors with layout [P,F], where partitioned SBUF inputs feed systolic-array rows and PSUM accumulates outputs.
- Each NeuronCore combines a 128×128 systolic array with vector and scalar engines for tensor operations and reductions.
- The Neuron Graph Compiler maps computation graphs to hardware instructions while determining tensor tiling, scheduling, and on-chip placement.It keeps active tensors in SBUF because transfers between device memory and SBUF are expensive.
- SBUF exposes partition and free axes, with the partition axis distributing tensor elements across 128 hardware partitions.The selected dimension determines available hardware parallelism.
- The compiler must choose layouts that satisfy architectural constraints while minimizing costly layout conversions.
B Proof of Theorem 4
The proof establishes that the dynamic program computes optimal layout-selection cost by induction over a nice tree decomposition. It also derives polynomial running time for fixed treewidth and layout count, with the optimum recovered at the root.
- DP table: The dynamic program tracks bag labelings and distinct consumer-layout requests for tensors produced at the bag frontier.Each state OPT(i, f, S) stores the minimum subtree cost for a bag labeling f and request sets S.
- Inductive computation: OPT is computed bottom-up using leaf, introduce, forget, and join transitions over the tree decomposition.Introduce nodes initialize configurations, forget nodes record input requests and finalize produced tensors, and join nodes combine compatible child states.
- Correctness: Correctness follows by induction because each decomposition node preserves the invariant that OPT equals the minimum cost over compatible labelings and request sets.The proof handles leaves, introductions, forgets, and joins, using contiguity and separation properties of nice decompositions.
- Root conclusion: At the root, the dynamic-programming entry OPT(r, ∅, ∅) equals the optimum of Layout Selection.The root bag and request state are empty, so the invariant applies to the complete graph.
- Runtime: The running time is polynomial in the graph size for fixed treewidth and maximum number of layouts.The analysis bounds state combinations and transitions per bag, yielding polynomial dependence on |V| when tw and M are fixed.