tilelang.cuda.language.loop¶

CUDA dialect of the loop constructs: the common loops plus CUDA hints.

Functions¶

Parallel(*extents[, coalesced_width, loop_layout, ...])

Construct a nested parallel loop, with CUDA lowering hints.

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

The unrolled For statement, with the CUDA unroll-factor pragma.

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

Alias of the CUDA dialect's unroll().

Module Contents¶

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

Construct a nested parallel loop, with CUDA lowering hints.

Same semantics as the common tilelang.language.loop.Parallel(). prefer_async requests the PTX cp.async rewrite for copies in this loop subtree even outside pipelined loops; it is a performance hint ignored by targets without async copy.

Parameters:

  • extents (int | PrimExpr) – Extents of the parallel loop nest.

  • coalesced_width (Optional[int]) – Width for coalesced memory access.

  • loop_layout (Optional[Fragment]) – Layout annotation for the parallel loop nest.

  • prefer_async (Optional[bool]) – When True, requests cp.async injection for this subtree; when False, forbids it. Lowered as the "parallel_prefer_async" annotation.

  • annotations (Optional[Dict[str, Any]]) – Additional loop annotations; values in it take precedence.

Return type:

tvm.tirx.script.builder.frame.ForFrame

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

The unrolled For statement, with the CUDA unroll-factor pragma.

Same semantics as the common tilelang.language.loop.unroll(). unroll_factor emits #pragma unroll N, which only the CUDA codegen honors; it is mutually exclusive with explicit.

Parameters:

  • start (PrimExpr) – Iteration range.

  • stop (PrimExpr) – Iteration range.

  • step (PrimExpr) – Iteration range.

  • explicit (bool) – Whether to explicitly unroll the loop at compile time.

  • unroll_factor (Optional[int]) – Partial unroll factor, lowered as the "pragma_unroll_factor" annotation.

  • annotations (Optional[Dict[str, Any]]) – Additional loop annotations; values in it take precedence.

Return type:

tvm.tirx.script.builder.frame.ForFrame

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

Alias of the CUDA dialect’s 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