feat(tensilelite): Add gfx1250 support for AdaptiveGemmNTAB

projects/hipblaslt/tensilelite/Tensile/KernelWriter.py

@@ -5736,6 +5736,11 @@ def _kernelBody(pack, packPre, nta, ntb):
       originalNta = tensorParametersA["NonTemporal"]
       originalNtb = tensorParametersB["NonTemporal"]
 
+      hasTH = self.states.asmCaps.get("HasTHModifier", False)
+      if hasTH:
+        originalThA = tensorParametersA["TemporalHint"]
+        originalThB = tensorParametersB["TemporalHint"]
+
       ntCombos = [[0, 0], [0, 4], [4, 0]]
       ntLabels = [Label("LoopBody_NTA{}_NTB{}".format(nta, ntb), "") for nta, ntb in ntCombos]
       ntLabelDone = Label("LoopBody_NTA_NTB_Done", "")
@@ -5799,6 +5804,9 @@ def _restoreNtabState():
         module.add(ntLabels[idx])
         tensorParametersA["NonTemporal"] = nta
         tensorParametersB["NonTemporal"] = ntb
+        if hasTH:
+          tensorParametersA["TemporalHint"] = 1 if nta else 0
+          tensorParametersB["TemporalHint"] = 1 if ntb else 0
         _kernelBody(pack, packPre, nta, ntb)
         # All paths but the last one need to skip the rest. Use long
         # branches everywhere because each kernelBody can easily exceed
@@ -5812,6 +5820,9 @@ def _restoreNtabState():
 
       tensorParametersA["NonTemporal"] = originalNta
       tensorParametersB["NonTemporal"] = originalNtb
+      if hasTH:
+        tensorParametersA["TemporalHint"] = originalThA
+        tensorParametersB["TemporalHint"] = originalThB
 
     if kernel["ExpertSchedulingMode"] > 0:
       module.add(SSetRegIMM32B32(dst=HWRegContainer(reg="26", value=[0,2]), src=0x0, comment="enable hardware dependency checking"))
@@ -9921,6 +9932,7 @@ def getTensorParameters(self, tP, kernel, itP, cM):
     tP["wtc"] = itP[cM].writeTileDimComponents                   # write vector components along tile dimension
     tP["idx"] = kernel["ProblemType"]["Index%d"%tP["tensorIdx"]] # index 0 is tile dimension belonging to A. Note 'idx' may not be in tP['ia'].
     tP["NonTemporal"] = kernel["NonTemporal%s"%cM]               # non-temporal read type
+    tP["TemporalHint"] = kernel.get("TemporalHint%s"%cM, 0)      # temporal-hint read type
     tP["shiftGR"] = 0 if (tP["bpeGR"] >= tP["bpeDS"]) else int(tP["glvw"] // 2 * (tP["bpeDS"] / self.states.bpr))  # Shift global read register for cvt spaces
     tP["bpeRatio"] = tP["bpeDS"] // tP["bpeGR"] if tP["bpeGR"] < tP["bpeDS"] else 1                                # g2lIdx multiplier
 

projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py

@@ -10341,7 +10341,7 @@ def globalReadGuardKBody(tP, optParams = None):
           dataType = kernel["ProblemType"]["MacDataType%s"%tcDataType if (tcDataType=='A' or tcDataType=='B') else "DataType"]
 
       isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal(
-          self.states.asmCaps, tP["NonTemporal"], _temporalHint(kernel, tc), _nonVolatile(kernel, tc))
+          self.states.asmCaps, tP["NonTemporal"], TemporalHint(tP["TemporalHint"]), _nonVolatile(kernel, tc))
       isLds = True if (kernel["DirectToLds%s"%tc] and not kernel["NonDTLTailLoop%s"%tc]) else False
       isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc]
       is16b = dataType.isHalf() or dataType.isBFloat16()
@@ -11276,7 +11276,7 @@ def globalReadBody(tP):
       bpe = tP["bpeGR"] if isAB else tP["bpe"]
       bpl = bpe * tP["glvw"]  # bytes per load
       isGlc, isSlc, isNT, scope, th, nv = decodeNonTemporal(
-          self.states.asmCaps, tP["NonTemporal"], _temporalHint(kernel, tc), _nonVolatile(kernel, tc))
+          self.states.asmCaps, tP["NonTemporal"], TemporalHint(tP["TemporalHint"]), _nonVolatile(kernel, tc))
       isLds = True if kernel["DirectToLds%s"%tc] else False
       isTr = (tc == "A" or tc == "B") and kernel["enableGLTr%s"%tc]
 

projects/hipblaslt/tensilelite/Tensile/Tests/common/gemm/gfx12/agntab_coverage_gfx1250.yaml

