发表机构
TU Wien; Amazon（维也纳技术大学; 亚马逊公司）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文首次形式化机器学习编译器的布局选择问题，提出组合优化模型与算法，实现多种现有策略的统一，在AI加速器编译器上的实验验证了其效果，可分离成本模型误差与搜索质量。
AI 中文摘要
现代机器学习编译器选择张量内存布局以在硬件约束下最小化执行成本。布局选择是全局的：一个算子在某一布局下可能最快，而其消费者则偏好另一布局，协调这些偏好需要显式的布局转换，而这可能损害模型性能。尽管布局选择具有实际重要性，但它缺乏形式化基础，因此当前编译器依赖临时启发式方法。本文首次对机器学习编译器中的布局选择进行形式化研究，将该问题表述为数据流图上的组合优化问题，最小化算子执行成本与每个张量转换成本之和。理论分析表明，即使对于仅包含二维张量矩阵乘法的程序，最优布局选择也具有计算难度。针对树宽有界的数据流图，设计了最优多项式时间算法；对于一般实例，给出了加权MaxSAT编码，可由现成求解器优化。该形式化统一了多种现有布局优化策略，包括XLA的布局分配、脉动阵列编译器中的分区维度选择，以及移动GPU优化器中的布局规划。我们在AI加速器的生产编译器中实现了该形式化，并在贪心启发式、编译器基于规则的策略和最优求解器下测量编译模型的执行时间。简单启发式在部分工作负载上使执行时间降低达5倍；在编译器成本模型准确时，求解器表现优于或匹配基于规则的策略；在数据移动复杂的工作负载上，求解器表现落后，而由于求解器精确最小化所述目标，该差距将成本模型误差与搜索质量隔离开，明确了编译器优化的实际发力方向。
英文摘要
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.
