# Tilelang Layout

> Use when working with TileLang layouts — the Layout/Fragment classes, layout inference and its cost model, the CuTe layout algebra and TileLang-to-CuTe conversion, swizzles, or when debugging layout-related passes and their annotations.

- Skill: `tile-ai/tilelang-layout` (Agent Skill, multi-file: 2 files)
- Install (CLI): `npx skillmds@latest add tile-ai/tilelang-layout`
- Raw SKILL.md: https://api.skillmd.com/api/skills/tile-ai/tilelang-layout/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML
- Author: tile-ai (https://skillmd.com/u/tile-ai)
- Updated: 2026-09-10
- Page: https://skillmd.com/skills/tile-ai/tilelang-layout

---


# TileLang Layout System

## Core types (`src/layout/layout.h`, Python `tilelang/layout/`)

**`tl::Layout`** maps logical coordinates to physical coordinates:
`input_size_` (logical shape) + `forward_index_` (one PrimExpr per output
axis). Expressions are written in process-wide placeholder Vars
`InputPlaceholder(i)` (printed `_i`, `_j`, ...). Key APIs:
`InputShape/OutputShape/GetForwardIndex/Forward(vars)/Inverse/
InverseWithLevel/Reshape/DetectInjective/DebugOutput`. `OutputShape()` is
DERIVED from the expressions' ranges by the Analyzer, not stored.

**`tl::Fragment : Layout`** adds a thread dimension for register buffers:
`forward_thread_` (PrimExpr, may reference `ReplicationPlaceholder()`,
printed `_rep`) and `replicate_size_`. `ThreadExtent()` =
`max(forward_thread)+1` under bound domains — also derived, not stored.
`thread_range_` (via `BindThreadRange`) offsets the thread space for warp
specialization; the forward map itself stays in normalized `[0, T)`.
Replication semantics: each logical point is held by `replicate_size`
threads; `Replicate/DeReplicate/CondenseReplicateVar` manipulate the rep
axis.

**Gotchas (verified, easy to trip on):**

- `FragmentNode::GetForwardVars()` PREPENDS `ReplicationPlaceholder()`
  when `replicate_size > 1` (`src/layout/layout.cc:1103`). The input
  placeholders are always the TRAILING `InputDim()` entries. Python code
  that zips `get_forward_vars()` against the shape positionally is wrong
  for replicated fragments.
- The canonical (thread, slot, rep) packing lives in
  `FragmentNode::InverseWithLevel` (`src/layout/layout.cc:1133`): rep
  becomes a trailing ORDINARY input dim of extent `ReplicateExtent()`
  (substituting the placeholder), thread a trailing output. The inverse
  then maps `(slot..., thread) -> (coords..., rep)` with rep last;
  `loop_partition.cc` consumes exactly this ordering and guards
  replica-zero for stores. Mirror this packing when you need a fragment
  as a plain multi-output layout.
- Layouts print via `DebugOutput()` / Python `repr`:
  `Fragment((2,) -> (2,), replicate: 256, thread: _rep, index: (_i,),
  thread_range: I.Range(0, 256))`. Structured fields are directly
  accessible from Python: `replicate_size`, `get_thread_size()`,
  `forward_thread`, `forward_index`, `thread_range`, `get_input_shape()`.

**Swizzles** (`tilelang/layout/swizzle.py`, `src/layout/swizzle_mode.*`)
are XOR-based shared-memory layouts. They are NOT expressible as strided
layouts; the CuTe side models them as a separate `Swizzle` functor
(`ComposedLayout`), never as (shape, stride) modes.

## CuTe layout algebra (`src/layout/cute_layout.{h,cc}`, `tilelang/layout/cute.py`)

A full CuTe implementation maintained for the MMA/TMA backends, exposed to
C++ (`namespace tvm::tl::cute`) and Python (`tilelang.layout.cute`):

- Layouts are `(shape, stride)` IntTuple trees, **COLUMN-major (first mode
  fastest)** — opposite of TileLang's row-major intuition. Multi-output
  codomains use `ScaledBasis` strides (`v@axis`, `E<i>`).
- Algebra: `Coalesce`, `RightInverse`, `LeftInverse`, `Composition`,
  `Complement`, `Cosize`, `LogicalDivide`, `Filter`, `Restrict`,
  `Layout::WithShape`, `Parse`/`Print` (exact CuTe spelling).
- Converters: `LayoutFromTileLang` (single flat output; multi-output
  tl::Layouts are serialized ROW-major, so outputs `[thread, slot]` yield
  `thread * slots + slot`), `LayoutFromTileLangHierarchical` (per-axis
  recovery onto basis axes), `ComposedLayoutFromTileLang` (swizzle
  recovery). All use probe-then-prove: numerically probe strides at
  one-hot points, then symbolically PROVE equivalence — a wrong recovery
  cannot slip through; failure returns `None`/nullopt.

**Failure conventions (must respect):** only the three `*FromTileLang`
converters return Optional. Every other algebra op ICHECK-crashes on its
preconditions (composition divisibility, non-constant extents in the
probe, complement non-injectivity, restrict rank...). Wrap calls when the
input is untrusted. `RightInverse` is PARTIAL: it silently drops stride-0
and non-const-stride modes and inverts only the maximal contiguous chain —
check bijectivity via `size(right_inverse(F)) == size(F)`.

**Conversion caveats:** the converters read only `GetForwardIndex()`; a
Fragment's `forward_thread` and replication are silently ignored unless
you pack them into a plain multi-output Layout first (see the canonical
packing above; substitute `ReplicationPlaceholder` with a trailing input
var — if it leaks into the probe, conversion gracefully returns None).

## Layout inference (`src/transform/layout_inference/`)

`tl.LayoutInference` assigns a layout to every fragment buffer and
parallel loop nest. Three levels: **strict** (annotations, MMA-imposed) →
**common** (BFS propagation through shared buffers) → **free** (per
connected component, try every member as inference root, keep the
cheapest attempt). Results are stored as IR annotations: `layout_map`
(Buffer -> Layout, on the SBlock) and `parallel_loop_layout` (Fragment,
on the outermost parallel For); `ParallelLoopLayoutValidator` enforces the
annotation contract.

"Cheapest" is a pluggable policy (`layout_cost_model.{h,cc}`):

- `tl.layout_cost_model="io-aware"` (opt-in): every
  fragment<->global copy and global-touching parallel loop is charged
  `max(bandwidth bytes, issue bytes)` under the attempt's layouts, scored
  symbolically on the CuTe algebra (pack -> `LayoutFromTileLang` ->
  `RightInverse` -> `Composition`; vector width read off coalesced modes;
  segments counted at warp/step granularity). Registers are the tiebreak.
  Statements outside the model are charged a conservative worst case — an
  attempt must never profit from opacity.
- `tl.layout_cost_model="register-count"` (default): register-slots-only
  ordering.
- The scoring formulas are guarded by the Python parity check
  `maint/layout_inference/run.py --cute` (symbolic scorer vs an
  independent NumPy enumeration oracle); keep `cute_model.py` in lockstep
  when changing `layout_cost_model.cc`.

Hardware geometry is parameterized: lane width via `MaxVectorLoadBits`
(`src/transform/loop_vectorize.h`, SHARED with the vectorizer so the
model's width beliefs match codegen), warp size from the target's
`thread_warp_size`, segment granularity 128B (see `BindMemoryGeometry`).

## Downstream consumers

- `loop_partition.cc`: partitions parallel loops per-thread via the
  fragment inverse; emits replica-zero guards for replicated stores.
- `loop_vectorize.cc`: plans vector widths (`GetVectorizeSize`,
  `IndicesCanVectorize`); the cost model mirrors its judgment.
- TMA/MMA lowering: `src/cuda/op/tma_layout.cc` and
  `producer_consumer_ws.cc` recover shared-buffer swizzles via
  `ComposedLayoutFromTileLang`; tcgen05/wgmma macro generators
  (`tilelang/cuda/intrinsics/macro/`) round-trip TMEM layouts through
  `to_tilelang`/`from_tilelang_hierarchical`.

## Tools

- **`maint/layout_inference/`** — the layout verification harness (see its
  README): golden layout snapshots per constructed case (`run.py`),
  lowered vector-width anchors (`--anchor`), and symbolic-vs-oracle parity
  (`--cute`, backed by `cute_model.py` + the numpy `oracle.py`). Extend it
  with a new case whenever you change inference, the cost model, or the
  converters; goldens are recorded then human-reviewed (`--record`).
- **Diagnostics**: the inference and cost-model passes log decisions via
  DLOG (debug builds); `tl.layout_visualization_enable` renders layouts.
- **Python inspection**: run `tl.transform.LayoutInference()` on a module
  and read the annotations (see `maint/layout_inference/common.py` for
  the extraction idiom), or `cute.Layout.from_tilelang(...)` to see a
  layout's (shape, stride) normal form.

## File map

| area | files |
|---|---|
| core types | `src/layout/layout.{h,cc}`, `tilelang/layout/{layout,fragment}.py` |
| swizzle | `src/layout/swizzle_mode.*`, `tilelang/layout/swizzle*.py` |
| CuTe algebra | `src/layout/cute_layout.{h,cc}`, `tilelang/layout/cute.py` |
| inference | `src/transform/layout_inference/layout_inference.cc` |
| cost model | `src/transform/layout_inference/layout_cost_model.{h,cc}` |
| MMA layouts | `src/layout/gemm_layouts.cc`, `src/layout/tcgen05_layout.*` |
| verification | `maint/layout_inference/` |