Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using F_2

Keren Zhou, Mario Lezcano Casado, Adam P. Goucher, Akhmed Rakhmati, Jeff Niu, Justin Lebar and 5 others

To run a matrix multiply fast on a GPU, you have to keep shuffling the numbers between the arrangement memory likes and the arrangement the tensor-core hardware demands — and today's compilers handle these arrangements one hard-coded case at a time, needing a separate hand-written conversion for every pair, which is buggy and explodes as new chips arrive; this work from the Triton team shows that every such arrangement is really just a matrix of ones and zeros, so all the conversions become plain linear algebra that composes automatically, wiping out the case explosion, fixing whole classes of bugs, and speeding real kernels.

The argument in one glance

Problem

GPU tensor layouts are handled one hard-coded case at a time

Fast tensor computation depends on tensor layouts — how a logical tensor maps to hardware registers, threads, warps, and memory. A GPU wants data one way for memory access and another for its matrix-multiply units, so data is constantly converted between layouts. But compilers hard-code a fixed set of layouts and a separate hand-written conversion for each pair.
What people did

Add each layout and conversion by hand

Support layouts as special compiler attributes, and for each new layout write the code for it and the routines to convert it to and from every other layout.
Why it fell short

The conversions explode quadratically and are buggy

The number of layout-to-layout conversions grows with the square of the number of layouts, so adding a layout means heavy compiler surgery. The manual work is error-prone — 12 percent of the bugs filed against Triton are layout-related — and, without first-class layouts, compilers leave data movement suboptimal, forcing experts to hand-optimize what the compiler cannot.
This paper's idea

Model every layout as a binary matrix (Linear Layouts)

Linear Layouts treats each tensor layout as a linear map over the two-element field — a matrix of ones and zeros acting on the bits of the hardware and tensor indices. Swizzling and broadcasting become XOR and AND on bit-vectors, and any layout conversion is just composing matrices (multiply, invert), eliminating the quadratic explosion. Integrated into Triton, it automatically derives optimal data movement, fixes many bugs, and delivers up to a 1.40 times speedup (1.07 times on average) across 265 real benchmarks.

Why this problem is worth solving

The speed of a GPU tensor kernel turns on a detail most people never see: exactly where each number physically lives.

A tensor is a logical array of numbers, but to compute with it on a GPU, every number must physically reside somewhere concrete — in a particular register of a particular thread within a particular warp, or in a spot in shared memory. The mapping from the logical tensor to these hardware resources is called the tensor layout, and it is decisive for performance. The reason is that different parts of the GPU want data arranged differently. When loading from memory, you want a coalesced layout, where adjacent threads read adjacent addresses, so the hardware can fetch a wide chunk in one go and maximize bandwidth. But the tensor-core units that do fast matrix multiplication — driven by special instructions on NVIDIA, AMD, and Intel GPUs — require their inputs distributed across threads and warps in a specific, different layout to compute correctly. So a kernel is constantly rearranging data between the layout that is good for memory and the layout that is good for compute, an operation called a layout conversion. To make matters harder, the required layouts differ across GPU generations (NVIDIA’s Ampere, Hopper, and Blackwell each differ), across vendors, and across data types, so the space of layouts and conversions a compiler must handle is large and growing.

The trouble is that today’s tools handle this explosion of layouts the worst possible way — by hand, one case at a time.

Deep-learning practitioners who need peak performance rely on vendor libraries like cuDNN and cuBLAS, but those cover only a limited set of standard operators; a custom operator from a new model falls outside them, forcing developers to write GPU kernels from scratch and wrestle with layouts directly. Deep-learning compilers such as Triton, TVM, and XLA do better by building layouts into their backends, but they support only a limited, fixed set of layouts and of conversions between them. Defining a new custom layout requires substantial changes to the compiler, and — this is the crux — a conversion routine must exist for each pair of layouts, so the number of conversions grows with the square of the number of layouts: a quadratic explosion. Every one of those routines is hand-written and error-prone; the paper notes that fully 12 percent of the bugs filed against Triton’s repository are layout-related. And because these compilers do not treat layouts as a first-class object to optimize, they often leave data movement suboptimal, so expert authors resort to hand-tuning conversions the compiler cannot express — as the authors of FlashAttention 3 did, manually using byte-permute and warp-shuffle instructions to bypass shared memory. What is missing is a general, systematic, robust way to represent layouts and their conversions.

What people did before, and where it broke

