tilelang.layout.partial_fragment

Wrapping the reducer partial-fragment layout.

Classes

PartialFragment

Layout of a reducer's per-thread partials.

Functions

make_fully_replicated_partial_fragment(shape, ...)

The wide plan: every participant holds one full-shape partial.

Module Contents

class tilelang.layout.partial_fragment.PartialFragment(shape, forward_fn=None, forward_thread_fn=None, replicate=1, forward_index_fn=None, combine=None)

Bases: tilelang.layout.fragment.Fragment

Layout of a reducer’s per-thread partials.

Same algebra as Fragment, but the replication coordinate enumerates addends that still await the finalize collective (combine lanes), not equal copies of a finished value. Inferred by layout inference for local.reducer buffers, or supplied by the user via T.annotate_layout({acc: PartialFragment(...)}) to pin the reducer’s physical plan.

static from_fragment(fragment)

Reinterpret a solved Fragment as per-thread partials.

Parameters:

fragment (tilelang.layout.fragment.Fragment)

Return type:

PartialFragment

as_post_collective()

The same algebraic map read as a plain Fragment (post-finalize).

Return type:

tilelang.layout.fragment.Fragment

tilelang.layout.partial_fragment.make_fully_replicated_partial_fragment(shape, thread_extent)

The wide plan: every participant holds one full-shape partial.

Return type:

PartialFragment