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

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

@@ -170,7 +170,8 @@ std::shared_ptr<GemmKernel> genGemmKernel(std::shared_ptr<SolutionParameters> ge
                 ShowValue(gemm->kernelType.scaleBMode));
 
     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.scaleAMode == Operations::ScaleMode::Separate)
     {
@@ -265,12 +266,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 +497,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 +530,14 @@ size_t workspaceRequired(std::shared_ptr<GemmKernel> gemm, const RocblasltContra
 
     auto runtimeArgs = commandArgs.runtimeArguments();
 
-    return gemm->commandKernel->scratchSpaceRequired(runtimeArgs);
+    // Sum scratch requirements for all policies
+    size_t total = 0;
+    for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
+    {
+        auto policy = static_cast<Operations::ScratchPolicy>(i);
+        total += gemm->commandKernel->scratchSpaceRequired(policy, runtimeArgs);
+    }
+    return total;
 }
 
 CommandArguments createCommandArguments(std::shared_ptr<GemmKernel>        gemm,
@@ -622,23 +636,53 @@ 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)
+    // Track allocated scratch memory for each policy
+    std::array<void*, static_cast<size_t>(Operations::ScratchPolicy::Count)> scratchPtrs = {};
+    std::array<size_t, static_cast<size_t>(Operations::ScratchPolicy::Count)> scratchSizes = {};
+
+    if(gemm->params->streamK)
     {
-        commandArgs.setArgument(
-            gemm->tagScratch, ArgumentType::Value, static_cast<unsigned char*>(prob.workspace));
+        auto runtimeArgs = commandArgs.runtimeArguments();
+        for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
+        {
+            auto policy = static_cast<Operations::ScratchPolicy>(i);
+            scratchSizes[i] = gemm->commandKernel->scratchSpaceRequired(policy, runtimeArgs);
+            if(scratchSizes[i] > 0)
+            {
+                commandArgs.setArgument(
+                    gemm->tagScratch.at(policy), ArgumentType::Value,
+                    static_cast<unsigned char*>(scratchPtrs[i]));
+            }
+        }
     }
 
     auto runtimeArgs = commandArgs.runtimeArguments();
 
     if(!gemm->commandKernel->matchesPredicates(runtimeArgs, LogLevel::Error))
     {
+        // Free allocated scratch memory before returning
+        for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
+        {
+            if(scratchPtrs[i] != nullptr)
+            {
+                AssertFatal(hipFree(scratchPtrs[i]), "Failed to free scratch memory" + ShowValue(i));
+            }
+        }
         return rocblaslt_status_invalid_value;
     }
 
     gemm->commandKernel->launchKernel(runtimeArgs, prob.stream);
+
+    // Free allocated scratch memory after kernel completes
+    for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
+    {
+        if(scratchPtrs[i] != nullptr)
+        {
+            AssertFatal(hipFree(scratchPtrs[i]), "Failed to free scratch memory" + ShowValue(i));
+        }
+    }
+
     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

@@ -69,6 +69,7 @@ struct SolutionIndexParameters
 {
     WorkGroupTileSize workgroupTile;
     bool              workgroupMapping;
+    bool              streamK;
 };
 
 int parametersToIndex(const SolutionIndexParameters& params);

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

@@ -196,13 +196,16 @@ std::shared_ptr<SolutionParameters>
         gemm->workgroupRemapXCC = true;
     }
 
-    // TODO: StreamK is not currently working with prefetching or workgroup mapping
-    if(gemm->streamK)
-    {
-        gemm->prefetch = false;
-        gemm->workgroupMappingDim = -1;
-        gemm->workgroupRemapXCC = false;
-    }
+    // Pass StreamK flag from solution index parameters
+    gemm->streamK = solutionIndexParameters.streamK;
+
+    // // StreamK is not currently working with prefetching or workgroup mapping
+    // if(gemm->streamK)
+    // {
+    //     gemm->prefetch = false;
+    //     gemm->workgroupMappingDim = -1;
+    //     gemm->workgroupRemapXCC = false;
+    // }
 
     return gemm;
 }

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

@@ -237,12 +237,21 @@ 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
+            size_t numTilesM = prob.m / wgt.m;
+            size_t numTilesN = prob.n / wgt.n;
+            size_t numTiles = numTilesM * numTilesN * prob.batch_count;
+            if(numTiles < analytical_hardware.N_CU)
+            {
+                params.back().streamK = true;
+            }
         }
     }
 
@@ -261,6 +270,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 +301,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