tilelang.ascend.language.gemm_op

Ascend dialect of the GEMM operators: the common ops plus Ascend hints.

Functions

gemm_blockscaled(A, B, C, SFA, SFB[, transpose_A, ...])

Ascend block-scaled MXFP8 GEMM, shadowing the common T.gemm_blockscaled.

gemm(A, B, C[, transpose_A, transpose_B, policy, ...])

TileLang GEMM operator, with the Ascend unit-flag hint.

Module Contents

tilelang.ascend.language.gemm_op.gemm_blockscaled(A, B, C, SFA, SFB, transpose_A=False, transpose_B=False, clear_accum=False, unit_flag_ctrl=None)

Ascend block-scaled MXFP8 GEMM, shadowing the common T.gemm_blockscaled.

SFA/SFB are required. For L1 A/B inputs they are the L1 scale-factor buffers, consumed directly by the fused L1 lowering. For L0 A/B inputs they are the MX scale-factor handles of the operand tiles (tilelang.ascend.language.alloc_l0a_sf() / alloc_l0b_sf()), loaded by a preceding T.copy(sf_l1, handle). Each handle must be allocated for the corresponding data buffer, select the same leading indices, and describe its compact trailing tile (zero origin and matching K). The MAD reads the slots implied by its A/B data addresses, so the operands here carry the read-region truth and select the block-scaled lowering.

Unlike the common surface, k_start and the sf_*_granularity_k knobs are implicit: the Ascend lowering derives the scale K offset from the SFA region slice, and MX scales cover 32 K elements per factor.

Parameters:

  • A (tilelang._typing.BufferLikeType)

  • B (tilelang._typing.BufferLikeType)

  • C (tilelang._typing.BufferLikeType)

  • SFA (tilelang._typing.BufferLikeType)

  • SFB (tilelang._typing.BufferLikeType)

  • transpose_A (bool)

  • transpose_B (bool)

  • clear_accum (bool)

  • unit_flag_ctrl (int | tvm.tirx.PrimExpr | None)

Return type:

tvm.tirx.PrimExpr

tilelang.ascend.language.gemm_op.gemm(A, B, C, transpose_A=False, transpose_B=False, policy=GemmWarpPolicy.Square, clear_accum=False, unit_flag_ctrl=None, annotations=None)

TileLang GEMM operator, with the Ascend unit-flag hint.

Same semantics as the common tilelang.language.gemm_op.gemm(). unit_flag_ctrl records the Cube unit-flag control on the tile op, which the Ascend lowering pairs with the following accumulator drain; None omits the annotation and lowers as 0.

On Ascend, L0 operand regions specify the effective MAD M/N/K. Their trailing matrix dimensions must start at zero and describe a compact tile. L0 allocations and producer copies may be padded for hardware alignment; for example, a transposed FP32 load can copy K32 while GEMM consumes A[:, :24] and B[:, :24]. Copy regions must cover the physical transfer. Allocation padding remains part of the storage budget.

Parameters:

  • A (BufferLikeType) – Input A, input B and output C.

  • B (BufferLikeType) – Input A, input B and output C.

  • C (BufferLikeType) – Input A, input B and output C.

  • transpose_A (bool) – Whether to transpose A. Defaults to False.

  • transpose_B (bool) – Whether to transpose B. Defaults to False.

  • policy (GemmWarpPolicy) – GEMM warp partition policy.

  • clear_accum (bool) – Whether to clear the accumulator.

  • unit_flag_ctrl (int | tirx.PrimExpr, optional) – Unit flag control for the instruction. None omits the annotation and lowers as 0.

  • annotations (Optional[dict]) – Additional annotations; values in it take precedence over the individual keywords.

Returns:

A handle to the GEMM operation.

Return type:

tirx.PrimExpr