tvm_tirx.backend.cuda.layout.tmem_datapath_layout

tvm_tirx.backend.cuda.layout.tmem_datapath_layout(datapath, rows, cols, sub_slab=0)

Return the TileLayout for a tcgen05 MMA datapath.

See PTX ISA §9.7.16.10.5 for the datapath enumeration. The returned layout is shape-compatible with a buffer of (rows, cols) and encodes the logical-row → physical-TMEM-lane mapping that the corresponding MMA writes to (and that the matching .16x*b / .32x32b atom expects to read).

Parameters:

• datapath (str) –

  PTX data path layout letter (all cta_group::2 layouts are the per-CTA view of the M-rows this CTA of the pair owns):

  • "A": M=256, .cta_group::2 — identity over 128 rows.

  • "B": M=128, .cta_group::2, dense A — 64 rows, column halves folded across the two 64-lane halves.

  • "C": M=128, .cta_group::2, sparse A — 64 rows scattered 16-per-warp, half datapath (same organization as F).

  • "D": M=128, .cta_group::1 — identity, full datapath.

  • "E": M=64, .ws — 64 rows, column halves folded (same organization as B).

  • "F": M=64, non-.ws — 64 rows scattered 16-per-warp, half datapath.

  • "G": M=32, .ws — 32 rows, column quarters folded across the four 32-lane warp slabs.

• rows (int) – Logical row count of the TMEM buffer. Must match the datapath’s M dimension: 128 for A/D, 64 for B/C/E/F, 32 for G.

• cols (int) – Logical column count. Datapath B requires an even count because its columns split into two equal lane halves.

• sub_slab (int) – For Layout F, select the lower (0) or upper (1) 16-lane half of each warp’s 32-lane TMEM partition. The upper half is useful as a 64-row read/write view of the high half-slab of a Layout D accumulator. Layouts D and B already span both halves and therefore only accept 0.

Returns:

TileLayout – Var-shape-compatible layout for (rows, cols).

Return type:

TileLayout