fix(hipblaslt): Load the RNDseed value into the sgpr when using software rounding

projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionSolution.hpp

@@ -318,6 +318,7 @@ namespace TensileLite
                 iot::mapOptional(io, "swizzleTensorA", s.swizzleTensorA);
                 iot::mapOptional(io, "swizzleTensorB", s.swizzleTensorB);
                 iot::mapOptional(io, "fusedGemmA2A", s.fusedGemmA2A);
+                iot::mapOptional(io, "stochasticRounding", s.stochasticRounding);
                 iot::mapOptional(io, "metadataLayout", s.metadataLayout);
                 // mxScaleFormat is mapped as optional so logic files that omit it
                 // (e.g. non-MX problems) deserialize cleanly with the default 0 = NoSwizzle.

projects/hipblaslt/tensilelite/tensilelite/Components/Signature.py

@@ -450,8 +450,8 @@ def __call__(self, writer) -> SignatureBase:
         # signature.offset counts from the very first arg including the common header.
         # The assembly loads these args with KernArgAddress already advanced past
         # that header by commonArgsSize, so subtract it.
+        commonArgsSize = userArgumentsInfo.commonArgsSize
         if not kernel["ProblemType"]["GroupedGemm"]:
-            commonArgsSize = userArgumentsInfo.commonArgsSize
             writer.states.batchOffsetDKernArgOffset = signature.offset - commonArgsSize
             signature.addArg("batchOffsetD", SVK.SIG_VALUE, "u64")
             writer.states.batchOffsetCKernArgOffset = signature.offset - commonArgsSize
@@ -483,6 +483,10 @@ def __call__(self, writer) -> SignatureBase:
             # that base = fusedA2AKernArgBase + fusedA2AKernArgLayout()[X].
             writer.states.fusedA2AKernArgBase = fusedBase - userArgumentsInfo.commonArgsSize
 
+        if kernel["ProblemType"]["StochasticRounding"] and not writer.states.asmCaps["v_prng_b32"]:
+            writer.states.rndSeedKernArgOffset = signature.offset - commonArgsSize
+            signature.addArg("RNDSeed", SVK.SIG_VALUE, "u32")
+
         activationType = ActivationType("all")
         for name in activationType.getAdditionalArgStringList():
             userArgumentsInfo.activationSize += userArgumentsInfo.actMaxSize

projects/hipblaslt/tensilelite/tensilelite/Contractions.py

@@ -107,7 +107,7 @@ class ProblemType:
                  'highPrecisionAccumulate', 'useInitialStridesAB', 'useInitialStridesCD', 'stridedBatched', 'groupedGemm',
                  'useGradient', 'activationType', 'activationArgLength', 'activationComputeDataType', 'activationNoGuard',
                  'sparse', 'f32XdlMathOp', 'supportDeviceUserArguments', 'outputAmaxD', 'swizzleTensorA', 'swizzleTensorB', 'metadataLayout',
-                 'mxBlockA', 'mxBlockB', 'mxTypeA', 'mxTypeB', 'mxScaleFormat', 'fusedGemmA2A']
+                 'mxBlockA', 'mxBlockB', 'mxTypeA', 'mxTypeB', 'mxScaleFormat', 'fusedGemmA2A', 'stochasticRounding']
     @classmethod
     def FromOriginalState(cls, d):
         indices = [None]*d['TotalIndices']
@@ -357,6 +357,8 @@ def FromOriginalState(cls, d):
         rv.metadataLayout = 0
         if 'MetadataLayout' in d:
             rv.metadataLayout = d['MetadataLayout']
+
+        rv.stochasticRounding = bool(d.get('StochasticRounding', False))
         return rv
 
     def __init__(self, freeIndices=None, batchIndices=None, boundIndices=None, aDims=None, bDims=None, cDims=None, dDims=None):

projects/hipblaslt/tensilelite/tensilelite/KernelWriter.py

@@ -383,6 +383,7 @@ class StateValues:
   batchOffsetCKernArgOffset: int         = 0
   batchOffsetAKernArgOffset: int         = 0
   batchOffsetBKernArgOffset: int         = 0
+  rndSeedKernArgOffset: int              = -1
   numSgprAlpha: int                      = 0 # For user arguments
   numSgprBeta: int                       = 0 # For user arguments
   numStoreSgprNames: List[str]           = field(init=False) # For post-loop kernel args

projects/hipblaslt/tensilelite/tensilelite/KernelWriterAssembly.py

