Stream-k initial support

.gitignore

@@ -36,12 +36,13 @@ tags
 .vscode
 
 # build-in-source directory and documentation artifacts
-build/
+build*/
 _build/
 _images/
 _static/
 _templates/
 _toc.yml
+Tensile.egg-info
 
 # matrices
 *.csr

clients/benchmarks/client_groupedgemm_fixed_mk.cpp

@@ -770,7 +770,7 @@ int test_hipblaslt(hipDataType                 in_datatype,
     }
 
     // Set User Preference attributes
-    uint64_t max_workspace_size = 32 * 1024 * 1024;
+    uint64_t max_workspace_size = 128 * 1024 * 1024;
     void*    d_workspace;
     CHECK_HIP_ERROR(hipMalloc(&d_workspace, max_workspace_size));
 

clients/benchmarks/client_sequence.cpp

@@ -248,7 +248,7 @@ class LayerConfigIOGeneralSettings
     uint32_t rotating           = 0; // Size in MB
     uint32_t cold_iters         = 1000;
     uint32_t iters              = 10;
-    int64_t  max_workspace_size = 32 * 1024 * 1024;
+    int64_t  max_workspace_size = 128 * 1024 * 1024;
     bool     graph_mode         = false;
 };
 

library/src/amd_detail/rocblaslt/src/tensile_host.cpp

