[rocRoller][hipblaslt]Use separate array for flags in rocRoller StreamK

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocroller/gemm.cpp

@@ -118,6 +118,10 @@ std::string genKernelName(std::shared_ptr<SolutionParameters> gemm)
     if(gemm->kernelType.scaleTypeB.mode != Operations::ScaleMode::None)
         rv << "SB_" << genScaleModeString(gemm->kernelType.scaleTypeB.mode) << "_";
 
+    if(gemm->streamK)
+    {
+        rv << "SK_";
+    }   
     rv << "WGT_";
     rocRoller::streamJoin(
         rv, std::vector{gemm->workgroupTile.m, gemm->workgroupTile.n, gemm->workgroupTile.k}, "x");
@@ -170,7 +174,8 @@ std::shared_ptr<GemmKernel> genGemmKernel(std::shared_ptr<SolutionParameters> ge
                 ShowValue(gemm->kernelType.scaleTypeB.mode));
 
     std::optional<Operations::OperationTag> tagTensorScaleA, tagLoadScaleA, tagBlockScaleA,
-        tagTensorScaleB, tagLoadScaleB, tagBlockScaleB, tagScratch, tagSKGrid, tagWGM;
+        tagTensorScaleB, tagLoadScaleB, tagBlockScaleB, tagSKGrid, tagWGM;
+    std::map<Operations::ScratchPolicy, Operations::OperationTag> tagScratch;
 
     if(gemm->kernelType.scaleTypeA.mode == Operations::ScaleMode::Separate)
     {
@@ -265,12 +270,19 @@ std::shared_ptr<GemmKernel> genGemmKernel(std::shared_ptr<SolutionParameters> ge
                                 DataDirection::ReadOnly,
                                 rocRoller::NUMWGS);
 
-        tagScratch = command->allocateTag();
-        command->allocateArgument(VariableType(DataType::UInt32, PointerType::PointerGlobal),
-                                *tagScratch,
-                                ArgumentType::Value,
-                                DataDirection::ReadWrite,
-                                rocRoller::SCRATCH);
+        // Create Scratch operations for each ScratchPolicy
+        for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
+        {
+            auto policy = static_cast<Operations::ScratchPolicy>(i);
+            auto tag = command->allocateTag();
+            command->addOperation(Operations::Scratch(tag, policy));
+            command->allocateArgument(VariableType(DataType::UInt32, PointerType::PointerGlobal),
+                                    tag,
+                                    ArgumentType::Value,
+                                    DataDirection::ReadWrite,
+                                    rocRoller::getScratchName(policy));
+            tagScratch[policy] = tag;
+        }
     }
 
     if(gemm->workgroupMappingDim != -1)
@@ -489,8 +501,7 @@ std::shared_ptr<GemmKernel> genGemmKernel(std::shared_ptr<SolutionParameters> ge
     if(tagTensorScaleB)
         gemmKernel->tagTensorScaleB = *tagTensorScaleB;
 
-    if(tagScratch)
-        gemmKernel->tagScratch = *tagScratch;
+    gemmKernel->tagScratch = tagScratch;
 
     if(tagSKGrid)
         gemmKernel->tagSKGrid = *tagSKGrid;
@@ -523,7 +534,9 @@ size_t workspaceRequired(std::shared_ptr<GemmKernel> gemm, const RocblasltContra
 
     auto runtimeArgs = commandArgs.runtimeArguments();
 
-    return gemm->commandKernel->scratchSpaceRequired(runtimeArgs);
+    // Only return scratch space for ScratchPolicy::None (uses prob.workspace)
+    return gemm->commandKernel->scratchSpaceRequired(
+        Operations::ScratchPolicy::None, runtimeArgs);
 }
 
 CommandArguments createCommandArguments(std::shared_ptr<GemmKernel>        gemm,
@@ -622,14 +635,34 @@ rocblaslt_status runGemmKernel(std::shared_ptr<GemmKernel>        gemm,
         }
         return rocblaslt_status_invalid_value;
     }
-
     auto commandArgs = createCommandArguments(gemm, prob, DEFAULT_WGM);
 
-    // Add scratch space
-    if(workSpaceRequired > 0)
+
+    if(gemm->params->streamK)
     {
-        commandArgs.setArgument(
-            gemm->tagScratch, ArgumentType::Value, static_cast<unsigned char*>(prob.workspace));
+        auto runtimeArgs = commandArgs.runtimeArguments();
+
+        // Use prob.workspace for ScratchPolicy::None
+        auto noneScratchSize = gemm->commandKernel->scratchSpaceRequired(
+            Operations::ScratchPolicy::None, runtimeArgs);
+        if(noneScratchSize > 0 && prob.workspace != nullptr)
+        {
+            commandArgs.setArgument(
+                gemm->tagScratch.at(Operations::ScratchPolicy::None),
+                ArgumentType::Value,
+                static_cast<unsigned char*>(prob.workspace));
+        }
+
+        // Use prob.Synchronizer for ScratchPolicy::ZeroedBeforeAndAfter
+        auto zeroedScratchSize = gemm->commandKernel->scratchSpaceRequired(
+            Operations::ScratchPolicy::ZeroedBeforeAndAfter, runtimeArgs);
+        if(zeroedScratchSize > 0 && prob.Synchronizer != nullptr)
+        {
+            commandArgs.setArgument(
+                gemm->tagScratch.at(Operations::ScratchPolicy::ZeroedBeforeAndAfter),
+                ArgumentType::Value,
+                static_cast<unsigned char*>(prob.Synchronizer));
+        }
     }
 
     auto runtimeArgs = commandArgs.runtimeArguments();
@@ -640,5 +673,6 @@ rocblaslt_status runGemmKernel(std::shared_ptr<GemmKernel>        gemm,
     }
 
     gemm->commandKernel->launchKernel(runtimeArgs, prob.stream);
+
     return rocblaslt_status_success;
 }

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocroller/include/gemm.hpp