Use vendor libraries for standard operations. Call cuDNN or cuBLAS to get peak performance on common operators. Where it broke. They cover only a limited set of operators; a custom operator forces writing kernels from scratch and dealing with layouts by hand.
Bake a fixed set of layouts into the compiler. Have compilers like Triton, TVM, and XLA support layouts as special attributes. Where it broke. Only a limited set of layouts and conversions is supported, and adding a new one requires heavy compiler changes.
Hand-write a conversion for each pair of layouts. Implement each layout-to-layout conversion routine individually. Where it broke. The number of conversions grows quadratically with the number of layouts, the manual work is error-prone (a large share of Triton’s bugs), and data movement is left suboptimal.
The gap. No framework represented all layouts and all their conversions in one general, composable, provably-correct formalism that a compiler could optimize automatically. Providing that, using linear algebra over a binary field, is what Linear Layouts does.

How Linear Layouts works

From hand-coded per-pair conversions to composable matrix algebra

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Turn a sprawl of special cases into one piece of mathematics. The heart of Linear Layouts is to stop treating each tensor layout as a bespoke special case and instead model every layout as a single kind of mathematical object: a linear map over the two-element field, a matrix of ones and zeros that acts on the bits of a hardware-resource index to produce the bits of a logical-tensor index — so swizzling and broadcasting become nothing more than XOR and AND operations on bit-vectors. Replacing ad-hoc, case-by-case reasoning about how the hardware behaves with one precise mathematical formalism is the same move the GPU memory-model work covered elsewhere on this site makes, capturing what the hardware does in an exact algebraic model so that a whole domain can be reasoned about uniformly instead of one situation at a time.

Compose conversions instead of enumerating them. Because every layout is a matrix, a conversion between any two layouts is simply the composition of their matrices — a matrix multiply and inverse — so the compiler computes arbitrary conversions on demand rather than storing a hand-written routine for each pair, which is exactly what dissolves the quadratic explosion that plagues layout systems. Generating a combinatorial space of cases from a small, general representation, rather than hand-writing every case, is the same strategy the relational-graph-network compiler work covered elsewhere on this site takes, producing what would be an unmanageable explosion of specialized kernels from a compact set of general templates instead.

Let the algebra reveal the fastest data movement, automatically. Having layouts as algebra lets the compiler search that structure for the best way to move data: a layout engine in Triton automatically derives swizzles that provably maximize read and write vectorization and minimize shared-memory bank conflicts, optimal warp shuffles, and generic lowerings of hardware instructions, so data movement is optimized generically rather than hand-tuned per kernel. Making the compiler automatically generate efficient GPU code and optimize data movement that experts otherwise tune by hand, cutting engineering effort, is the same aim the deep-learning-training compiler work covered elsewhere on this site pursues, turning laborious hand-optimization of GPU kernels into automatic, systematic code generation.

Following a layout conversion through Linear Layouts

Converting a memory layout into a tensor-core layout

A trace of Linear Layouts computing and optimizing one conversion, following the paper's design.

Step 1 of 5

Load in a memory-friendly layout

A kernel loads a tensor from global memory in a coalesced, bandwidth-friendly layout, but the tensor-core matrix-multiply instruction needs the data in a different, specific layout.

Two layouts, good for two different purposes, must be reconciled.

Why it works

Layouts are the key to tensor performance. Because a GPU forces data into different arrangements for memory access and for compute → tensors must be constantly converted between layouts → so how layouts and their conversions are handled determines performance.
One algebra replaces a sprawl of special cases. Because every layout can be written as a binary matrix acting on index bits → the messy, case-by-case handling collapses into linear algebra → so any layout in the family plugs in and is treated uniformly.
Composition kills the quadratic explosion. Because a conversion between two layouts is just the composition of their matrices → arbitrary conversions are computed by matrix multiply and inverse rather than hand-written per pair → so the number-of-layouts-squared explosion of code disappears.
Algebraic structure exposes the best data movement. Because swizzling and broadcasting are XOR and AND on bit-vectors within the algebra → the compiler can search that structure for the conversion that maximizes vectorization and avoids bank conflicts → so optimal hardware primitives are derived generically rather than hand-tuned.
First-class layouts are more correct and less work. Because layouts and conversions are computed from a formalism instead of coded case-by-case → whole classes of layout bugs vanish and backend effort drops → so the compiler is both more robust and easier to extend.

What it costs (trade-offs and scope)

The layout family is those expressible as linear maps. Linear Layouts covers any layout that can be represented as a permutation of indices or via swizzling — a broad and practically important family — but a layout that cannot be expressed as a linear map over the two-element field falls outside the framework.
The speedups are modest on average. The improvement is up to 1.40 times but averages 1.07 times across 265 benchmarks; much of the value is in correctness, robustness, and engineering effort rather than a uniformly large speedup, and common access patterns were already well served by the legacy heuristics.
It is a compiler-backend change tied to Triton. The work is implemented in and integrated with Triton’s GPU backend. Its benefits are realized through Triton (though the mathematical framework is general); other compilers would need their own integration.
The abstraction has a learning curve. Reasoning about layouts as binary matrices and bit-level linear algebra is powerful but unfamiliar, so adopting or extending the framework requires thinking in these terms rather than in concrete layout diagrams.
Optimality is within the modeled objectives. The automatic swizzling provably maximizes vectorization and minimizes bank conflicts, which are the right proxies for data-movement efficiency, but they are specific objectives; the true best kernel could depend on factors outside them.
Hardware coverage depends on modeling each target. Supporting a GPU’s special units means expressing their required layouts in the formalism and lowering to their intrinsics; the approach makes this systematic, but each new hardware family still must be brought into the model.

