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.
Figures & tables
Figure 1. Two layout assignments for the computation in (a). Layouts are row-major (RM) or column-major (CM); dimensions are in storage order. (b) All tensors remain in RM. (c) Inputs are transposed to CM, and the output tensor R is transposed back. Three panels: a small tensor program and two layout assignments for it Panel~(a) lists a program func f(A,B) computing C=matMul(A,B), then D=reduce(C,dim=0), then R=add(D,B), returning~R. Panels~(b) and~(c) show the same dataflow graph twice, with operators as boxes, tensors as labelled arrows, and each tensor's element grid drawn beside it in storage order. In~(b) every tensor is row-major: inputs A[4×8] and B[8×3] feed matMul, producing C[4×3]; reduce in its RM form produces D[1×3]; add combines D with B (broadcasting D over the rows of B) to give R[8×3], and no conversion is needed. In~(c) the inputs arrive transposed as A⊤[8×4] and B⊤[3×8], so matMul yields C⊤[3×4], the cheaper CM form of reduce yields D⊤[3×1], and add yields R⊤[3×8]; a final transpose operator, shaded differently from the compute operators, converts R⊤ back to R[8×3].
Figure 2. Local versus global layout selection on a graph of Conv and BatchNorm operators feeding a shared Add . Layouts are NHWC or NCHW ; inserted Transpose operators are shown explicitly. Legend plus two layout assignments of a fan-out/fan-in graph, one with transposes and one without The leftmost panel is a legend of four operator boxes: convolution (prefers NHWC), batch norm (prefers NCHW), add (accepts NHWC or NCHW), and transpose. The two remaining panels each show the same graph: a single Input node fanning out to a row of alternating convolution and batch-norm operators (an ellipsis marks the omitted middle ones), which all fan back in to one shared Add that feeds Output. In the locally-optimal assignment~(b), each operator takes its own preferred layout, so the convolutions are NHWC and the batch-norms NCHW; because the shared Add runs in NHWC, a transpose box is inserted on the outgoing edge of every batch-norm. In the globally optimized assignment~(c), every operator including the batch-norms runs in NHWC, so all edges reach the Add directly and no transpose boxes appear at all.
Figure 3. Reduction from Odd Cycle Transversal to Layout Selection (Theorem 3 ). (a) An undirected graph H with a triangle {u,v,w} . (b) The corresponding ML program: each vertex becomes a tensor with two layouts {A,B} , and each edge becomes a matmul requiring its inputs in different layouts (cost 0 ) or infeasible (cost ∞ ). The triangle forms an odd cycle of opposite-layout constraints that cannot be two-colored with one layout per tensor; at least one tensor must be read in both layouts, incurring a transpose cost of 1 . A four-vertex graph and the layout-selection program built from it Panel~(a) draws an undirected graph on vertices u, v, w, x with five edges: {u,v}, {v,w}, {u,w}, {u,x} and {w,x}. Vertices u, v, w form a triangle, that is, an odd cycle, and u, w, x form a second one. Panel~(b) lists the constructed program func f(): first one source operator per vertex, producing tensors tu, tv, tw, tx, each with layout drawn from {A,B}; then one matmul per edge, ruv=matmul(tu,tv), rvw=matmul(tv,tw), ruw=matmul(tu,tw), rux=matmul(tu,tx) and rwx=matmul(tw,tx), each annotated with the edge it encodes. Two trailing comments record the costs: a matmul costs 0 if and only if its two input layouts differ, and the transpose cost between A and B is 1 in either direction.
Figure 4. Treewidth-based layout selection applied to the dataflow graph from Figure 1 . Each node in the tree decomposition is labeled by type (leaf, introduce, forget, or join), and annotated with its bag contents. The algorithm of Theorem 4 traverses this decomposition bottom-up: introduce nodes record operator configurations, forget nodes charge operator and transpose costs, and the join node combines two subtrees. A four-operator dataflow graph and a nice tree decomposition of it with bag contents Panel~(a) redraws the dataflow graph of Figure~\ref{fig:layout-example} with abbreviated names: input tensors A and B feed the matmul mm, whose result feeds the reduction red, whose result feeds add; B also feeds add directly by a long edge on the right, and add produces R. Panel~(b) draws a nice tree decomposition of width~2, so every bag holds at most three elements. The root is a leaf-labelled bag {B}; below it an introduce `add`'' step gives $\{B,\texttt{add}\}$, which is a join node with two children. The left branch goes $\{B,\texttt{add}\}$, then forget B, introduce add'' to {R,add}, then forget $R$'' to the leaf $\{\texttt{add}\}$. The right branch goes $\{B,\texttt{add}\}$, then introduce red'' to {B,add,red}, then forget `add`, introduce `mm`'' to $\{B,\texttt{red},\texttt{mm}\}$, then forget red, introduce A'' to {A,B,mm}, then forget $A$'' to $\{B,\texttt{mm}\}$, then forget mm'' to the leaf {B}.
Figure 5. RQ1 (Practical Impact of Layout Selection): Relative model performance (as fraction of best known result) on Trainium instances. Two cumulative-distribution plots of relative model performance, over all models and over dense transformers only Both panels plot model performance on hardware as a percentage of the best known result (y, 0 to 100) against the percentage of models (x, 0 to 100), as one non-decreasing curve per algorithm: \alg{Rule-Based} in green, \alg{Greedy} in purple and \alg{Local} in red, with a dashed horizontal line at 100%. Higher and earlier curves are better. In panel~(a), over all models, the green \alg{Rule-Based} curve starts near 71% and rises steeply, exceeding 95% within the first few percent of models and staying essentially flat at 100% thereafter. The purple \alg{Greedy} and red \alg{Local} curves start below 20%, remain under 30% for the first tenth of models, then climb gradually, with \alg{Greedy} consistently above \alg{Local} by roughly ten percentage points through the middle of the range; both converge to the green curve only past about 60% of models. Panel~(b), restricted to dense transformers, shows the same green curve but a much smaller spread: \alg{Greedy} and \alg{Local} nearly coincide, rise past 80% within the first tenth of models, and track just below \alg{Rule-Based} for the rest of the range.
Figure 6. RQ2 (Solver-based Algorithms in Practice): Most real-world dataflow graphs exhibit low treewidth, yet Treewidth exceeds the memory budget on 8% of instances, whereas MaxSAT remains practical. A bar chart of the treewidth distribution and a cumulative runtime plot for the two exact algorithms Panel~(a) is a bar chart of the percentage of dataflow graphs (y, 0 to 60) by treewidth (x, 1 to 8), with each bar labelled: 5.3% at width~1, 54.2% at~2, 25.2% at~3, 7.6% at~4, 4.6% at~5, 1.5% at~6, and 0.8% at each of~7 and~8; the distribution is thus concentrated on widths 2 and~3, which together account for about four fifths of the graphs. Panel~(b) plots runtime in seconds on a logarithmic y~axis from 0.001 to 600 against the percentage of models (x, 0 to 100), with one non-decreasing curve for \alg{Treewidth} in blue and one for \alg{MaxSAT} in orange. The blue \alg{Treewidth} curve sits roughly one order of magnitude below the orange \alg{MaxSAT} curve across the whole middle of the range, staying under a tenth of a second for about 70% of models, but then jumps vertically to the 600-second ceiling at about 90% of models, marking the instances that exhaust the memory budget. The orange \alg{MaxSAT} curve rises more smoothly, stays below ten seconds for about 90% of models, and reaches the same ceiling only in the last few percent.
Type
Model
Local
Greedy
Rule-Based
Solver-Based
Dense Transformers
ALBERT-large
53.4
45.8
44.7
45.3
BERT-base
19.2
19.3
18.5
18.5
BERT-large
29.0
29.6
28.7
28.7
BART-base
2.4
2.6
2.4
2.4
DistilBART
2.2
2.3
2.2
2.1
Electra-large
23.2
23.7
23.0
23.0
Table 1. Model execution time (millions of cycles). Lower is better. The best result per model, and every result within 1% of it, is bold.
Figure 7. RQ4 (Cost Model Fidelity): Algorithms achieve near-optimal objective values under the cost model, yet predicted costs can diverge from actual hardware execution time. A cumulative plot of optimality gaps and a heatmap of predicted cost against measured performance Panel~(a) plots the optimality gap in percent (y, 0 to 15) against the percentage of models (x, 0 to 100), as one non-increasing curve per algorithm: \alg{Solver-Based} in blue lies flat on 0% by definition; \alg{Greedy} in purple starts at about 0.5% and is indistinguishable from zero past the first tenth of models; \alg{Local} in red starts near 12% and falls below 1% by about a quarter of models; \alg{Rule-Based} in green starts near 7% and decays most slowly of all, staying above the other heuristics from roughly 30% of models onward. All curves are therefore within a few percent of the optimum over most of the suite. Panel~(b) is a 5×5 heatmap whose x~axis is the normalized predicted layout cost from Min to Max and whose y~axis is the normalized measured model performance on hardware from Min to Max, with cell shading giving the frequency of algorithm-model points and a dashed diagonal marking perfect agreement. Two dark cells sit on the diagonal, at the lowest-cost/best-performance corner and the highest-cost/worst-performance corner, but two further clusters of mass sit off it: a moderately dark cell at lowest predicted cost with worst measured performance, and a shaded cell at highest predicted cost with best measured performance. The interior cells are all pale, so the distribution is concentrated in the corners rather than along the diagonal.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 8. The NeuronCore tensor engine: a 128×128 systolic array connected to two on-chip SRAMs. The state buffer (SBUF) stores input tensors partitioned across 128 partitions, each feeding one row of the array. The partial sum buffer (PSUM) accumulates output results. A tensor (red) has layout [P,F]. Illustration adapted from ( 5 ) . Block diagram of the tensor engine with the SBUF and PSUM on-chip SRAMs A schematic with three rectangular blocks. The tensor engine, shaded green, is a square block annotated 128~rows along its left edge and 128~columns along its top edge. To its right sits the state buffer SBUF, shaded blue and annotated 128~partitions along its right edge; an arrow runs leftward from SBUF into the tensor engine, one partition feeding each row of the array. Below the tensor engine sits the partial sum buffer PSUM, also shaded blue and annotated 128~partitions along its bottom edge, with an arrow running downward from the tensor engine into it. A tensor resident in SBUF is highlighted in red to show its layout \layout{P,F}, with its first dimension spread across the partitions and its second dimension laid out along the free axis within each partition.
Tensor program optimization is essential for modern machine learning systems, but its search space is enormous. Existing auto-schedulers reduce measurement cost with learned cost models, yet they usually evaluate each candidate as a static code snapshot, ignoring the schedule trajectory that produced it. This makes them insensitive to action dependencies and vulnerable to superficial code variations. We propose a world-model-inspired evaluator that models schedule evaluation as action-conditioned latent dynamics over program states. Starting from the initial program, it rolls out scheduling actions in a continuous latent space with a lightweight transition model, avoiding expensive AST mutation and repeated code encoding. The final dynamic representation is combined with action and hardware features to rank candidates. Implemented in TVM AutoScheduler, our method improves representative-subgraph latency over Ansor by 1.37× on GPU and 1.54× on CPU under the same 64-trial budget. It also matches Ansor-10K within 2.2% geometric mean using 10× fewer measurements, and accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61×/3.67× geometric mean.
Haolin Pan, Lianghong Huang, Xvlin Zhou +2
Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Hangzhou, China · Institute of Software, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.
Optimized kernels such as FlashAttention and FlashDecoding are crucial for accelerating today's large models. Most of them are handwritten by experts because existing ML compilers cannot match their efficiency. Producing such kernels requires fusing computations with multiple reductions, which requires both algebraic transformation of the computation graph and operator scheduling of the transformed graph. Unfortunately, searching the two jointly yields a space too large to navigate. We propose Cleave, an ML compiler built on symbolic decoupling: Cleave discovers transformations by performing superoptimization on a graph with symbolic shapes, and then schedules each resulting graph on concrete shapes. Representing shapes as symbols makes equivalence checking cheap and lets a new Split operator, with a symbolic split count, parallelize along a reduction dimension. Cleave's scheduler fuses graphs with multiple reductions through iterative tiling and horizontal fusion. Evaluation on common LLM subgraphs shows that Cleave generates kernels up to 2.8x faster than the best baseline (1.6x on average) and reduces compilation time by 5.9x on average compared to Mirage. For dynamic workloads captured from production serving traces, Cleave compiles each operator once and achieves geometric mean speedups of 1.4x and 1.7x over FlashInfer's handwritten FA2 and FA3 backends. Cleave's code is available at: https://github.com/nyu-systems/cleave