[FEATURE][TIRx][CUDA] Support TMEM datapath B

docs/tirx/layout.rst

@@ -269,6 +269,7 @@ TMEM datapath layouts
     accum = tmem_datapath_layout("D", 128, cols)
     lower = tmem_datapath_layout("F", 64, cols, sub_slab=0)
     upper = tmem_datapath_layout("F", 64, cols, sub_slab=1)
+    paired = tmem_datapath_layout("B", 64, cols)
 
 Layout D maps logical row ``r`` directly to ``TLane = r`` and spans both
 16-lane halves of every warp's 32-lane TMEM partition. Layout F maps its 64
@@ -287,6 +288,32 @@ upper-half aliases of the same 128-row Layout D allocation. The
 ``row=0`` or ``row=16`` instruction. Layout D already occupies both halves,
 so a nonzero ``sub_slab`` is rejected.
 
+Layout B is the per-CTA accumulator placement for an M=64
+``tcgen05.mma.cta_group::2`` operation. PTX describes the two-CTA operation
+as M=128, while each CTA owns a logical ``(64, N)`` tile. Its columns split
+across the two 64-lane halves:
+
+.. math::
+
+   \mathrm{TLane}
+   = r + 64\left\lfloor\frac{c}{N/2}\right\rfloor,
+   \qquad
+   \mathrm{TCol} = c \bmod (N/2).
+
+Thus the tile occupies all 128 lanes and ``N/2`` tensor-memory columns, the
+same physical footprint as Layout D ``(128, N/2)``. ``N`` must be even and
+Layout B does not accept ``sub_slab``. Its register image uses the existing
+fragment API:
+
+.. code-block:: python
+
+    frag = T.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
+    Tx.wg.copy_async(frag[:, :], paired_accumulator[:, :])
+    T.ptx.tcgen05.wait.ld()
+
+The logical ``(64, N)`` fragment is one physical ``.32x32b`` transfer over
+all 128 lanes; each thread owns ``N/2`` contiguous fp32 registers.
+
 Beyond GPU registers
 ~~~~~~~~~~~~~~~~~~~~~~
 

docs/tirx/native_basics/cuda/buffers.rst

@@ -379,7 +379,7 @@ tensor as a view at a column offset, and one warp frees it at the end:
         T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
         T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group)
 
-You manage the column offsets and the ``tmem_layout`` (a datapath D/F layout)
+You manage the column offsets and the ``tmem_layout`` (a datapath D/F/B layout)
 yourself. This is exactly the sequence the pool below emits.
 
 Pool
@@ -393,7 +393,9 @@ bump-allocation, and the datapath layout:
     tmem_addr = pool.alloc((1,), "uint32")          # pool = the kernel's smem pool
     tmem_pool = T.TMEMPool(pool, total_cols=512, cta_group=cta_group,
                            tmem_addr=tmem_addr)
-    acc = tmem_pool.alloc((CTA_M, 512), "float32")  # allocated_addr set for you
+    # Choose the layout required by the instruction that consumes the buffer:
+    acc = tmem_pool.alloc((CTA_M, 512), "float32")  # Layout D when CTA_M=128
+    # acc = tmem_pool.alloc((64, N), "float32", datapath="B")  # cta_group=2
     tmem_pool.commit()                               # emits tcgen05.alloc (one warp)
     # ... use acc ...
     tmem_pool.dealloc()                              # emits tcgen05.dealloc (one warp)

docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst

@@ -59,10 +59,12 @@ lowering:
      - ``(tmem, local)`` or ``(local, tmem)`` — exactly one side is tensor memory
    * - register layout
      - matched against a ``tcgen05_atom_layout`` (``.16x64b`` / ``.16x128b`` /
-       ``.16x256b``) for the fast path; otherwise the ``.32x32b`` fallback
+       ``.16x256b``) for the fast path; otherwise the ``.32x32b`` fallback.
+       Layout B uses the special ``.32x32b`` logical ``(64, N)`` image
    * - tmem datapath
