tilelang.language.loop¶

The language interface for tl programs.

Attributes¶

Serial

Unroll

Functions¶

Parallel(*extents[, coalesced_width])

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])

unroll(start[, stop, step, explicit, unroll_factor, ...])

The unrolled For statement.

Module Contents¶

tilelang.language.loop.Parallel(*extents, coalesced_width=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.

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[tir.PrimExpr]) – The list of dominators.

  • wave_size (int) – The wave size.

  • index (int) – The tile index in one wave.

  • group_size (tir.PrimExpr) – The group size.

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. if num_stages is 0, pipeline will not be enabled.

  • order (list[int] | None)

  • stage (list[int] | None)

  • sync (list[list[int]] | None)

  • group (list[list[int]] | None)

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

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

  • start (tvm.tir.PrimExpr)

  • stop (tvm.tir.PrimExpr | None)

  • step (tvm.tir.PrimExpr | None)

  • annotations (dict[str, Any] | None)

Return type:

tvm.script.ir_builder.tir.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¶
tilelang.language.loop.Unroll¶