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

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

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

@@ -53,6 +53,7 @@
 #include <rocRoller/KernelGraph/RegisterTagManager_fwd.hpp>
 #include <rocRoller/KernelGraph/ScopeManager_fwd.hpp>
 #include <rocRoller/KernelOptions.hpp>
+#include <rocRoller/Operations/Scratch_fwd.hpp>
 #include <rocRoller/ScheduledInstructions_fwd.hpp>
 #include <rocRoller/Scheduling/Scheduling_fwd.hpp>
 #include <rocRoller/Utilities/Random_fwd.hpp>
@@ -127,18 +128,23 @@ namespace rocRoller
         void setKernel(AssemblyKernelPtr);
 
         /**
-         * @brief Returns an expression representing how much scratch space is required (in bytes)
+         * @brief Allocate scratch space for the specified scratch policy.
          *
-         * @return Expression::ExpressionPtr
+         * @param policy The scratch policy to allocate for
+         * @param size Number of bytes requested
+         * @return Expression::ExpressionPtr The offset before this allocation
          */
-        Expression::ExpressionPtr getScratchAmount() const;
+        Expression::ExpressionPtr allocateScratch(Operations::ScratchPolicy policy,
+                                                  Expression::ExpressionPtr size);
 
         /**
-         * @brief Allocate more scratch space
+         * @brief Returns an expression representing how much scratch space is required (in bytes)
+         *        for the specified scratch policy.
          *
-         * @param size Number of bytes requested
+         * @param policy The scratch policy to query
+         * @return Expression::ExpressionPtr
          */
-        void allocateScratch(Expression::ExpressionPtr size);
+        Expression::ExpressionPtr getScratchAmount(Operations::ScratchPolicy policy) const;
 
         /**
          * @brief Get register scope manager.
@@ -168,14 +174,15 @@ namespace rocRoller
         std::array<std::shared_ptr<Register::Allocator>, static_cast<size_t>(Register::Type::Count)>
             m_allocators;
 
-        std::shared_ptr<Scheduling::IObserver>     m_observer;
-        AssemblyKernelPtr                          m_kernel;
-        std::shared_ptr<ArgumentLoader>            m_argLoader;
-        std::shared_ptr<ScheduledInstructions>     m_instructions;
-        std::shared_ptr<MemoryInstructions>        m_mem;
-        LabelAllocatorPtr                          m_labelAllocator;
-        std::shared_ptr<LDSAllocator>              m_ldsAllocator;
-        Expression::ExpressionPtr                  m_scratchAllocator;
+        std::shared_ptr<Scheduling::IObserver> m_observer;
+        AssemblyKernelPtr                      m_kernel;
+        std::shared_ptr<ArgumentLoader>        m_argLoader;
+        std::shared_ptr<ScheduledInstructions> m_instructions;
+        std::shared_ptr<MemoryInstructions>    m_mem;
+        LabelAllocatorPtr                      m_labelAllocator;
+        std::shared_ptr<LDSAllocator>          m_ldsAllocator;
+        std::array<Expression::ExpressionPtr, static_cast<size_t>(Operations::ScratchPolicy::Count)>
+                                                   m_scratchSizes;
         std::shared_ptr<CopyGenerator>             m_copier;
         std::shared_ptr<BranchGenerator>           m_brancher;
         std::shared_ptr<CrashKernelGenerator>      m_crasher;

shared/rocroller/lib/include/rocRoller/KernelGraph/CoordinateGraph/Dimension.hpp

@@ -153,8 +153,11 @@ namespace rocRoller
              *
              * @param size How many elements make up the User dimension.
              * @param offset Location of data within the scratch space
+             * @param argName Name of the argument for this scratch space
              */
-            User(Expression::ExpressionPtr size, Expression::ExpressionPtr offset);
+            User(Expression::ExpressionPtr size,
+                 Expression::ExpressionPtr offset,
+                 std::string const&        argName);
 
             std::string name() const override;
         };

shared/rocroller/lib/include/rocRoller/KernelGraph/Utils.hpp

@@ -319,11 +319,15 @@ namespace rocRoller
          *
          * @param size
          * @param varType