-     - classified ``D`` (M=128 identity) or ``F`` (M=64 scattered); an F layout
-       also selects the lower or upper 16-lane sub-slab of each warp partition
+     - classified ``D`` (M=128 identity), ``F`` (M=64 scattered), or ``B``
+       (per-CTA M=64, ``cta_group=2`` column split); an F layout also selects
+       the lower or upper 16-lane sub-slab of each warp partition
 
 Demonstration program
 ----------------------
@@ -128,6 +130,10 @@ fragment spans two 16-row slabs, so the warps issue the atom twice
            shape=shape, num=num_eff,
            row=(sub_slab + slab) * 16, col=col_off_32b)
 
+Layout B is routed before the ordinary atom matching. Its logical
+``(64, N)`` fragment is emitted as one physical ``.32x32b.x{N/2}``
+instruction over all 128 lanes.
+
 The dispatch emits **no** wait — the caller issues ``tcgen05.wait.ld()`` /
 ``wait.st()`` (as in the demo).
 
@@ -166,6 +172,36 @@ The lower view emits ``row=0`` and the upper view emits ``row=16`` for
 ``.16x64b``, ``.16x128b``, and ``.16x256b`` atoms. Layout D has 128 rows and
 already spans both sub-slabs, so it only accepts ``sub_slab=0``.
 
+Layout B readback
+-----------------
+
+Layout B is produced by an M=64-per-CTA ``tcgen05.mma`` with
+``cta_group=2``. Its logical ``(64, N)`` columns are split into two
+``N/2`` halves: the low half uses physical TMEM lanes 0–63, and the high
+half uses lanes 64–127. It therefore occupies all 128 lanes and ``N/2``
+``TCol`` values.
+
+Use the public allocation and fragment APIs together:
+
+.. code-block:: python
+
+    accumulator = tmem_pool.alloc((64, N), "float32", datapath="B")
+    frag = T.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
+
+    Tx.wg.copy_async(frag[:, :], accumulator[:, :])
+    T.ptx.tcgen05.wait.ld()
+
+    # The inverse direction emits tcgen05.st with the same physical image.
+    Tx.wg.copy_async(accumulator[:, :], frag[:, :])
+    T.ptx.tcgen05.wait.st()
+
+This is a single ``tcgen05.{ld,st}.32x32b.x{N/2}`` issue. ``N`` must be
+even, ``N/2`` must be a valid PTX ``num``, and the fragment is fp32-only.
+The first implementation intentionally requires the full logical
+``(64, N)`` region: a partial logical-column slice is not contiguous after
+the two-way lane split. A conventional ``.16x*b`` fragment is rejected with
+an actionable error.
+
 Generated TIRx IR
 -----------------
 
@@ -207,7 +243,9 @@ How inputs change the algorithm
    * - direction
      - ``tmem → local`` → ``tcgen05.ld``; ``local → tmem`` → ``tcgen05.st`` (same
        shape/num logic)
-   * - datapath D vs F
+   * - datapath D vs F vs B
      - ``D`` (M=128) covers all 128 rows; an M=128 ``.16x*b`` copy issues two slabs
        (``row = 0`` / ``row = 16``). ``F`` (M=64) scatters rows to lanes and
-       its layout selects one issue at ``row = 0`` or ``row = 16``
+       its layout selects one issue at ``row = 0`` or ``row = 16``. ``B``
+       (per-CTA M=64, ``cta_group=2``) splits N across two 64-lane halves and
+       emits one logical-64-row ``.32x32b`` image

docs/tirx/tile_primitives/gemm_async.rst

@@ -57,9 +57,9 @@ A single predicate — single-thread or warp scope:
        ``float8_e4m3fn`` / ``float4_e2m1fn`` with ``SFA`` / ``SFB`` scale factors in
        tmem; accumulator always ``float32``
    * - shape
