Batch the non-grouped Conv GEMM with gemmStridedBatched

core/inc/SOFIE/ROperator.hxx

@@ -81,6 +81,7 @@ public:
    virtual std::string Header() { return "";}
    virtual std::string GetFusableOutputTensorName() { return "";}
    virtual std::string GetBlasConfig() { return ""; }
+   virtual bool UsesBatchedGemm() { return false; }
    virtual void UpdateFusableTensorName(std::string, const std::function<void(const std::string&)>& removal_func){ return;};
 
    // Elementwise kernel fusion interface

core/inc/SOFIE/ROperator_Conv.hxx

@@ -874,22 +874,22 @@ public:
       // Step 3 + 4: Im2Col then GEMM — structure differs for grouped vs non-grouped
       // -----------------------------------------------------------------------
       if (fAttrGroup == 1) {
-         // Non-grouped: single im2col per batch, then GEMM
-         out << SP << SP << "// Step 3: im2col\n";
+         // batched gives each sample its own _xcol slice, otherwise all reuse slice 0
+         out << SP << SP << "std::size_t const col_offset = fBatchedGemm ? n * " << colElements << "u : 0u;\n\n";
          out << SP << SP << "{\n";
          out << SP << SP << SP << "auto const elementsPerThread_im2col = Vec::all(static_cast<Idx>(1));\n";
          out << SP << SP << SP << "auto const elementsPerGrid_im2col   = Vec::all(Idx{" << colElements << "});\n";
          out << SP << SP << SP << "auto const workDiv_im2col = sofie_workdiv(elementsPerGrid_im2col);\n";
          out << SP << SP << SP << "alpaka::exec<Acc>(queue, workDiv_im2col, im2colKernel_" << opName
             << ", alpaka::getPtrNative(deviceBuf_" << fNX << ") + x_offset"
-            << ", alpaka::getPtrNative(deviceBuf_" << imcol << ")"
+            << ", alpaka::getPtrNative(deviceBuf_" << imcol << ") + col_offset"
             << ", static_cast<Idx>(" << colElements << "));\n";
-         out << SP << SP << SP << "alpaka::wait(queue);\n";
+         out << SP << SP << SP << "if (!fBatchedGemm) alpaka::wait(queue);\n";
          out << SP << SP << "}\n\n";
 
          if (!fNB.empty()) {
                size_t biasElements = gemm_n * gemm_m;
-               out << SP << SP << "// Step 4a: broadcast bias into output slice\n";
+               out << SP << SP << "// broadcast bias into this sample's output slice\n";
                out << SP << SP << "{\n";
                out << SP << SP << SP << "auto const elementsPerThread_bias = Vec::all(static_cast<Idx>(1));\n";
                out << SP << SP << SP << "auto const elementsPerGrid_bias   = Vec::all(Idx{" << biasElements << "});\n";
@@ -898,24 +898,19 @@ public:
                   << ", alpaka::getPtrNative(deviceBuf_" << fNB << ")"
                   << ", alpaka::getPtrNative(deviceBuf_" << fNY << ") + out_offset"
                   << ", static_cast<Idx>(" << biasElements << "));\n";
-               out << SP << SP << SP << "alpaka::wait(queue);\n";
+               out << SP << SP << SP << "if (!fBatchedGemm) alpaka::wait(queue);\n";
                out << SP << SP << "}\n\n";
-               out << SP << SP << "// Step 4b: GEMM beta=1 accumulates onto bias-initialised output\n";
-               out << SP << SP << "blas.matmul('n', 'n', "
-                  << gemm_m << ", " << gemm_n << ", " << gemm_k
-                  << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << imcol << ")"
-                  << ", alpaka::getPtrNative(deviceBuf_" << convK << ")"
-                  << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << fNY << ") + out_offset);\n\n";
-         } else {
-               out << SP << SP << "// Step 4: GEMM beta=0 (no bias)\n";
-               out << SP << SP << "blas.matmul('n', 'n', "
-                  << gemm_m << ", " << gemm_n << ", " << gemm_k
-                  << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << imcol << ")"
-                  << ", alpaka::getPtrNative(deviceBuf_" << convK << ")"
-                  << ", 0.0f, alpaka::getPtrNative(deviceBuf_" << fNY << ") + out_offset);\n\n";
          }
-         // Wait for GEMM to finish before next batch overwrites the shared _xcol buffer.
-         out << SP << SP << "alpaka::wait(queue);\n\n";
+
+         out << SP << SP << "if (!fBatchedGemm) {\n";
+         out << SP << SP << SP << "blas.matmul('n', 'n', "
+            << gemm_m << ", " << gemm_n << ", " << gemm_k
+            << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << imcol << ")"
+            << ", alpaka::getPtrNative(deviceBuf_" << convK << ")"
+            << ", " << (fNB.empty() ? "0.0f" : "1.0f")
+            << ", alpaka::getPtrNative(deviceBuf_" << fNY << ") + out_offset);\n";
+         out << SP << SP << SP << "alpaka::wait(queue);\n";
+         out << SP << SP << "}\n\n";
 
       } else {
          // Grouped convolution: im2col and GEMM per group with group-adjusted input pointer.
@@ -970,6 +965,20 @@ public:
       }
 
       out << SP << "}\n"; // end batch loop
