tilelang.rocm.language.gemm_op

ROCm dialect of T.gemm: the common GEMM plus ROCm-specific knobs.

Functions

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

TileLang GEMM operator for ROCm.

Module Contents

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

TileLang GEMM operator for ROCm.

Same semantics as the common tilelang.language.gemm_op.gemm(). k_pack packs multiple matrix-core operations along K in the MFMA/WMMA lowering (CDNA/RDNA); it is a performance knob with no counterpart on other targets.

Parameters:

  • A (BufferLikeType, i.e. Buffer | BufferLoad | BufferRegion, or Var) – Input buffer A.

  • B (BufferLikeType) – Input buffer B.

  • C (BufferLikeType) – Output buffer 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.

  • k_pack (int) – Number of packed matrix cores along K. Must be 1 or 2. Defaults to 1.

  • annotations (Optional[dict]) – Additional annotations.

Returns:

A handle to the GEMM operation.

Return type:

tirx.Call