@@ -32,8 +32,11 @@
 #include <rocRoller/CommandSolution.hpp>
 #include <rocRoller/Expression.hpp>
 #include <rocRoller/KernelGraph/CoordinateGraph/Dimension.hpp>
+#include <rocRoller/Operations/Scratch_fwd.hpp>
 #include <rocRoller/TensorDescriptor.hpp>
 
+#include <map>
+
 /**
  * @brief GemmKernel
  *
@@ -57,7 +60,7 @@ struct GemmKernel
     rocRoller::Operations::OperationTag tagTensorScaleA;
     rocRoller::Operations::OperationTag tagTensorScaleB;
 
-    rocRoller::Operations::OperationTag tagScratch;
+    std::map<rocRoller::Operations::ScratchPolicy, rocRoller::Operations::OperationTag> tagScratch;
     rocRoller::Operations::OperationTag tagSKGrid;
     rocRoller::Operations::OperationTag tagWGM;
 

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocroller/include/solution_selection.hpp

@@ -71,6 +71,7 @@ struct SolutionIndexParameters
 {
     WorkGroupTileSize workgroupTile;
     bool              workgroupMapping;
+    bool              streamK;
 
     auto operator<=>(const SolutionIndexParameters& other) const = default;
 };

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocroller/parameter_selection.cpp

@@ -37,6 +37,7 @@ std::string SolutionParameters::toString() const
     result << "MachineInstruction:" << machineInstruction.m << "x" << machineInstruction.n << "x"
            << machineInstruction.k << std::endl;
     result << "WorkgroupSize:" << workgroupSizeX << "x" << workgroupSizeY << std::endl;
+    result << "StreamK: " << streamK << std::endl;
     result << "LoadA: " << loadPathA << std::endl;
     result << "LoadB: " << loadPathB << std::endl;
     result << "LDS Usage";
@@ -196,10 +197,12 @@ std::shared_ptr<SolutionParameters>
         gemm->workgroupRemapXCC = true;
     }
 
-    // TODO: StreamK is not currently working with prefetching or workgroup mapping
+    // Pass StreamK flag from solution index parameters
+    gemm->streamK = solutionIndexParameters.streamK;
+
+    // StreamK is not currently working with workgroup mapping due to register pressure
     if(gemm->streamK)
     {
-        gemm->prefetch = false;
         gemm->workgroupMappingDim = -1;
         gemm->workgroupRemapXCC = false;
     }

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocroller/solution_selection.cpp

@@ -237,12 +237,22 @@ std::vector<SolutionIndexParameters> chooseSolutionIndexParameters(
                    || !std::has_single_bit(static_cast<uint>(wgt.n))))
                 continue;
 
-            params.push_back({wgt, true});
+            params.push_back({wgt, true, false});
 
             if (prob.k < USE_WORKGROUP_MAPPING_K_SIZE)
             {
                 params.back().workgroupMapping = false;
             }
+
+            // Enable StreamK when number of output tiles < number of CUs and not f6 data type
+            size_t numTilesM = prob.m / wgt.m;
+            size_t numTilesN = prob.n / wgt.n;
+            size_t numTiles = numTilesM * numTilesN * prob.batch_count;
+            auto isF6 = (kernelType.typeA == rocRoller::DataType::FP6 || kernelType.typeA == rocRoller::DataType::BF6 || kernelType.typeB == rocRoller::DataType::FP6 || kernelType.typeB == rocRoller::DataType::BF6);
+            if(numTiles < analytical_hardware.N_CU && !isF6)
+            {
+                params.back().streamK = true;
+            }
         }
     }
 
@@ -261,6 +271,8 @@ int parametersToIndex(const SolutionIndexParameters& params)
     result |= ((params.workgroupTile.m / REQUIRED_MULTIPLE_M_N) << pos);
     pos += MAX_BITS_WORKGROUPTILE_M;
     result |= ((params.workgroupMapping ? 1 : 0) << pos);
+    pos += 1;
+    result |= ((params.streamK ? 1 : 0) << pos);
 
     // Set top bit indicating it is a rocRoller index
     result |= (1 << 31);
@@ -290,6 +302,8 @@ SolutionIndexParameters indexToParameters(int index)
         = ((index >> pos) & mask(MAX_BITS_WORKGROUPTILE_M)) * REQUIRED_MULTIPLE_M_N;
     pos += MAX_BITS_WORKGROUPTILE_M;
     result.workgroupMapping = (index >> pos) & 1;
+    pos += 1;
+    result.streamK = (index >> pos) & 1;
 
     return result;
 }

shared/rocroller/client/include/client/GEMMSolution.hpp

@@ -29,6 +29,7 @@
 #include "GEMMParameters.hpp"
 
 #include <rocRoller/Operations/CommandArgument_fwd.hpp>
+#include <rocRoller/Operations/Scratch_fwd.hpp>
 
 using namespace rocRoller;
 
@@ -61,7 +62,8 @@ namespace rocRoller
                 {
                 }
 
-                virtual Operations::OperationTag getScratchTag() const
+                virtual Operations::OperationTag
+                    getScratchTag(Operations::ScratchPolicy scratchPolicy) const
                 {
                     return {};
                 }

shared/rocroller/client/include/client/StreamKGEMMSolution.hpp

@@ -41,14 +41,16 @@ namespace rocRoller
         {
             class StreamKGEMMSolution : public DataParallelGEMMSolution
             {
-                Operations::OperationTag m_scratchTag, m_numWGsTag;
+                std::map<Operations::ScratchPolicy, Operations::OperationTag> m_scratchTags;
+                Operations::OperationTag                                      m_numWGsTag;
 
             public:
                 using DataParallelGEMMSolution::DataParallelGEMMSolution;
 
-                Operations::OperationTag getScratchTag() const override
+                Operations::OperationTag
+                    getScratchTag(Operations::ScratchPolicy scratchPolicy) const override
                 {
-                    return m_scratchTag;
+                    return m_scratchTags.at(scratchPolicy);
                 }
 
             protected:
@@ -63,15 +65,31 @@ namespace rocRoller
                                                                DataDirection::ReadOnly,
                                                                rocRoller::NUMWGS);
 
-                    m_scratchTag = command->allocateTag();
-                    command->addOperation(rocRoller::Operations::Scratch(
-                        m_scratchTag, Operations::ScratchPolicy::None));
+                    // Create a scratch operation for tile data
+                    m_scratchTags[Operations::ScratchPolicy::None] = command->allocateTag();
+                    command->addOperation(
+                        Operations::Scratch(m_scratchTags[Operations::ScratchPolicy::None],
+                                            Operations::ScratchPolicy::None));
                     command->allocateArgument(
                         VariableType(DataType::UInt32, PointerType::PointerGlobal),
-                        m_scratchTag,
+                        m_scratchTags[Operations::ScratchPolicy::None],
                         ArgumentType::Value,
                         DataDirection::ReadWrite,
-                        rocRoller::SCRATCH);
+                        getScratchName(Operations::ScratchPolicy::None));
+
+                    // Create a scratch operation for flags
+                    m_scratchTags[Operations::ScratchPolicy::ZeroedBeforeAndAfter]
+                        = command->allocateTag();
+                    command->addOperation(Operations::Scratch(
+                        m_scratchTags[Operations::ScratchPolicy::ZeroedBeforeAndAfter],
+                        Operations::ScratchPolicy::ZeroedBeforeAndAfter));
+                    command->allocateArgument(
+                        VariableType(DataType::UInt32, PointerType::PointerGlobal),
+                        m_scratchTags[Operations::ScratchPolicy::ZeroedBeforeAndAfter],
+                        ArgumentType::Value,
+                        DataDirection::ReadWrite,
+                        getScratchName(Operations::ScratchPolicy::ZeroedBeforeAndAfter));
+
                     return command;
                 }
 

shared/rocroller/client/src/gemm.cpp

@@ -25,6 +25,7 @@
  *******************************************************************************/
 
 #include "rocRoller/Serialization/YAML.hpp"
