首次为机器学习编译器的张量布局选择建立形式化框架,提升模型执行效率。
Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers

- 将布局选择建模为数据流图上的组合优化问题,兼顾计算与转换开销。
- 在有界树宽图上设计多项式时间最优算法,在一般情况用MaxSAT求解。
- 实测显示简单启发式方法性能下降最高达5倍,验证了形式化方法的有效性。
现代机器学习编译器通过选择张量内存布局来最小化执行成本,但布局选择是全局性的:一个算子在某种布局下最快,其下游却可能偏好另一种布局,需显式转换,而转换会损害性能。尽管重要,当前布局选择缺乏形式基础,依赖经验启发式。本文首次对机器学习编译器中的布局选择进行形式化研究,将其建模为数据流图上的组合优化问题,目标是最小化算子执行成本与每张量转换成本之和。理论分析表明,即使仅含二维张量矩阵乘法,最优布局选择也属计算困难问题。针对有界树宽数据流图,设计了多项式时间最优算法;对一般情况,提出可由现成求解器优化的加权MaxSAT编码。该框架统一了XLA布局分配、流式阵列编译器的分区维度选择及移动GPU优化器的布局规划等策略。我们在一款AI加速器生产编译器中实现该形式化,对比贪心启发式、规则驱动策略与最优求解器的执行时间。简单启发式在部分工作负载上使执行时间恶化高达5倍;当编译器成本模型准确时,求解器表现不劣于或优于规则策略;在复杂数据移动工作负载上稍逊,但因求解器精确最小化目标,该差距揭示了成本模型误差而非搜索质量的问题,明确了编译器优化的实际价值所在。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。