Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers (arxiv.org)

A new study has formalized the process of layout selection for machine learning (ML) compilers, addressing a longstanding issue where layout choices can significantly impact execution costs. Current compilers rely on ad-hoc heuristics that can lead to inefficiencies, especially when operators have conflicting layout preferences, necessitating costly conversions. By framing layout selection as a combinatorial optimization problem over dataflow graphs, the researchers have proposed both an optimal polynomial-time algorithm for specific graph types and a weighted MaxSAT encoding for more general scenarios, marking a significant advancement in how tensor memory layouts are managed. This formalization not only unifies various existing layout optimization strategies—such as those used in XLA and mobile GPU optimizers—but also enhances model performance by minimizing execution time effectively. In practical tests with a compiler for an AI accelerator, the researchers found that simple heuristics could degrade performance by up to five times compared to optimal approaches. Their findings demonstrate that when the compiler's cost model is accurate, the formally optimized solutions outperformed traditional methods, highlighting critical areas where improved compiler strategies can lead to tangible gains in AI workload execution.