This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 58b71d78cf [FIX][TIRx][CUDA] Fix tcgen05 register fragment layouts
(#20068)
58b71d78cf is described below
commit 58b71d78cf28b49ba559f0790e770b5130394e49
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Jul 29 14:06:01 2026 -0400
[FIX][TIRx][CUDA] Fix tcgen05 register fragment layouts (#20068)
## Motivation and context
`tcgen05.ld/st` with a `.16x*b` atom accesses one 16-row half-slab from
each 32-row TMEM partition owned by a warp. Across a four-warp
warpgroup, the two physical half-slabs are:
| View | Physical TMEM lanes | PTX row immediate |
|---|---|---|
| lower half | `0..15, 32..47, 64..79, 96..111` | `0` |
| upper half | `16..31, 48..63, 80..95, 112..127` | `16` |
Layout D for an M=128 accumulator occupies both halves. Layout F exposes
an M=64 logical tile over one half, which is useful both for a native
M=64 accumulator and for reading either 64-row half of an existing
Layout D accumulator.
Before this PR, `tmem_datapath_layout("F", 64, cols)` could describe
only the lower half. The copy dispatcher classified a TMEM buffer only
as `"D"` or `"F"` and every M=64 `.16x*b` operation started at `row=0`.
As a result, there was no layout-preserving way to create a recognized
64-row view of the upper half of a Layout D accumulator: a normal Layout
F view still addressed the lower half, while a hand-written `+16@TLane`
layout was not recognized by the dispatcher.
The half-slab selection belongs in the buffer layout because layout is
the source of truth for physical placement in TIRx. It should not be an
out-of-band `copy_async` option. This PR therefore records the selection
in Layout F, carries it through datapath classification, and derives the
PTX row immediate from it.
A related invariant is that a default M=128 TMEM allocation must be
structurally identical to named Layout D. The dispatcher recognizes
datapaths structurally, so keeping a separate hand-written default
layout creates an unnecessary drift risk. This PR makes the default call
the public Layout D factory directly.
Finally, the register-side `tcgen05_atom_layout` must agree with the PTX
mapping from a logical `(row, col)` to `(laneid, wid_in_wg, register)`.
A self-consistent load/store round trip is not enough to prove that
mapping: raw PTX can move the same bits back even when the logical
layout label is wrong. Elementwise dispatch does consume that label, so
this PR adds direct mapping and elementwise compilation fences for all
`.16x*b` atom families. These are coverage additions; the production
atom-layout construction itself is unchanged here.
## Changes
- Add `sub_slab={0,1}` to `tmem_datapath_layout("F", ...)`.
- Encode the upper view as a `+16@TLane` offset and reject invalid
selectors, including nonzero selectors for Layout D.
- Classify TMEM layouts as `(datapath, sub_slab)` and emit `.16x*b` with
`row=(sub_slab + slab) * 16`.
- Preserve the existing M=128 behavior: Layout D with a 128-row `.16x*b`
fragment still emits two operations at rows `0` and `16`.
- Build the default M=128 TMEM layout through `tmem_datapath_layout("D",
...)`.
- Document the physical lane mapping and supported datapath/atom
combinations.
- Add direct atom-layout mapping coverage and a warpgroup elementwise
regression.
## Testing
- Static Layout F checks cover every logical row for both `sub_slab=0`
and `sub_slab=1`.
- B200 readback tests populate one Layout D accumulator, then verify
that lower and upper Layout F views reproduce the two corresponding
register halves for `.16x64b`, `.16x128b`, and `.16x256b`.
- Negative tests cover invalid sub-slab values and incompatible
datapath/atom pairings.
- Direct `(row, col) -> (laneid, wid_in_wg, register)` sweeps cover
supported `.16x*b` shapes and repetitions.
- Warpgroup elementwise codegen verifies that an atom-layout fragment
canonicalizes and slices correctly.
- Changed-files pre-commit checks.
---
docs/tirx/layout.rst | 31 +++++
.../tile_primitives/copy_async/tcgen05_ldst.rst | 57 +++++++--
python/tvm/backend/cuda/lang/alloc_pool.py | 7 ++
.../tile_primitive/copy_async/tcgen05_ldst.py | 75 ++++++++----
python/tvm/tirx/layout.py | 32 +++--
.../cuda/copy_async/test_tmem_16xnb.py | 136 ++++++++++++++++++++-
.../tile_primitive/cuda/elementwise/test_binary.py | 37 +++++-
7 files changed, 326 insertions(+), 49 deletions(-)
diff --git a/docs/tirx/layout.rst b/docs/tirx/layout.rst
index 690e03ef97..d8bff61f62 100644
--- a/docs/tirx/layout.rst
+++ b/docs/tirx/layout.rst
@@ -256,6 +256,37 @@ MMAs and double-buffering. So the one ``TileLayout`` model
expresses both the
accumulator (a pure placement, no replica) and its scale factors (a replicated,
routed placement) in the same tensor-memory address space.
+TMEM datapath layouts
+~~~~~~~~~~~~~~~~~~~~~
+
+``tmem_datapath_layout`` provides the canonical row placement for supported
+``tcgen05`` MMA datapaths:
+
+.. code-block:: python
+
+ from tvm.tirx.layout import tmem_datapath_layout
+
+ 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)
+
+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
+logical rows according to
+
+.. math::
+
+ \mathrm{TLane}
+ = 32\left\lfloor\frac{r}{16}\right\rfloor
+ + 16\,\mathrm{sub\_slab}
+ + (r \bmod 16).
+
+Thus F with ``sub_slab=0`` and ``sub_slab=1`` can describe lower- and
+upper-half aliases of the same 128-row Layout D allocation. The
+``tcgen05_ldst`` copy dispatch recognizes the layout and emits the matching
+``row=0`` or ``row=16`` instruction. Layout D already occupies both halves,
+so a nonzero ``sub_slab`` is rejected.
+
Beyond GPU registers
~~~~~~~~~~~~~~~~~~~~~~
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
index ee5510d2c7..d34147f60e 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
@@ -61,8 +61,8 @@ lowering:
- matched against a ``tcgen05_atom_layout`` (``.16x64b`` / ``.16x128b`` /
``.16x256b``) for the fast path; otherwise the ``.32x32b`` fallback
* - tmem datapath
- - classified ``D`` (M=128 identity) or ``F`` (M=64 scattered) — sets how
- fragment rows map to lanes
+ - 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
Demonstration program
----------------------
@@ -111,21 +111,61 @@ is a store (``tcgen05.st``).
fp32 columns) and the ``num`` count. If nothing matches it falls back to
``.32x32b`` and probes ``num ∈ {1, 2, 4, 8, …}`` against the column width.
-**3. Issue per datapath slab.** For an M=128 ``.16x*b`` copy the fragment
spans two
-16-row slabs, so the warps issue the atom twice (``row = 0`` and ``row =
16``); the
-``.32x32b`` path covers M=128 in a single issue (``row = 0``):
+**3. Issue per datapath slab.** For an M=128 Layout D ``.16x*b`` copy the
+fragment spans two 16-row slabs, so the warps issue the atom twice
+(``row = 0`` and ``row = 16``). An M=64 Layout F copy issues once, at
+``row = 0`` for ``sub_slab=0`` or ``row = 16`` for ``sub_slab=1``. The
+``.32x32b`` path covers M=128 in a single issue at ``row = 0``:
.. code-block:: python
op = T.ptx.tcgen05.ld if load else T.ptx.tcgen05.st
- for slab in range(n_slabs): # 1 for .32x32b / M=64; 2 for
.16x*b M=128
+ classified = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
+ datapath, sub_slab = classified if classified is not None else (None, 0)
+ for slab in range(n_slabs): # 1 for M=64; 2 for M=128
op(tmem_buf.allocated_addr[0],
*[local_32b[reg_base + i] for i in range(regs_eff)],
- shape=shape, num=num_eff, row=slab * 16, col=col_off_32b)
+ shape=shape, num=num_eff,
+ row=(sub_slab + slab) * 16, col=col_off_32b)
The dispatch emits **no** wait — the caller issues ``tcgen05.wait.ld()`` /
``wait.st()`` (as in the demo).
+Selecting the upper F sub-slab
+-------------------------------
+
+``sub_slab`` is part of the tensor-memory layout rather than an option on
+``copy_async``. This keeps physical TMEM occupation explicit and lets two
+64-row buffers alias the lower and upper halves of one 128-row allocation:
+
+.. code-block:: python
+
+ from tvm.tirx.layout import tmem_datapath_layout
+
+ lower = T.decl_buffer(
+ (64, cols),
+ "float32",
+ scope="tmem",
+ allocated_addr=tmem_addr[0],
+ layout=tmem_datapath_layout("F", 64, cols, sub_slab=0),
+ )
+ upper = T.decl_buffer(
+ (64, cols),
+ "float32",
+ scope="tmem",
+ allocated_addr=tmem_addr[0],
+ layout=tmem_datapath_layout("F", 64, cols, sub_slab=1),
+ )
+
+ Tx.wg.copy_async(lower_frag, lower)
+ T.ptx.tcgen05.wait.ld()
+ Tx.wg.copy_async(upper_frag, upper)
+ T.ptx.tcgen05.wait.ld()
+
+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``.
+
Generated TIRx IR
-----------------
@@ -169,4 +209,5 @@ How inputs change the algorithm
shape/num logic)
* - datapath D vs F
- ``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
+ (``row = 0`` / ``row = 16``). ``F`` (M=64) scatters rows to lanes and
+ its layout selects one issue at ``row = 0`` or ``row = 16``
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py
b/python/tvm/backend/cuda/lang/alloc_pool.py
index c635999b86..c7ce8942a6 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -54,6 +54,13 @@ _POOL_UNSET = object()
def _default_tmem_layout(rows, cols):
+ # Keep the default M=128 layout structurally identical to Layout D. The
+ # tcgen05 copy dispatch recognizes default-allocated accumulators by that
+ # structure, so single-source the named datapath instead of duplicating it.
+ if rows == 128:
+ from tvm.tirx.layout import tmem_datapath_layout
+
+ return tmem_datapath_layout("D", rows, cols)
return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
index 081ea5a772..b3842303bb 100644
--- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
@@ -82,14 +82,17 @@ def _match_tcgen05_atom_layout(buf):
def _classify_tmem_datapath(tmem_buf):
- """Return ``"D"`` / ``"F"`` if ``tmem_buf.layout`` matches a known tcgen05
- datapath (PTX ISA §9.7.16.10.5), else ``None``.
+ """Return ``(datapath, sub_slab)`` if ``tmem_buf.layout`` matches a known
+ tcgen05 datapath (PTX ISA §9.7.16.10.5), else ``None``.
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")``.
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.
"""
if tmem_buf.layout is None:
return None
@@ -99,16 +102,24 @@ def _classify_tmem_datapath(tmem_buf):
cand = tmem_datapath_layout("D", 128, tmem_buf.shape[1]).canonicalize()
try:
tvm.ir.assert_structural_equal(buf_layout, cand)
- return "D"
+ return ("D", 0)
except (AssertionError, ValueError):
return None
if rows == 64:
- cand = tmem_datapath_layout("F", 64, tmem_buf.shape[1]).canonicalize()
- try:
- tvm.ir.assert_structural_equal(buf_layout, cand)
- return "F"
- except (AssertionError, ValueError):
- return None
+ # 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
+ # out-of-band copy_async option.
+ for sub_slab in (0, 1):
+ cand = tmem_datapath_layout(
+ "F", 64, tmem_buf.shape[1], sub_slab=sub_slab
+ ).canonicalize()
+ try:
+ tvm.ir.assert_structural_equal(buf_layout, cand)
+ return ("F", sub_slab)
+ except (AssertionError, ValueError):
+ continue
+ return None
return None
@@ -150,10 +161,13 @@ def _check_tmem_layout_for_atom(tmem_buf, atom_kind,
frag_rows):
M=128 variants, 64 for ``.16x*b`` M=64). If the buffer's layout is
unrecognized (i.e. it isn't Layout D or Layout F), the dispatch falls
back to the structural assertions below.
+
+ Returns ``(datapath, sub_slab)`` for a recognized compatible layout.
"""
- datapath = _classify_tmem_datapath(tmem_buf)
- if datapath is None:
+ classified = _classify_tmem_datapath(tmem_buf)
+ if classified is None:
return None
+ datapath, sub_slab = classified
allowed = _TMEM_ATOM_COMPAT.get((datapath, atom_kind, frag_rows), False)
if not allowed:
raise ValueError(
@@ -163,7 +177,7 @@ def _check_tmem_layout_for_atom(tmem_buf, atom_kind,
frag_rows):
f"buffer was allocated via tmem_pool.alloc(..., "
f"datapath={datapath!r})."
)
- return datapath
+ return (datapath, sub_slab)
def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) ->
PrimFunc | None:
@@ -195,10 +209,10 @@ def copy_tmem_local_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -> P
# Try the .16x* (M=64) path first by structural-matching the register-side
# layout against ``tcgen05_atom_layout(instr_shape, (64, K), dtype)``. The
- # TMEM-side layout is the standard (128, W):(1@TLane, 1@TCol); the M=64
- # fragment lives at lanes 0..15 of each warp's accessible slab (per PTX
- # 9.7.16.8.1), so each warp issues with row_offset=0 and collectively the
- # 4 warps cover all 64 rows.
+ # An M=64 TMEM-side Layout F fragment lives in either lanes 0..15 or
+ # 16..31 of each warp's accessible slab (per PTX 9.7.16.8.1). The layout
+ # selects the half-slab, and the four warps collectively cover all 64
+ # logical rows.
atom_match = _match_tcgen05_atom_layout(local_buf)
if atom_match is not None:
@@ -307,15 +321,14 @@ def _emit_16xnb_path(
"""``.16x*b`` fragment path using ``tcgen05.{ld,st}.<shape>.x<num>`` (one
of ``.16x64b``, ``.16x128b``, ``.16x256b``).
- Each of the warpgroup's 4 warps issues the atom with ``row_offset=0`` to
- cover lanes 0..15 of its 32-lane TMEM partition (one 16-row slab); the
- four warps collectively span M=64 rows. When ``frag_rows == 128`` the
- dispatch emits a second issue with ``row_offset=16`` to also cover lanes
- 16..31 of each warp's partition, doubling the fragment's row coverage to
- M=128. The two atoms share the same column footprint; the layout factory
- surfaces the combined per-thread register vector with the second slab's
- regs in the high half of the m-axis (so the dispatch can split regs
- contiguously between the two PTX calls).
+ For an M=64 Layout F fragment, each warp issues the atom at the lower or
+ upper 16-lane half selected by the TMEM layout, and the four warps
+ collectively span 64 rows. For an M=128 Layout D fragment, the dispatch
+ emits both ``row_offset=0`` and ``row_offset=16`` issues. The two atoms
+ share the same column footprint; the layout factory surfaces the combined
+ per-thread register vector with the second slab's regs in the high half of
+ the m-axis (so the dispatch can split regs contiguously between the two
+ PTX calls).
"""
# Per-atom column footprint in fp32 columns:
# .16x64b → 2N .16x128b → 4N .16x256b → 8N
@@ -344,7 +357,15 @@ def _emit_16xnb_path(
# warp partition rule and the atom's lane access pattern are baked into
# the hardware); the layout classification just keeps the buffer's
# logical row indexing in sync with the physical TMEM occupation.
- datapath = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
+ classified = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
+ datapath, sub_slab = classified if classified is not None else (None, 0)
+ # A .16x*b issue covers one 16-lane half-slab. A 64-row fragment issues
+ # once; a 128-row fragment issues for both halves. ``sub_slab`` comes from
+ # the TMEM layout and shifts the first issue to the upper half.
+ assert sub_slab + n_slabs <= 2, (
+ f".16x*b sub_slab={sub_slab} with frag_rows={frag_rows} exceeds the 2 "
+ "sub-slabs of each warp's 32-lane TMEM partition"
+ )
if datapath == "F":
# Layout F: buffer shape (64, W), scattered row→lane.
@@ -425,7 +446,7 @@ def _emit_16xnb_path(
op(
tmem_buf.allocated_addr[0],
*[local_32b[local_reg_base + reg_base + i] for i in
range(regs_eff)],
- shape=shape, num=num_eff, row=slab * 16, col=col_off_32b,
+ shape=shape, num=num_eff, row=(sub_slab + slab) * 16,
col=col_off_32b,
)
# fmt: on
return impl
diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py
index 6b1aac6be3..3557e35367 100644
--- a/python/tvm/tirx/layout.py
+++ b/python/tvm/tirx/layout.py
@@ -593,7 +593,8 @@ __all__ += ["tcgen05_atom_layout", "tmem_datapath_layout",
"wg_local_layout"]
# Supported today:
# - ``"D"``: M=128, ``.cta_group::1``, full datapath. Identity row→lane.
# - ``"F"``: M=64, non-``.ws``, half datapath (4x1 lane utilization).
-# Logical row r → physical lane (r // 16) * 32 + (r % 16).
+# Logical row r → physical lane
+# (r // 16) * 32 + sub_slab * 16 + (r % 16).
#
# Layouts A / B / C / E / G are reserved for future expansion.
@@ -601,7 +602,7 @@ __all__ += ["tcgen05_atom_layout", "tmem_datapath_layout",
"wg_local_layout"]
_TMEM_DATAPATH_ROWS = {"D": 128, "F": 64}
-def tmem_datapath_layout(datapath: str, rows: int, cols: int) -> "TileLayout":
+def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int =
0) -> "TileLayout":
"""Return the ``TileLayout`` for a tcgen05 MMA datapath.
See PTX ISA §9.7.16.10.5 for the datapath enumeration. The returned
@@ -621,6 +622,12 @@ def tmem_datapath_layout(datapath: str, rows: int, cols:
int) -> "TileLayout":
dimension: 128 for D, 64 for F.
cols : int
Logical column count.
+ 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. Layout D already spans both halves and therefore only
+ accepts ``0``.
Returns
-------
@@ -637,21 +644,32 @@ def tmem_datapath_layout(datapath: str, rows: int, cols:
int) -> "TileLayout":
raise ValueError(
f"tmem_datapath_layout: datapath={datapath!r} expects
rows={expected}, got {rows}"
)
+ if sub_slab not in (0, 1):
+ raise ValueError(f"tmem_datapath_layout: sub_slab must be 0 or 1, got
{sub_slab}")
tlane = Axis.get("TLane")
tcol = Axis.get("TCol")
if datapath == "D":
- # M=128, identity row→lane: row r ∈ [0, 128) → physical lane r.
+ # M=128, identity row→lane: row r ∈ [0, 128) → physical lane r. D
+ # already spans both 16-lane sub-slabs of every warp partition.
+ if sub_slab != 0:
+ raise ValueError(
+ "tmem_datapath_layout: datapath='D' (M=128) already spans both
"
+ "sub-slabs; sub_slab must be 0"
+ )
return TileLayout(S[(rows, cols) : (1 @ tlane, 1 @ tcol)])
# Layout F: M=64 scattered. Logical row r = wid * 16 + intra (wid ∈ [0,4),
- # intra ∈ [0,16)) → physical lane wid * 32 + intra, i.e.
- # ``r // 16`` is the warp selector and ``r % 16`` is the within-slab lane.
+ # intra ∈ [0,16)) → physical lane wid * 32 + sub_slab * 16 + intra.
# ``TileLayout`` decomposes a scalar row index via ``SplitCoord``
# (src/tirx/ir/layout/utils.cc), which uses row-major ordering: with
# shape ``(s0, s1)`` the FIRST iter receives ``coord // s1`` (the high
# bits) and the SECOND receives ``coord % s1`` (the low bits). So we
# pin the warp selector to iter 0 (extent 4, TLane stride 32) and the
- # within-slab lane to iter 1 (extent 16, TLane stride 1).
- return TileLayout(S[(4, 16, cols) : (32 @ tlane, 1 @ tlane, 1 @ tcol)])
+ # within-slab lane to iter 1 (extent 16, TLane stride 1), then shift the
+ # TMEM lane offset by 16 for the upper sub-slab.
+ spec = S[(4, 16, cols) : (32 @ tlane, 1 @ tlane, 1 @ tcol)]
+ if sub_slab:
+ spec = spec + (16 @ tlane)
+ return TileLayout(spec)
def wg_local_layout(cols, rows=128):
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py
index ff9f20c6a6..31c159dc94 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py
@@ -41,6 +41,7 @@ import pytest
import tvm
import tvm.testing
+from tvm.backend.cuda.lang.alloc_pool import _default_tmem_layout
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
@@ -362,28 +363,44 @@ def _next_pow2(x: int) -> int:
# even if the row→lane formula doesn't match PTX's actual behavior. This
# test bypasses compilation and checks the layout's ``apply`` method
# directly against ``_frag_row_to_tmem_lane`` for every M=64 logical row.
-def test_tmem_datapath_layout_F_row_to_lane_mapping():
[email protected]("sub_slab", [0, 1])
+def test_tmem_datapath_layout_F_row_to_lane_mapping(sub_slab):
"""Layout F: every logical row r ∈ [0, 64) must land at physical TMEM
- lane ``(r // 16) * 32 + (r % 16)`` — the canonical scatter that the
- ``.16x*b`` M=64 PTX accesses (warp i on lanes ``i * 32 .. i * 32 + 15``).
+ lane ``(r // 16) * 32 + sub_slab * 16 + (r % 16)``. The default lower
+ sub-slab is the canonical M=64 scatter; sub-slab 1 is its upper-half view.
"""
cols = 32
- layout = tmem_datapath_layout("F", 64, cols)
+ layout = tmem_datapath_layout("F", 64, cols, sub_slab=sub_slab)
for r in range(64):
for c in [0, 1, 7, 16, 31]:
# Use ``apply(coord, shape=[64, cols])`` so (r, c) gets flattened
# row-major before SplitCoord into the shard iters.
axis_values = layout.apply(r, c, shape=[64, cols])
- expected_lane = (r // 16) * 32 + (r % 16)
+ expected_lane = (r // 16) * 32 + sub_slab * 16 + (r % 16)
assert int(axis_values["TLane"]) == expected_lane, (
f"(r={r}, c={c}) mapped to TLane {int(axis_values['TLane'])}, "
- f"expected {expected_lane} (= (r//16)*32 + (r%16))"
+ f"expected {expected_lane} "
+ f"(= (r//16)*32 + {sub_slab}*16 + (r%16))"
)
assert int(axis_values["TCol"]) == c, (
f"(r={r}, c={c}) mapped to TCol {int(axis_values['TCol'])},
expected {c}"
)
+def test_tmem_datapath_layout_rejects_invalid_sub_slab():
+ with pytest.raises(ValueError, match="sub_slab must be 0 or 1"):
+ tmem_datapath_layout("F", 64, 8, sub_slab=2)
+ with pytest.raises(ValueError, match="already spans both sub-slabs"):
+ tmem_datapath_layout("D", 128, 8, sub_slab=1)
+
+
+def test_default_tmem_layout_M128_matches_datapath_D():
+ """The default accumulator layout and named Layout D must not drift."""
+ tvm.ir.assert_structural_equal(
+ _default_tmem_layout(128, 32), tmem_datapath_layout("D", 128, 32)
+ )
+
+
@pytest.mark.parametrize("shape", ["16x64b", "16x128b", "16x256b"])
@pytest.mark.parametrize("rep", [1, 2, 4])
def test_tcgen05_atom_layout_apply_matches_decompose_fp32(shape, rep):
@@ -435,6 +452,113 @@ def test_tmem_datapath_layout_D_row_to_lane_mapping():
)
[email protected]
[email protected](not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
[email protected]("shape,rep", [("16x256b", 4), ("16x128b", 4),
("16x64b", 8)])
+def test_tcgen05_16xnb_sub_slab_view_read(shape, rep):
+ """F sub-slab views split a known-correct M=128 read into low/high
halves."""
+ dtype = "float32"
+ cols = _COL_FACTOR_FP32[shape] * rep
+ regs128 = _REGS_FACTOR[shape] * rep * 2
+ regs64 = _REGS_FACTOR[shape] * rep
+ tmem_cols = max(32, _next_pow2(cols))
+ atom128 = tcgen05_atom_layout(shape, (128, cols), dtype)
+ atom64 = tcgen05_atom_layout(shape, (64, cols), dtype)
+ layout_d = tmem_datapath_layout("D", 128, tmem_cols)
+ layout_f0 = tmem_datapath_layout("F", 64, tmem_cols, sub_slab=0)
+ layout_f1 = tmem_datapath_layout("F", 64, tmem_cols, sub_slab=1)
+
+ @T.prim_func
+ def kernel(A_ptr: T.handle, B128_ptr: T.handle, B0_ptr: T.handle, B1_ptr:
T.handle) -> None:
+ A = T.match_buffer(A_ptr, (128, regs128), dtype)
+ B128 = T.match_buffer(B128_ptr, (128, regs128), dtype)
+ B0 = T.match_buffer(B0_ptr, (128, regs64), dtype)
+ B1 = T.match_buffer(B1_ptr, (128, regs64), dtype)
+ T.device_entry()
+ warp_id = T.warp_id([4])
+ T.cta_id([2])
+ wg_id = T.warpgroup_id([1])
+ T.warp_id_in_wg([4])
+ T.lane_id([32])
+ tid = T.thread_id([128])
+ tmem_addr = T.alloc_shared([1], "uint32")
+ if wg_id == 0:
+ if warp_id == 0:
+ T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=tmem_cols,
cta_group=1)
+ T.tvm_storage_sync("shared")
+ tmem_d = T.decl_buffer(
+ (128, tmem_cols),
+ dtype,
+ scope="tmem",
+ allocated_addr=tmem_addr[0],
+ layout=layout_d,
+ )
+ tmem_f0 = T.decl_buffer(
+ (64, tmem_cols),
+ dtype,
+ scope="tmem",
+ allocated_addr=tmem_addr[0],
+ layout=layout_f0,
+ )
+ tmem_f1 = T.decl_buffer(
+ (64, tmem_cols),
+ dtype,
+ scope="tmem",
+ allocated_addr=tmem_addr[0],
+ layout=layout_f1,
+ )
+ source = T.alloc_local((regs128,), dtype)
+ for i in range(regs128):
+ source[i] = A[tid, i]
+ T.cuda.cta_sync()
+ Tx.wg.copy_async(tmem_d[0:128, 0:cols], source.view(128, cols,
layout=atom128))
+ T.ptx.tcgen05.wait.st()
+ T.cuda.cta_sync()
+
+ full = T.alloc_local((regs128,), dtype)
+ Tx.wg.copy_async(full.view(128, cols, layout=atom128),
tmem_d[0:128, 0:cols])
+ T.ptx.tcgen05.wait.ld()
+ T.cuda.cta_sync()
+ for i in range(regs128):
+ B128[tid, i] = full[i]
+
+ lower = T.alloc_local((regs64,), dtype)
+ Tx.wg.copy_async(lower.view(64, cols, layout=atom64),
tmem_f0[0:64, 0:cols])
+ T.ptx.tcgen05.wait.ld()
+ T.cuda.cta_sync()
+ for i in range(regs64):
+ B0[tid, i] = lower[i]
+
+ upper = T.alloc_local((regs64,), dtype)
+ Tx.wg.copy_async(upper.view(64, cols, layout=atom64),
tmem_f1[0:64, 0:cols])
+ T.ptx.tcgen05.wait.ld()
+ T.cuda.cta_sync()
+ for i in range(regs64):
+ B1[tid, i] = upper[i]
+
+ if warp_id == 0:
+ T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+ T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=tmem_cols,
cta_group=1)
+
+ target = tvm.target.Target("cuda")
+ with target:
+ mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target,
tir_pipeline="tirx")
+ source_np = tvm.testing.generate_random_array(dtype, (128, regs128))
+
+ def run_and_check():
+ dev = tvm.cuda(0)
+ source_dev = tvm.runtime.tensor(source_np, dev)
+ full_dev = tvm.runtime.tensor(np.zeros((128, regs128), dtype), dev)
+ lower_dev = tvm.runtime.tensor(np.zeros((128, regs64), dtype), dev)
+ upper_dev = tvm.runtime.tensor(np.zeros((128, regs64), dtype), dev)
+ mod(source_dev, full_dev, lower_dev, upper_dev)
+ full_np = full_dev.numpy()
+ np.testing.assert_array_equal(lower_dev.numpy(), full_np[:,
:regs64])
+ np.testing.assert_array_equal(upper_dev.numpy(), full_np[:,
regs64:])
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
# Negative tests: the datapath/atom pairing matrix in ``tcgen05_ldst.py``
# must reject mismatched combinations. We construct a Layout F TMEM buffer
# (64 rows, scattered) and try to read it with a ``.16x*b`` M=128 atom,
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
index 06a1dd68a0..2156c24ea9 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
@@ -24,7 +24,7 @@ import tvm.testing
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
-from tvm.tirx.layout import S, TileLayout, wg_local_layout
+from tvm.tirx.layout import S, TileLayout, tcgen05_atom_layout, wg_local_layout
@pytest.mark.parametrize(
@@ -866,5 +866,40 @@ def test_binary_add_f16_scalar_fallback_dispatch():
)
+def test_mul_tcgen05_16x256b_atom_warpgroup_dispatch():
+ """A split laneid/wid_in_wg register atom must canonicalize before
slicing."""
+ rows, cols = 64, 64
+ regs_per_thread = 32
+ atom_layout = tcgen05_atom_layout("16x256b", (rows, cols), "float32")
+
+ @T.prim_func
+ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+ A = T.match_buffer(A_ptr, (128, regs_per_thread), "float32")
+ B = T.match_buffer(B_ptr, (128, regs_per_thread), "float32")
+ T.device_entry()
+ T.cta_id([1])
+ T.warpgroup_id([1])
+ T.warp_id_in_wg([4])
+ T.lane_id([32])
+ tid = T.thread_id_in_wg([128])
+ src = T.alloc_buffer((rows, cols), "float32", scope="local",
layout=atom_layout)
+ dst = T.alloc_buffer((rows, cols), "float32", scope="local",
layout=atom_layout)
+ src_local = src.local(regs_per_thread)
+ dst_local = dst.local(regs_per_thread)
+ for i in T.serial(regs_per_thread):
+ src_local[i] = A[tid, i]
+ Tx.wg.mul(dst, src, T.float32(3.0))
+ for i in T.serial(regs_per_thread):
+ B[tid, i] = dst_local[i]
+
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+ with target:
+ mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target,
tir_pipeline="tirx")
+ src = mod.mod.imports[0].inspect_source()
+ assert re.search(r"mul\.[a-z]+\.ftz\.f32x2", src) or re.search(
+ r"tvm_builtin_ptx_mul_packed_", src
+ ), f"expected packed mul_f32x2 for tcgen05 atom; got:\n{src[:2000]}"
+
+
if __name__ == "__main__":
tvm.testing.main()
