diff --git a/core/inc/SOFIE/ROperator.hxx b/core/inc/SOFIE/ROperator.hxx index aa203521..a2999251 100644 --- a/core/inc/SOFIE/ROperator.hxx +++ b/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& removal_func){ return;}; // Elementwise kernel fusion interface diff --git a/core/inc/SOFIE/ROperator_Conv.hxx b/core/inc/SOFIE/ROperator_Conv.hxx index 5e305e64..1570ff1f 100644 --- a/core/inc/SOFIE/ROperator_Conv.hxx +++ b/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(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(queue, workDiv_im2col, im2colKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNX << ") + x_offset" - << ", alpaka::getPtrNative(deviceBuf_" << imcol << ")" + << ", alpaka::getPtrNative(deviceBuf_" << imcol << ") + col_offset" << ", static_cast(" << 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(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(" << 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 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; diff --git a/core/src/RModel_ALPAKA.cxx b/core/src/RModel_ALPAKA.cxx index 9e0e84c8..82baac1f 100644 --- a/core/src/RModel_ALPAKA.cxx +++ b/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"; diff --git a/test/alpaka/TestAlpakaConv.cxx b/test/alpaka/TestAlpakaConv.cxx index 56c1b0c7..ec3606ba 100644 --- a/test/alpaka/TestAlpakaConv.cxx +++ b/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 input(100); + std::iota(input.begin(), input.end(), 0.0f); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{input.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < input.size(); ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{input.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sizeof(ConvBatch4_ExpectedOutput::correct) / sizeof(float)})); + + { + SOFIE_ConvBatch4::Session 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(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;