tilelang.layout.fragment

Wrapping Layouts.

Classes

Fragment

A Fragment layout object that encapsulates iteration variables (forward_vars),

Functions

make_swizzled_layout(buffer)

Module Contents

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

Bases: tilelang.layout.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.

property thread

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

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.

repeat(repeats, repeat_on_thread=False, lower_dim_first=True)

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)

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

condense_rep_var()

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

map_forward_thread(indices)

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

__repr__()

String representation of the Fragment for debugging and logging.

Returns:

A string showing the thread dimension and the index dimension.

Return type:

str

tilelang.layout.fragment.make_swizzled_layout(buffer)

Parameters:

buffer (tvm.tir.Buffer)