tilelang.layout.swizzle¶

Wrapping Layouts.

Functions¶

`make_swizzled_layout`(buffer[, k_major, allow_pad])
`make_volta_swizzled_layout`(buffer[, is_a, k_inner])
`make_wgmma_swizzled_layout`(buffer[, continuity, k_major])
`make_tcgen05mma_swizzled_layout`(buffer[, continuity, ...])
`make_full_bank_swizzled_layout`(*args)
`make_half_bank_swizzled_layout`(*args)
`make_quarter_bank_swizzled_layout`(*args)
`make_linear_layout`(buffer_or_load_or_region)	Create a row-major linear layout for any dimension.
`make_gemm_fragment_8x8`()	Create a standard 8x8 GEMM fragment layout for ldmatrix/stmatrix.
`make_gemm_fragment_8x8_transposed`()	Create a transposed 8x8 GEMM fragment layout for ldmatrix/stmatrix.
`make_fully_replicated_layout_fragment`(buffer, threads)	Create a fully replicated layout for a fragment buffer.

Module Contents¶

tilelang.layout.swizzle.make_swizzled_layout(buffer, k_major=True, allow_pad=True)¶

Parameters:

buffer (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion)
k_major (bool)
allow_pad (bool)

tilelang.layout.swizzle.make_volta_swizzled_layout(buffer, is_a=True, k_inner=True)¶

Parameters:

buffer (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion)
is_a (bool)
k_inner (bool)

tilelang.layout.swizzle.make_wgmma_swizzled_layout(buffer, continuity=None, k_major=True)¶

Parameters:

buffer (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion)
continuity (int)
k_major (bool)

tilelang.layout.swizzle.make_tcgen05mma_swizzled_layout(buffer, continuity=None, k_major=True)¶

Parameters:

buffer (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion)
continuity (int)
k_major (bool)

tilelang.layout.swizzle.make_full_bank_swizzled_layout(*args)¶

Parameters:: args – buffer/BufferLoad/BufferRegion or (stride, continuous, element_size)

Examples

make_full_bank_swizzled_layout(buffer) make_full_bank_swizzled_layout(stride, continuous, element_size)

tilelang.layout.swizzle.make_half_bank_swizzled_layout(*args)¶

Parameters:: args – buffer/BufferLoad/BufferRegion or (stride, continuous, element_size)

Examples

make_half_bank_swizzled_layout(buffer) make_half_bank_swizzled_layout(stride, continuous, element_size)

tilelang.layout.swizzle.make_quarter_bank_swizzled_layout(*args)¶

Parameters:: args – buffer/BufferLoad/BufferRegion or (stride, continuous, element_size)

Examples

make_quarter_bank_swizzled_layout(buffer) make_quarter_bank_swizzled_layout(stride, continuous, element_size)

tilelang.layout.swizzle.make_linear_layout(buffer_or_load_or_region)¶

Create a row-major linear layout for any dimension.

Parameters:: buffer_or_load_or_region (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion) – Buffer, BufferLoad, or BufferRegion
Returns:: A row-major linear layout
Return type:: Layout

tilelang.layout.swizzle.make_gemm_fragment_8x8()¶

Create a standard 8x8 GEMM fragment layout for ldmatrix/stmatrix.

This layout matches the warp-level matrix multiplication pattern used in tensor cores.

Returns:: An 8x8 fragment layout
Return type:: Fragment

tilelang.layout.swizzle.make_gemm_fragment_8x8_transposed()¶

Create a transposed 8x8 GEMM fragment layout for ldmatrix/stmatrix.

This layout is the transposed version of make_gemm_fragment_8x8, useful for different access patterns in matrix operations.

Returns:: A transposed 8x8 fragment layout
Return type:: Fragment

tilelang.layout.swizzle.make_fully_replicated_layout_fragment(buffer, threads)¶

Create a fully replicated layout for a fragment buffer.

A fully replicated fragment means all threads hold identical copies of the entire buffer. This is useful for index buffers or masks that need to be accessed uniformly across all threads.

Parameters:

buffer (tvm.tir.Buffer | tvm.tir.BufferLoad | tvm.tir.BufferRegion) – Buffer, BufferLoad, or BufferRegion to get shape information
threads (int) – Number of threads (replicate extent)

Returns:

A fully replicated layout where each thread has a complete copy

Return type:

Fragment

Example

>>> C_local = T.alloc_fragment((2,), T.float32)
>>> layout = make_fully_replicated_layout_fragment(C_local, 256)
>>> T.annotate_layout({C_local: layout})