@@ -1113,9 +1113,6 @@ def defineVariableSgprs(self, kernel):
         self.states.nonPostLoopSgpr.append("PackKForV2")
         self.states.nonPostLoopSgpr.append("PackKForV3")
 
-    if kernel["ProblemType"]["StochasticRounding"] and not self.states.asmCaps["v_prng_b32"]:
-      module.add(self.defineSgpr("RNDSeed", 1))
-
     if kernel["_UseSgprForGRO"]:
       needFirstSgprOffset = kernel["DirectToLdsA"] and kernel["UseInstOffsetForGRO"]
       numberOfSgpr = self.states.a.numVgprGlobalReadOffsets if needFirstSgprOffset else (self.states.a.numVgprGlobalReadOffsets-1)
@@ -2683,6 +2680,11 @@ def getKernelArgLoadModule(self, kernel, sgprStartIdx, numsOfLoad, preloadNum):
         sgprOffset += 4
         self.argLoader.setOffset(sgprOffset)
 
+    if self.states.rndSeedKernArgOffset >= 0:
+        kernelArgs.add(self.argLoader.loadKernArg(
+            "RNDSeed", "KernArgAddress",
+            sgprOffset=hex(self.states.rndSeedKernArgOffset), dword=1))
+
     # Batch offset arguments for general batched mode are loaded on-demand from
     # their kernarg byte offsets. Those offsets are recorded accurately by
     # the signature builder (see Signature.py) into self.states.batchOffset*KernArgOffset,
@@ -2966,7 +2968,10 @@ def defineAndResources(self, kernel, tPA, tPB, tPM):
             self.sgprPool.checkIn(tmpSgpr)
             break
           self.states.preloadGuard.append(tmpSgpr)
-
+
+      if kernel["ProblemType"]["StochasticRounding"] and not self.states.asmCaps["v_prng_b32"]:
+        module.add(self.defineSgpr("RNDSeed", 1))
+
       ### temp sgpr for groupedgemm ###
       # can be start from sgpr_preload_end
       sgprNumsOfGemm = self.sgprPool.checkOut(1, tag="defineAndResources_sgprNumsOfGemm", preventOverflow=False)

projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_stochastic_rounding_library.py

@@ -0,0 +1,62 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""The library problem type must keep the YAML StochasticRounding value.
+
+Kernel generation reads the logic YAML directly. Verify StochasticRounding
+value is preserved and propagated to the host.
+"""
+
+import pytest
+
+from tensilelite.Common.Utilities import state
+from tensilelite.Contractions import ProblemType
+
+pytestmark = pytest.mark.unit
+
+
+def generic_gemm(**overrides) -> ProblemType:
+    """A single-batch GEMM dict that FromOriginalState accepts."""
+    source = {
+        "OperationType": "GEMM",
+        "DataType": "H",
+        "DestDataType": "H",
+        "ComputeDataType": "S",
+        "TransposeA": False,
+        "TransposeB": True,
+        "Batched": True,
+        "UseBeta": True,
+        "TotalIndices": 4,
+        "NumIndicesC": 3,
+        "IndicesFree": [0, 1],
+        "IndicesBatch": [2],
+        "IndicesSummation": [3],
+        "IndexAssignmentsA": [0, 3, 2],
+        "IndexAssignmentsB": [3, 1, 2],
+        "ComplexConjugateA": False,
+        "ComplexConjugateB": False,
+        "ActivationComputeDataType": "S",
+    }
+    source.update(overrides)
+    problem = ProblemType.FromOriginalState(source)
+    # Solution.FromOriginalState sets this after ProblemType.FromOriginalState.
+    problem.mxScaleFormat = 0
+    return problem
+
+
+def test_library_state_has_yaml_stochastic_rounding():
+    """The library state has the YAML StochasticRounding value."""
+    assert "stochasticRounding" in ProblemType.StateKeys
+
+    enabled = state(generic_gemm(StochasticRounding=True))
+    disabled = state(generic_gemm(StochasticRounding=False))
+    omitted = state(generic_gemm())
+
+    assert enabled["stochasticRounding"] is True
+    assert disabled["stochasticRounding"] is False
+    assert omitted["stochasticRounding"] is False
+
+
+def test_problem_type_defaults_stochastic_rounding_off():
+    """A problem type built without the YAML key does not append RNDSeed."""
+    assert generic_gemm().stochasticRounding is False