+#include <algorithm>
 #include <filesystem>
 #include <fstream>
 #include <string>
@@ -344,14 +345,18 @@ namespace rocRoller::Client::GEMMClient
         auto runtimeArgs = commandArgs.runtimeArguments();
 
         // Note: the lifetime of deviceScratch needs to exceed kernel executions
-        std::shared_ptr<uint8_t> deviceScratch;
+        std::shared_ptr<uint8_t>
+            deviceScratch[static_cast<size_t>(Operations::ScratchPolicy::Count)];
+
+        for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
         {
-            auto scratchSpaceRequired = commandKernel->scratchSpaceRequired(runtimeArgs);
+            auto policy               = static_cast<Operations::ScratchPolicy>(i);
+            auto scratchSpaceRequired = commandKernel->scratchSpaceRequired(policy, runtimeArgs);
             if(scratchSpaceRequired > 0)
             {
-                deviceScratch = make_shared_device<uint8_t>(scratchSpaceRequired, 0);
+                deviceScratch[i] = make_shared_device<uint8_t>(scratchSpaceRequired, 0);
                 commandArgs.setArgument(
-                    gemm->getScratchTag(), ArgumentType::Value, deviceScratch.get());
+                    gemm->getScratchTag(policy), ArgumentType::Value, deviceScratch[i].get());
             }
         }
 
