From 02046f68fde9c7a509fdfc131509c523de2fcf09 Mon Sep 17 00:00:00 2001 From: James Newling Date: Wed, 21 Oct 2026 12:22:58 -0700 Subject: [PATCH] fix(tensilelite): match 16-bit TN CMS prefetch and loop load order The gfx950 MT256x160x64 TN custom schedule issues B before A, but its prologue prefetched A before B. Partial VMEM waits in the first loop and the no-global-load loop can leave required B DirectToLds writes pending at LDS reads. Set the existing SwapGlobalReadOrder property and exchange schedule stream keys so both prefetches issue B before A while retaining both SIMD loop bodies and every wait/barrier. Extend the emitted-assembly dependency regression to BF16 and FP16 with ordinary and StreamK scheduling. JIRA ID : AIHPBLAS-5228 --- .../tensilelite/Components/CustomSchedule.py | 9 ++-- .../unit/test_data/bf16_tn_dtl_prefetch.yaml | 48 +++++++++++++++++ .../Tests/unit/test_tf32_dtl_prefetch.py | 53 +++++++++++++++---- 3 files changed, 96 insertions(+), 14 deletions(-) create mode 100644 projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_data/bf16_tn_dtl_prefetch.yaml diff --git a/projects/hipblaslt/tensilelite/tensilelite/Components/CustomSchedule.py b/projects/hipblaslt/tensilelite/tensilelite/Components/CustomSchedule.py index 0eed0cdb313f..676c291c8bf3 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/Components/CustomSchedule.py +++ b/projects/hipblaslt/tensilelite/tensilelite/Components/CustomSchedule.py @@ -1884,6 +1884,9 @@ def _get_schedule_256x160x64_16bit(kernel, useLDSTr, TLDS): nglshift = nllshift = 13 opt1 = ScheduleInfo(2, numMfma, optSchedule, syncCode, nglshift, nllshift) elif isTN(kernel) and (not useLDSTr) and TLDS==1: + # The partial VMEM waits need B before A in the prefetch as well as + # the loop. GRA/GRB name the first/second streams when swapping. + kernel["SwapGlobalReadOrder"] = 1 syncTable = [ -1, SWaitCnt(dscnt=3, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1 (partial) before starting main loop"), 4, SWaitCnt(dscnt=0+2, vlcnt=-1, vscnt=-1, comment="Wait for prior LRA1 (complete) for the remaining main loop"), @@ -1909,11 +1912,11 @@ def _get_schedule_256x160x64_16bit(kernel, useLDSTr, TLDS): 'LRA0' : [[13,15,18,21,24,26,28,30], [13,16,19,22,25,27,29,31]], - # GRB must not start before the SYNC at idx 15 (LRB0 completion). - 'GRB' : [[14,14, 17,17, 20,20, 23,23, 26,26], + # B loads follow the barrier that completes LRB0. + 'GRA' : [[14,14, 17,17, 20,20, 23,23, 26,26], [15,15, 18,18, 21,21, 24,24, 27,27]], # Buffer loads. - 'GRA' : [[40,40, 43,43, 46,46, 49,49, 59,59, 62,62, 65,65, 67,67], + 'GRB' : [[40,40, 43,43, 46,46, 49,49, 59,59, 62,62, 65,65, 67,67], [41,41, 44,44, 47,47, 57,57, 60,60, 63,63, 66,66, 68,68]], # Prefetch next iteration. # Need 5 local reads for B (MIWaveTileB=5). diff --git a/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_data/bf16_tn_dtl_prefetch.yaml b/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_data/bf16_tn_dtl_prefetch.yaml new file mode 100644 index 000000000000..1afc86e3b066 --- /dev/null +++ b/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_data/bf16_tn_dtl_prefetch.yaml @@ -0,0 +1,48 @@ +# Copyright Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +GlobalParameters: + SyncsPerBenchmark: 0 + NumElementsToValidate: 0 + NumWarmups: 0 + NumBenchmarks: 1 + Device: 0 +BenchmarkProblems: +- - OperationType: GEMM + DataType: B + DestDataType: B + ComputeDataType: S + HighPrecisionAccumulate: true + TransposeA: true + TransposeB: false + UseBeta: true + Batched: true + - InitialSolutionParameters: null + BenchmarkCommonParameters: + - KernelLanguage: [Assembly] + ForkParameters: + - MatrixInstruction: + - [16, 16, 32, 1, 1, 8, 5, 2, 2] + - DepthU: [64] + - PrefetchGlobalRead: [2] + - PrefetchLocalRead: [1] + - ScheduleIterAlg: [3] + - ExpandPointerSwap: [true] + - TransposeLDS: [1] + - LDSTrInst: [false] + - LocalReadVectorWidth: [8] + - GlobalReadVectorWidthA: [8] + - GlobalReadVectorWidthB: [8] + - DirectToLds: [1] + - LdsPadA: [16] + - LdsPadB: [32] + - LdsBlockSizePerPadA: [1024] + - LdsBlockSizePerPadB: [1024] + - 1LDSBuffer: [0] + - StaggerU: [0] + - SourceSwap: [true] + - UseCustomMainLoopSchedule: [1] + BenchmarkJoinParameters: null + BenchmarkFinalParameters: + - ProblemSizes: + - Exact: [256, 160, 1, 192] diff --git a/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_tf32_dtl_prefetch.py b/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_tf32_dtl_prefetch.py index cf09ae83a934..c1e017aa0b3c 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_tf32_dtl_prefetch.py +++ b/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_tf32_dtl_prefetch.py @@ -14,21 +14,22 @@ pytestmark = pytest.mark.unit -def _assert_prefetch_ready(assembly, loads_a): +def _assert_prefetch_ready(assembly, loads_a, loads_b=4): # Keep the last prefetched set pending, just as the nonzero prologue wait # permits. The first loop iteration must retire it before next-bank reads. prefetch = assembly.split("// PGR=2 but only 1 loop", 2)[2] prefetch = prefetch.split("label_skipPGR2_1:", 1)[0] load_pattern = r"^buffer_load_.*vgprGlobalReadOffset([AB]).*\blds\b" initial = re.findall(load_pattern, prefetch, re.MULTILINE) - assert len(initial) == loads_a + 4 - assert f"s_waitcnt vmcnt({loads_a + 4})" in assembly + assert len(initial) == loads_a + loads_b + assert f"s_waitcnt vmcnt({loads_a + loads_b})" in assembly # Test both SIMD schedules. Unlike the steady-state CMS validator, seed # the queue from the actual prologue, not a previous main-loop iteration. paths = re.findall( r"^label_LoopBeginL_\d+:\n(.*?)(?=^s_cbranch_scc0 label_LoopBeginL_\d+)", - assembly, re.MULTILINE | re.DOTALL, + assembly, + re.MULTILINE | re.DOTALL, ) assert len(paths) == 2 for body in paths: @@ -49,9 +50,10 @@ def _assert_prefetch_ready(assembly, loads_a): read = re.match(r"ds_read_.*v\[vgprLocalReadAddr([AB])", line) if read and read[1] in swapped: tensor = read[1] - assert (tensor, True) not in pending, ( - f"prefetched {tensor} is still pending at {line}: {pending}" - ) + assert ( + tensor, + True, + ) not in pending, f"prefetched {tensor} is still pending at {line}: {pending}" checked.add(tensor) assert checked == {"A", "B"} @@ -63,10 +65,12 @@ def test_tf32_dtl_prefetch_matches_loop_waits(tile_m, stream_k, tmp_path): (Path(__file__).parent / "test_data" / "tf32_dtl_prefetch.yaml").read_text() ) params = config["BenchmarkProblems"][0][1] - params["ForkParameters"].extend([ - {"MatrixInstruction": [[16, 16, 32, 1, 1, tile_m // 32, 4, 2, 2]]}, - {"StreamK": [stream_k]}, - ]) + params["ForkParameters"].extend( + [ + {"MatrixInstruction": [[16, 16, 32, 1, 1, tile_m // 32, 4, 2, 2]]}, + {"StreamK": [stream_k]}, + ] + ) path = tmp_path / "tf32.yaml" path.write_text(yaml.safe_dump(config)) results = emit_kernels_from_config(path, limit=1, arch="gfx950") @@ -74,3 +78,30 @@ def test_tf32_dtl_prefetch_matches_loop_waits(tile_m, stream_k, tmp_path): _, assembly, error = results[0] assert error == 0 _assert_prefetch_ready(assembly, tile_m // 32) + + +@pytest.mark.parametrize("dtype", ["B", "H"]) +@pytest.mark.parametrize("tile_processing_strategy", ["None", "StreamK"]) +@pytest.mark.parametrize("schedule_gr_over_barrier", [0, 1]) +def test_16bit_tn_256x160_dtl_prefetch_matches_loop_waits( + dtype, tile_processing_strategy, schedule_gr_over_barrier, tmp_path +): + config = yaml.safe_load( + (Path(__file__).parent / "test_data" / "bf16_tn_dtl_prefetch.yaml").read_text() + ) + problem, params = config["BenchmarkProblems"][0] + problem["DataType"] = dtype + problem["DestDataType"] = dtype + params["ForkParameters"].extend( + [ + {"TileProcessingStrategy": [tile_processing_strategy]}, + {"ScheduleGROverBarrier": [schedule_gr_over_barrier]}, + ] + ) + path = tmp_path / "tn.yaml" + path.write_text(yaml.safe_dump(config)) + results = emit_kernels_from_config(path, limit=1, arch="gfx950") + assert len(results) == 1 + _, assembly, error = results[0] + assert error == 0 + _assert_prefetch_ready(assembly, loads_a=8, loads_b=5)