@@ -1886,7 +1886,7 @@ rocblaslt_status makeArgument(rocblaslt_handle             handle,
             if(useUserArgs)
             {
                 data->kernels = solution->solveGroupedGemmGPU(
-                    data->problem.gemms, data->inputs, nullptr, workspace, stream);
+                    data->problem.gemms, data->inputs, *hardware, nullptr, workspace, stream);
             }
             else
             {
@@ -2016,7 +2016,7 @@ rocblaslt_status getDeviceUserArgumentsValuesFromContractionProblem(rocblaslt_ha
         {
             std::shared_ptr<TensileDataGroupedGemm> data
                 = std::static_pointer_cast<TensileDataGroupedGemm>(gemmData);
-            auto  solution = library->getSolutionByIndex(data->algoIndex);
+            auto  solution = library->getSolutionByIndex(*hardware, data->algoIndex);
             auto& problem  = data->problem.gemms[0];
             if(problem.activationComputeType() == Tensile::DataType::Float)
             {
@@ -2077,7 +2077,7 @@ rocblaslt_status runKernelFromNewDeviceUserArguments(rocblaslt_handle       hand
             for(auto& it : data->kernels)
             {
                 uint8_t* arg      = it.args.rawdata();
-                auto     solution = library->getSolutionByIndex(data->algoIndex);
+                auto     solution = library->getSolutionByIndex(*hardware, data->algoIndex);
                 if(solution->internalArgsSupport.useUniversalArgs)
                 {
                     if(deviceUserArgs != nullptr)
@@ -2144,11 +2144,11 @@ rocblaslt_status runKernelFromDeviceUserArguments(rocblaslt_handle             h
         // don't overwrite data->algoIndex = *solutionIndex; here
         if(gemmType == rocblaslt::RocGemmType::ROCBLASLT_GROUPED_GEMM)
         {
-            auto solution = library->getSolutionByIndex(*solutionIndex);
+            auto solution = library->getSolutionByIndex(*hardware, *solutionIndex);
             std::shared_ptr<TensileDataGroupedGemm> data
                 = std::static_pointer_cast<TensileDataGroupedGemm>(gemmData);
             auto kernel = solution->solveGroupedGemmGPU(
-                data->problem.gemms, data->inputs, deviceUserArgs, workspace, stream);
+                data->problem.gemms, data->inputs,*hardware, deviceUserArgs, workspace, stream);
             status = hip2RocStatus(adapter->launchKernels(kernel, stream, nullptr, nullptr));
         }
         else
@@ -2187,7 +2187,8 @@ void _convertToHeuristicResultArray(
     rocblaslt_matmul_heuristic_result                           heuristicResultsArray[],
     int*                                                        returnAlgoCount,
     size_t                                                      maxWorkSpaceBytes,
-    const Tensile::ContractionProblemGemm&                      problem)
+    const Tensile::ContractionProblemGemm&                      problem,
+    const Tensile::Hardware&                                    hardware)
 {
     *returnAlgoCount = std::min((int)solutions.size(), requestedAlgoCount);
     for(size_t i = 0; i < *returnAlgoCount; i++)
@@ -2199,7 +2200,7 @@ void _convertToHeuristicResultArray(
         heuristicResultsArray[i].algo.max_workspace_bytes = maxWorkSpaceBytes;
         heuristicResultsArray[i].algo.fallback            = false;
         heuristicResultsArray[i].state                    = rocblaslt_status_success;
-        heuristicResultsArray[i].workspaceSize = solution->requiredWorkspaceSize(problem);
+        heuristicResultsArray[i].workspaceSize = solution->requiredWorkspaceSize(problem, hardware);
     }
     for(size_t i = *returnAlgoCount; i < requestedAlgoCount; i++)
     {
@@ -2296,7 +2297,8 @@ rocblaslt_status getBestSolutions(RocblasltContractionProblem const& prob,
                                    heuristicResultsArray,
                                    returnAlgoCount,
                                    maxWorkSpaceBytes,
-                                   data->problem);
+                                   data->problem,
+                                   *hardware);
 
     return rocblaslt_status_success;
 }
@@ -2367,7 +2369,7 @@ rocblaslt_status getAllSolutions(MyProblem&
         heuristicResults[i].algo.fallback            = false;
         heuristicResults[i].state                    = rocblaslt_status_success;
         if constexpr(std::is_same<MyProblem, Tensile::ContractionProblemGemm>::value)
-            heuristicResults[i].workspaceSize = solution->requiredWorkspaceSize(prob);
+            heuristicResults[i].workspaceSize = solution->requiredWorkspaceSize(prob, *hardware);
         else
             heuristicResults[i].workspaceSize = 0;
         i++;
@@ -2452,7 +2454,7 @@ rocblaslt_status
     for(auto index : solutionIndex)
     {
         isOutOfBound  = isOutOfBound && (index > lastSolutionIndex);
-        auto solution = library->getSolutionByIndex(index);
+        auto solution = library->getSolutionByIndex(*hardware, index);
         if(!solution)
             continue;
         rocblaslt_matmul_heuristic_result result;
@@ -2546,7 +2548,7 @@ rocblaslt_status isSolutionSupported(rocblaslt_handle       handle,
         }
         else
         {
-            *workspaceSizeInBytes = solution->requiredWorkspaceSize(tensile_prob);
+            *workspaceSizeInBytes = solution->requiredWorkspaceSize(tensile_prob, *hardware);
         }
     }
     else if constexpr(std::is_same<MyProblem, Tensile::ContractionProblemGroupedGemm>::value)
@@ -2581,7 +2583,7 @@ rocblaslt_status isSolutionSupported(rocblaslt_handle       handle,
 
         bool isSupported  = true;
         bool isNormalGemm = true;
-        auto problemWs    = solution->requiredWorkspaceSizeGroupedGemm(tensile_prob.gemms);
+        auto problemWs    = solution->requiredWorkspaceSizeGroupedGemm(tensile_prob.gemms, *hardware);
         for(int i = 0; i < tensile_prob.gemms.size(); i++)
         {
             tensile_prob.gemms[i].setWorkspaceSize(algo->max_workspace_bytes);
@@ -2730,7 +2732,8 @@ rocblaslt_status getBestSolutions(rocblaslt_handle       handle,
                                        heuristicResults.data(),
                                        &returnAlgoCount,
                                        workspaceBytes,
-                                       data->problem);
+                                       data->problem,
+                                       *hardware);
     }
     else if(gemmType == rocblaslt::RocGemmType::ROCBLASLT_GROUPED_GEMM)
     {
@@ -2755,7 +2758,8 @@ rocblaslt_status getBestSolutions(rocblaslt_handle       handle,
                                        heuristicResults.data(),
                                        &returnAlgoCount,
                                        workspaceBytes,
-                                       data->problem.gemms[0]);
+                                       data->problem.gemms[0],
+                                       *hardware);
     }
 
     return rocblaslt_status_success;
@@ -2811,6 +2815,10 @@ std::string getSolutionNameFromData(rocblaslt_handle             handle,
     int  wgm           = 0;
     int  solutionIndex = -1;
 
+    std::shared_ptr<Tensile::Hardware> hardware;
+
+    hardware     = Tensile::hip::GetDevice(*deviceProp);
+
     if(gemmType == rocblaslt::RocGemmType::ROCBLASLT_GEMM)
     {
         std::shared_ptr<TensileDataGemm> data = std::static_pointer_cast<TensileDataGemm>(gemmData);
@@ -2828,7 +2836,7 @@ std::string getSolutionNameFromData(rocblaslt_handle             handle,
     }
     if(solutionIndex == -1)
         return "";
-    auto        solution       = library->getSolutionByIndex(solutionIndex);
+    auto        solution       = library->getSolutionByIndex( *hardware, solutionIndex);
     std::string modifiedString = "";
     if(gsu != solution->sizeMapping.globalSplitU && gsu != 0)
     {
@@ -2853,9 +2861,10 @@ std::string getKernelNameFromAlgoIndex(rocblaslt_handle handle, const rocblaslt_
     std::shared_ptr<hipDeviceProp_t>                                                 deviceProp;
 
     auto adapter = get_library_and_adapter(&library, &deviceProp, handle->device);
-
+    std::shared_ptr<Tensile::Hardware>                                               hardware;
+    hardware = Tensile::hip::GetDevice(*deviceProp);
     int* solutionIndex = (int*)algo.data;
-    auto solution      = library->getSolutionByIndex(*solutionIndex);
+    auto solution      = library->getSolutionByIndex(*hardware, *solutionIndex);
     return solution->kernelName;
 }
 
@@ -2865,9 +2874,10 @@ std::string getSolutionNameFromAlgoIndex(rocblaslt_handle handle, const rocblasl
     std::shared_ptr<hipDeviceProp_t>                                                 deviceProp;
 
     auto adapter = get_library_and_adapter(&library, &deviceProp, handle->device);
-
+    std::shared_ptr<Tensile::Hardware>                                               hardware;
+    hardware = Tensile::hip::GetDevice(*deviceProp);
     int* solutionIndex = (int*)algo.data;
-    auto solution      = library->getSolutionByIndex(*solutionIndex);
+    auto solution      = library->getSolutionByIndex(*hardware , *solutionIndex);
     return solution->solutionName;
 }
 

tensilelite/Tensile/AsmStoreState.py

@@ -39,7 +39,7 @@ class StoreState:
     # the generation of the store code.
     ##############################################################################
     class StoreConstConfig:
-        def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic):
+        def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic, isWorkspace=False):
             self.gwvw = gwvw
             self.lsu = kernel["LocalSplitU"]
 
@@ -76,14 +76,18 @@ def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic):
                 self.numElementsPerBatchLimitedBySgprs = 1 # dummy value
                   #assert self.numElementsPerBatchLimitedBySgprs > 0, "numElementsPerBatchLimitedBySgprs=0 for %s"%self.kernelName
 
+            bpeC = kernelWriter.states.bpeCexternal
+            if isWorkspace:
+                bpeC = kernelWriter.states.bpeCinternal
+
             if atomic:
                 # flat atomics have another VGPR to allow different data for return#
                 regsPerElement = 2 if kernel["BufferStore"] else (3 + 1) # + 1 for alignment
                 # The atomic loop processes multiple elements in single instruction
                 # so will use VGPR from consec elements? TODO
-                self.numVgprsPerDataPerVI = (1.0 * regsPerElement * kernelWriter.states.bpeCexternal) / kernelWriter.states.bpr
+                self.numVgprsPerDataPerVI = (1.0 * regsPerElement * bpeC) / kernelWriter.states.bpr
             elif beta:
-                self.numVgprsPerDataPerVI = (1.0 * kernelWriter.states.bpeCexternal) / kernelWriter.states.bpr
+                self.numVgprsPerDataPerVI = (1.0 * bpeC) / kernelWriter.states.bpr
             else:
                 self.numVgprsPerDataPerVI = 0.0
 
@@ -104,7 +108,7 @@ def __init__(self, kernelWriter, kernel, ss, gwvw, edge, beta, atomic):
             self.lsuStartVgprOffset = 0
 
     # StoreState constructor:
-    def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements):
+    def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements, isWorkspace=False):
         self.kernelWriter = kernelWriter
         self.kernel = kernel
         self.lsu = kernel["LocalSplitU"]
@@ -158,7 +162,7 @@ def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements):
             if not atomic and len(kernel["PackedC1IndicesX"]) == 1:
                 self.optSrdIncForRow = 1
 
-        if kernel["StoreRemapVectorWidth"]:
+        if kernel["StoreRemapVectorWidth"] and not isWorkspace:
             self.optSrdIncForRow = 1
 
         if kernel["ProblemType"]["UseInitialStridesCD"]:
@@ -173,7 +177,7 @@ def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements):
         assert (not (self.optSingleColVgpr and self.optSharedColVgpr))
 
 
-        self.cfg = self.StoreConstConfig(kernelWriter, kernel, self, gwvw, edge, beta, atomic)
+        self.cfg = self.StoreConstConfig(kernelWriter, kernel, self, gwvw, edge, beta, atomic, isWorkspace)
 
         # Use to detect new rows:
         self.lastCoordOffset1 = 0
@@ -262,9 +266,22 @@ def __init__(self, kernelWriter, kernel, gwvw, edge, beta, atomic, elements):
         # For detecting when we are running first batch
         self.firstBatch = True
 
+        # TODO code migrated from globalWrite, comments outdated with respect to features like activation
+        # how many vgprs are needed for zero elements
+        # 2 for addressC in vgpr for addition - already checked out
+        # 2 for coord0,1 of thread - already checked out
+        # 2 for tmp - already checked out
+
+        # 5 = how many vgprs are needed per element (flat)
+        #  - 2 for addr
+        #  - 3 for GLOBAL_OFFSET_C calculation (can overlap below, therefore max)
+        #  - if beta gwvw*rpe for new value
+        #  - if atomic 2*rpe for old and cmp values
+
         # Calculate numVgprsPerElement
         # print("numVgprsPerAddr=%u, numVgprsPerDataPerVI=%u, numVgprPerValuC=%u"%(self.cfg.numVgprsPerAddr, self.cfg.numVgprsPerDataPerVI, self.cfg.numVgprPerValuC))
         self.numVgprsPerElement = self.cfg.numVgprPerValuC*gwvw + self.cfg.numVgprsPerAddr + int(ceil(self.cfg.numVgprsPerDataPerVI * gwvw))
+        # TODO STREAM-K Update numVgprsPerElement for workspace
         if kernel["GroupLoadStore"] and kernel["ProblemType"]["UseBeta"]:
             self.numVgprsPerElement += self.cfg.numVgprsPerAddr
         if kernel["ProblemType"]["UseE"] and (kernel["GlobalSplitU"] == 1):
@@ -377,7 +394,7 @@ def getStoreElementsInfoForBatch(self, kernel, batchElements):
 
         return self.elementCoord0, self.elementCoord1
 
-    def setupStoreElementsForBatch(self, kernel, gwvw, batchElements, batchElementSgprs, isOptNLL, factorDim):
+    def setupStoreElementsForBatch(self, kernel, gwvw, batchElements, batchElementSgprs, isOptNLL, factorDim, isWorkspace=False):
 
         self.elementAddr              = []
         self.elementDataE             = []
@@ -394,6 +411,9 @@ def setupStoreElementsForBatch(self, kernel, gwvw, batchElements, batchElementSg
         #generate elementCoord0/1
         self.getStoreElementsInfoForBatch(kernel, batchElements)
         kw = self.kernelWriter
+        dataType = kernel["ProblemType"]["DestDataType"]
+        if isWorkspace:
+            dataType = kernel["ProblemType"]["ComputeDataType"]
 
         biasVgprMap = {}
         scaleAVecVgprMap = {}
@@ -514,7 +534,7 @@ def setupStoreElementsForBatch(self, kernel, gwvw, batchElements, batchElementSg
                 if self.cfg.halfDataRegPerVI:
                     # TODO- check (H,H,H,H,S,S)
                     if kernel["ProblemType"]["HighPrecisionAccumulate"] and \
-                       (kernel["ProblemType"]["DestDataType"].isBFloat16() or kernel["ProblemType"]["DestDataType"].isHalf()):
+                       (dataType.isBFloat16() or dataType.isHalf()):
                         data = kw.vgprPool.checkOutAligned(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw), \
                               int(ceil(int(2*self.cfg.numVgprsPerDataPerVI*self.cfg.gwvw))), "writeBatch-data for ei=%u and ei=%u"%(elementIdx,elementIdx+1), preventOverflow=not isOptNLL)
                     else: