tilelang.language.loop¶

Loop related language interfaces in TileLang.

Functions¶

`Parallel`(*extents[, coalesced_width, loop_layout, ...])	Tools to construct nested parallel for loop.
`Persistent`(domain, wave_size, index[, group_size])	Tools to construct persistent for loop.
`Pipelined`(start[, stop, num_stages, order, stage, ...])	Tools to construct pipelined for loop.
`serial`(start[, stop, step, annotations])	The serial For statement.
`unroll`(start[, stop, step, explicit, unroll_factor, ...])	The unrolled For statement.
`Serial`(start[, stop, step, annotations])	Alias of T.serial.
`Unroll`(start[, stop, step, explicit, unroll_factor, ...])	Alias of T.unroll.
`vectorized`(start[, stop, annotations])	The vectorized For statement.
`Vectorized`(start[, stop, annotations])	Alias of T.vectorized.

Module Contents¶

tilelang.language.loop.Parallel(*extents, coalesced_width=None, loop_layout=None, prefer_async=None, annotations=None)¶

Tools to construct nested parallel for loop.: This can be used to create element-wise tensor expression.

Parameters:

extents (PrimExpr) – The extents of the iteration.
coalesced_width (Optional[int]) – The coalesced width of the parallel loop.
loop_layout (to the outermost generated loop only. If you omit) – A layout annotation for the parallel loop nest, expressed as a T.Fragment. When provided, it is attached as the "parallel_loop_layout" annotation on the outermost parallel loop. For a k-dimensional T.Parallel(...) nest, the fragment’s InputDim must equal k.
prefer_async (Optional[bool]) – Optional hint for PTX async-copy rewrite in this parallel loop subtree. When set to True, it requests cp.async injection even outside pipelined loops. False/None keeps default behavior. Internally lowered as loop annotation "parallel_prefer_async".
annotations (Optional[Dict[str, Any]]) – Optional user-provided loop annotations attached to the outermost generated parallel loop. For example: {"parallel_async_without_async_commit_wait": True}.
constraints (Notes on layout)
---------------------------
during (TileLang validates parallel loop layout annotations)
ParallelLoopLayoutValidator. (tl.transform.LayoutInference with)
are (The key constraints)
after (- Every parallel loop must be covered by a layout annotation) – layout inference. For a nested parallel nest, this annotation must live on the outermost loop; inner parallel loops must not carry the layout annotation themselves.
k (- For a nest depth of) – InputDim == k.
satisfy (the layout must) – InputDim == k.
loop (- Violations (missing annotation on the outermost) – inner loops, or mismatched InputDim) cause a compilation error.
on (outermost loop can manage its inner nest. Therefore the layout is placed) – inner loops, or mismatched InputDim) cause a compilation error.
Rationale (inner loops cannot control/annotate their outer loops, while the)
on
region. (the outermost loop so lowering passes can rewrite the entire)
easy (To make this)
loop_layout
loop_layout
the (compiler will try to infer a valid layout and attach it during)
the
pass. (LayoutInference)

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tilelang.language.loop.Persistent(domain, wave_size, index, group_size=8)¶

Tools to construct persistent for loop.

Parameters:

domain (List[tirx.PrimExpr]) – The list of dominators.
wave_size (int) – The wave size.
index (int) – The tile index in one wave.
group_size (tirx.PrimExpr) – The group size.

Return type:

tvm.tirx.script.builder.frame.ForFrame

tilelang.language.loop.Pipelined(start, stop=None, num_stages=0, order=None, stage=None, sync=None, group=None)¶

Tools to construct pipelined for loop.

Parameters:

start (PrimExpr) – The minimum value of iteration.
stop (PrimExpr) – The maximum value of iteration.
num_stages (int) – The max number of buffer used between pipeline producers and consumers. For compiler-inferred pipelines, if num_stages is 0, pipeline will not be enabled. For manually scheduled pipelines, prefer specifying order and stage without num_stages; the pipeline depth is inferred as max(stage) + 1.
order (Optional[List[int]]) – Optional manual emission order for scheduled statements in the loop body. This list should describe executable pipeline statements such as copies, fills, GEMMs, reductions, stores, waits, or commits.
stage (Optional[List[int]]) – Optional manual pipeline stage for each scheduled statement in the loop body. The list is aligned with order and follows the same statement counting rule.
sync (Optional[List[List[int]]]) – Optional synchronization metadata for manual pipeline lowering.
group (Optional[List[List[int]]]) – Optional producer grouping metadata for manual pipeline lowering.

Return type:

tvm.tirx.script.builder.frame.ForFrame

Notes

Replayable scalar Bind statements created by local aliases are not scheduled pipeline operations and should not consume entries in order or stage. For example, in the body below only the copy and store need annotation entries; base is replayed automatically at each use:

for i in T.Pipelined(n, order=[1, 0], stage=[0, 1]):
    base: T.int32 = i * block
    T.copy(A[base], shared)
    T.copy(shared, B[base])

A replayable Bind may also read a buffer that is not written by the pipeline body, such as idx = Ids[i]. If a Bind reads a buffer that is written inside the same pipeline body, it is kept as a scheduled statement because the load has a pipeline dependency and cannot be freely replayed.

Older code that included replayable scalar Bind entries in order/stage is accepted for compatibility, but those entries are ignored by the pipeline passes. New code should annotate only the scheduled statements. Avoid setting num_stages together with manual order/stage unless you intentionally need an explicit depth override.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

Parameters:

start (tvm.tirx.PrimExpr)
stop (tvm.tirx.PrimExpr | None)
num_stages (int)
order (list[int] | None)
stage (list[int] | None)
sync (list[list[int]] | None)
group (list[list[int]] | None)

tilelang.language.loop.serial(start, stop=None, step=None, *, annotations=None)¶

The serial For statement.

Parameters:

start (PrimExpr) – The minimum value of iteration.
stop (PrimExpr) – The maximum value of iteration.
step (PrimExpr) – The step size of the iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tilelang.language.loop.unroll(start, stop=None, step=None, *, explicit=False, unroll_factor=None, annotations=None)¶

The unrolled For statement.

Parameters:

start (PrimExpr) – The minimum value of iteration.
stop (PrimExpr) – The maximum value of iteration.
step (PrimExpr) – The step size of the iteration.
explicit (bool) – Whether to explicitly unroll the loop.
unroll_factor (int) – The unroll factor of the loop.
annotations (Dict[str, Any]) – The optional annotations of the For statement.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tilelang.language.loop.Serial(start, stop=None, step=None, *, annotations=None)¶

Alias of T.serial.

Parameters:

start (tvm.tirx.PrimExpr)
stop (tvm.tirx.PrimExpr | None)
step (tvm.tirx.PrimExpr | None)
annotations (dict[str, Any] | None)

Return type:

tvm.tirx.script.builder.frame.ForFrame

tilelang.language.loop.Unroll(start, stop=None, step=None, *, explicit=False, unroll_factor=None, annotations=None)¶

Alias of T.unroll.

Parameters:

start (tvm.tirx.PrimExpr)
stop (tvm.tirx.PrimExpr | None)
step (tvm.tirx.PrimExpr | None)
explicit (bool)
unroll_factor (int | None)
annotations (dict[str, Any] | None)

Return type:

tvm.tirx.script.builder.frame.ForFrame

tilelang.language.loop.vectorized(start, stop=None, *, annotations=None)¶

The vectorized For statement.

Parameters:

start (PrimExpr) – The minimum value of iteration.
stop (PrimExpr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tilelang.language.loop.Vectorized(start, stop=None, *, annotations=None)¶

Alias of T.vectorized.

Parameters:

start (tvm.tirx.PrimExpr)
stop (tvm.tirx.PrimExpr | None)
annotations (dict[str, Any] | None)

Return type:

tvm.tirx.script.builder.frame.ForFrame