diff --git a/projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionSolution.hpp b/projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionSolution.hpp index 7c4fe03ec921..0cb8d89d0d43 100644 --- a/projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionSolution.hpp +++ b/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. diff --git a/projects/hipblaslt/tensilelite/tensilelite/Components/Signature.py b/projects/hipblaslt/tensilelite/tensilelite/Components/Signature.py index d511afa30183..e4aa4af1750e 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/Components/Signature.py +++ b/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 diff --git a/projects/hipblaslt/tensilelite/tensilelite/Contractions.py b/projects/hipblaslt/tensilelite/tensilelite/Contractions.py index 2d7fbb8c6f64..2404101c2359 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/Contractions.py +++ b/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): diff --git a/projects/hipblaslt/tensilelite/tensilelite/KernelWriter.py b/projects/hipblaslt/tensilelite/tensilelite/KernelWriter.py index 4059edc4cefa..7db542a1229c 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/KernelWriter.py +++ b/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 diff --git a/projects/hipblaslt/tensilelite/tensilelite/KernelWriterAssembly.py b/projects/hipblaslt/tensilelite/tensilelite/KernelWriterAssembly.py index 9890a66b4c8d..4e51780fcdf0 100644 --- a/projects/hipblaslt/tensilelite/tensilelite/KernelWriterAssembly.py +++ b/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) diff --git a/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_stochastic_rounding_library.py b/projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_stochastic_rounding_library.py new file mode 100644 index 000000000000..059436d8ef95 --- /dev/null +++ b/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