+         * @param policy The scratch policy to use for allocation
          * @param context
          * @return User
          */
-        rocRoller::KernelGraph::CoordinateGraph::User newScratchCoordinate(
-            Expression::ExpressionPtr size, VariableType varType, ContextPtr context);
+        rocRoller::KernelGraph::CoordinateGraph::User
+            newScratchCoordinate(Expression::ExpressionPtr size,
+                                 VariableType              varType,
+                                 Operations::ScratchPolicy policy,
+                                 ContextPtr                context);
 
         /**
          * @brief Replace operation with a new operation.

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

@@ -33,6 +33,7 @@
 #include <string>
 
 #include <rocRoller/AssertOpKinds_fwd.hpp>
+#include <rocRoller/Operations/Scratch_fwd.hpp>
 #include <rocRoller/Utilities/EnumBitset.hpp>
 #include <rocRoller/Utilities/Settings_fwd.hpp>
 
@@ -49,6 +50,12 @@ namespace rocRoller
     const std::string NUMWGS  = "numWGs";
     const std::string WGM     = "WGM";
 
+    // Helper to get scratch argument name for a specific policy
+    inline std::string getScratchName(Operations::ScratchPolicy policy)
+    {
+        return rocRoller::SCRATCH + "_" + Operations::toString(policy);
+    }
+
     class KernelOptions
     {
     public:

shared/rocroller/lib/source/CommandSolution.cpp

@@ -646,9 +646,10 @@ namespace rocRoller
         return m_context;
     }
 
-    size_t CommandKernel::scratchSpaceRequired(RuntimeArguments const& args) const
+    size_t CommandKernel::scratchSpaceRequired(Operations::ScratchPolicy policy,
+                                               RuntimeArguments const&   args) const
     {
-        auto amount = m_context->getScratchAmount();
+        auto amount = m_context->getScratchAmount(policy);
 
         auto times = evaluationTimes(amount);
         AssertFatal(times[Expression::EvaluationTime::Translate]

shared/rocroller/lib/source/Context.cpp

@@ -47,8 +47,12 @@
 namespace rocRoller
 {
     Context::Context()
-        : m_scratchAllocator(Expression::literal(0u))
     {
+        // Initialize scratch sizes for each policy with zero
+        for(size_t i = 0; i < static_cast<size_t>(Operations::ScratchPolicy::Count); ++i)
+        {
+            m_scratchSizes[i] = Expression::literal(0u);
+        }
     }
 
     ContextPtr Context::ForDefaultHipDevice(std::string const&   kernelName,
@@ -287,14 +291,18 @@ namespace rocRoller
         m_kernel = assemblyKernel;
     }
 
-    Expression::ExpressionPtr Context::getScratchAmount() const
+    Expression::ExpressionPtr Context::allocateScratch(Operations::ScratchPolicy policy,
+                                                       Expression::ExpressionPtr size)
     {
-        return m_scratchAllocator;
+        auto idx            = static_cast<size_t>(policy);
+        auto currentOffset  = m_scratchSizes[idx];
+        m_scratchSizes[idx] = simplify(m_scratchSizes[idx] + size);
+        return currentOffset;
     }
 
-    void Context::allocateScratch(Expression::ExpressionPtr size)
+    Expression::ExpressionPtr Context::getScratchAmount(Operations::ScratchPolicy policy) const
     {
-        m_scratchAllocator = simplify(m_scratchAllocator + size);
+        return m_scratchSizes[static_cast<size_t>(policy)];
     }
 
     void Context::scheduleCopy(Instruction const& inst)

shared/rocroller/lib/source/KernelGraph/CoordinateGraph/Dimension.cpp

@@ -117,9 +117,11 @@ namespace rocRoller
             return name() + stag;
         }
 
-        User::User(Expression::ExpressionPtr size, Expression::ExpressionPtr offset)
+        User::User(Expression::ExpressionPtr size,
+                   Expression::ExpressionPtr offset,
+                   std::string const&        argName)
             : BaseDimension(size, Expression::literal(1u), offset)
-            , argumentName(rocRoller::SCRATCH)
+            , argumentName(argName)
         {
         }