tilelang.ascend.language.allocate¶

Ascend on-chip buffer allocation: L1 (CBuf) and the L0A/L0B/L0C Cube scopes.

Functions¶

alloc_shared(shape, dtype[, scope])

Allocate a UB buffer.

alloc_l1(shape, dtype[, scope])

Allocate an L1 buffer (__cbuf__) on Ascend NPU.

alloc_l0a(shape, dtype[, scope])

Allocate an L0A buffer (__ca__) on Ascend NPU.

alloc_l0b(shape, dtype[, scope])

Allocate an L0B buffer (__cb__) on Ascend NPU.

alloc_l0c(shape, dtype[, scope, layout])

Allocate an L0C buffer (__cc__) on Ascend NPU.

alloc_l0a_sf(buf[, sf_dtype, sf_shape])

Return the MX scale-factor handle of an L0A data tile.

alloc_l0b_sf(buf[, sf_dtype, sf_shape])

Return the MX scale-factor handle of an L0B data tile.

Module Contents¶

tilelang.ascend.language.allocate.alloc_shared(shape, dtype, scope='shared.dyn')¶

Allocate a UB buffer.

Same surface as the common T.alloc_shared minus its CUDA-only bool workaround (bool buffers there downgrade to the static “shared” scope because the smem-merge pass cannot handle bool): Ascend UB handles bool buffers in the requested scope directly.

Parameters:

  • shape (tilelang._typing.ShapeType)

  • dtype (tilelang._typing.DType)

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l1(shape, dtype, scope='shared.l1')¶

Allocate an L1 buffer (__cbuf__) on Ascend NPU.

Parameters:

  • shape (tilelang._typing.ShapeType) – Buffer shape

  • dtype (tilelang._typing.DType) – Data type

  • scope – Memory scope. Defaults to “shared.l1”

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l0a(shape, dtype, scope='shared.l0a')¶

Allocate an L0A buffer (__ca__) on Ascend NPU.

Layout is deferred to the consuming gemm operation, which infers K-major or MN-major from its transpose flags.

Parameters:

  • shape (tilelang._typing.ShapeType)

  • dtype (tilelang._typing.DType)

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l0b(shape, dtype, scope='shared.l0b')¶

Allocate an L0B buffer (__cb__) on Ascend NPU.

Layout is deferred to the consuming gemm operation, which infers K-major or MN-major from its transpose flags.

Parameters:

  • shape (tilelang._typing.ShapeType)

  • dtype (tilelang._typing.DType)

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l0c(shape, dtype, scope='shared.l0c', layout=True)¶

Allocate an L0C buffer (__cc__) on Ascend NPU.

Parameters:

  • shape (tilelang._typing.ShapeType) – Buffer shape

  • dtype (tilelang._typing.DType) – Data type

  • scope – Memory scope. Defaults to “shared.l0c”

  • layout (bool) – Whether to annotate the fixed Ascend L0C accumulator layout.

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l0a_sf(buf, sf_dtype='uint16', sf_shape=None)¶

Return the MX scale-factor handle of an L0A data tile.

The handle (scope shared.l0a.sf) materializes no storage: the hardware keys a tile’s scale slots to the tile’s own address. Load scales with T.copy(sf_l1, handle) and pass the handle to T.gemm_blockscaled as SFA. Allocation permanently binds this handle to buf; GEMM verifies the binding and leading indices. Scale slots are sticky, so one scale load may serve several data loads (see testing/ascend/language/test_tilelang_ascend_mx_sf_slots.py).

The default shape preserves all leading dimensions and assumes trailing untransposed (rows, K) matrix dimensions with one scale per 32 K elements packed into sf_dtype (uint16 = one pair per 64 K elements); pass sf_shape explicitly for transposed tiles. Allocate one SF handle per data buffer and reuse that handle.

Parameters:

  • buf (tvm.tirx.buffer.Buffer)

  • sf_dtype (tilelang._typing.DType)

  • sf_shape (tilelang._typing.ShapeType | None)

Return type:

tvm.tirx.buffer.Buffer

tilelang.ascend.language.allocate.alloc_l0b_sf(buf, sf_dtype='uint16', sf_shape=None)¶

Return the MX scale-factor handle of an L0B data tile.

See alloc_l0a_sf(); identical semantics for the L0B slot shadow.

Parameters:

  • buf (tvm.tirx.buffer.Buffer)

  • sf_dtype (tilelang._typing.DType)

  • sf_shape (tilelang._typing.ShapeType | None)

Return type:

tvm.tirx.buffer.Buffer