@@ -457,6 +462,24 @@ namespace rocRoller::Client::GEMMClient
             auto [correct, rnorm] = validate<A, B, C, D>(
                 hostA, hostB, hostC, hostD, hostScaleA, hostScaleB, problemParams, arch);
 
+            // Verify ZeroedBeforeAndAfter scratch is all zeros after kernel
+            auto zeroedIdx = static_cast<size_t>(Operations::ScratchPolicy::ZeroedBeforeAndAfter);
+            if(deviceScratch[zeroedIdx])
+            {
+                auto zeroedSize = commandKernel->scratchSpaceRequired(
+                    Operations::ScratchPolicy::ZeroedBeforeAndAfter, runtimeArgs);
+                std::vector<uint8_t> zeroedResult(zeroedSize);
+                AssertFatal(hipMemcpy(zeroedResult.data(),
+                                      deviceScratch[zeroedIdx].get(),
+                                      zeroedSize,
+                                      hipMemcpyDeviceToHost)
+                            == (hipError_t)HIP_SUCCESS);
+                AssertFatal(
+                    std::all_of(
+                        zeroedResult.begin(), zeroedResult.end(), [](uint8_t v) { return v == 0; }),
+                    "ZeroedBeforeAndAfter scratch should be all zeros after kernel execution");
+            }
+
             result.checked = true;
             result.correct = correct;
             result.rnorm   = rnorm;

shared/rocroller/lib/include/rocRoller/CommandSolution.hpp

@@ -307,13 +307,17 @@ namespace rocRoller
 
         /**
          * @brief Returns the total number of bytes required for scratch space
+         *        for the specified scratch policy.
          *
-         * If this value is greather than 0, the user is required to allocate this
+         * If this value is greater than 0, the user is required to allocate this
          * amount of device memory and pass it into the kernel.
          *
+         * @param policy The scratch policy to query
+         * @param args The runtime arguments
          * @return size_t
          */
-        size_t scratchSpaceRequired(RuntimeArguments const& args) const;
+        size_t scratchSpaceRequired(Operations::ScratchPolicy policy,
+                                    RuntimeArguments const&   args) const;
 
         /**
          * @brief Returns the workgroup size