fix(hipblaslt): wait for LDS operands in handwritten gfx950 final loops

projects/hipblaslt/CHANGELOG.md

@@ -46,6 +46,7 @@ Full documentation for hipBLASLt is available at [rocm.docs.amd.com/projects/hip
 
 ### Resolved issues
 
+* Fixed premature LDS reads when handwritten gfx950 BF16 and FP16 TN kernels skip the second global prefetch at `K=64`.
 * Fixed `hipblaslt-bench` using C's batch stride for D and computing its CPU reference with the wrong layout when C and D have different leading dimensions or batch strides.
 * Fixed output-amax accumulation omitting packed-store values and returning zero when C/D scaling is disabled. Invalid Stream-K or split-reduction combinations with output-amax are rejected during solution validation.
 * Fixed GEMM output scaling reading C/D scale values before their scalar memory loads completed.

projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s

@@ -1412,7 +1412,10 @@ s_xor_b32 s[sgprLocalWriteAddrA], s[sgprSwapA], s[sgprLocalWriteAddrA] // swap R
 /* local write swap b */
 s_xor_b32 s[sgprLocalWriteAddrB], s[sgprSwapB], s[sgprLocalWriteAddrB] // swap Red Blk SGPR
 s_cmp_eq_u32 s[sgprLoopCounterL], 0x1              // PGR=2 but only 1 loop
-s_cbranch_scc1 label_skipPGR2                      // PGR=2 but only 1 loop
+s_cbranch_scc0 label_PrefetchGlobalRead2           // Issue the second prefetch when present.
+s_waitcnt vmcnt(0)                                // Complete the only prefetch before LDS reads.
+s_branch label_skipPGR2
+label_PrefetchGlobalRead2:
 s_mov_b32 m0, s[sgprLocalWriteAddrA]               // m0 <- LDS write address
 
 buffer_load_dwordx4 v[vgprGlobalReadOffsetA+0], s[sgprSrdA:sgprSrdA+3], 0 offen offset:0, lds // G -> Reg 0_0_0_0
@@ -2376,6 +2379,7 @@ s_cbranch_scc1 label_OptNLL_End                    // jump if edges required
 
 
 /*  mfmaIndex:0  */
+s_waitcnt lgkmcnt(0) // Complete LDS reads before the optimized final loop.
 v_mfma_f32_16x16x32_bf16 acc[0:3], v[vgprValuB_X0_I0+0+0+0:vgprValuB_X0_I0+0+0+0+3], v[vgprValuA_X0_I0+0+0+0:vgprValuA_X0_I0+0+0+0+3], acc[0:3] // left value = acc[0+0:3+0]
 ds_read_b128 v[vgprValuA_X1_I0+0:vgprValuA_X1_I0+0+3], v[vgprLocalReadAddrA] offset:64 // L -> Reg lro=32 swapByteOffset=0 ti=256 vIdx=0 eIdx=0 rIdx=0 oIdx=0 buffer=1 iui=0
 

projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950.s

@@ -1436,7 +1436,10 @@ s_xor_b32 s[sgprLocalWriteAddrA], s[sgprSwapA], s[sgprLocalWriteAddrA] // swap R
 /* local write swap b */
 s_xor_b32 s[sgprLocalWriteAddrB], s[sgprSwapB], s[sgprLocalWriteAddrB] // swap Red Blk SGPR
 s_cmp_eq_u32 s[sgprLoopCounterL], 0x1              // PGR=2 but only 1 loop
-s_cbranch_scc1 label_skipPGR2                      // PGR=2 but only 1 loop
+s_cbranch_scc0 label_PrefetchGlobalRead2           // Issue the second prefetch when present.
+s_waitcnt vmcnt(0)                                // Complete the only prefetch before LDS reads.
+s_branch label_skipPGR2
+label_PrefetchGlobalRead2:
 s_mov_b32 m0, s[sgprLocalWriteAddrA]               // m0 <- LDS write address
 
 buffer_load_dwordx4 v[vgprGlobalReadOffsetA+0], s[sgprSrdA:sgprSrdA+3], 0 offen offset:0, lds // G -> Reg 0_0_0_0
@@ -2684,6 +2687,7 @@ s_cbranch_scc1 label_OptNLL_End                    // jump if edges required
 
 
 /*  mfmaIndex:0  */
+s_waitcnt lgkmcnt(0) // Complete LDS reads before the optimized final loop.
 v_mfma_f32_16x16x32_bf16 acc[0:3], v[vgprValuB_X0_I0+0+0+0:vgprValuB_X0_I0+0+0+0+3], v[vgprValuA_X0_I0+0+0+0:vgprValuA_X0_I0+0+0+0+3], acc[0:3] // left value = acc[0+0:3+0]
 ds_read_b128 v[vgprValuA_X1_I0+0:vgprValuA_X1_I0+0+3], v[vgprLocalReadAddrA] offset:64 // L -> Reg lro=32 swapByteOffset=0 ti=256 vIdx=0 eIdx=0 rIdx=0 oIdx=0 buffer=1 iui=0
 

projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_HHS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s

@@ -1696,7 +1696,10 @@ s_xor_b32 s[sgprLocalWriteAddrA], s[sgprSwapA], s[sgprLocalWriteAddrA] // swap R
 /* local write swap b */
 s_xor_b32 s[sgprLocalWriteAddrB], s[sgprSwapB], s[sgprLocalWriteAddrB] // swap Red Blk SGPR
 s_cmp_eq_u32 s[sgprLoopCounterL], 0x1              // PGR=2 but only 1 loop
-s_cbranch_scc1 label_skipPGR2                      // PGR=2 but only 1 loop
+s_cbranch_scc0 label_PrefetchGlobalRead2           // Issue the second prefetch when present.
+s_waitcnt vmcnt(0)                                // Complete the only prefetch before LDS reads.
+s_branch label_skipPGR2
+label_PrefetchGlobalRead2:
 
 s_mov_b32 m0, s[sgprLocalWriteAddrA]               // m0 <- LDS write address
 buffer_load_dwordx4 v[vgprGlobalReadOffsetA+0], s[sgprSrdA:sgprSrdA+3], 0 offen offset:0, lds // G -> Reg 0_0_0_0
@@ -2942,6 +2945,7 @@ s_cbranch_scc1 label_OptNLL_End                    // jump if edges required
 
 
 /*  mfmaIndex:0  */
+s_waitcnt lgkmcnt(0) // Complete LDS reads before the optimized final loop.
 v_mfma_f32_16x16x32_f16 acc[0:3], v[vgprValuB_X0_I0+0+0+0:vgprValuB_X0_I0+0+0+0+3], v[vgprValuA_X0_I0+0+0+0:vgprValuA_X0_I0+0+0+0+3], acc[0:3] // left value = acc[0+0:3+0]
 ds_read_b128 v[vgprValuA_X1_I0+0:vgprValuA_X1_I0+0+3], v[vgprLocalReadAddrA] offset:64 // L -> Reg lro=32 swapByteOffset=0 ti=256 vIdx=0 eIdx=0 rIdx=0 oIdx=0 buffer=1 iui=0
 

projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_CustomKernels_final_loop_wait.py

@@ -0,0 +1,89 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+"""Check LDS operand readiness on the handwritten gfx950 optimized final path."""
+
+from collections import deque
+import re
+
+import pytest
+
+from Tensile.CustomKernels import getCustomKernelContents
+
+pytestmark = pytest.mark.unit
+
+KERNELS = [
+    f"Custom_Cijk_Alik_Bljk_{dtype}_BH_MT256x256x64_MI16x16x1_UserArgs_shortname{variant}_gfx950"
+    for dtype, variant in [("BBS", 0), ("BBS", 1), ("HHS", 0)]
+]
+
+
+def _pending_final_loop_operands(source):
+    """Follow the beta=0, alpha=1, GSU=1, full-tile fallthrough path.
+
+    Start at the preceding LDS drain, track the subsequent reads, and check
+    their MFMA consumers across the final-loop boundary. These kernels use
+    only ds_read_b128 for LGKM operations in this interval. FIFO retirement
+    accounts for gfx950's 15 outstanding LGKM operations as well as waits.
+    This is a focused dependency check, not a general assembly interpreter.
+    """
+    definitions = dict(re.findall(r"^\.set\s+(\w+),\s*([^/\n]+)", source, re.M))
+
+    def value(expression):
+        # Register expressions in these kernels are sums of constants and
+        # .set symbols. Fail on unsupported syntax instead of evaluating code.
+        terms = expression.strip().split("+")
+        result = 0
+        for term in terms:
+            term = term.strip()
+            if re.fullmatch(r"0x[0-9a-fA-F]+|[0-9]+", term):
+                result += int(term, 0)
+            else:
+                assert re.fullmatch(r"\w+", term), term
+                result += value(definitions[term])
+        return result
+
+    def registers(operand):
+        bounds = operand.split(":")
+        first, last = value(bounds[0]), value(bounds[-1])
+        return set(range(first, last + 1))
+
+    boundary = source.index("label_toPGR1:")
+    start = source.rindex("s_waitcnt lgkmcnt(0)", 0, boundary)
+    end = source.index("label_OptNLL_End:", boundary)
+    # All branches selecting the optimized path fall through. Its final
+    # branch exits to the epilogue after the compute instructions checked here.
+    lines = source[start:end].splitlines()
+    pending = deque()
+    violations = []
+    reads = mfmas = 0
+    for line in lines:
+        instruction = line.split("//", 1)[0].strip()
+        assert not instruction.startswith(("s_load_", "ds_write_")), instruction
+        wait = re.match(r"s_waitcnt\s+lgkmcnt\((\d+)\)", instruction)
+        if wait:
+            while len(pending) > int(wait[1]):
+                pending.popleft()
+        elif instruction.startswith("ds_read_"):
+            assert instruction.startswith("ds_read_b128 "), instruction
+            destination = re.match(r"ds_read_b128\s+v\[([^\]]+)\]", instruction)
+            assert destination, instruction
+            pending.append(registers(destination[1]))
+            reads += 1
+            if len(pending) > 15:
+                pending.popleft()
+        elif instruction.startswith("v_mfma_"):
+            operands = re.findall(r"v\[([^\]]+)\]", instruction)
+            assert len(operands) == 2, instruction
+            used = registers(operands[0]) | registers(operands[1])
+            outstanding = set().union(*pending)
+            if used & outstanding:
+                violations.append((sorted(used & outstanding), instruction))
+            mfmas += 1
+    assert reads >= 16 and mfmas >= 128, "The expected compute path was not checked"
+    return violations
+
+
+@pytest.mark.parametrize("kernel", KERNELS)
+def test_optimized_final_loop_consumes_completed_lds_reads(kernel):
+    violations = _pending_final_loop_operands(getCustomKernelContents(kernel))
+    assert not violations, f"{kernel}: pending LDS operands at {violations}"

projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_custom_tn_prefetch_waits.py

@@ -0,0 +1,90 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Check DirectToLDS completion before the handwritten gfx950 TN local reads.
+
+The PGR2 waits count the younger second prefetch. On the one-DepthU path,
+that prefetch is skipped, so the first prefetch needs its own completion
+before the barrier makes its LDS writes visible to the local reads.
+"""
+
+import re
+
+import pytest
+
+from Tensile.resources import custom_kernel_text
+
+pytestmark = pytest.mark.unit
+
+KERNELS = [
+    f"Custom_Cijk_Alik_Bljk_{dtype}_BH_MT256x256x64_MI16x16x1_UserArgs_shortname{variant}_gfx950"
+    for dtype, variant in [("BBS", 0), ("BBS", 1), ("HHS", 0)]
+]
+
+
+def _prefetch_hazards(source, loop_count):
+    # No VM wait precedes the PGR2 branch: the first A/B prefetch is pending.
+    start = source.index("/* prefetch: global -> local */")
+    branch = source.index("s_cmp_eq_u32 s[sgprLoopCounterL], 0x1", start)
+    prefix = source[start:branch]
+    end = source.index("label_openLoopL:", branch)
+    code = [line.split("//", 1)[0].strip() for line in source[branch:end].splitlines()]
+    labels = {line[:-1]: i for i, line in enumerate(code) if line.endswith(":")}
+    pending, issued, completed, visible = [], set(), set(), set()
+    prefetch_counts = {"A": 0, "B": 0}
+
+    def issue(instruction):
+        match = re.search(r"s\[sgprSrd([AB]):", instruction)
+        assert match and "lds" in instruction
+        tensor = match.group(1)
+        index = prefetch_counts[tensor]
+        event = (tensor, index // 8, index % 8)
+        prefetch_counts[tensor] += 1
+        pending.append(event)
+        issued.add(event)
+
+    for instruction in prefix.splitlines():
+        if instruction.startswith("buffer_load_"):
+            issue(instruction)
+        assert not re.search(r"s_waitcnt.*vmcnt", instruction)
+    assert prefetch_counts == {"A": 8, "B": 8}
+    ip, scc, hazards, reads = 0, False, [], 0
+    while ip < len(code):
+        instruction = code[ip]
+        if instruction.startswith("s_cmp_eq_u32 s[sgprLoopCounterL], 0x1"):
+            scc = loop_count == 1
+        elif instruction.startswith("s_cbranch_scc"):
+            operation, target = instruction.split()
+            if scc == operation.endswith("1"):
+                ip = labels[target]
+                continue
+        elif instruction.startswith("s_branch "):
+            ip = labels[instruction.split()[1]]
+            continue
+        elif instruction.startswith("buffer_load_"):
+            issue(instruction)
+        elif instruction.startswith("s_waitcnt "):
+            match = re.search(r"vmcnt\((\d+)\)", instruction)
+            if match:
+                threshold = int(match.group(1))
+                retire = max(0, len(pending) - threshold)
+                completed.update(pending[:retire])
+                del pending[:retire]
+        elif instruction.startswith("s_barrier"):
+            visible.update(completed)
+        elif instruction.startswith("ds_read_b128 "):
+            tensor = re.search(r"v\[vgprLocalReadAddr([AB])\]", instruction).group(1)
+            first_prefetch = {event for event in issued if event[0:2] == (tensor, 0)}
+            if not first_prefetch <= visible:
+                hazards.append((ip, tensor, sorted(first_prefetch - visible)))
+            reads += 1
+        ip += 1
+    assert reads == 16
+    assert prefetch_counts == {"A": 8 * min(loop_count, 2), "B": 8 * min(loop_count, 2)}
+    return hazards
+
+
+@pytest.mark.parametrize("kernel", KERNELS)
+@pytest.mark.parametrize("loop_count", [1, 2])
+def test_first_prefetch_is_visible_before_local_reads(kernel, loop_count):
+    assert not _prefetch_hazards(custom_kernel_text(kernel), loop_count)