tilelang.layout.fragment module#

Wrapping Layouts.

class tilelang.layout.fragment.Fragment(shape, forward_fn=None, forward_thread_fn=None, replicate=1, forward_index_fn=None)#

Bases: Layout

A Fragment layout object that encapsulates iteration variables (forward_vars), thread iteration variables (forward_thread), and index transformations (forward_index). This class supports replication (thread_replicate) and index mapping for fine-grained control over multi-dimensional data layouts.

condense_rep_var() → Fragment#

Condense or fold the replicate variable into the existing iteration space. This operation may be used to reduce dimensionality if the replicate variable is no longer needed as a separate dimension.

Returns:: A new Fragment where the replicate variable is condensed.
Return type:: Fragment

get_thread_size()#: Returns the extent (range size) of the thread dimension. If the Fragment was replicated over threads, this will reflect the number of threads.

map_forward_thread(indices: List[PrimExpr]) → PrimExpr#

Get the thread mapping expression for a given set of argument indices.

Parameters:: indices (list of PrimExpr) – Indices for which to compute the thread mapping.
Returns:: The computed thread expression for the provided indices.
Return type:: PrimExpr

repeat(repeats, repeat_on_thread: bool = False, lower_dim_first: bool = True) → Fragment#

Returns a new Fragment that repeats the iteration space a given number of times.

Parameters:

repeats (int) – Number of times to repeat.
repeat_on_thread (bool, optional) – If set, the repeat will happen on the thread dimension.
lower_dim_first (bool, optional) – If set to True, repeat on lower dimensions first.

Returns:

A new Fragment with the repeated iteration space.

Return type:

Fragment

replicate(replicate: int) → Fragment#

Replicate the Fragment across a new thread dimension.

Parameters:: replicate (int) – The replication factor or number of threads.
Returns:: A new Fragment with an additional replicate dimension.
Return type:: Fragment

property thread#: Returns the forward_thread (IterVar) of the Fragment, representing the thread dimension or mapping.

tilelang.layout.fragment.make_swizzled_layout(buffer: Buffer)#