tilelang.ascend.language.annotations¶

Ascend auto-schedule annotations: multi-buffer version control and memory capacity overrides consumed by the Z3 scheduler.

Functions¶

annotate_buffer_versions(buffer_versions_map)

Control automatic multi-buffer version counts and indexing modes.

annotate_manual_multi_buffer(*items)

Declare buffers that the user has manually multi-buffered.

annotate_unlimit_memory(*scopes)

Remove the memory capacity limit for the given scopes during auto-scheduling.

Module Contents¶

tilelang.ascend.language.annotations.annotate_buffer_versions(buffer_versions_map)¶

Control automatic multi-buffer version counts and indexing modes.

When using T.Persistent with num_stages > 1, the Z3 auto-scheduler computes optimal buffer version counts for each on-chip buffer. Each mapping value may be one of:

  • 1: opt out of multi-buffer eligibility, including explicit owner claims;

  • num_versions >= 2: use a fixed version count and automatic mode;

  • (num_versions, mode): use a fixed count and explicit mode;

  • mode: select a mode while leaving the version count to the scheduler.

mode is "auto", "iteration", or "counter". "iteration" uses the affine flattened-loop index. "counter" uses a monotonic local counter that advances after a completed active buffer epoch. "auto" keeps the affine fast path when safe and selects a counter for conditionally executed or non-affine epochs.

A fixed count is useful when you want to guarantee double-buffering for a critical buffer regardless of the Z3 solver’s choice.

{buf: 1} removes the storage (including aliases) from multi-buffer eligibility, overriding inferred or explicit owner claims and preserving ordinary dependencies instead of owner-exclusion dependencies. Overriding explicit multi_buffer_eligible claims emits a warning once per storage. Use this when the eligibility heuristic misses reads of previous data. To keep eligibility with one version, explicitly specify a mode: {buf: (1, "auto")}.

AutoSchedule consumes this map and replaces it with the selected version counts, including versions chosen automatically by the solver.

Example

>>> @T.prim_func
... def my_kernel(x: T.Tensor((M, K), T.float16),
...               w: T.Tensor((N, K), T.float16),
...               y: T.Tensor((M, N), T.float16)):
...     with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx:
...         fixed = T.alloc_shared((block_M, block_K), T.float16)
...         inferred = T.alloc_shared((block_M, block_K), T.float16)
...         T.annotate_buffer_versions({
...             fixed: (2, "counter"),
...             inferred: "iteration",
...         })
...         with T.Persistent(T.ceildiv(M, block_M), num_stages=4) as i:
...             ...
Parameters:

buffer_versions_map (dict)

tilelang.ascend.language.annotations.annotate_manual_multi_buffer(*items)¶

Declare buffers that the user has manually multi-buffered.

Unlike annotate_buffer_versions() (which asks the Z3 auto-scheduler to version a buffer for you), this annotation tells AutoSchedule that the buffer is already multi-buffered by hand: the user writes the version index into every access themselves (e.g. buf[l % 2] / buf[(l + 1) % 2] ping-pong), and the pass only needs to analyze dependencies and emit the set_flag / wait_flag synchronization.

The per-buffer version count N is used as the ring size / the upper bound on the cross-iteration distance the analysis searches for — it does NOT reshape or reallocate the buffer and does NOT assume the version index lives in dimension 0. Each positional argument is either:

  • a bare buffer: N defaults to its leading dimension buffer.shape[0] (which must then be a compile-time constant) — the common T.alloc_shared((N, ...)) ping-pong case; or

  • a {buffer: N} dict: gives N explicitly, decoupled from the shape. Use this when the buffer has no dedicated leading version dim — e.g. a flattened buffer multi-buffered by disjoint sub-ranges — or when you want a version count that differs from shape[0].

Example

>>> @T.prim_func
... def my_kernel(...):
...     with T.Kernel(...) as core_id:
...         buf = T.alloc_shared((2, N), T.bfloat16)     # leading version dim
...         flat = T.alloc_shared((3 * N,), T.bfloat16)  # flat, sub-ranges
...         T.annotate_manual_multi_buffer(buf, {flat: 3})
...         for l in T.serial(L):
...             ...  # buf[l % 2] ping-pong; flat[(l % 3) * N : ...] triple-buffer
tilelang.ascend.language.annotations.annotate_unlimit_memory(*scopes)¶

Remove the memory capacity limit for the given scopes during auto-scheduling.

When unlimit’d, a scope’s capacity is effectively unbounded so Z3 will not constrain multi-buffer versions based on that memory pool. Use this when you know a particular scope has spare headroom.

Valid scopes: "shared", "shared.l1", "shared.l0c", "shared.l0a", "shared.l0b".

Example

>>> @T.prim_func
... def my_kernel(x: T.Tensor((M, K), T.float16),
...               w: T.Tensor((N, K), T.float16),
...               y: T.Tensor((M, N), T.float16)):
...     with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx:
...         T.annotate_unlimit_memory("shared")
...         ...
Parameters:

scopes (str)