@@ -0,0 +1,103 @@
+TestParameters:
+  marks: [skip-gfx900, skip-gfx906, skip-gfx908, skip-gfx90a, skip-gfx940, skip-gfx941, skip-gfx942, skip-gfx950, skip-gfx1010, skip-gfx1011, skip-gfx1012, skip-gfx1030, skip-gfx1100, skip-gfx1101, skip-gfx1102, skip-gfx1200, skip-gfx1201]
+
+GlobalParameters:
+  NumElementsToValidate: -1
+  NumWarmups: 1
+  MaxEnqueuesPerSync: 1
+  SleepPercent: 0
+  PrintSolutionRejectionReason: True
+  DataInitTypeA: 23
+  DataInitTypeB: 24
+  DataInitTypeBeta: 0
+  DataInitTypeAlpha: 1
+  CSVExportWinner: True
+  CSVMergeSameProblemID: True
+  Device: 0
+  MaxWorkspaceSize: 134217728
+  KernelTime: True
+  RotatingBufferSize: 512
+  Architecture: gfx1250
+  CodeObjectVersion: "4"
+
+.AdaptiveBase: &adaptive_base
+  CustomKernels:
+  InitialSolutionParameters:
+  BenchmarkCommonParameters:
+    - KernelLanguage: ["Assembly"]
+  ForkParameters:
+    - MatrixInstruction:
+      - [16, 16, 32, 1, 1, 1, 1, 1, 1] # gfx1250 WMMA bf16 16x16x32
+    - WavefrontSize: [32]
+    - DepthU: [32]
+    - PrefetchGlobalRead: [1]
+    - PrefetchLocalRead: [1]
+    - TransposeLDS: [-1]
+    - WorkGroupMapping: [1]
+    - 1LDSBuffer: [-1]
+    - GlobalSplitU: [1]
+    - GlobalSplitUAlgorithm: ["MultipleBuffer"]
+    # ============ StinkyTofu scheduling + AGNTAB/AG multi-body under test ============
+    - ScheduleIterAlg: [0, 1, 2, 3, 4]
+    - AdaptiveGemm: [0, 1]
+    - AdaptiveGemmNTAB: [0, 1]
+    - InitCIterWmma: [0, 1]
+  BenchmarkJoinParameters:
+  BenchmarkFinalParameters:
+    - ProblemSizes:
+      - Exact: [128,    128, 1, 128]   # baseline   -> NTA0_NTB0
+      - Exact: [32768,  128, 1, 192]   # NTA4 path  -> NTA4_NTB0
+      - Exact: [128,  16384, 1, 256]   # NTB4 path  -> NTA0_NTB4
+
+BenchmarkProblems:
+  # =================================================================
+  # AGNTAB x AG x StinkyTofu coverage -- 4 transpose x 3 sizes
+  # Focus: exercise the multi-body main loop + multi-store epilogue through StinkyTofu.
+  # =================================================================
+  - # NN adaptive
+    - OperationType: GEMM
+      DataType: b
+      DestDataType: b
+      ComputeDataType: s
+      HighPrecisionAccumulate: True
+      TransposeA: False
+      TransposeB: False
+      UseBeta: True
+      Batched: True
+    - <<: *adaptive_base
+
+  - # TN adaptive
+    - OperationType: GEMM
+      DataType: b
+      DestDataType: b
+      ComputeDataType: s
+      HighPrecisionAccumulate: True
+      TransposeA: True
+      TransposeB: False
+      UseBeta: True
+      Batched: True
+    - <<: *adaptive_base
+
+  - # NT adaptive
+    - OperationType: GEMM
+      DataType: b
+      DestDataType: b
+      ComputeDataType: s
+      HighPrecisionAccumulate: True
+      TransposeA: False
+      TransposeB: True
+      UseBeta: True
+      Batched: True
+    - <<: *adaptive_base
+
+  - # TT adaptive
+    - OperationType: GEMM
+      DataType: b
+      DestDataType: b
+      ComputeDataType: s
+      HighPrecisionAccumulate: True
+      TransposeA: True
+      TransposeB: True
+      UseBeta: True
+      Batched: True
+    - <<: *adaptive_base

projects/hipblaslt/tensilelite/rocisa/test/test_mubuf.py

@@ -28,7 +28,7 @@
 import rocisa
 from rocisa.code import Module, SignatureBase
 from rocisa.container import MUBUFModifiers, sgpr, vgpr
-from rocisa.enum import CacheScope
+from rocisa.enum import CacheScope, TemporalHint, NonVolatile
 from rocisa.instruction import (
     BufferAtomicAddF32,
     BufferLoadB32,
@@ -156,6 +156,116 @@ def test_mubuf_scope_modifiers_stinkytofu(_mubuf_scope_asm):
     )
 
 