+
+      if (fAttrGroup == 1) {
+         std::string convBeta = fNB.empty() ? "0.0f" : "1.0f";
+         out << SP << "if (fBatchedGemm) {\n";
+         out << SP << SP << "alpaka::wait(queue);\n";
+         out << SP << SP << "blas.gemmStridedBatched('n', 'n', "
+            << gemm_m << ", " << gemm_n << ", " << gemm_k << ", 1.0f, "
+            << "alpaka::getPtrNative(deviceBuf_" << imcol << "), " << gemm_m << ", " << colElements << ", "
+            << "alpaka::getPtrNative(deviceBuf_" << convK << "), " << gemm_k << ", 0, "
+            << convBeta << ", alpaka::getPtrNative(deviceBuf_" << fNY << "), "
+            << gemm_m << ", " << gemm_n * gemm_m << ", " << bsize << ");\n";
+         out << SP << SP << "alpaka::wait(queue);\n";
+         out << SP << "}\n";
+      }
       return out.str();
    }
 
@@ -978,7 +987,39 @@ public:
    std::vector<std::string> GetBlasRoutines() override { return { std::string("Gemm"), std::string("Axpy") }; }
 
 
+   bool UsesBatchedGemm() override { return fAttrGroup == 1; }
+
+   std::string GenerateInitCode_GPU_ALPAKA() override {
+      if (fAttrGroup != 1) return "";
+      if (fShapeX.empty() || fShapeW.empty() || fShapeY.empty()) return "";
+
+      size_t bsize      = fShapeX[0].dim;
+      size_t oDepth     = (fDim > 2) ? fShapeY[2].dim    : 1;
+      size_t oHeight    = (fDim > 1) ? fShapeY[fDim].dim : 1;
+      size_t oWidth     = fShapeY[fDim + 1].dim;
+      size_t kernelSize = fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2];
+      size_t colElements = fShapeW[1] * kernelSize * oDepth * oHeight * oWidth;
+
+      // _xcol is declared for one sample, batching needs a slice per sample
+      std::stringstream out;
+      out << SP << "if (!fBatchedGemm) {\n";
+      out << SP << SP << "blas.addLayoutConfig(" << BlasLayoutConfig() << ");\n";
+      if (bsize > 1) {
+         out << SP << "} else {\n";
+         out << SP << SP << "deviceBuf_" << imcol << " = alpaka::allocBuf<" << fType
+             << ", size_t>(devAcc, Ext1D::all(Idx{" << bsize * colElements << "}));\n";
+      }
+      out << SP << "}\n";
+      return out.str();
+   }
+
    std::string GetBlasConfig(){
+      // batched path is legacy cuBLAS, the fallback registers its own layout
+      if (fAttrGroup == 1) return "";
+      return BlasLayoutConfig();
+   }
+
+   std::string BlasLayoutConfig(){
       size_t oDepth_  = (fDim > 2) ? fShapeY[2].dim    : 1;
       size_t oHeight_ = (fDim > 1) ? fShapeY[fDim].dim : 1;
       size_t oWidth_  = fShapeY[fDim + 1].dim;

core/src/RModel_ALPAKA.cxx

@@ -695,6 +695,14 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() {
       fGC += RModelProfilerGPU::GenerateSessionMembers();
    }
 
+   bool hasBatchedGemm = false;
+   for (size_t id = 0; id < fOperators.size(); id++) {
+      if (fSkipOperators.count(id)) continue;
+      if (fOperators[id]->UsesBatchedGemm()) hasBatchedGemm = true;
+   }
+   if (hasBatchedGemm)
+      fGC += "\nbool fBatchedGemm = false;\n";
+
    // Session constructor
    if (fUseSession) {
       std::string sessionName = "\n\nSession";
@@ -719,8 +727,13 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() {
             fGC += "        size_t " + p.first + " = " + p.second;
          }
       }
+      if (hasBatchedGemm)
+         fGC += ",\n        bool batchedGemm = true";
       fGC += ") {\n";
-
+
+      if (hasBatchedGemm)
+         fGC += SP + "fBatchedGemm = batchedGemm;\n";
+
       GenerateTemporaryInitializedTensorContainers_GPU_ALPAKA();
       if (fUseWeightFile) {
          fGC += "\n//--- reading weights from file\n";

test/alpaka/TestAlpakaConv.cxx

@@ -459,6 +459,40 @@ TEST_F(SofieAlpakaTest, ConvBatch4)
       EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i;
 }
 
+TEST_F(SofieAlpakaTest, ConvBatch4NoBatchedGemm)
+{
+   constexpr float TOLERANCE = DEFAULT_TOLERANCE;
+
+   std::vector<float> input(100);
+   std::iota(input.begin(), input.end(), 0.0f);
+
+   auto input_h = alpaka::allocBuf<float, Idx>(host, Ext1D::all(Idx{input.size()}));
+   float* input_ptr = reinterpret_cast<float*>(alpaka::getPtrNative(input_h));
+   for (Idx i = 0; i < input.size(); ++i) input_ptr[i] = input[i];
+
+   auto input_d = alpaka::allocBuf<float, Idx>(device, Ext1D::all(Idx{input.size()}));
+   alpaka::memcpy(queue, input_d, input_h);
+   alpaka::wait(queue);
+
+   auto result_h = alpaka::allocBuf<float, Idx>(host, Ext1D::all(Idx{sizeof(ConvBatch4_ExpectedOutput::correct) / sizeof(float)}));
+
+   {
+      SOFIE_ConvBatch4::Session<alpaka::TagGpuCudaRt> session("ConvBatch4_FromONNX_GPU_ALPAKA.dat", false);
+      auto result = session.infer(input_d);
+      alpaka::wait(queue);
+      cudaDeviceSynchronize();
+      alpaka::memcpy(queue, result_h, result);
+      alpaka::wait(queue);
+   }
+
+   float* res_ptr = reinterpret_cast<float*>(alpaka::getPtrNative(result_h));
+   float* correct = ConvBatch4_ExpectedOutput::correct;
+   constexpr size_t nOut_batch4_noBatched = sizeof(ConvBatch4_ExpectedOutput::correct) / sizeof(float);
+
+   for (size_t i = 0; i < nOut_batch4_noBatched; ++i)
+      EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i;
+}
+
 TEST_F(SofieAlpakaTest, ConvBatch8)
 {
    constexpr float TOLERANCE = DEFAULT_TOLERANCE;