tilelang.language.annotations¶

Annotation helpers exposed on the TileLang language surface.

Attributes¶

WSID

Functions¶

`use_swizzle`(panel_size[, order, enable])	Annotate a kernel to use a specific threadblock swizzle pattern.
`annotate_layout`(layout_map)	Annotate the layout of the buffer.
`annotate_safe_value`(safe_value_map)	Annotate the safe value of the buffer.
`annotate_l2_hit_ratio`(l2_hit_ratio_map)	Annotate the L2 hit ratio of the buffer.
`annotate_min_blocks_per_sm`(n)	Annotate the minimum number of thread blocks per SM (multiprocessor).
`annotate_restrict_buffers`(*buffers)	Mark the given buffer parameters as non-restrict.
`annotate_ws_schedule`(schedule)	Attach a warp-specialization schedule to the kernel.
`ws_op`(op_id)	Manually annotate the enclosed statement(s) with a stable

Module Contents¶

tilelang.language.annotations.WSID = 'tl.ws_op_id'¶

tilelang.language.annotations.use_swizzle(panel_size, order='row', enable=True)¶

Annotate a kernel to use a specific threadblock swizzle pattern.

Parameters:

panel_size (int)
order (str)
enable (bool)

tilelang.language.annotations.annotate_layout(layout_map)¶

Annotate the layout of the buffer.

Parameters:: layout_map (dict)

tilelang.language.annotations.annotate_safe_value(safe_value_map)¶

Annotate the safe value of the buffer.

Parameters:: safe_value_map (dict)

tilelang.language.annotations.annotate_l2_hit_ratio(l2_hit_ratio_map)¶

Annotate the L2 hit ratio of the buffer.

Parameters:: l2_hit_ratio_map (dict)

tilelang.language.annotations.annotate_min_blocks_per_sm(n)¶

Annotate the minimum number of thread blocks per SM (multiprocessor).

When set, this value is passed as the second argument of __launch_bounds__(maxThreadsPerBlock, minBlocksPerMultiprocessor) in the generated CUDA kernel. A larger value hints the compiler to limit register usage so that more blocks can reside on each SM simultaneously, which can improve occupancy at the cost of potentially more register spilling.

Example

>>> @T.prim_func
... def my_kernel(...):
...     with T.Kernel(...):
            T.annotate_min_blocks_per_sm(2)
...         ...

Parameters:: n (int)

tilelang.language.annotations.annotate_restrict_buffers(*buffers)¶

Mark the given buffer parameters as non-restrict.

This annotation tells codegen to omit the __restrict__ qualifier for the specified kernel buffer parameters. Use this when two (or more) buffers may alias, for example overlapping slices from the same base tensor.

Example

>>> @T.prim_func
... def buggy_kernel(x: T.Tensor((N,), T.float32),
...                  y: T.Tensor((N,), T.float32)):
...     T.annotate_restrict_buffers(x, y)
...     with T.Kernel(N, threads=32) as pid:
...         y[pid] = x[pid] + 1

tilelang.language.annotations.annotate_ws_schedule(schedule)¶

Attach a warp-specialization schedule to the kernel.

schedule is a typed WSSchedule object describing how to transform the straight-line kernel into a warp-specialized one: warp roles, pipelines (full/empty barrier pairs protecting multi-versioned buffers), and per-role instruction sequences per loop scope. It is materialized by the MaterializeWSSchedule pass; see examples/aws/gemm.py for a complete example.

tilelang.language.annotations.ws_op(op_id)¶

Manually annotate the enclosed statement(s) with a stable warp-specialization op or scope id.

Wraps the body in an AttrStmt with key WSID. This is the annotation of last resort, for statement forms that cannot carry annotations= themselves. The enclosed statements — a scalar Bind, or several statements such as an inlined scheduler method — become ONE opaque op:

>>> with T.ws_op("rescale_vote"):
...     should_rescale = T.any_sync(scale_shared[tid] < 1.0)
>>> with T.ws_op("sched_next"):
...     sched.next_tile()

With a scope id the wrapper encloses a while loop and opens that scope (serial scope loops carry the id in their own annotations):

>>> with T.ws_op("loop_wave"):
...     while sched.valid():
...         ...

Tile ops and loops pass annotations={T.WSID: ...} directly instead.

Parameters:: op_id (str)