+@pytest.fixture(scope="module")
+def _mubuf_th_asm() -> str:
+    mod = Module("mubuf_th_modifiers")
+    mod.add(
+        BufferLoadB32(
+            dst=vgpr(13),
+            vaddr=vgpr(33),
+            saddr=sgpr(64, 4),
+            soffset=sgpr(47),
+            mubuf=MUBUFModifiers(offen=True, scope=CacheScope.SCOPE_CU, th=TemporalHint.TH_NT),
+        )
+    )
+    mod.add(
+        BufferStoreB32(
+            src=vgpr(12),
+            vaddr=vgpr(32),
+            saddr=sgpr(60, 4),
+            soffset=sgpr(46),
+            mubuf=MUBUFModifiers(offen=True, isStore=True, scope=CacheScope.SCOPE_CU, th=TemporalHint.TH_NT),
+        )
+    )
+    mod.setParent()
+
+    sig = SignatureBase(
+        kernelName="mubuf_th_modifiers",
+        kernArgsVersion=1,
+        codeObjectVersion="4",
+        groupSegmentSize=0,
+        sgprWorkGroup=(1, 1, 0),
+        vgprWorkItem=0,
+        flatWorkGroupSize=64,
+        numSgprPreload=0,
+    )
+
+    st = rocisa.toStinkyTofuModule(
+        mod, _ISA, "mubuf_th_modifiers", signature=sig, options={"OptLevel": 0}
+    )
+    st.runOptimizationPipeline()
+    return st.emitAssembly()
+
+
+def test_mubuf_temporal_hint_modifiers_stinkytofu(_mubuf_th_asm):
+    assert re.search(
+        r"buffer_load_b32 v13, v33, s\[64:67\], s47 offen offset:0 scope:SCOPE_CU th:TH_LOAD_NT",
+        _mubuf_th_asm,
+    )
+    assert re.search(
+        r"buffer_store_b32 v12, v32, s\[60:63\], s46 offen offset:0 scope:SCOPE_CU th:TH_STORE_NT",
+        _mubuf_th_asm,
+    )
+
+
+@pytest.fixture(scope="module")
+def _mubuf_nv_asm() -> str:
+    mod = Module("mubuf_nv_modifiers")
+    mod.add(
+        BufferLoadB32(
+            dst=vgpr(13),
+            vaddr=vgpr(33),
+            saddr=sgpr(64, 4),
+            soffset=sgpr(47),
+            mubuf=MUBUFModifiers(offen=True, scope=CacheScope.SCOPE_CU, nv=NonVolatile.NV),
+        )
+    )
+    mod.add(
+        BufferStoreB32(
+            src=vgpr(12),
+            vaddr=vgpr(32),
+            saddr=sgpr(60, 4),
+            soffset=sgpr(46),
+            mubuf=MUBUFModifiers(
+                offen=True,
+                isStore=True,
+                scope=CacheScope.SCOPE_CU,
+                th=TemporalHint.TH_NT,
+                nv=NonVolatile.NV,
+            ),
+        )
+    )
+    mod.setParent()
+
+    sig = SignatureBase(
+        kernelName="mubuf_nv_modifiers",
+        kernArgsVersion=1,
+        codeObjectVersion="4",
+        groupSegmentSize=0,
+        sgprWorkGroup=(1, 1, 0),
+        vgprWorkItem=0,
+        flatWorkGroupSize=64,
+        numSgprPreload=0,
+    )
+
+    st = rocisa.toStinkyTofuModule(
+        mod, _ISA, "mubuf_nv_modifiers", signature=sig, options={"OptLevel": 0}
+    )
+    st.runOptimizationPipeline()
+    return st.emitAssembly()
+
+
+def test_mubuf_non_volatile_modifiers_stinkytofu(_mubuf_nv_asm):
+    assert re.search(
+        r"buffer_load_b32 v13, v33, s\[64:67\], s47 offen offset:0 scope:SCOPE_CU nv",
+        _mubuf_nv_asm,
+    )
+    assert re.search(
+        r"buffer_store_b32 v12, v32, s\[60:63\], s46 offen offset:0 scope:SCOPE_CU th:TH_STORE_NT nv",
+        _mubuf_nv_asm,
+    )
+
+
 def test_rocisa_load_null_soffset_adds_offen():
     inst = BufferLoadB32(
         dst=vgpr(13),

shared/stinkytofu/include/stinkytofu/ir/asm/StinkyModifiers.hpp

@@ -177,6 +177,17 @@ inline TemporalHint parseTemporalHint(std::string_view th) {
     return TemporalHint::TH_NONE;
 }
 
+enum class NonVolatile : uint8_t { NV_NONE = 0, NV = 1 };
+
+inline std::string_view toString(NonVolatile nv) {
+    return nv == NonVolatile::NV ? "nv" : "";
+}
+
+// Inverse of toString(): assembly token -> enum.
+inline NonVolatile parseNonVolatile(std::string_view nv) {
+    return nv == "nv" ? NonVolatile::NV : NonVolatile::NV_NONE;
+}
+
 // 9-bit DPP permutation control selector (matches the hardware dpp_ctrl field).
 // Three encoding shapes:
 //   singleton     — the named value IS the encoding   (e.g. ROW_MIRROR = 0x140)
@@ -404,7 +415,7 @@ struct MUBUFModifiers : public TypedModifier<MUBUFModifiers> {
                    bool nt = false, bool lds = false, bool isStore = false,
                    bool hasMUBUFConst = false, bool hasGLCModifier = false,
                    bool hasSC0Modifier = false, MUBUFScope scope = MUBUFScope::SCOPE_NONE,
-                   TemporalHint th = TemporalHint::TH_NONE)
+                   TemporalHint th = TemporalHint::TH_NONE, NonVolatile nv = NonVolatile::NV_NONE)
         : TypedModifier<MUBUFModifiers>(),
           offset12(offset12),
           offen(offen),
