tilelang.ascend.language.annotations¶

Ascend storage placement, multi-buffer versions and scheduling capacity annotations.

Functions¶

annotate_layout(layout_map)

Annotate each Buffer view, preserving aliases with distinct Cube layouts.

annotate_buffer_addresses(buffer_addresses)

Place on-chip buffers at exact byte addresses within their memory pools.

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_layout(layout_map)¶

Annotate each Buffer view, preserving aliases with distinct Cube layouts.

Views share storage, not a layout conversion. A producer must write the representation consumed through each view; annotating a view never moves or transposes existing values. Canonical Cube views sharing storage must have the same dtype and padded physical size.

Each memory view keeps its explicit or operation-required layout. Layouts never propagate to other memory views sharing the same storage; annotate every view that needs a non-default interpretation. Unannotated views use normal inference or their ordinary dense/strided addressing. Fragment aliases must always agree on thread ownership.

Parameters:

layout_map (dict)

tilelang.ascend.language.annotations.annotate_buffer_addresses(buffer_addresses)¶

Place on-chip buffers at exact byte addresses within their memory pools.

Supports UB, L1, L0A, L0B, L0C and SS in the Ascend scheduled pipeline. Each placed storage must also have an explicit positive version count in annotate_buffer_versions(), including 1 for unversioned storage. A mode without a count, inferred versions, and manual multi-buffer metadata alone do not satisfy this requirement. Annotation order does not matter.

Addresses are non-negative Python integers, aligned to 32 bytes for UB/SS, 512 for L1/L0A/L0B, and 1024 for L0C. A placement includes all physical versions and layout padding; views of one storage share the same base. Bound L0 scale-factor handles follow their data buffer’s placement.

Partial maps are allowed. Other allocations in a placed pool use the first aligned free interval, without storage reuse. Overlap is rejected even for disjoint lifetimes or positive conflict hints. Different pools and private AIV UBs have independent addresses; common UB storage reserves both UBs. Placement does not change synchronization or ownership.

Every allocation in a placed pool must have a positive static physical size and fit its hardware capacity, including the 32 KiB UB reservation for SIMT. Pools without placements retain their ordinary allocation behavior.

Example

>>> T.annotate_buffer_versions({left: 2, accum: 1})
>>> T.annotate_buffer_addresses({left: 262144, accum: 1024})
Parameters:

buffer_addresses (dict)

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", "counter", or "write_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.

"write_counter" requires an explicit version count and advances on the first executed writing Task in each owner-loop iteration; later writes use the same version, and iterations without writes retain the preceding value. Users guarantee initialized reads and that each update writes the new value before reading it; use multi_buffer_eligible to override inferred owners when needed.

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)