triton.experimental.gluon.language.amd.cdna5.make_partitioned_dot_layouts

triton.experimental.gluon.language.amd.cdna5.make_partitioned_dot_layouts()

Create partitioned shared memory layouts and WMMA layout for a CDNA5 GEMM in order to avoid LDS partition conflicts.

Parameters:

  • block_m – M dimension tile size of the shared operand buffer. Must be at least four logical WMMA instruction tiles because the M dimension is split into 2 partitions and each partition must be at least 2 instructions wide.

  • block_n – N dimension tile size of the shared operand buffer. Must be at least two logical WMMA instruction tiles because the N dimension is split into 2 partitions and each partition must be at least 1 instruction wide.

  • original_layout_a – PaddedSharedLayout for operand A. Shape is [block_m, block_k] when not transposed (K contiguous) and [block_k, block_m] when transposed (M contiguous).

  • original_layout_b – PaddedSharedLayout for operand B. Shape is [block_k, block_n] when not transposed (N contiguous) and [block_n, block_k] when transposed (K contiguous).

  • num_warps – Number of warps per CTA. Currently must be 4 or 8.

  • instr_shape – Physical WMMA instruction shape as [M, N, K].

  • a_transposed – Whether A is transposed in shared memory, i.e. M is the contiguous axis instead of K.

  • b_transposed – Whether B is transposed in shared memory, i.e. K is the contiguous axis instead of N.

  • slice_m – M dimension of a dot operation after slicing. Defaults to block_m (unsliced).

  • slice_n – N dimension of a dot operation after slicing. Defaults to block_n (unsliced).

  • transposed – Whether the returned WMMA result layout is transposed. Its logical instruction tile is N x M when true and M x N when false. Defaults to true.

Returns:

A tuple (shared_layout_a, shared_layout_b, wmma_layout). The wmma_layout is sized for slice_m x slice_n; the two shared layouts partition the full block_m / block_n.