The evaluation

Three ways to handle tensor layouts for GPU kernels

Property	Vendor libraries	Legacy compiler layouts	Linear Layouts
Supports custom and new layouts
Layout-to-layout conversions	Not applicable
Correctness and robustness	High for supported ops
Automatic data-movement optimization	Internal to the library
Outcome	Peak but inflexible	Suboptimal and buggy	Up to 1.40x, average 1.07x, fewer bugs

Three ways to handle tensor layouts for GPU kernels, as framed by the paper. Vendor libraries (cuDNN, cuBLAS) are fast but cover only standard operators. Legacy compiler layout systems (as in Triton before this work) bake in a limited set of layouts and hand-write a conversion routine for each pair, so conversions grow with the square of the number of layouts, adding a layout requires heavy changes, and the manual code is error-prone (12 percent of Triton's bugs are layout-related) with suboptimal data movement. Linear Layouts models every layout as a matrix of ones and zeros over the two-element field acting on index bits, so swizzling and broadcasting are XOR and AND operations, and any conversion is computed by composing matrices (multiply and inverse) — eliminating the quadratic explosion. Integrated into Triton as a layout engine that automatically chooses and propagates layouts, it derives provably optimal swizzles (maximizing vectorization, minimizing shared-memory bank conflicts) and warp shuffles. It improves correctness, fixes many pre-existing Triton bugs, reduces backend engineering effort, and yields up to a 1.40 times speedup, averaging 1.07 times across 265 real-world benchmark cases.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

Linear Layouts is a beautiful example of the right mathematical abstraction dissolving a messy engineering problem. GPU tensor kernels live or die by tensor layouts — how a logical tensor maps onto registers, threads, warps, and memory — and because memory access and tensor-core compute want different layouts, kernels constantly convert between them, across a space of layouts that keeps growing with every GPU generation, vendor, and data type. Handling this the old way, with a hand-written conversion routine for each pair of layouts, produces a quadratic explosion of error-prone code, and indeed a large fraction of Triton’s bugs were layout-related. The insight is that every one of these layouts is really a linear map over the two-element field — a binary matrix acting on index bits — so swizzling and broadcasting are just XOR and AND, and any conversion is a matrix composition. That single reframing collapses the case explosion into linear algebra: conversions are computed on demand, any layout in the family plugs in, and the algebraic structure can be searched to derive provably optimal data movement, including swizzles that maximize vectorization and avoid bank conflicts. Integrated into Triton — a widely-used, real, multi-vendor GPU compiler — it improves correctness, fixes many bugs, reduces backend engineering effort, and speeds real kernels, up to 1.40 times and 1.07 times on average across 265 benchmarks. Read the scope precisely. The framework covers layouts expressible as index permutations or swizzles, a broad but not universal family; the average speedup is modest, with much of the value in correctness and maintainability rather than raw performance; it is realized through Triton’s backend specifically; and reasoning in binary matrices has a learning curve. The durable idea is that when a systems problem has devolved into an unmanageable pile of special cases, the highest-leverage move is often to find the mathematical structure underneath it — here, linear algebra over a binary field — because the right abstraction turns hand-written case analysis into composition, and case analysis into proof.

Background

Why finding the underlying math tames systems complexity

It helps to see Linear Layouts within a recurring pattern: some of the most durable advances in systems come not from a faster algorithm but from recognizing that a tangle of special cases is really one clean mathematical object in disguise, at which point the special cases become instances of a general theory and the hardest operations become routine composition. Here the tangle is the combinatorial mess of GPU tensor layouts and their pairwise conversions, and the underlying object is a matrix over the two-element field; once seen that way, the quadratic explosion of hand-written conversions becomes matrix multiplication, and finding the best data movement becomes a search over algebraic structure. The tensor-computation compiler work covered elsewhere on this site reflects the same instinct, treating the mapping of tensor computation onto hardware as a principled abstraction to be optimized rather than a bag of tricks, so that performance follows from the structure rather than from manual tuning. The common thread is that in performance-critical systems the growth of hardware variety and workload complexity tends to produce case-by-case sprawl, which is both slow to extend and a breeding ground for bugs, and that the antidote is a representation general enough to absorb the variety and precise enough to compute and prove over. When such a representation is found, it pays three ways at once: it removes redundant engineering, it eliminates whole classes of errors, and it exposes optimizations that were invisible in the case-by-case view. As accelerators proliferate and their layouts multiply, the ability to name the algebra underneath them is likely to be a decisive tool for keeping their software both correct and fast.

If you remember only three things

1. Fast GPU tensor computation depends on tensor layouts — the mapping from a logical tensor to physical hardware resources (which register of which thread in which warp, or which spot in shared memory holds each element). Layouts are decisive because a GPU wants data coalesced for memory bandwidth but arranged differently for its tensor-core matrix-multiply units, so kernels constantly convert between a memory-friendly and a compute-friendly layout, and the required layouts differ across GPU generations, vendors, and data types. Deep-learning compilers handle this poorly: they support a fixed set of layouts as special attributes and hand-write a conversion routine for each pair, so the number of conversions grows with the square of the number of layouts, adding a layout requires heavy compiler changes, and the manual code is error-prone — 12 percent of the bugs filed against Triton are layout-related.
2. Linear Layouts models every tensor layout as a linear map over the two-element field: a matrix of ones and zeros that acts on the bits of a hardware-resource index to produce the bits of the logical-tensor index. In this formalism, complex arrangements like swizzling and broadcasting are just XOR and AND operations on bit-vectors, and any layout-to-layout conversion is computed by composing the two layouts' matrices (matrix multiplication and inverse) — so conversions are generated on demand rather than hand-written per pair, eliminating the quadratic explosion. Integrated into Triton's GPU backend as a layout engine that automatically chooses and propagates layouts for any operation, it introduces algorithms for provably optimal swizzling (maximizing read/write vectorization, minimizing shared-memory bank conflicts), optimal warp-shuffle generation, and generic lowering of hardware intrinsics.
3. Across 265 real-world benchmark cases, Linear Layouts yields up to a 1.40 times speedup (1.07 times on average) over legacy Triton, while improving correctness by fixing many pre-existing layout bugs and reducing engineering effort in the compiler backend. Scope: it covers layouts expressible as index permutations or swizzles (broad but not universal), the average speedup is modest with much of the value in correctness and maintainability, it is realized through Triton specifically, and reasoning in binary matrices has a learning curve. The durable lesson is that when a systems problem has devolved into a pile of hand-written special cases, the highest-leverage move is often to find the mathematical structure underneath it — here, linear algebra over a binary field — because the right abstraction turns case-by-case code into composition and case analysis into proof.

Check yourself

Check your understanding

1. What is a tensor layout, and why does it matter for GPU performance?
2. Why does the traditional way of handling layouts scale so badly?
3. What is the key idea of Linear Layouts?
4. Beyond raw speed, what does integrating Linear Layouts into Triton achieve, and what are its results?

Sources

The paper is “Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using F_2” by Keren Zhou, Mario Lezcano Casado, Adam P. Goucher, Akhmed Rakhmati, Jeff Niu, Justin Lebar, Pawel Szczerbuk, Peter Bell, Phil Tillet, Thomas Raoux, and Zahi Moudallal, published at ASPLOS 2026, DOI 10.1145/3760250.3762221. The authors are affiliated with OpenAI and George Mason University.

This explainer is based on the full text of the paper, obtained from the authors’ arXiv version (arXiv:2505.23819), which corresponds to the CC-BY version of record. The problem framing (tensor layouts as mappings between logical tensors and GPU hardware resources; the need to convert between coalesced memory-friendly layouts and tensor-core compute layouts; the proliferation of layouts across GPU generations, vendors, and data types; the limitations of vendor libraries and of deep-learning compilers that support a fixed set of layouts with hand-written per-pair conversions leading to a quadratic explosion; the report that 12 percent of Triton’s bugs are layout-related; and the example of FlashAttention 3 hand-optimizing conversions), the Linear Layouts approach (modeling every layout as a linear map over the two-element field, a binary matrix acting on the bits of the indices; expressing swizzling and broadcasting as XOR and AND on bit-vectors; composing arbitrary conversions via matrix multiplication and inverse to eliminate the quadratic explosion; and integration into Triton’s GPU backend as a layout engine that automatically chooses and propagates layouts), the novel algorithms (automatic optimal swizzling discovery that provably maximizes read/write vectorization and minimizes bank conflicts, automatic optimal warp-shuffle generation, and generic lowering of hardware intrinsics), and the results (up to a 1.40 times speedup with an average of 1.07 times across 265 real-world benchmark cases, alongside improved correctness through fixing many pre-existing Triton bugs and reduced backend engineering effort) are all drawn directly from the paper. Background explanations of GPUs, threads and warps, tensor cores, shared memory and bank conflicts, coalesced access, and linear algebra over a binary field are standard and provided to make the paper’s contribution understandable.