-     - ``M ∈ {64, 128}`` (×2 for cta_group=2); ``N`` divisible by 8 (cta_group=1) or
-       16 (cta_group=2); ``K`` divisible by ``MMA_K`` = 16 (f16/bf16) / 32 (fp8) /
-       64 (fp4)
+     - per-CTA ``M ∈ {64, 128}``; ``N`` divisible by 8 (cta_group=1) or 16
+       (cta_group=2); ``K`` divisible by ``MMA_K`` = 16 (f16/bf16) / 32 (fp8) /
+       64 (fp4). With cta_group=2, the CTA pair covers twice the per-CTA M
    * - cta_group
      - ``1`` (one CTA) or ``2`` (two CTAs split the operand)
 
@@ -126,6 +126,39 @@ instruction descriptor is encoded at runtime. As with the other async ops, the
 dispatch emits **no** completion — the caller's ``tcgen05.commit`` + mbarrier wait
 close it.
 
+Accumulator datapaths and readback
+----------------------------------
+
+The accumulator layout must match the MMA's row placement:
+
+* Layout D is the M=128 identity placement.
+* Layout F is the single-CTA M=64 scattered placement.
+* Layout B is the per-CTA M=64 placement for ``cta_group=2``. Its logical
+  N columns split across physical lane halves 0–63 and 64–127, so it
+  occupies all 128 lanes and ``N/2`` tensor-memory columns.
+
+Allocate and read a Layout B result as follows:
+
+.. code-block:: python
+
+    accumulator = tmem_pool.alloc((64, N), "float32", datapath="B")
+    Tx.gemm_async(
+        accumulator[:, :],
+        A_smem[:, :],
+        B_smem[:, :],
+        dispatch="tcgen05",
+        cta_group=2,
+    )
+
+    frag = T.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")
+    Tx.wg.copy_async(frag[:, :], accumulator[:, :])
+    T.ptx.tcgen05.wait.ld()
+
+The fragment is a logical ``(64, N)`` view of one physical
+``.32x32b`` transfer over all 128 lanes. The gemm write-side layout and
+the allocation/readback layout are produced by the same
+``tmem_datapath_layout("B", 64, N)`` factory.
+
 Generated TIRx IR
 -----------------
 
@@ -167,7 +200,8 @@ How inputs change the algorithm
        addresses and a runtime-encoded instruction descriptor
    * - cta_group
      - ``1`` → one CTA, ``M ∈ {64, 128}``; ``2`` → two CTAs split the operand,
-       ``M ∈ {128, 256}`` and half the per-CTA N
+       each with per-CTA ``M ∈ {64, 128}`` and half the B rows. The
+       per-CTA M=64 output uses Layout B
    * - M / N / K extents
      - set the ``(mi, ni, ki)`` unrolled loop counts; K iterations accumulate into
        the same tmem accumulator

python/tvm/backend/cuda/lang/alloc_pool.py

@@ -323,7 +323,8 @@ def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None
             Explicit ``TileLayout``. Mutually exclusive with ``datapath``.
         datapath : str | None
             Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath,
