[BugFix][Hopper][Blackwell] Read the WGMMA/UMMA K-panel stride from the layout instead of the operand extent

testing/python/cuda/test_tilelang_cuda_tcgen05_operand_layout.py

@@ -367,3 +367,33 @@ def plain_2d(m, k):
         reference.is_k_major,
     )
     assert staged.slice_byte_offset - reference.slice_byte_offset == 2 * m_extent * k_extent * 2
+
+
+def test_tcgen05_descriptor_k_panel_stride_follows_the_buffer():
+    """An MN slice keeps the buffer's K-panel spacing, not the slice's.
+
+    ``block_K`` past one swizzle atom stores the operand as (MN, (atom, panels)),
+    so stepping to the next K panel moves by ``buffer_MN * atom`` elements. That
+    step must be read off the layout: reconstructing it from the operand extent
+    (``mn_extent * atom``) halves the stride for a half-height slice and makes
+    every K atom past the first read the wrong panel.
+    """
+
+    # Two M atoms tall so the slice origin stays on the 128-row atom grid that
+    # validate_tcgen05_operand_regions requires, and two K panels wide.
+    m_extent, k_extent = 256, 128
+    atom_elems = 64  # 128B swizzle atom over bfloat16
+
+    def plain_2d(m, k):
+        return m * atom_elems + k % atom_elems + (k // atom_elems) * m_extent * atom_elems
+
+    layout = _swizzle_128b((m_extent, k_extent), plain_2d)
+    buffer = tvm.tirx.decl_buffer((m_extent, k_extent), "bfloat16", scope="shared")
+
+    whole = compute_umma_descriptor(layout, buffer, False, region=_ranges((0, 0), (m_extent, k_extent)))
+    sliced = compute_umma_descriptor(layout, buffer, False, region=_ranges((128, 0), (128, k_extent)))
+
+    assert whole.k_panel_elems == m_extent * atom_elems
+    # Same spacing for the slice, even though its MN extent is half as large.
+    assert sliced.k_panel_elems == m_extent * atom_elems
+    assert sliced.slice_byte_offset == 128 * atom_elems * 2

testing/python/cuda/test_tilelang_cuda_wgmma_operand_layout.py

@@ -0,0 +1,71 @@
+"""Unit tests for the WGMMA/UMMA K-panel stride decode.
+
+``decode_k_panel_elems`` is the single source of truth for how far an atom offset
+steps to reach the next K swizzle-atom panel. Reconstructing that step from the
+operand's MN extent is wrong for an operand that is a slice of a wider buffer, so
+the decode must either return the layout's own stride or refuse — never fall back
+silently, which is the wrong-code path the decode exists to remove.
+"""
+
+import pytest
+
+import tilelang.testing  # noqa: F401  (registers the test main)
+from tilelang.layout import SwizzleMode, cute
+from tilelang.cuda.intrinsics.macro.wgmma_macro_generator import decode_k_panel_elems
+
+ATOM = 64  # elements per 128B swizzle atom at 16-bit dtypes
+B128 = SwizzleMode.SWIZZLE_128B
+NONE = SwizzleMode.NONE
+
+
+def _k_mode(shape, stride):
+    return cute.make_layout(shape, stride)
+
+
+def test_canonical_multi_panel_returns_the_layout_stride():
+    """(atom, panels):(1, panel_stride) -- the step is read off the layout."""
+    assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 8192]), True, B128, ATOM, "WGMMA") == 8192
+    assert decode_k_panel_elems(_k_mode([ATOM, 4], [1, 8192]), True, B128, ATOM, "WGMMA") == 8192
+    # The stride is the buffer's, so a taller buffer gives a larger step for the
+    # same operand shape -- exactly what a sliced operand needs.
+    assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 16384]), True, B128, ATOM, "WGMMA") == 16384
+
+
+def test_single_panel_and_mn_major_need_no_step():
+    """None is only returned where ``ki // k_atom_size`` is always zero."""
+    # Flat K mode: one atom wide, so there is no second panel to step to.
+    assert decode_k_panel_elems(_k_mode(ATOM, 1), True, B128, ATOM, "WGMMA") is None
+    # MN-major operands never reach the panel term.
+    assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 8192]), False, B128, ATOM, "WGMMA") is None
+    # Unswizzled layouts have no panels at all.
+    assert decode_k_panel_elems(_k_mode(128, 1), True, NONE, 128, "WGMMA") is None
+
+
+def test_panels_first_ordering_is_rejected_not_misread():
+    """A (panels, atom) K mode must not have its trailing atom stride read as the step.
+
+    Taking ``stride[-1]`` without checking that the leading sub-mode is the
+    contiguous atom would return 1 here — a silently wrong panel step.
+    """
+    with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"):
+        decode_k_panel_elems(_k_mode([2, ATOM], [8192, 1]), True, B128, ATOM, "WGMMA")
+
+
+def test_multi_panel_without_a_constant_stride_asserts():
+    """Refusing beats falling back to the extent-based reconstruction."""
+    # Flat but two atoms wide: no panel sub-mode to read a stride from.
+    with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"):
+        decode_k_panel_elems(_k_mode(2 * ATOM, 1), True, B128, ATOM, "UMMA")
+    # Three sub-modes: panel spacing is not expressible as one scalar step.
+    with pytest.raises(AssertionError, match=r"spans 4 swizzle-atom panels"):
+        decode_k_panel_elems(_k_mode([ATOM, 2, 2], [1, 8192, 16384]), True, B128, ATOM, "WGMMA")
+
+
+def test_partial_atom_extent_still_counts_as_two_panels():
+    """A K extent that is not a whole number of atoms must not read as single-panel."""
+    with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"):
+        decode_k_panel_elems(_k_mode(ATOM + 16, 1), True, B128, ATOM, "WGMMA")
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()

testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py

@@ -0,0 +1,178 @@
+"""A shared operand that is a slice of a wider buffer must match the same GEMM
+written with one buffer per slice.
+
+Regression coverage for the WGMMA K-panel stride: a K-major operand whose
+``block_K`` spans more than one swizzle atom is stored as (MN, (atom, panels)),
+and stepping from one K panel to the next moves by the *buffer's* MN extent. The
+offset formulas used to reconstruct that step from the operand's own extent,
+which is smaller for a slice, so ``T.gemm(A_s, B_s[0:64, :], ...)`` silently read
+the wrong panel for every ``ki`` past the first atom. ``block_K == 64`` (bf16 =
+one 128B atom, single panel) never exercised the term and stayed correct.
+"""
+
+import pytest
+import torch
+
+import tilelang
+import tilelang.language as T
+import tilelang.testing
+
+ACCUM = T.float32
+
+
+def _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced):
+    """Split one GEMM into two half-width ones over a shared operand.
+
+    ``operand`` selects which operand is split ("A" splits M, "B" splits N),
+    ``transposed`` gives that operand's ``transpose_`` flag (which decides whether
+    the split axis is the buffer's leading or trailing one), and ``sliced`` picks
+    the shape under test (one buffer consumed as two slices) versus the reference
+    (one buffer per half).
+    """
+    trans_a = operand == "A" and transposed
+    trans_b = operand == "B" and transposed
+    a_shape = (K, M) if trans_a else (M, K)
+    b_shape = (N, K) if trans_b else (K, N)
+    split_full = block_N if operand == "B" else block_M
+    rest = split_full - split
+
+    def a_tile(extent):
+        return (block_K, extent) if trans_a else (extent, block_K)
+
+    def b_tile(extent):
+        return (extent, block_K) if trans_b else (block_K, extent)
+
+    @T.prim_func
+    def main(
+        A: T.Tensor(a_shape, dtype),
+        B: T.Tensor(b_shape, dtype),
+        C: T.Tensor((M, N), dtype),
+    ):
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            C0 = T.alloc_fragment((split, block_N) if operand == "A" else (block_M, split), ACCUM)
+            C1 = T.alloc_fragment((rest, block_N) if operand == "A" else (block_M, rest), ACCUM)
+            T.clear(C0)
+            T.clear(C1)
+
+            if operand == "B":
+                A_s = T.alloc_shared(a_tile(block_M), dtype)
+                if sliced:
+                    B_s = T.alloc_shared(b_tile(block_N), dtype)
+                else:
+                    B_s0 = T.alloc_shared(b_tile(split), dtype)
+                    B_s1 = T.alloc_shared(b_tile(rest), dtype)
+            else:
+                B_s = T.alloc_shared(b_tile(block_N), dtype)
+                if sliced:
+                    A_s = T.alloc_shared(a_tile(block_M), dtype)
+                else:
+                    A_s0 = T.alloc_shared(a_tile(split), dtype)
+                    A_s1 = T.alloc_shared(a_tile(rest), dtype)
+
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=0):
+                if operand == "B":
+                    if trans_a:
+                        T.copy(A[ko * block_K, by * block_M], A_s)
+                    else:
+                        T.copy(A[by * block_M, ko * block_K], A_s)
+                    if sliced:
+                        if trans_b:
+                            T.copy(B[bx * block_N, ko * block_K], B_s)
+                            T.gemm(A_s, B_s[0:split, :], C0, trans_a, trans_b)
+                            T.gemm(A_s, B_s[split:block_N, :], C1, trans_a, trans_b)
+                        else:
+                            T.copy(B[ko * block_K, bx * block_N], B_s)
+                            T.gemm(A_s, B_s[:, 0:split], C0, trans_a, trans_b)
+                            T.gemm(A_s, B_s[:, split:block_N], C1, trans_a, trans_b)
+                    else:
+                        if trans_b:
+                            T.copy(B[bx * block_N, ko * block_K], B_s0)
+                            T.copy(B[bx * block_N + split, ko * block_K], B_s1)
+                        else:
+                            T.copy(B[ko * block_K, bx * block_N], B_s0)
+                            T.copy(B[ko * block_K, bx * block_N + split], B_s1)
+                        T.gemm(A_s, B_s0, C0, trans_a, trans_b)
+                        T.gemm(A_s, B_s1, C1, trans_a, trans_b)
+                else:
+                    if trans_b:
+                        T.copy(B[bx * block_N, ko * block_K], B_s)
+                    else:
+                        T.copy(B[ko * block_K, bx * block_N], B_s)
+                    if sliced:
+                        if trans_a:
+                            T.copy(A[ko * block_K, by * block_M], A_s)
+                            T.gemm(A_s[:, 0:split], B_s, C0, trans_a, trans_b)
+                            T.gemm(A_s[:, split:block_M], B_s, C1, trans_a, trans_b)
+                        else:
+                            T.copy(A[by * block_M, ko * block_K], A_s)
+                            T.gemm(A_s[0:split, :], B_s, C0, trans_a, trans_b)
+                            T.gemm(A_s[split:block_M, :], B_s, C1, trans_a, trans_b)
+                    else:
+                        if trans_a:
+                            T.copy(A[ko * block_K, by * block_M], A_s0)
+                            T.copy(A[ko * block_K, by * block_M + split], A_s1)
+                        else:
+                            T.copy(A[by * block_M, ko * block_K], A_s0)
+                            T.copy(A[by * block_M + split, ko * block_K], A_s1)
+                        T.gemm(A_s0, B_s, C0, trans_a, trans_b)
+                        T.gemm(A_s1, B_s, C1, trans_a, trans_b)
+
+            if operand == "B":
+                T.copy(C0, C[by * block_M, bx * block_N])
+                T.copy(C1, C[by * block_M, bx * block_N + split])
+            else:
+                T.copy(C0, C[by * block_M, bx * block_N])
+                T.copy(C1, C[by * block_M + split, bx * block_N])
+
+    return main
+
+
+def run_sliced_operand(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed):
+    torch_dtype = getattr(torch, dtype)
+    trans_a = operand == "A" and transposed
+    trans_b = operand == "B" and transposed
+
+    torch.manual_seed(0)
+    a = torch.randn((K, M) if trans_a else (M, K), device="cuda", dtype=torch_dtype)
+    b = torch.randn((N, K) if trans_b else (K, N), device="cuda", dtype=torch_dtype)
+    ref = (a.T if trans_a else a).float() @ (b.T if trans_b else b).float()
+
+    outputs = {}
+    for sliced in (False, True):
+        kernel = tilelang.compile(
+            _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced),
+            out_idx=[2],
+            target="cuda",
+            pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
+        )
+        outputs[sliced] = kernel(a, b).float()
+
+    # The two shapes issue the same instructions over the same values, so they
+    # agree bit-for-bit; the fp32 reference only bounds the accumulation error.
+    torch.testing.assert_close(outputs[True], outputs[False], rtol=0, atol=0)
+    torch.testing.assert_close(outputs[True], ref, rtol=1e-2, atol=1)
+
+
+@tilelang.testing.requires_cuda
+@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
+@pytest.mark.parametrize("operand", ["A", "B"])
+@pytest.mark.parametrize("transposed", [False, True])
+@pytest.mark.parametrize("block_K", [64, 128, 256])
+def test_gemm_sliced_shared_operand(operand, transposed, block_K):
+    """block_K 128/256 put more than one swizzle atom on the K axis."""
+    run_sliced_operand(
+        M=256,
+        N=256,
+        K=512,
+        block_M=128,
+        block_N=128,
+        block_K=block_K,
+        dtype="bfloat16",
+        split=64,
+        operand=operand,
+        transposed=transposed,
+    )
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()