# Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

arXiv: https://arxiv.org/abs/2608.21555 · HTML: https://mlsanity.com/paper/2608.21555/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 66 references that are in the map.
- First posted on arXiv 21 Aug 2026.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks](https://mlsanity.com/paper/1808.07412/) (2018)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [A Survey on Compiler Autotuning using Machine Learning](https://mlsanity.com/paper/1801.04405/) (2018)
- [Optimizing CNN Model Inference on CPUs](https://mlsanity.com/paper/1809.02697/) (2018)
- [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://mlsanity.com/paper/2308.13490/) (2023)
- [SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile](https://mlsanity.com/paper/2404.13528/) (2024)
- [oneDNN Graph Compiler: A Hybrid Approach for High-Performance Deep Learning Compilation](https://mlsanity.com/paper/2301.01333/) (2023)
- [Hardware Acceleration for Neural Networks: A Comprehensive Survey](https://mlsanity.com/paper/2512.23914/) (2025)
- [FlexSA: Flexible Systolic Array Architecture for Efficient Pruned DNN Model Training](https://mlsanity.com/paper/2004.13027/) (2020)
- [ArrayFlex: A Systolic Array Architecture with Configurable Transparent Pipelining](https://mlsanity.com/paper/2211.12600/) (2022)
- [Autotuning OpenCL Workgroup Size for Stencil Patterns](https://mlsanity.com/paper/1511.02490/) (2015)
- [Improved Basic Block Reordering](https://mlsanity.com/paper/1809.04676/) (2018)
- [Joint Program and Layout Transformations to enable Convolutional Operators on Specialized Hardware based on Constraint Programming](https://mlsanity.com/paper/2104.04731/) (2021)
- [Modeling Layout Abstractions Using Integer Set Relations](https://mlsanity.com/paper/2511.10374/) (2025)

## Abstract (arXiv)

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× 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.