fix(tensilelite): match 16-bit TN CMS prefetch order (AIHPBLAS-5228)

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).

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]

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,14 +65,43 @@ 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")
     assert len(results) == 1
     _, 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)