diff --git a/projects/hipblaslt/CHANGELOG.md b/projects/hipblaslt/CHANGELOG.md index acaa2a339664..a5a16dcb628a 100644 --- a/projects/hipblaslt/CHANGELOG.md +++ b/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. diff --git a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s b/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s index edd84ab0f857..d63073a0ea64 100644 --- a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s +++ b/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 diff --git a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950.s b/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950.s index 1474d0fe0873..cb37cd885def 100644 --- a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950.s +++ b/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 diff --git a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_HHS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s b/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_HHS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s index fa9cc71de9be..2a8be6b0d345 100644 --- a/projects/hipblaslt/tensilelite/Tensile/CustomKernels/tensile/Custom_Cijk_Alik_Bljk_HHS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname0_gfx950.s +++ b/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 diff --git a/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_CustomKernels_final_loop_wait.py b/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_CustomKernels_final_loop_wait.py new file mode 100644 index 000000000000..d14cdaaa132e --- /dev/null +++ b/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}" diff --git a/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_custom_tn_prefetch_waits.py b/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_custom_tn_prefetch_waits.py new file mode 100644 index 000000000000..bd2af5bd622a --- /dev/null +++ b/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)