-            ``"F"`` for M=64 non-``.ws`` scattered). When provided, the buffer's
+            ``"F"`` for M=64 non-``.ws`` scattered, or ``"B"`` for per-CTA
+            M=64 ``.cta_group::2`` "2x2"). When provided, the buffer's
             layout is derived from ``tmem_datapath_layout(datapath, *shape)``
             so the row index reflects the *physical* TMEM lane occupation
             (PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b``

python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py

@@ -88,11 +88,13 @@ def _classify_tmem_datapath(tmem_buf):
     Layout D (M=128, identity row→lane) is the default returned by
     ``_default_tmem_layout``. Layout F (M=64 non-``.ws``, scattered) is the
     explicit opt-in produced by ``tmem_pool.alloc(..., datapath="F")``.
+    Layout B is the per-CTA M=64, ``.cta_group::2`` "2x2" placement produced
+    by ``tmem_pool.alloc(..., datapath="B")``.
     The dispatch uses this to pair each ``.16x*b`` / ``.32x32b`` atom with a
     compatible layout — see ``_check_tmem_layout_for_atom``.
 
-    ``sub_slab`` is always 0 for Layout D. For Layout F it selects the lower
-    (0) or upper (1) 16-lane half of each warp's 32-lane partition.
+    ``sub_slab`` is always 0 for Layout D and B. For Layout F it selects the
+    lower (0) or upper (1) 16-lane half of each warp's 32-lane partition.
     """
     if tmem_buf.layout is None:
         return None
@@ -106,6 +108,15 @@ def _classify_tmem_datapath(tmem_buf):
         except (AssertionError, ValueError):
             return None
     if rows == 64:
+        # Layout B splits N into two N/2 column halves, together spanning all
+        # 128 lanes. Its structure is disjoint from Layout F; try it first.
+        if int(tmem_buf.shape[1]) % 2 == 0:
+            cand = tmem_datapath_layout("B", 64, tmem_buf.shape[1]).canonicalize()
+            try:
+                tvm.ir.assert_structural_equal(buf_layout, cand)
+                return ("B", 0)
+            except (AssertionError, ValueError):
+                pass
         # Layout F may occupy either 16-lane half of each warp's 32-lane
         # partition. The layout carries that choice as a +16 TLane offset;
         # thread it through to the PTX row immediate instead of adding an
@@ -142,13 +153,21 @@ def _classify_tmem_datapath(tmem_buf):
 #                                |           |   high slab (row=16) is garbage
 #   F (M=64 scatter)x .32x32b    | no       | F only utilizes 16 of each
 #                                |           |   warp's 32 lanes
+#   B (M=64 2x2)    x bare atom  | no       | B splits N into two N/2
+#                                |           |   lane-halves; use the dedicated
+#                                |           |   logical (64, N) Layout B image
+#                                |           |   handled by
+#                                |           |   _emit_datapath_b_path
 _TMEM_ATOM_COMPAT = {
     ("D", "32x32b", 128): True,
     ("D", "16x*b", 64): True,
     ("D", "16x*b", 128): True,
     ("F", "32x32b", 128): False,
     ("F", "16x*b", 64): True,
     ("F", "16x*b", 128): False,
+    ("B", "32x32b", 128): False,
+    ("B", "16x*b", 64): False,
+    ("B", "16x*b", 128): False,
 }
 
 
@@ -207,6 +226,20 @@ def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> P
     elem_per_32b = 32 // elem_size
     assert len(local_buf.shape) == len(tmem_buf.shape) == 2
 
+    # Datapath B is identified from the TMEM side before probing ordinary
+    # register atoms. A logical (64, N) Layout B tile is physically a
+    # (128, N/2) .32x32b transfer and therefore needs its dedicated reshape.
+    if _classify_tmem_datapath(tmem_buf) == ("B", 0):
+        return _emit_datapath_b_path(
+            direction=direction,
+            tmem_buf=tmem_buf,
+            local_buf=local_buf,
+            tmem_region=tmem_region,
+            local_region=local_region,
+            elem_per_32b=elem_per_32b,
+            analyzer=analyzer,
+        )
+
     # Try the .16x* (M=64) path first by structural-matching the register-side
     # layout against ``tcgen05_atom_layout(instr_shape, (64, K), dtype)``. The
     # An M=64 TMEM-side Layout F fragment lives in either lanes 0..15 or
@@ -452,6 +485,78 @@ def impl():
     return impl
 
 
+def _emit_datapath_b_path(
+    *, direction, tmem_buf, local_buf, tmem_region, local_region, elem_per_32b, analyzer
+) -> PrimFunc:
+    """Read or write a Layout B (per-CTA M=64, ``.cta_group::2``) accumulator.
+
+    Layout B splits a logical ``(64, N)`` tile into two ``N/2`` column halves
+    over physical lanes 0..63 and 64..127. It is therefore transferred as one
+    physical ``.32x32b`` ``(128, N/2)`` register file, re-labeled by
+    ``tcgen05_atom_layout("32x32b", (64, N), "float32")``.
+    """
+    if elem_per_32b != 1:
+        raise ValueError(
+            "datapath B readback expects an fp32 fragment (32-bit cells), got "
+            f"dtype={local_buf.dtype!r}"
+        )
+    if int(local_buf.shape[0]) != 64:
+        raise ValueError(
+            "datapath B (.cta_group::2 M=64) fragment must be (64, N); a "
+            f"128-row .32x32b fragment reads the wrong region. Got rows={local_buf.shape[0]}. "
+            "Allocate it with T.alloc_tcgen05_ldst_frag('32x32b', (64, N), 'float32')."
+        )
+
+    n_cols = int(local_buf.shape[1])
+    n_half = n_cols // 2
+    expected_local = tcgen05_atom_layout("32x32b", (64, n_cols), local_buf.dtype).canonicalize()
+    try:
+        tvm.ir.assert_structural_equal(local_buf.layout.canonicalize(), expected_local)
+    except (AssertionError, ValueError) as err:
+        raise ValueError(
+            "datapath B (.cta_group::2 M=64) requires a matching Layout B "
+            "register fragment. Allocate it with "
+            "T.alloc_tcgen05_ldst_frag('32x32b', (64, N), 'float32'); a "
+            ".16x*b fragment reads the wrong physical lanes and columns. "
+            f"(fragment layout mismatch: {err})"
+        ) from err
+
+    # A partial logical-column slice is not contiguous after the two-way lane
+    # split. Keep this first implementation deliberately strict and transfer
+    # the complete logical tile on both sides.
+    tmem_st, tmem_extent = get_st_extent(tmem_region)
+    local_st, local_extent = get_st_extent(local_region)
+    if not (
+        analyzer.can_prove_equal(tmem_st[0], 0)
+        and analyzer.can_prove_equal(tmem_st[1], 0)
+        and analyzer.can_prove_equal(tmem_extent[0], 64)
+        and analyzer.can_prove_equal(tmem_extent[1], n_cols)
+    ):
+        raise ValueError("datapath B copy must cover the full (64, N) TMEM buffer")
+    if not (
+        analyzer.can_prove_equal(local_st[0], 0)
+        and analyzer.can_prove_equal(local_st[1], 0)
+        and analyzer.can_prove_equal(local_extent[0], 64)
+        and analyzer.can_prove_equal(local_extent[1], n_cols)
+    ):
+        raise ValueError("datapath B copy must cover the full (64, N) register fragment")
+
+    op = T.ptx.tcgen05.ld if direction == "tmem2local" else T.ptx.tcgen05.st
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        local_storage = local_buf.view(n_half, layout=TileLayout(S[n_half]))
+        local_32b = local_storage.view("uint32")
+        op(
+            tmem_buf.allocated_addr[0],
+            *[local_32b[i] for i in range(n_half)],
+            shape="32x32b", num=n_half, row=0, col=0,
+        )
+    # fmt: on
+    return impl
+
+
 # === Variant: copy_async/tmem<->local (priority=10) ===
 #
 # When: one buffer is in tmem (tensor memory, Blackwell SM100+) and the other

python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py

@@ -776,8 +776,10 @@ def _atom_off(dim):
     # the full row range, the slice layout structurally matches Layout F over
     # (M=64, N) — assert against that base instead of the Layout D identity.
     if is_2x2:
-        N_half = N // 2
-        base = TileLayout(S[(M, 2, N_half) : (1 @ TLane, 64 @ TLane, 1 @ TCol)])
+        # Layout B (per-CTA M=64, .cta_group::2 "2x2"). Single-source the
+        # write layout with allocation and readback so the two sides cannot
+        # drift.
+        base = tmem_datapath_layout("B", M, N)
     elif (
         M == 64
         and int(C_buffer.shape[0]) == 64