Modern GPUs compute and move multi-dimensional tensors in hardware, but the languages that reach those instructions (CUDA C++ and PTX assembly) still speak in flat one-dimensional buffers and scalar thread indices — so the fastest instructions, which demand intricate mappings of data to specific thread groups, are accessible only through index-arithmetic contortions that even experts struggle to write and no existing tensor IR can express; Graphene (NVIDIA) makes both data AND threads first-class hierarchically-tileable tensors, so an optimized kernel becomes a mapping between data tiles and thread tiles that decomposes step by step down to the hardware instructions themselves — expressive enough to exactly match cuBLAS at peak, to fuse kernels no library offers (up to 2.39× over chained library calls), and to edge out NVIDIA's own hand-written MLPerf attention kernels.
The argument in one glance
-
Problem
The hardware became tensor-shaped; the software stayed flat
Driven by deep learning, NVIDIA GPUs grew instructions that operate on whole multi-dimensional tensors: Volta's Tensor Cores multiply small matrices in hardware, and later architectures added instructions like ldmatrix that move entire tiles through the memory hierarchy in one shot. These instructions are MANDATORY for peak performance — replacing ldmatrix with simpler equivalent moves costs up to 17% in matrix-multiply kernels. But they are exposed only in PTX (NVIDIA's virtual assembly), reachable only through CUDA C++ with inline assembly, and CUDA C++ has no notion of a multi-dimensional tensor — of data, or of threads.
-
What people did
Three escape hatches, all in use today
The ecosystem routes around the gap in three ways. Call vendor libraries (cuBLAS, cuDNN) and let NVIDIA's experts write the PTX. Expose tensor instructions as built-in operations inside compiler IRs (TVM's tensorize, MLIR's GPU dialect). Or offer a friendly surface language — OpenAI's Triton — and let heavyweight compiler passes do the lowering.
-
Why it fell short
Punt, under-express, or bury
Libraries just relocate the problem — someone still writes the CUDA/PTX — and no library exists for novel fused operations, which is where the big wins now live. Built-in-op IRs describe tensors as dimension-and-stride integer lists, which structurally CANNOT express what the fast instructions demand: swizzled memory layouts that dodge shared-memory bank conflicts, and mappings onto specific thread groups — ldmatrix wants a warp reshaped into 2×2 groups of eight threads, Volta's Tensor Core instruction wants "quad-pairs" of eight NON-CONTIGUOUS threads (threads 0–3 joined with 16–19). And DSL compilers like Triton hide exactly this knowledge inside instruction-selection passes, so extending them requires intimate knowledge of both the GPU and the compiler's internals.
-
This paper's idea
Make threads tensors too, and let the kernel be a mapping between tiles
Graphene (NVIDIA) is a low-level IR where data AND threads are both first-class tensors that tile hierarchically — a tile is just a smaller nested tensor, and a dimension's size and stride can themselves be tuples, which is enough to express every swizzled layout and non-contiguous thread group real kernels use. A warp is a tensor of 32 threads; reshape it into 2×2 groups of 8 exactly the way you tile a matrix. Computations are "specs" — declared mappings from input tensors to output tensors under a thread configuration — that decompose hierarchically until every leaf matches an "atomic spec," the declared signature of an actual hardware instruction. Code generation is then just printing. Graphene exactly matches cuBLAS at practical peak on Volta and Ampere, fuses kernels no library provides (multi-layer MLP up to 2.39× over chained cuBLASLt calls, LSTM cell 1.82×), slightly beats NVIDIA's hand-written MLPerf BERT attention kernels, and speeds up real Transformer inference by up to 59%.
Why this problem is worth solving
Every deep-learning framework, compiler, and hand-tuned
kernel ultimately funnels through one narrow doorway: the
hardware operations, exposed
only in PTX (NVIDIA’s virtual assembly language), that
compute or move whole multi-dimensional tiles at once —
Tensor Core matrix-multiply-accumulate instructions, and
data movers like ldmatrix, which uses all 32 threads of a
warp cooperatively to pull four 8×8 matrices from shared
memory into the threads’ registers in a single
instruction
CUDA C++ can only express these contracts as scalar index
arithmetic on flat one-dimensional buffers — modulo-and-divide
soup plus address-space casts plus inline assembly — and the
paper’s opening figure shows the result: a conceptually simple
picture (tile the warp, tile the data, match them up) becomes
nineteen lines of obfuscated indexing that experts find hard
to write and compilers find hard to reason about. Meanwhile
memory arrangements beyond
row-major and column-major, where elements are interleaved in
patterns designed so that simultaneous accesses from
different threads land in different shared-memory banks. A
bank can serve one thread at a time; conflicting accesses
serialize and stall the pipeline. Peak-performance kernels
routinely store intermediate tiles in such layouts — and IRs
that describe a tensor as one stride per dimension cannot
represent them, because one stride per dimension forces every
dimension to be traversed wholly before the next
What people did before, and where it broke
- Vendor libraries (cuBLAS, cuBLASLt, cuDNN). The default answer, and the performance reference. But this only relocates the authorship problem to library engineers writing CUDA/PTX by hand — and libraries enumerate fixed menus. The fastest technique in modern inference is fusing operator chains into single kernels, and no library can pre-build every fusion a model might need.
- Tensor instructions as compiler built-ins (TVM’s tensorize, MLIR’s GPU dialect, AMOS). These describe what an instruction computes, but not the thread arrangements it demands — none represents executing threads explicitly, so quad-pairs and ldmatrix-style mappings are out of reach. Their stride-list tensor descriptions also cannot state swizzled layouts.
- High-level DSLs with heavyweight lowering (Triton). Programmers write friendly tile-level Python; instruction selection happens inside complex transformation passes. The knowledge this paper cares about exists there — but buried in compiler internals, extensible only by people who know both the GPU architecture and the compiler deeply.
- Fireiron (the direct ancestor). Introduced decomposable specifications — but only for matrix multiply, and only for two-dimensional matrices.
- Distributed-training mesh frameworks (Mesh TensorFlow, GShard, GSPMD). Pioneered mapping data dimensions onto meshes of devices — the inspiration for threads-as-tensors — but their meshes are flat, and GPU compute is inherently hierarchical (grids of blocks of warps of threads).
How Graphene works
Data tensors, thread tensors, and specs that decompose to instructions
-
Tensors with hierarchical shapes new
Every tensor has a name, shape (dimensions AND strides), element type, and memory space (global, shared, or registers). The two recursions do all the work. First, a dimension's size and stride can be TUPLES — multiple strides per dimension — which expresses interleaved and swizzled layouts far beyond row/column-major, while logical 2-D coordinates keep working unchanged: declare the layout once at allocation, never touch an access. Second, an element type can itself be a shape — that IS a tiled tensor, outer shape arranging the tiles, inner shape arranging elements within each tile.
-
Tiling as the universal verb new
A. tile([2],[4]) carves a 4×8 tensor into 2×4 tiles — and tile sizes are themselves tiny 1-D tensors WITH STRIDES, so tile([2:2],[4]) builds tiles from every OTHER row: non-contiguous tiles, in one or both dimensions. That expressiveness is not exotic — it is exactly what ldmatrix's data-to-thread contract requires. Partial tiles (sizes that do not divide evenly) are handled by overapproximation plus predicated accesses; shapes can be symbolic (M, N) for dynamic networks.
-
Logical thread groups: the compute hierarchy is a tensor new
A warp is #w: [32].thread — a tensor whose scalar type is thread rather than fp16. Tile and reshape it with the SAME operations as data: four groups of eight, rearranged 2×2, is two lines. Volta's quad-pairs — eight non-contiguous threads, 0–3 with 16–19 — fall out as a single non-contiguous tiling [(4,2):(1,16)]. No built-in hierarchy zoo: quad-pairs appeared with Volta and vanished in later architectures, so hard-coding them into an IR is exactly wrong; logical thread groups express any arrangement, past or future, and the scalar index expressions are derived automatically at code generation.
-
Specs: computations as declared tensor mappings new
A spec states outputs ← Spec ⟨blocks, threads⟩ (inputs), optionally followed by an implementation in braces. Eight built-ins cover GPU computation: Move, MatMul, unary/binary pointwise, Reduction, Shfl (data exchange within thread groups), Init, Allocate — plus a generic Spec for anything else, including arbitrary fusions, whose meaning is entirely its decomposition.
-
Hierarchical decomposition down to atomic specs new
A kernel-level spec (say, a device-wide GEMM) is implemented by decomposing into finer specs on TILES of data and thread tensors — block-level multiplies, warp-level moves — until every remaining leaf matches an ATOMIC spec: a pre-declared signature pairing tensor shapes and thread arrangement with a hardware instruction. A Move of one fp32 scalar per thread is ld.global.u32; a warp-wide Move of shared-memory rows into 2×2-tiled register pairs IS ldmatrix; MatMul on quad-pairs with 4×1 and 1×4 fp16 register tiles is Volta's mma.m8n8k4. The instruction's brutal contract becomes a type signature.
-
Code generation is printing
Because the IR states every tile, layout, and thread assignment explicitly, there is no clever lowering left to do: control flow prints as CUDA C++, atomic specs print as their instruction (inline PTX where needed), and tensor accesses compile into automatically simplified scalar index arithmetic. Graphene is generated from a Python API today — too verbose to hand-write — and is positioned as a TARGET language for ML compilers like XLA, TVM, or Triton.
The idea
Name the thing every fast kernel is secretly about. The paper’s diagnosis is that an optimized GPU kernel is not really a program over flat buffers — it is a mapping between a hierarchy of data tiles and a hierarchy of thread tiles, crossed with layout choices at each staging point. CUDA C++ has no vocabulary for any of those three, so the mapping is smeared across scalar index expressions where neither humans nor compilers can see it. Graphene reifies all three — tiles as nested tensors, layouts as hierarchical strides, thread arrangements as tensors of threads — and once the mapping is a first-class object, expressing an optimized kernel stops being contortion and becomes bookkeeping: tile the data, tile the threads the same way, and state which tile goes to whom.
Symmetry is the trick: threads tile exactly like data. The genuinely novel move is representing the compute hierarchy itself as a tensor, so one set of operations — tile, reshape, index — serves both sides of the mapping. This kills two birds. Expressiveness: any thread arrangement any instruction has ever demanded (contiguous octets, 2×2 group grids, non-contiguous quad-pairs) is a tiling expression, not a built-in. Durability: architectures churn — quad-pairs existed for exactly one generation — and an IR with hard-coded hierarchies churns with them, while an IR where arrangements are ordinary expressions absorbs the next architecture without surgery.
Put the intelligence in the representation, so the compiler can be dumb. Triton and its kin accept an impoverished input language and compensate with brilliant transformation passes — which makes the compiler the bottleneck for extension. Graphene inverts the design: the IR is expressive enough to state the finished optimization explicitly, so code generation degenerates to printing, and adding support for a new instruction means declaring one atomic spec — its tensor shapes, thread arrangement, and opcode — not writing a compiler pass. The evaluation’s exact-match with cuBLAS is the proof that nothing essential was lost in translation: the IR can say everything the library authors’ hand-written PTX says.
One instruction, three languages
Following ldmatrix from contract to code
ldmatrix: one warp moves a 16×16 fp16 tile (four 8×8 matrices) from shared memory into the registers of its 32 threads. Skipping it costs up to 17% in GEMM.
Step 1 of 5
The contract the hardware dictates
Conceptually the picture is clean: tile the warp into four groups of eight threads, arranged 2×2. Each group owns one 8×8 matrix in shared memory. Each thread reads one row of its group's matrix — and after the instruction fires, each thread holds eight values in registers: two adjacent values from each of the four matrices. Reads and receives differ: what a thread TOUCHES in shared memory is not what it ENDS UP with.
the mapping is the instruction — deviate and you get garbage, not an error fast kernels lean on it heavily; simpler equivalent moves cost up to 17%
Why it works
- An optimized kernel is at heart a mapping between tiled data and tiled threads → representing both sides as the same kind of object (decomposable tensors) makes the mapping directly statable → what CUDA C++ smears into index arithmetic becomes visible, checkable structure.
- Dimensions whose sizes and strides are tuples subsume row-major, column-major, interleaved, and swizzled layouts in one notation → layout is declared once at allocation while logical coordinates stay fixed → bank-conflict-dodging layouts stop contaminating every access site in the kernel.
- Thread arrangements are tiling expressions, not built-ins → any contract any instruction demands (including non-contiguous quad-pairs) is expressible, and architecture churn does not force IR redesign → one IR spans Volta, Ampere, and futures.
- Decomposition terminates at atomic specs whose signatures are hardware instructions → the IR states the finished optimization explicitly, so code generation is printing plus index simplification → extending the system is declaring signatures, not authoring compiler passes.
- Fusion is just a generic spec with a decomposition → intermediates stay in shared memory instead of round-tripping through global memory → speedups no library can offer, because no library can enumerate every fusion.
What it costs
- Verbosity by design. Graphene is a target language, not a surface language — humans generate it from a Python API, and its redundant shape annotations make it impractical to write by hand.
- It represents optimizations; it does not find them. The cuBLAS-matching GEMM borrowed cuBLAS’s own tile sizes; choosing tilings, layouts, and fusion strategies remains the job of experts or future search-based compilers layered on top.
- Per-architecture atomic-spec tables. Every instruction variant needs its declared signature per target; the economy is that these are declarations, not passes, but the catalog must still be curated.
- NVIDIA-shaped. The memory spaces, hierarchy, and evaluated instructions are CUDA’s; the IR’s ideas are general, but the artifact targets NVIDIA GPUs (evaluated on Volta V100 and Ampere RTX A6000, pre-Hopper).
- Integration is future work. At publication Graphene stands beside, not inside, XLA/TVM/Triton; the envisioned role as their shared backend — and comparisons against Triton-generated code — are not in the paper.
The evaluation
Two hypotheses, tested on V100 (Volta) and RTX A6000 (Ampere) with clocks locked by the Nsight Compute profiler: can Graphene express the fastest known single kernels, and can it beat the state of the art where libraries end?
Three routes to fast GPU tensor code
| Property | Vendor libraries + hand-written CUDA/PTX | Tensor compilers and IRs (TVM, MLIR, Triton) | Graphene |
|---|---|---|---|
| Who writes the data-to-thread mapping | |||
| Swizzled / interleaved layouts | Hand-managed index arithmetic at every access | ||
| Peak single-kernel performance (GEMM) | cuBLAS — the practical peak reference | Typically relies on libraries or falls short on tensor-instruction control | |
| Fusions beyond the library menu | Only what ships: cuBLASLt covers GEMM + bias/ReLU epilogues; deeper chains mean one kernel per op | Possible in principle; in practice bounded by what passes can lower | |
| Attention (FMHA) and end-to-end | — |
Check yourself
Check your understanding
1. Existing tensor IRs describe a tensor as a list of dimensions and one stride per dimension. Predict what this makes structurally impossible to express — and how Graphene's shape notation escapes.
2. Graphene represents a warp as a tensor of 32 threads and tiles it exactly like a matrix. Why is this symmetry — threads as tensors — the load-bearing novelty rather than a notational convenience?
3. Triton takes friendly input and does the hard work in compiler passes; Graphene's code generator is deliberately dumb — nearly a printer. What justifies calling the dumb backend the better engineering bet?
4. Graphene's biggest speedups are not on GEMM (where it exactly TIES cuBLAS) but on fusions: 2.39× on multi-layer MLP, 1.82× on LSTM, wins on FMHA. Predict where fusion's advantage comes from — and where it must disappear.
Sources
- ACM DL: Graphene: An IR for Optimized Tensor Computations on GPUs — ASPLOS 2023 (open access)