@@ -417,7 +428,8 @@ struct MUBUFModifiers : public TypedModifier<MUBUFModifiers> {
           hasGLCModifier(hasGLCModifier),
           hasSC0Modifier(hasSC0Modifier),
           scope(scope),
-          th(th) {}
+          th(th),
+          nv(nv) {}
 
     int offset12;
     uint32_t offen : 1;
@@ -431,6 +443,7 @@ struct MUBUFModifiers : public TypedModifier<MUBUFModifiers> {
     uint32_t hasSC0Modifier : 1;
     MUBUFScope scope;
     TemporalHint th;
+    NonVolatile nv;
 };
 
 // Carries just the cache scope token for SOPP-format memory fences such as

shared/stinkytofu/src/conversion/rocisa/ToStinkyTofuUtils.cpp

@@ -130,16 +130,28 @@ stinkytofu::TemporalHint convertTemporalHint(rocisa::TemporalHint th) {
     }
 }
 
+stinkytofu::NonVolatile convertNonVolatile(rocisa::NonVolatile nv) {
+    switch (nv) {
+        case rocisa::NonVolatile::NV_NONE:
+            return stinkytofu::NonVolatile::NV_NONE;
+        case rocisa::NonVolatile::NV:
+            return stinkytofu::NonVolatile::NV;
+        default:
+            return stinkytofu::NonVolatile::NV_NONE;
+    }
+}
+
 stinkytofu::MUBUFModifiers convertMUBUFModifiers(const rocisa::MUBUFModifiers& rocMod,
                                                  const std::map<std::string, int>& asmCaps) {
     bool hasMUBUFConst = asmCaps.count("HasMUBUFConst") && asmCaps.at("HasMUBUFConst");
     bool hasGLCModifier = asmCaps.count("HasGLCModifier") && asmCaps.at("HasGLCModifier");
     bool hasSC0Modifier = asmCaps.count("HasSC0Modifier") && asmCaps.at("HasSC0Modifier");
     stinkytofu::MUBUFScope scope = convertMUBUFScope(rocMod.scope);
     stinkytofu::TemporalHint th = convertTemporalHint(rocMod.th);
+    stinkytofu::NonVolatile nv = convertNonVolatile(rocMod.nv);
     return stinkytofu::MUBUFModifiers(rocMod.offen, rocMod.offset12, rocMod.glc, rocMod.slc,
                                       rocMod.nt, rocMod.lds, rocMod.isStore, hasMUBUFConst,
-                                      hasGLCModifier, hasSC0Modifier, scope, th);
+                                      hasGLCModifier, hasSC0Modifier, scope, th, nv);
 }
 
 /// Returns true when vaddr is the MUBUF "off" keyword.
@@ -1105,8 +1117,7 @@ static std::shared_ptr<StinkyAsmModule> toStinkyTofuModule(
         }
         if (const auto* subMod = dynamic_cast<const rocisa::Module*>(item.get())) {
             if (subMod->name == "loopBody") {
-                loopBodyIdx = i;
-                break;
+                loopBodyIdx = i;  // keep updating loop bodies
             }
         }
     }
@@ -1172,7 +1183,8 @@ static std::shared_ptr<StinkyAsmModule> toStinkyTofuModule(
     static const std::string kScope = "expertScheduleMode2";
 
     // Traverse top-level items, injecting the loopWithPrefetch group name
-    // for items in the detected prefetch region [pgrStartIdx, loopBodyIdx].
+    // for items in the detected prefetch region [pgrStartIdx, loopBodyIdx]
+    // (spans all main-loop bodies when present).
     for (int i = 0; i < static_cast<int>(module.itemList.size()); ++i) {
         const auto& item = module.itemList[i];
         const bool inPGR = hasPGR && (i >= pgrStartIdx && i <= loopBodyIdx);