Layout Visualization¶

TileLang provides two related ways to inspect data layouts:

tilelang.tools.plot_layout plots a T.Layout or T.Fragment that you already have.
Layout inference visualization prints fragment layouts discovered by the compiler and can save a plot for each inferred two-dimensional fragment.

Use the first workflow while designing a mapping and the second when checking the mapping selected for buffers in a compiled kernel.

Installation¶

PNG, PDF, and SVG output requires Matplotlib. Install TileLang’s visualization extra:

pip install "tilelang[vis]"

Text-only compiler output does not create a Matplotlib figure.

Plot a Layout Directly¶

This example plots a 4 by 4 transpose mapping:

import tilelang.language as T
from tilelang.tools import plot_layout

transpose = T.Layout([4, 4], lambda i, j: (j, i))

plot_layout(
    transpose,
    save_directory="./tmp",
    name="transpose_4x4",
    formats="png",
)

The call writes ./tmp/transpose_4x4.png and prints the saved path. In the default input view, each grid cell represents an input position and its label is the flattened output position. Set view="output" to put output positions on the grid and label them with their source coordinates.

For a T.Fragment, cells are colored by thread and labeled with both the thread ID (T) and thread-local register index (L):

plot_layout(fragment, name="mma_load", formats="pdf")

`plot_layout` API¶

plot_layout(
    layout,
    save_directory="./tmp",
    name="layout",
    colormap=None,
    verbose=False,
    formats="pdf",
    view="input",
    grid_shape=None,
)

Parameter	Meaning
`layout`	A `T.Layout` or `T.Fragment`. Other objects raise `TypeError`.
`save_directory`	Directory created for output files. Defaults to `./tmp`.
`name`	Base filename without an extension.
`colormap`	Matplotlib colormap name. The defaults are `Spectral` for `T.Layout` and `RdPu` for `T.Fragment`.
`verbose`	Print each mapping while building the plot.
`formats`	A string: `pdf`, `png`, `svg`, `all`, or a comma-separated combination such as `png,svg`.
`view`	For `T.Layout`, use `input` or `output`. It is not used for `T.Fragment`.
`grid_shape`	For the input view of a `T.Layout`, override the display grid with `(rows, columns)`. Its product must equal the number of input elements.

plot_layout returns None. With formats="all", it writes PDF, PNG, and SVG files. A higher-dimensional T.Layout is flattened to a two-dimensional grid: all dimensions except the last form the row coordinate, and the last dimension forms the column coordinate.

Visualize Inferred Layouts¶

Enable the compiler pass through pass_configs on a JIT kernel:

import tilelang
import tilelang.language as T


@tilelang.jit(
    pass_configs={
        tilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_ENABLE: True,
        tilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_FORMATS: "txt,svg",
    }
)
def kernel(A, block_M, block_N):
    M, N = T.const("M, N")
    A: T.Tensor((M, N), T.float16)
    B = T.empty((M, N), T.float16)

    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
        fragment = T.alloc_fragment((block_M, block_N), T.float16)
        T.copy(A[by * block_M, bx * block_N], fragment)
        T.copy(fragment, B[by * block_M, bx * block_N])

    return B

Compile or invoke the kernel normally. For each inferred T.Fragment found in a block’s layout_map annotation, the pass prints output like:

C_local inferred layout:
  Shape: [32, 32] -> [8]
  Thread: <thread-index expression>
  Index:  [<local-index expression>]
  Replicate:  1

Image files are written under ./tmp with names derived from the buffer, such as C_local_layout.svg.

TL_LAYOUT_VISUALIZATION_ENABLE defaults to False. When it is enabled, TL_LAYOUT_VISUALIZATION_FORMATS accepts txt, png, pdf, svg, all, or a comma-separated combination. Omitting the formats setting selects text-only output. The compiler prints the textual mapping whenever visualization is enabled; txt does not add an image format.

Limitations¶

Direct T.Fragment plotting requires a two-dimensional input shape and single-valued thread and local-index mappings.
Compiler-generated images are limited to inferred fragments with two-dimensional input shapes. Other inferred fragment shapes are still printed, but image generation is skipped with a warning.
Plots enumerate layout elements in Python. Large shapes can produce large figures and take substantial time and memory.
The output-space view infers its shape from mapped coordinates. Mappings that are non-bijective or sparse may overwrite a cell or leave empty cells, so the plot should be treated as a diagnostic view rather than layout validation.