tilelang.layout.fragment

Wrapping Layouts.

類別

Fragment

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

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.

參數:

• 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.

回傳:

A new Fragment with the repeated iteration space.

回傳型別:

Fragment

replicate(replicate)

Replicate the Fragment across a new thread dimension.

參數:

replicate (int) -- The replication factor or number of threads.

回傳:

A new Fragment with an additional replicate dimension.

回傳型別:

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.

回傳:

A new Fragment where the replicate variable is condensed.

回傳型別:

Fragment

map_forward_thread(indices)

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

參數:

indices (list of PrimExpr) -- Indices for which to compute the thread mapping.

回傳:

The computed thread expression for the provided indices.

回傳型別:

PrimExpr

__repr__()

String representation of the Fragment for debugging and logging.

回傳:

A string showing the thread dimension and the index dimension.

回傳型別:

str

is_equal(other)

Check if the current fragment is equal to another fragment.

參數:

other (Fragment)

回傳型別:

bool