diff --git a/check_style.sh b/check_style.sh index 22a56e4d..91ae9ccf 100644 --- a/check_style.sh +++ b/check_style.sh @@ -1,14 +1,17 @@ #!/bin/bash set -e -# Directories -SRC_DIR="./include" -TEST_DIR="./tests" +# Directories (SOFIE's own layout: source lives under core/parsers/utils, +# tests under test/ - unlike sofieBLAS this script was originally written for, +# SOFIE has no top-level include/ or tests/ directory) +SRC_DIRS="./core ./parsers ./utils" +TEST_DIR="./test" echo "📝 Discovering source/header files..." -FILES=$(find "$SRC_DIR" "$TEST_DIR" \ +FILES=$(find $SRC_DIRS "$TEST_DIR" \ -path "$TEST_DIR/build" -prune -o \ + -path "$TEST_DIR/input_models/references" -prune -o \ -type f \( \ -name '*.cpp' -o -name '*.cc' -o -name '*.cxx' -o \ -name '*.h' -o -name '*.hpp' -o -name '*.hxx' -o -name '*.hh' \ @@ -31,7 +34,7 @@ done echo "🔍 Running clang-tidy..." for file in $FILES; do echo "Linting $file" - clang-tidy "$file" --extra-arg=-std=c++20 -- -I"$SRC_DIR" || true + clang-tidy "$file" --extra-arg=-std=c++20 -- -I./core/inc -I./parsers/inc -I./utils || true done echo "✅ Formatting and linting complete." diff --git a/core/inc/SOFIE/RModel.hxx b/core/inc/SOFIE/RModel.hxx index a31fa5b8..1eb7324e 100644 --- a/core/inc/SOFIE/RModel.hxx +++ b/core/inc/SOFIE/RModel.hxx @@ -8,6 +8,10 @@ namespace SOFIE { +/// Session member holding the construction-time value of a shape parameter (N -> fN); +/// shared by the CPU and GPU session generators +std::string GetMemberNameForDimShape(std::string name); + class RModel final : public RModel_Base { friend class RModelProfiler; @@ -60,7 +64,7 @@ private: std::vector opIndices; ///< consecutive op indices forming this group std::string inputTensor; ///< input tensor name of the first op std::string outputTensor; ///< output tensor name of the last op - size_t numElements = 0; + std::string iterationLengthExpr; ///< element count: literal for static tensors, runtime expression for dynamic bool isFused() const { return opIndices.size() > 1; } std::string suffix() const { std::string s; @@ -211,6 +215,12 @@ public: // generate the infer function signature for inference on ALPAKA. If isdecl= false generate the calling infer function // used to infer the sub-graphs std::string GenerateInferSignature_GPU_ALPAKA(bool isdecl = true); + /// iterate the GPU infer argument sequence in order: onParam(param) for each dynamic + /// shape parameter at the position of the input introducing it, onInput(name) for each + /// input tensor; single source of the argument-order convention shared by the + /// signatures, the _infer_impl call sites and the Session constructor + void ForEachInferArg_GPU_ALPAKA(const std::function &onParam, + const std::function &onInput) const; // generate the _infer_impl signature using ViewPlainPtr types instead of Buf types std::string GenerateImplSignature_GPU_ALPAKA(bool isdecl = true); diff --git a/core/inc/SOFIE/RModelProfilerGPU.hxx b/core/inc/SOFIE/RModelProfilerGPU.hxx index bc4aab26..4140ee66 100644 --- a/core/inc/SOFIE/RModelProfilerGPU.hxx +++ b/core/inc/SOFIE/RModelProfilerGPU.hxx @@ -35,7 +35,7 @@ public: static std::string GenerateMemoryReport(const MemoryInfo &info); static std::string GenerateBeginInferCode(); - static std::string GenerateOperatorCode(ROperator &op, size_t op_idx); + static std::string GenerateOperatorCode(ROperator &op, size_t op_idx, const std::vector &dynParamNames); static std::string GenerateEndInferCode(); RModelProfilerGPU() = delete; diff --git a/core/inc/SOFIE/ROperator.hxx b/core/inc/SOFIE/ROperator.hxx index adccb308..af6910d8 100644 --- a/core/inc/SOFIE/ROperator.hxx +++ b/core/inc/SOFIE/ROperator.hxx @@ -77,6 +77,10 @@ public: virtual void Initialize(RModel&) = 0; virtual std::string Generate(std::string OpName) = 0; //expect unique opName for each operator within the same RModel virtual std::string Generate_GPU_ALPAKA(std::string OpName){ return "";} //expect unique opName for each operator within the same RModel + //dynParamNames: the model's shape parameters in infer-argument order + virtual std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector & /*dynParamNames*/) { + return Generate_GPU_ALPAKA(OpName); + } // generate initialization code for session constructor virtual std::string GenerateInitCode() { return "";} virtual std::string GenerateInitCode_GPU_ALPAKA() { return "";}; @@ -87,6 +91,9 @@ public: // generate session data members specific to operator virtual std::string GenerateSessionMembersCode(std::string /*opName*/) { return ""; } virtual std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) { return ""; } + virtual std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector & /*dynParamNames*/) { + return Generate_GPU_Kernel_ALPAKA(opName); + } virtual std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string /*opName*/) { return ""; } virtual std::string Header() { return "";} virtual std::string GetFusableOutputTensorName() { return "";} diff --git a/core/inc/SOFIE/ROperator_BasicBinary.hxx b/core/inc/SOFIE/ROperator_BasicBinary.hxx index 9a1a963e..69f22d1a 100644 --- a/core/inc/SOFIE/ROperator_BasicBinary.hxx +++ b/core/inc/SOFIE/ROperator_BasicBinary.hxx @@ -276,6 +276,9 @@ public: << ConvertShapeToString(fShapeY) << std::endl; } // we convert non-dim shapes to Dim shapes + // (fShapeA/fShapeB were rank-padded to fShapeY by MultidirectionalBroadcastShape) + fDimShapeA = ConvertShapeToDim(fShapeA); + fDimShapeB = ConvertShapeToDim(fShapeB); fDimShapeY = ConvertShapeToDim(fShapeY); } } else { @@ -451,97 +454,63 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) { + // fDimShapeA/fDimShapeB are rank-padded to fDimShapeY by Initialize, so a broadcast + // dim is a literal 1. An input is either fully broadcast (scalar, all dims 1), not + // broadcast (same shape as the output), or partially broadcast - only the last case + // needs per-dim index decomposition in the kernel. + bool IsPartialBroadcast(const std::vector &shape) const { + return ConvertDimShapeToLength(shape) != "1" && !UTILITY::AreSameShape(shape, fDimShapeY); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; + if (fDimShapeY.empty()) + throw std::runtime_error("SOFIE Operator Basic Binary called to Generate without being initialized first"); + + const std::size_t D = fDimShapeY.size(); + auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY); + auto stridesA = UTILITY::ComputeStrideFromShape(fDimShapeA); + auto stridesB = UTILITY::ComputeStrideFromShape(fDimShapeB); + bool isAScalar = ConvertDimShapeToLength(fDimShapeA) == "1"; + bool isBScalar = ConvertDimShapeToLength(fDimShapeB) == "1"; + bool isAPartialBroadcast = IsPartialBroadcast(fDimShapeA); + bool isBPartialBroadcast = IsPartialBroadcast(fDimShapeB); std::string op; op = "\n//------ "+opName+"_"+BinaryOperatorTrait::Name()+"_KERNEL_ALPAKA\n"; op += SP + "struct Binary"+opName+BinaryOperatorTrait::Name()+"Kernel {\n"; op += SP + SP + "template\n"; - op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * A, T const * B, T * C) const {\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * A, T const * B, T * C"; + for (auto &p : dynParamNames) + op += ", std::size_t const " + p; + op += ", std::size_t const totalElements) const {\n"; op += SP + SP + SP + "auto idx = alpaka::getIdx(acc)[0];\n"; - op += SP + SP + SP + "if (idx < " + std::to_string(ConvertShapeToLength(fShapeY)) + ") {\n"; - auto stridesA = UTILITY::ComputeStrideFromShape(fShapeA); - auto stridesB = UTILITY::ComputeStrideFromShape(fShapeB); - - for(size_t id_s = 0; id_s < stridesA.size(); ++id_s){ - if(fShapeA[id_s] == 1) - stridesA[id_s] = 0; - } - - for(size_t id_s = 0; id_s < stridesB.size(); ++id_s){ - if(fShapeB[id_s] == 1) - stridesB[id_s] = 0; - } - - auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); - - // --- Fast-path index simplifications --- - // Check whether A is broadcast (all strides zero → single element) - bool isAScalar = true; - for (const auto& s : stridesA) { if (s != 0) { isAScalar = false; break; } } - // Check whether B is broadcast (all strides zero → single element) - bool isBScalar = true; - for (const auto& s : stridesB) { if (s != 0) { isBScalar = false; break; } } - // Check whether A has the same contiguous layout as Y (no broadcasting) - bool isAContiguous = (fShapeA.size() == fShapeY.size()); - if (isAContiguous) { - for (size_t i = 0; i < fShapeA.size(); ++i) - if (fShapeA[i] != fShapeY[i]) { isAContiguous = false; break; } - } - // Check whether B has the same contiguous layout as Y (no broadcasting) - bool isBContiguous = (fShapeB.size() == fShapeY.size()); - if (isBContiguous) { - for (size_t i = 0; i < fShapeB.size(); ++i) - if (fShapeB[i] != fShapeY[i]) { isBContiguous = false; break; } - } - - std::string flattened_index_A = ""; - std::string flattened_index_B = ""; - - if (isAScalar) { - // A is a single broadcast value - flattened_index_A = "0"; - } else if (isAContiguous) { - // A and Y have identical shapes → direct index - flattened_index_A = "idx"; - } else { - // General broadcast case: decompose idx into per-dim coords - std::string temp = "idx"; - for (size_t id_s = 0; id_s < fShapeA.size(); ++id_s) { - auto strideY = stridesY[id_s]; - auto strideA = stridesA[id_s]; - std::string coord = "(int)(" + temp + " / " + std::to_string(strideY) + ")"; - flattened_index_A += coord + " * " + std::to_string(strideA) + " + "; - temp = temp + " - (" + coord + " * " + std::to_string(strideY) + ")"; + op += SP + SP + SP + "if (idx < totalElements) {\n"; + + // decompose idx into output coords, skipping broadcast dims (stride 0) + if (isAPartialBroadcast || isBPartialBroadcast) { + op += SP + SP + SP + SP + "std::size_t remaining = idx;\n"; + op += SP + SP + SP + SP + "std::size_t coord;\n"; + if (isAPartialBroadcast) op += SP + SP + SP + SP + "std::size_t idxA = 0;\n"; + if (isBPartialBroadcast) op += SP + SP + SP + SP + "std::size_t idxB = 0;\n"; + for (std::size_t d = 0; d < D; d++) { + std::string sY = "(" + stridesY[d].GetVal() + ")"; + op += SP + SP + SP + SP + "coord = remaining / " + sY + ";\n"; + if (d + 1 < D) + op += SP + SP + SP + SP + "remaining -= coord * " + sY + ";\n"; + if (isAPartialBroadcast && fDimShapeA[d].GetVal() != "1") + op += SP + SP + SP + SP + "idxA += coord * (" + stridesA[d].GetVal() + ");\n"; + if (isBPartialBroadcast && fDimShapeB[d].GetVal() != "1") + op += SP + SP + SP + SP + "idxB += coord * (" + stridesB[d].GetVal() + ");\n"; } - if (!flattened_index_A.empty()) - flattened_index_A.erase(flattened_index_A.size() - 3); } - - if (isBScalar) { - // B is a single broadcast value - flattened_index_B = "0"; - } else if (isBContiguous) { - // B and Y have identical shapes → direct index - flattened_index_B = "idx"; - } else { - // General broadcast case - std::string temp = "idx"; - for (size_t id_s = 0; id_s < fShapeB.size(); ++id_s) { - auto strideY = stridesY[id_s]; - auto strideB = stridesB[id_s]; - std::string coord = "(int)(" + temp + " / " + std::to_string(strideY) + ")"; - flattened_index_B += coord + " * " + std::to_string(strideB) + " + "; - temp = temp + " - (" + coord + " * " + std::to_string(strideY) + ")"; - } - if (!flattened_index_B.empty()) - flattened_index_B.erase(flattened_index_B.size() - 3); - } - - op += "C[idx] = " + BinaryOperatorTrait::Op("A["+flattened_index_A+"]", "B["+flattened_index_B+"]") + ";\n"; - op += "}\n}\n};\n"; + std::string indexA = isAScalar ? "0" : (isAPartialBroadcast ? "idxA" : "idx"); + std::string indexB = isBScalar ? "0" : (isBPartialBroadcast ? "idxB" : "idx"); + op += SP + SP + SP + SP + "C[idx] = " + BinaryOperatorTrait::Op("A["+indexA+"]", "B["+indexB+"]") + ";\n"; + op += SP + SP + SP + "}\n"; + op += SP + SP + "}\n"; + op += SP + "};\n"; return op; } @@ -552,7 +521,7 @@ public: return SP + "Binary"+OpName+BinaryOperatorTrait::Name()+"Kernel binary" + OpName + "Kernel;\n"; } - std::string Generate_GPU_ALPAKA(std::string OpName) { + std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; @@ -567,7 +536,10 @@ public: out << SP << "auto const workDiv_" << fNY << " = sofie_workdiv(elementsPerGrid_" << fNY << ");\n"; out << SP << "auto task_" << OpName << " = alpaka::createTaskKernel(workDiv_" << fNY << ", binary" << OpName << "Kernel, alpaka::getPtrNative(deviceBuf_" << fNA - << "), alpaka::getPtrNative(deviceBuf_" << fNB << "), alpaka::getPtrNative(deviceBuf_" << fNY << "));\n"; + << "), alpaka::getPtrNative(deviceBuf_" << fNB << "), alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << length << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << OpName << ");\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_BasicNary.hxx b/core/inc/SOFIE/ROperator_BasicNary.hxx index 928ab1c2..1947f53b 100644 --- a/core/inc/SOFIE/ROperator_BasicNary.hxx +++ b/core/inc/SOFIE/ROperator_BasicNary.hxx @@ -145,6 +145,11 @@ public: fNBroadcastedInputs.emplace_back("tensor_" + fNInputs[i]); } } + // rank-pad each input shape to fShapeY; the GPU emitters read them directly + for (auto &s : fShapeInputs) { + auto y = fShapeY; + UTILITY::MultidirectionalBroadcastShape(s, y); + } fType = ConvertTypeToString(model.GetTensorType(fNInputs[0])); } @@ -185,6 +190,109 @@ public: return out.str(); } + std::string GetGPUCombine(const std::string& acc_v, const std::string& val) const { + if (Op == EBasicNaryOperator::Max) + return acc_v + " = (" + acc_v + " > " + val + ") ? " + acc_v + " : " + val + ";"; + if (Op == EBasicNaryOperator::Min) + return acc_v + " = (" + acc_v + " < " + val + ") ? " + acc_v + " : " + val + ";"; + return acc_v + " = " + acc_v + " + " + val + ";"; // Sum and Mean both accumulate + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { + OpName = "op_" + OpName; + size_t nIn = fNInputs.size(); + const std::size_t D = fShapeY.size(); + auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); + auto isOne = [](const Dim &d) { return d.GetVal() == "1"; }; + std::vector isScalar(nIn), isContiguous(nIn), isPartial(nIn); + std::vector> strides(nIn); + bool anyPartial = false; + for (size_t i = 0; i < nIn; i++) { + isScalar[i] = ConvertDimShapeToLength(fShapeInputs[i]) == "1"; + isContiguous[i] = UTILITY::AreSameShape(fShapeInputs[i], fShapeY); + isPartial[i] = !isScalar[i] && !isContiguous[i]; + if (isPartial[i]) { + strides[i] = UTILITY::ComputeStrideFromShape(fShapeInputs[i]); + anyPartial = true; + } + } + std::string op; + op += "\n//------ BASICNARY_KERNEL_ALPAKA\n"; + op += SP + "struct BasicNaryKernel_" + OpName + " {\n"; + op += SP + SP + "template(acc)[0];\n"; + op += SP + SP + SP + "if (idx >= n) return;\n"; + + // decompose idx into output coords, skipping broadcast dims (stride 0) + if (anyPartial) { + op += SP + SP + SP + "std::size_t remaining = idx;\n"; + op += SP + SP + SP + "std::size_t coord;\n"; + for (size_t i = 0; i < nIn; i++) + if (isPartial[i]) + op += SP + SP + SP + "std::size_t idx" + std::to_string(i) + " = 0;\n"; + for (std::size_t d = 0; d < D; d++) { + std::string sY = "(" + stridesY[d].GetVal() + ")"; + op += SP + SP + SP + "coord = remaining / " + sY + ";\n"; + if (d + 1 < D) + op += SP + SP + SP + "remaining -= coord * " + sY + ";\n"; + for (size_t i = 0; i < nIn; i++) { + if (!isPartial[i] || isOne(fShapeInputs[i][d])) continue; + op += SP + SP + SP + "idx" + std::to_string(i) + " += coord * (" + strides[i][d].GetVal() + ");\n"; + } + } + } + auto index = [&](size_t i) -> std::string { + if (isContiguous[i]) return "idx"; + if (isScalar[i]) return "0"; + return "idx" + std::to_string(i); + }; + op += SP + SP + SP + "TOut v = static_cast(in0[" + index(0) + "]);\n"; + for (size_t i = 1; i < nIn; i++) + op += SP + SP + SP + "{ TOut w = static_cast(in" + std::to_string(i) + "[" + index(i) + "]); " + GetGPUCombine("v", "w") + " }\n"; + if (Op == EBasicNaryOperator::Mean) + op += SP + SP + SP + "v = v / static_cast(" + std::to_string(nIn) + ");\n"; + op += SP + SP + SP + "out[idx] = v;\n"; + op += SP + SP + "}\n"; + op += SP + "};\n"; + return op; + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string OpName) override { + OpName = "op_" + OpName; + return SP + "BasicNaryKernel_" + OpName + " basicNaryKernel_" + OpName + ";\n"; + } + + std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { + if (fShapeY.empty()) + throw std::runtime_error("SOFIE BasicNary Op called to Generate without being initialized first"); + OpName = "op_" + OpName; + std::stringstream out; + std::string length = ConvertDimShapeToLength(fShapeY); + out << "\n//------ BASICNARY_GPU_ALPAKA\n"; + out << SP << "auto const elementsPerGrid_" << OpName << " = Vec::all(Idx{" << length << "});\n"; + out << SP << "auto const workDiv_" << OpName << " = sofie_workdiv(elementsPerGrid_" << OpName << ");\n"; + out << SP << "auto task_" << OpName << " = alpaka::createTaskKernel(workDiv_" << OpName + << ", basicNaryKernel_" << OpName; + for (auto &in : fNInputs) + out << ", alpaka::getPtrNative(deviceBuf_" << in << ")"; + out << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << length << "));\n"; + out << SP << "alpaka::enqueue(queue, task_" << OpName << ");\n"; + return out.str(); + } + std::vector GetStdLibs() override {return { std::string("cmath") }; } }; diff --git a/core/inc/SOFIE/ROperator_BatchNormalization.hxx b/core/inc/SOFIE/ROperator_BatchNormalization.hxx index 8bc3b3d5..05b0717b 100644 --- a/core/inc/SOFIE/ROperator_BatchNormalization.hxx +++ b/core/inc/SOFIE/ROperator_BatchNormalization.hxx @@ -28,14 +28,11 @@ private: std::string fNMean; std::string fNVar; std::string fNY; + std::string fNFusedScale; // scale/sqrt(var+eps) fused over channels, shape [C] EActivationType fActivation; - std::vector fShapeX; - std::vector fShapeScale; - std::vector fShapeB; - std::vector fShapeMean; - std::vector fShapeVar; - std::vector fShapeY; + std::vector fShapeX; + std::vector fShapeY; std::string fType; @@ -54,6 +51,8 @@ public: fInputTensorNames = { fNX }; fOutputTensorNames = { fNY }; + fNFusedScale = fNScale + "_fused_inv_std_dev"; + if(std::is_same::value){ fType = "float"; } @@ -69,171 +68,100 @@ public: return {out}; } - std::vector> ShapeInference(std::vector> input) override { - if (input.size() != 5 ) { - throw - std::runtime_error("SOFIE BatchNormalization Op Shape inference need 5 input tensors"); - } - for(size_t i = 0; i < input.size(); i++) { - if (input[i].size() != 4) { - throw - std::runtime_error("SOFIE BatchNormalization Op Shape inference only accept tensor with 4 dimensions"); - } - } - - auto ret = input; - return ret; - } - void Initialize(RModel& model) override { if (!model.CheckIfTensorAlreadyExist(fNX)) { - throw - std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNX + " fnx is not found in model"); + throw std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNX + " is not found in model"); } if (!model.CheckIfTensorAlreadyExist(fNScale)) { - throw - std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNScale + " fns is not found in model"); + throw std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNScale + " is not found in model"); } - if (!model.CheckIfTensorAlreadyExist(fNB)) { - throw - std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNB + " fnb is not found in model"); + if (!model.CheckIfTensorAlreadyExist(fNB)) { + throw std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNB + " is not found in model"); } if (!model.CheckIfTensorAlreadyExist(fNMean)) { - throw - std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNMean + " fnm is not found in model"); + throw std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNMean + " is not found in model"); } if (!model.CheckIfTensorAlreadyExist(fNVar)) { - throw - std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNVar + " fnv is not found in model"); + throw std::runtime_error("SOFIE BatchNormalization op Input Tensor " + fNVar + " is not found in model"); } - fShapeX = model.GetTensorShape(fNX); - - if (fShapeX.size() < 2 || fShapeX.size() > 4) { - throw - std::runtime_error("SOFIE BatchNormalization Op input tensor " + fNX + " fnx has wrong shape : " + ConvertShapeToString(fShapeX)); + fShapeX = model.GetDimTensorShape(fNX); + if (fShapeX.size() < 2 || fShapeX.size() > 4) { + throw std::runtime_error("SOFIE BatchNormalization Op input tensor " + fNX + + " has wrong shape : " + ConvertDimShapeToString(fShapeX)); } - fShapeScale = model.GetTensorShape(fNScale); - fShapeB = model.GetTensorShape(fNB); - fShapeMean = model.GetTensorShape(fNMean); - fShapeVar = model.GetTensorShape(fNVar); fShapeY = fShapeX; model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); - if (fShapeB.size() == 1) { - // Broadcast scale, bias, input_mean and input_var to shape_X - auto original_B = model.GetInitializedTensorData(fNB); - auto original_S = model.GetInitializedTensorData(fNScale); - auto original_M = model.GetInitializedTensorData(fNMean); - auto original_V = model.GetInitializedTensorData(fNVar); - size_t batchSize = fShapeX[0]; - size_t channels = fShapeX[1]; - size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1; - size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1; - size_t n = batchSize * channels * height * width; - if (fType == "float") { - float *original_bias = static_cast(original_B.get()); - float *original_scale = static_cast(original_S.get()); - float *original_mean = static_cast(original_M.get()); - float *original_var = static_cast(original_V.get()); - float *new_bias = new float[n]; - float *new_scale = new float[n]; - float *new_mean = new float[n]; - float *new_var = new float[n]; - size_t bs = 0, ch = 0, h = 0, w = 0; - for (ch = 0; ch < channels; ch++) { - for (h = 0; h < height; h++) { - for (w = 0; w < width; w++) { - new_bias[bs * channels * height * width + ch * height * width + h * width + w] = original_bias[ch]; - new_scale[bs * channels * height * width + ch * height * width + h * width + w] = - original_scale[ch]; - new_mean[bs * channels * height * width + ch * height * width + h * width + w] = original_mean[ch]; - new_var[bs * channels * height * width + ch * height * width + h * width + w] = original_var[ch]; - } - } - } - size_t Batchoffset = channels * height * width; - for (bs = 1; bs < batchSize; bs++) { - std::copy(new_bias, new_bias + Batchoffset, new_bias + (bs * Batchoffset)); - std::copy(new_scale, new_scale + Batchoffset, new_scale + (bs * Batchoffset)); - std::copy(new_mean, new_mean + Batchoffset, new_mean + (bs * Batchoffset)); - std::copy(new_var, new_var + Batchoffset, new_var + (bs * Batchoffset)); - } - //// new_var =1. / sqrt(input_var + fepsilon) - for (size_t i = 0; i < n; i++) { - new_var[i] = 1. / sqrt(new_var[i] + fepsilon); - new_scale[i] *= new_var[i]; // include var in new scale - } - std::vector new_bias_shape = {batchSize, channels, height, width}; - std::shared_ptr new_bias_ptr(new_bias, std::default_delete()); - std::shared_ptr new_scale_ptr(new_scale, std::default_delete()); - std::shared_ptr new_mean_ptr(new_mean, std::default_delete()); - std::shared_ptr new_var_ptr(new_var, std::default_delete()); - model.UpdateInitializedTensor(fNB, model.GetTensorType(fNB), new_bias_shape, new_bias_ptr); - model.UpdateInitializedTensor(fNScale, model.GetTensorType(fNScale), new_bias_shape, new_scale_ptr); - model.UpdateInitializedTensor(fNMean, model.GetTensorType(fNMean), new_bias_shape, new_mean_ptr); - model.UpdateInitializedTensor(fNVar, model.GetTensorType(fNVar), new_bias_shape, new_var_ptr); - fShapeB = model.GetTensorShape(fNB); - fShapeScale = model.GetTensorShape(fNScale); - fShapeMean = model.GetTensorShape(fNMean); - fShapeVar = model.GetTensorShape(fNVar); + auto original_S = model.GetInitializedTensorData(fNScale); + auto original_V = model.GetInitializedTensorData(fNVar); + auto shape_S = model.GetTensorShape(fNScale); + if (shape_S.size() != 1) { + throw std::runtime_error("SOFIE BatchNormalization 'scale' tensor must be 1D (per-channel)."); + } + size_t channels = shape_S[0]; + + if (fType == "float") { + float *original_scale_ptr = static_cast(original_S.get()); + float *original_var_ptr = static_cast(original_V.get()); + float *fused_scale_data = new float[channels]; + for (size_t i = 0; i < channels; i++) { + fused_scale_data[i] = original_scale_ptr[i] / std::sqrt(original_var_ptr[i] + fepsilon); } + std::shared_ptr fused_scale_ptr(fused_scale_data, std::default_delete()); + model.AddInitializedTensor(fNFusedScale, model.GetTensorType(fNScale), {channels}, fused_scale_ptr); } } - std::string Generate(std::string OpName) override { - OpName = "op_" + OpName; + std::string Generate(std::string opName) override { + opName = "op_" + opName; if (fShapeX.empty()){ throw std::runtime_error("SOFIE Batch Normalization called to Generate without being initialized first"); } std::stringstream out; - //// Batch Norm op - size_t batchSize = fShapeX[0]; - size_t channels = fShapeX[1]; - size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1; - size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1; - size_t n = batchSize * channels * height * width; - - //// copy X into Y - out << "\n\n//---- BatchNorm\n"; - out << SP << "constexpr int " << OpName << "_N =" << batchSize * channels * height * width << ";\n"; - out << SP << "constexpr int "< 2) { + auto spatialShape = fShapeX; + spatialShape.erase(spatialShape.begin(), spatialShape.begin() + 2); + spatial_dim = ConvertDimShapeToLength(spatialShape); + } - if(fActivation == EActivationType::RELU){ - out << SP << "for (int id = 0; id < " << ConvertShapeToLength(fShapeY) << " ; id++){\n"; - out << SP << SP << "tensor_" << fNY << "[id] = ((tensor_" << fNY << "[id] > 0 )? tensor_" << fNY << "[id] : 0);\n"; - out << SP << "}\n"; + out << "\n\n//---- BatchNorm" << (fActivation == EActivationType::RELU ? " + ReLU " : " ") << opName << "\n"; + out << SP << "{\n"; + out << SP << " size_t i = 0;\n"; + out << SP << " for (size_t n = 0; n < " << batchSize << "; ++n) {\n"; + out << SP << " for (size_t c = 0; c < " << channels << "; ++c) {\n"; + out << SP << " const float mean_val = tensor_" << fNMean << "[c];\n"; + out << SP << " const float fused_scale_val = tensor_" << fNFusedScale << "[c];\n"; + out << SP << " const float bias_val = tensor_" << fNB << "[c];\n"; + out << SP << " for (size_t sp = 0; sp < " << spatial_dim << "; ++sp) {\n"; + out << SP << " float val = (tensor_" << fNX << "[i] - mean_val) * fused_scale_val + bias_val;\n"; + if (fActivation == EActivationType::RELU) { + out << SP << " tensor_" << fNY << "[i] = (val > 0.0f) ? val : 0.0f;\n"; + } else { + out << SP << " tensor_" << fNY << "[i] = val;\n"; } + out << SP << " i++;\n"; + out << SP << " }\n"; + out << SP << " }\n"; + out << SP << " }\n"; + out << SP << "}\n"; + return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeX.empty()) throw std::runtime_error("SOFIE BatchNormalization called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeY); - + std::string channels = fShapeX[1].GetVal(); + std::string spatial_dim = ConvertDimShapeToLength(std::vector(fShapeX.begin() + 2, fShapeX.end())); //"1" for rank-2 input std::string kname = "BatchNormKernel_" + opName; std::string op; op = "\n//------ BATCHNORM_KERNEL_ALPAKA\n"; @@ -242,10 +170,12 @@ public: op += SP + SP + "ALPAKA_FN_ACC void operator()(\n"; op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ X,\n"; - op += SP + SP + SP + "T const* __restrict__ scale,\n"; + op += SP + SP + SP + "T const* __restrict__ fused_scale,\n"; op += SP + SP + SP + "T const* __restrict__ bias,\n"; op += SP + SP + SP + "T const* __restrict__ mean,\n"; op += SP + SP + SP + "T* __restrict__ Y,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -253,14 +183,12 @@ public: op += SP + SP + SP + "auto const grid_thread_extent = alpaka::getWorkDiv(acc)[0];\n\n"; op += SP + SP + SP + "for (std::size_t i = global_thread_idx; i < totalElements; i += grid_thread_extent) {\n"; - - op += SP + SP + SP + SP + "T val = (X[i] - mean[i]) * scale[i] + bias[i];\n"; - + op += SP + SP + SP + SP + "std::size_t const c = (i / (" + spatial_dim + ")) % (" + channels + ");\n"; + op += SP + SP + SP + SP + "T val = (X[i] - mean[c]) * fused_scale[c] + bias[c];\n"; if (fActivation == EActivationType::RELU) op += SP + SP + SP + SP + "Y[i] = val > static_cast(0) ? val : static_cast(0);\n"; else op += SP + SP + SP + SP + "Y[i] = val;\n"; - op += SP + SP + SP + "}\n"; op += SP + SP + "}\n"; op += SP + "};\n"; @@ -274,34 +202,38 @@ public: return SP + kname + " batchNormKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeX.empty()) throw std::runtime_error("SOFIE BatchNormalization called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeY); + std::string totalElements = ConvertDimShapeToLength(fShapeY); std::string kname = "batchNormKernel_" + opName; + std::string dynArgs; + for (auto &p : dynParamNames) dynArgs += ", static_cast(" + p + ")"; + std::stringstream out; out << "\n//------ BATCHNORM_GPU_ALPAKA\n"; out << SP << "auto const elementsPerThread_" << fNY << " = Vec::all(static_cast(1));\n"; out << SP << "auto const elementsPerGrid_" << fNY << " = Vec::all(Idx{" << totalElements << "});\n"; out << SP << "auto const workDiv_" << fNY << " = sofie_workdiv(elementsPerGrid_" << fNY << ");\n"; - + out << SP << "auto task_" << fNY << " = alpaka::createTaskKernel(workDiv_" << fNY << ", " << kname - << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNScale << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNB << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNMean << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNFusedScale << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNB << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNMean << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << dynArgs << ", static_cast(" << totalElements << "));\n"; out << SP <<"alpaka::enqueue(queue, task_" << fNY << ");\n"; - + return out.str(); } - std::vector GetBlasRoutines() override { return { std::string("Copy"), std::string("Axpy") }; } + std::vector GetBlasRoutines() override { return {}; } }; }//SOFIE diff --git a/core/inc/SOFIE/ROperator_Comparision.hxx b/core/inc/SOFIE/ROperator_Comparision.hxx index 1e02d534..fbe6e450 100644 --- a/core/inc/SOFIE/ROperator_Comparision.hxx +++ b/core/inc/SOFIE/ROperator_Comparision.hxx @@ -18,35 +18,35 @@ struct ComparisionTrait{}; template struct ComparisionTrait { static const std::string Name() { return "Equal"; } - static std::string Op(const std::string & t1, const std::string t2) { return t1 + " == " + t2 + " ? true : false "; } + static std::string Op(const std::string & t1, const std::string t2) { return t1 + " == " + t2; } static bool Result(T v1, T v2) { return v1 == v2;} }; template struct ComparisionTrait { static const std::string Name() { return "Less"; } - static std::string Op(const std::string & t1, const std::string t2) { return t1 + " < " + t2 + " ? true : false "; } + static std::string Op(const std::string & t1, const std::string t2) { return t1 + " < " + t2; } static bool Result(T v1, T v2) { return v1 < v2;} }; template struct ComparisionTrait { static const std::string Name() { return "LessOrEqual"; } - static std::string Op(const std::string & t1, const std::string t2) { return t1 + " <= " + t2 + " ? true : false "; } + static std::string Op(const std::string & t1, const std::string t2) { return t1 + " <= " + t2; } static bool Result(T v1, T v2) { return v1 <= v2;} }; template struct ComparisionTrait { static const std::string Name() { return "Greater"; } - static std::string Op(const std::string & t1, const std::string t2) { return t1 + " > " + t2 + " ? true : false "; } + static std::string Op(const std::string & t1, const std::string t2) { return t1 + " > " + t2; } static bool Result(T v1, T v2) { return v1 > v2;} }; template struct ComparisionTrait { static const std::string Name() { return "GreaterOrEqual"; } - static std::string Op(const std::string & t1, const std::string t2) { return t1 + " >= " + t2 + " ? true : false " ; } + static std::string Op(const std::string & t1, const std::string t2) { return t1 + " >= " + t2 ; } static bool Result(T v1, T v2) { return v1 >= v2;} }; @@ -54,18 +54,18 @@ template class ROperator_Comparision final : public ROperator{ private: - bool fIsModelOutput = false; std::string fNX1; std::string fNX2; std::string fNY; std::vector fShapeX1; std::vector fShapeX2; + std::vector fDimShapeX1; + std::vector fDimShapeX2; std::vector fShapeY; - std::string fNBroadcastedX1; - std::string fNBroadcastedX2; + std::vector fDimShapeY; ETensorType fTensorType1 = ETensorType::UNDEFINED; ETensorType fTensorType2 = ETensorType::UNDEFINED; - bool fBroadcast = false; + int fBroadcastFlag = 0; public: @@ -77,145 +77,314 @@ public: fOutputTensorNames = { fNY }; } + // type of output given input std::vector TypeInference(std::vector input) override { return input; } + // shape of output tensors given input tensors std::vector> ShapeInference(std::vector> input) override { - auto ret = input; + auto ret = input; // return vector size 1 with first input return ret; } void Initialize(RModel& model) override { + // input must be a graph input, or already initialized intermediate tensor if (!model.CheckIfTensorAlreadyExist(fNX1)){ throw std::runtime_error(std::string("SOFIE Comparision Op Input Tensor ") + fNX1 + "is not found in model"); } if (!model.CheckIfTensorAlreadyExist(fNX2)) { throw std::runtime_error(std::string("SOFIE Comparision Op Input Tensor ") + fNX2 + "is not found in model"); } - fShapeX1 = model.GetTensorShape(fNX1); - fShapeX2 = model.GetTensorShape(fNX2); + if (model.IsDynamicTensor(fNX1)) + fDimShapeX1 = model.GetDynamicTensorShape(fNX1); + else { + fShapeX1 = model.GetTensorShape(fNX1); + fDimShapeX1 = ConvertShapeToDim(fShapeX1); + } + if (model.IsDynamicTensor(fNX2)) + fDimShapeX2 = model.GetDynamicTensorShape(fNX2); + else { + fShapeX2 = model.GetTensorShape(fNX2); + fDimShapeX2 = ConvertShapeToDim(fShapeX2); + } fTensorType1 = model.GetTensorType(fNX1); fTensorType2 = model.GetTensorType(fNX2); - bool broadcast = !UTILITY::AreSameShape(fShapeX1, fShapeX2); - if (broadcast) { - // ONNX comparison ops support multidirectional broadcasting (numpy semantics): - // both inputs can be broadcast to the common output shape. - auto ret = UTILITY::MultidirectionalBroadcastShape(fShapeX1, fShapeX2); - fShapeY = ret.second; - bool broadcastX1 = !UTILITY::AreSameShape(fShapeX1, fShapeY); - bool broadcastX2 = !UTILITY::AreSameShape(fShapeX2, fShapeY); - if (broadcastX1) { - if (model.IsInitializedTensor(fNX1)) { - auto data = model.GetInitializedTensorData(fNX1); - std::shared_ptr broadcastedData( - UTILITY::UnidirectionalBroadcast(static_cast(data.get()), fShapeX1, fShapeY), - std::default_delete()); - model.UpdateInitializedTensor(fNX1, model.GetTensorType(fNX1), fShapeY, broadcastedData); - fShapeX1 = fShapeY; - } else { - fNBroadcastedX1 = "Broadcasted" + fNX1; - model.AddIntermediateTensor(fNBroadcastedX1, model.GetTensorType(fNX1), fShapeY); + // case of non dynamic tensors + if (!fShapeX1.empty() && !fShapeX2.empty()) { + bool broadcastX1 = false; + bool broadcastX2 = false; + if (UTILITY::AreSameShape(fShapeX1, fShapeX2)) { + // no broadcast needed + fShapeY = fShapeX1; + } else { + // Y is the common shape of A and B + fShapeY = UTILITY::MultidirectionalBroadcastShape(fShapeX1, fShapeX2).second; + broadcastX1 = !UTILITY::AreSameShape(fShapeX1, fShapeY); + broadcastX2 = !UTILITY::AreSameShape(fShapeX2, fShapeY); + } + // keep the Dim copies in sync with the (possibly rank-padded) static shapes, + // the GPU emitters read them directly + fDimShapeX1 = ConvertShapeToDim(fShapeX1); + fDimShapeX2 = ConvertShapeToDim(fShapeX2); + + // analyze case of constant tensors or shape tensors (which have known shapes but data as Dim values + // normal case with non-dynamic tensor is also here + T *data1 = nullptr; + T *data2 = nullptr; + std::unique_ptr broadcastedData1; + std::unique_ptr broadcastedData2; + // data for shape tensors + std::vector shapeData1; + std::vector shapeData2; + size_t length = ConvertShapeToLength(fShapeY); + bool *outData = new bool[length]; + if (model.IsInitializedTensor(fNX1)) { + data1 = static_cast(model.GetInitializedTensorData(fNX1).get()); + if (broadcastX1) { + broadcastedData1 = std::unique_ptr( + UTILITY::UnidirectionalBroadcast(data1, fShapeX1, fShapeY)); + data1 = broadcastedData1.get(); } + + } else if (model.IsShapeTensor(fNX1)) { + shapeData1 = model.GetShapeTensorValues(fNX1); } - if (broadcastX2) { - if (model.IsInitializedTensor(fNX2)) { - auto data = model.GetInitializedTensorData(fNX2); - std::shared_ptr broadcastedData( - UTILITY::UnidirectionalBroadcast(static_cast(data.get()), fShapeX2, fShapeY), - std::default_delete()); - model.UpdateInitializedTensor(fNX2, model.GetTensorType(fNX2), fShapeY, broadcastedData); - fShapeX2 = fShapeY; - } else { - fNBroadcastedX2 = "Broadcasted" + fNX2; - model.AddIntermediateTensor(fNBroadcastedX2, model.GetTensorType(fNX2), fShapeY); + if (model.IsInitializedTensor(fNX2)) { + data2 = static_cast(model.GetInitializedTensorData(fNX2).get()); + if (broadcastX2) { + broadcastedData2 = std::unique_ptr( + UTILITY::UnidirectionalBroadcast(data2, fShapeX2, fShapeY)); + data2 = broadcastedData2.get(); } + } else if (model.IsShapeTensor(fNX2)) { + shapeData2 = model.GetShapeTensorValues(fNX2); + } + if (data1 && data2) { + fIsOutputConstant = true; + for (size_t i = 0; i < length; i++) + outData[i] = ComparisionTrait::Result(data1[i], data2[i]); + model.AddConstantTensor(fNY, fShapeY, outData); + if (model.Verbose()) + std::cout << ComparisionTrait::Name() << " op ---> " << fNY << " " + << ConvertShapeToString(fShapeY) << " : " << ConvertValuesToString(length, outData) + << std::endl; + } else if ((data1 || !shapeData1.empty()) && (data2 || !shapeData2.empty())) { + fIsOutputConstant = true; + if (data1 && !data2) { + // data 1 is constant and data2 is shape + for (size_t i = 0; i < length; i++) { + if (shapeData2[i].isParam) { + if (shapeData2[i].dim == size_t(-1) || data1[i] > 0) { + fIsOutputConstant = false; + break; + } else { + // assume a comparison is done with .dim = 0 + shapeData2[i].dim = 0; + } + } + outData[i] = ComparisionTrait::Result(data1[i], static_cast(shapeData2[i].dim)); + } + } else if (!data1 && data2) { + // data 1 is shape and dat2 is constant + for (size_t i = 0; i < length; i++) { + if (shapeData1[i].isParam) { + if (shapeData1[i].dim == size_t(-1) || data2[i] > 0) { + fIsOutputConstant = false; + break; + } else { + // assume a comparison is done with .dim = 0 + shapeData1[i].dim = 0; + } + } + outData[i] = ComparisionTrait::Result(static_cast(shapeData1[i].dim), data2[i]); + } + } else if (!shapeData1.empty() && !shapeData2.empty()) { + // both data1 and data2 are shape tensors + for (size_t i = 0; i < length; i++) { + if (!shapeData1[i].isParam && !shapeData2[i].isParam) { + outData[i] = ComparisionTrait::Result(shapeData1[i].dim, shapeData2[i].dim); + } else if (shapeData1[i].isParam && shapeData2[i].isParam) { + if (shapeData1[i].param == shapeData2[i].param) + outData[i] = ComparisionTrait::Result(1, 1); // comparison of two equal value + else { + fIsOutputConstant = false; + break; + } + } else { + fIsOutputConstant = false; + break; + } + } + } + if (fIsOutputConstant) { + model.AddConstantTensor(fNY, fShapeY, outData); + if (model.Verbose()) + std::cout << ComparisionTrait::Name() << " op ---> " << fNY << " " + << ConvertShapeToString(fShapeY) << " : " << ConvertValuesToString(length, outData) + << " (constant) " << std::endl; + } + } + delete[] outData; + // case of non constant output (no constant or shape tensors) + if (!fIsOutputConstant && !fShapeY.empty()) { + model.AddIntermediateTensor(fNY, ETensorType::BOOL, fShapeY); + fDimShapeY = ConvertShapeToDim(fShapeY); + if (model.Verbose()) + std::cout << ComparisionTrait::Name() << " op ---> " << fNY << " " + << ConvertShapeToString(fShapeY) << std::endl; } } else { - fShapeY = fShapeX1; - } - if (model.IsInitializedTensor(fNX1) && model.IsInitializedTensor(fNX2)) { - fIsOutputConstant = true; - auto data1 = static_cast(model.GetInitializedTensorData(fNX1).get()); - auto data2 = static_cast(model.GetInitializedTensorData(fNX2).get()); - size_t length = ConvertShapeToLength(fShapeY); - bool * outData = new bool[length]; - for (size_t i = 0; i < length; i++) - outData[i] = ComparisionTrait::Result(data1[i], data2[i]); - model.AddConstantTensor(fNY, fShapeY, outData); - if (model.Verbose()) - std::cout << ComparisionTrait::Name() << " op ---> " << fNY << " " << ConvertShapeToString(fShapeY) << " : " - << ConvertValuesToString(length,outData) << std::endl; - delete [] outData; - } else { - model.AddIntermediateTensor(fNY, ETensorType::BOOL, fShapeY); + // case of dynamic tensors + // case A or B have dynamic shapes. We need to broadcast if shape are not same + auto ret = UTILITY::MultidirectionalBroadcastShape(fDimShapeX1, fDimShapeX2); + fBroadcastFlag = ret.first; + fDimShapeY = ret.second; + // case of all parametric shapes and MultiDirectionalBroadcastShape return the max of the 2 + // need to do before we declare the output tensor shape and the broadcasted ones + if (ret.first & 4) { + // check if one of the parameter is an input dimension + // define function to find this + auto IsInputDimParam = [&](const std::string &p) { + auto inputNames = model.GetInputTensorNames(); + for (auto &input : inputNames) { + for (auto &i_s : model.GetDimTensorShape(input)) { + if (i_s.isParam && i_s.param == p) + return true; + } + } + return false; + }; + for (size_t i = 0; i < fDimShapeY.size(); i++) { + auto &s = fDimShapeY[i]; + if (s.isParam && s.param.find("std::max") != std::string::npos) { + if (IsInputDimParam(fDimShapeX1[i].param)) { + // case dim is 1 we indicate that the input parameter is equal to 1 + if (fDimShapeX1[i].dim != 1) + s = fDimShapeX1[i]; + else + s = fDimShapeX2[i]; + } else if (IsInputDimParam(fDimShapeX2[i].param)) { + if (fDimShapeX2[i].dim != 1) + s = fDimShapeX2[i]; + else + s = fDimShapeX1[i]; + } + } + } + } + + model.AddIntermediateTensor(fNY, ETensorType::BOOL, fDimShapeY); + if (model.Verbose()) { + std::cout << ComparisionTrait::Name() << " : " << fNX1 << " " << ConvertDimShapeToString(fDimShapeX1) << " , " + << fNX2 << " " << ConvertDimShapeToString(fDimShapeX2) << " --> " + << fNY << " " << ConvertDimShapeToString(fDimShapeY) << std::endl; + model.PrintIntermediateTensors(); + } } - const auto & outputTensorNames = model.GetOutputTensorNames(); - fIsModelOutput = false; - if (std::find(outputTensorNames.begin(), outputTensorNames.end(), fNY) != outputTensorNames.end()) - fIsModelOutput = true; } - std::string Generate(std::string OpName) override { + std::string Generate(std::string opName) override { if (fIsOutputConstant) return ""; - OpName = "op_" + OpName; - if (fShapeY.empty()) { + opName = "op_" + opName; + + if (fDimShapeY.empty()) { throw std::runtime_error("SOFIE Comparision Op called to Generate without being initialized first"); } std::stringstream out; - out << SP << "\n//------ " << ComparisionTrait::Name() << "\n"; - size_t length = ConvertShapeToLength(fShapeY); - if (!fNBroadcastedX1.empty()) { - std::string type1 = ConvertTypeToString(fTensorType1); - out << SP << "// Broadcasting uninitialized tensor " << fNX1 << "\n"; - out << SP << "{\n"; - out << SP << SP << type1 << "* data = SOFIE::UTILITY::UnidirectionalBroadcast<" << type1 << ">(tensor_" << fNX1 << ", " << ConvertShapeToString(fShapeX1) << ", " << ConvertShapeToString(fShapeY) << ");\n"; - out << SP << SP << "std::copy(data, data + " << length << ", tensor_" << fNBroadcastedX1 << ");\n"; - out << SP << SP << "delete[] data;\n"; - out << SP << "}\n"; + out << SP << "\n//------ " << ComparisionTrait::Name() << " " << opName + << " --> " << ConvertShapeToString(fShapeY) << "\n"; + + // need to add check if tensors are compatible as in binary operator + + // use same code as Binary operator + auto stridesA = UTILITY::ComputeStrideFromShape(fDimShapeX1); + auto stridesB = UTILITY::ComputeStrideFromShape(fDimShapeX2); + auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY); + + std::string compute_idx_X1, compute_idx_X2, compute_idx_Y; + if (fDimShapeX1.empty() || + std::all_of(fDimShapeX1.begin(), fDimShapeX1.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) { + compute_idx_X1 = "0"; + } else { + for (size_t i = 0; i < fDimShapeX1.size(); ++i) { + if (fDimShapeX1[i].dim == 1 || fDimShapeX1[i].GetVal() == "1") + continue; + compute_idx_X1 += "idx_" + std::to_string(i + (fDimShapeY.size() - fDimShapeX1.size())); + if (stridesA[i].GetVal() != "1") + compute_idx_X1 += " * " + stridesA[i].GetVal(); + compute_idx_X1 += " + "; + } + // remove last 3 character " + " + for (int j = 0; j < 3; j++) + compute_idx_X1.pop_back(); + } + if (fDimShapeX2.empty() || + std::all_of(fDimShapeX2.begin(), fDimShapeX2.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) { + compute_idx_X2 = "0"; + } else { + for (size_t i = 0; i < fDimShapeX2.size(); ++i) { + if (fDimShapeX2[i].dim == 1 || fDimShapeX2[i].GetVal() == "1") + continue; + compute_idx_X2 += "idx_" + std::to_string(i + (fDimShapeY.size() - fDimShapeX2.size())); + if (stridesB[i].GetVal() != "1") + compute_idx_X2 += " * " + stridesB[i].GetVal(); + compute_idx_X2 += " + "; + } + // remove last 3 character " + " + for (int j = 0; j < 3; j++) + compute_idx_X2.pop_back(); } - if (!fNBroadcastedX2.empty()) { - std::string type2 = ConvertTypeToString(fTensorType2); - out << SP << "// Broadcasting uninitialized tensor " << fNX2 << "\n"; - out << SP << "{\n"; - out << SP << SP << type2 << "* data = SOFIE::UTILITY::UnidirectionalBroadcast<" << type2 << ">(tensor_" << fNX2 << ", " << ConvertShapeToString(fShapeX2) << ", " << ConvertShapeToString(fShapeY) << ");\n"; - out << SP << SP << "std::copy(data, data + " << length << ", tensor_" << fNBroadcastedX2 << ");\n"; - out << SP << SP << "delete[] data;\n"; - out << SP << "}\n"; + int nloop = 0; + if (fDimShapeY.empty() || + std::all_of(fDimShapeY.begin(), fDimShapeY.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) { + compute_idx_Y = "0"; + } else { + for (size_t i = 0; i < fDimShapeY.size(); ++i) { + if (fDimShapeY[i].dim != 1 && fDimShapeY[i].GetVal() != "1") { + nloop++; + for (int j = 0; j < nloop; j++) out << SP; + out << "for (size_t idx_" << i << " = 0; idx_" << i << " < " << fDimShapeY[i] + << "; ++idx_" << i << "){\n"; + compute_idx_Y += "idx_" + std::to_string(i); + if (stridesY[i].GetVal() != "1") + compute_idx_Y += " * " + stridesY[i].GetVal(); + compute_idx_Y += " + "; + } + } + // remove last 3 characters " + " + for (int j = 0; j < 3; j++) + compute_idx_Y.pop_back(); } - const std::string& nameX1 = fNBroadcastedX1.empty()? fNX1 : fNBroadcastedX1; - const std::string& nameX2 = fNBroadcastedX2.empty()? fNX2 : fNBroadcastedX2; - out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n"; - out << SP << SP << "fTensor_" << fNY << "[id] = " << ComparisionTrait::Op( "tensor_" + nameX1 + "[id]" , "tensor_" + nameX2 + "[id]") << " ;\n"; - out << SP << "}\n"; - if (!fIsModelOutput) - out << SP << "const std::vector & tensor_" << fNY << " = fTensor_" << fNY << ";\n"; + for (int j = 0; j < nloop + 1; j++) out << SP; + out << "tensor_" << fNY << "[" << compute_idx_Y << "] = " + << ComparisionTrait::Op( "tensor_" + fNX1 + "[" + compute_idx_X1 + "]" , + "tensor_" + fNX2 + "[" + compute_idx_X2 + "]") << " ;\n"; + + + for (int i = nloop; i > 0; i--) { + for (int j = 0; j < i; j++) out << SP; + out << "}\n"; + } + + return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; - if (fShapeY.empty()) + if (fDimShapeY.empty()) throw std::runtime_error("SOFIE Comparision Op called to Generate without being initialized first"); - const std::size_t D = fShapeY.size(); - std::size_t totalElements = ConvertShapeToLength(fShapeY); - - std::vector shapeX1_padded(D, 1); - std::vector shapeX2_padded(D, 1); - { - size_t off1 = D - fShapeX1.size(); - for (size_t i = 0; i < fShapeX1.size(); ++i) - shapeX1_padded[off1 + i] = fShapeX1[i]; - size_t off2 = D - fShapeX2.size(); - for (size_t i = 0; i < fShapeX2.size(); ++i) - shapeX2_padded[off2 + i] = fShapeX2[i]; - } + // fDimShapeX1/fDimShapeX2 are rank-padded to fDimShapeY by Initialize + const std::size_t D = fDimShapeY.size(); - auto stridesX1 = UTILITY::ComputeStrideFromShape(shapeX1_padded); - auto stridesX2 = UTILITY::ComputeStrideFromShape(shapeX2_padded); - auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); + auto stridesX1 = UTILITY::ComputeStrideFromShape(fDimShapeX1); + auto stridesX2 = UTILITY::ComputeStrideFromShape(fDimShapeX2); + auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY); + + auto isOne = [](const Dim &d){ return d.GetVal() == "1"; }; std::string type1 = ConvertTypeToString(fTensorType1); std::string type2 = ConvertTypeToString(fTensorType2); @@ -231,6 +400,8 @@ public: op += SP + SP + SP + type1 + " const* __restrict__ x1,\n"; op += SP + SP + SP + type2 + " const* __restrict__ x2,\n"; op += SP + SP + SP + "uint8_t* __restrict__ output,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -239,32 +410,28 @@ public: op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - for (std::size_t d = 0; d < D; ++d) { - op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(stridesY[d]) + "u) % " - + std::to_string(fShapeY[d]) + "u;\n"; - } + EmitOutputCoordsFromThreadIdx(op, SP + SP + SP + SP, stridesY, fDimShapeY); op += "\n"; op += SP + SP + SP + SP + "std::size_t const x1_idx =\n"; for (std::size_t d = 0; d < D; ++d) { - if (shapeX1_padded[d] == 1) + if (isOne(fDimShapeX1[d])) op += SP + SP + SP + SP + SP + "0u"; else op += SP + SP + SP + SP + SP + "out_" + std::to_string(d) - + " * " + std::to_string(stridesX1[d]) + "u"; + + " * (" + stridesX1[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } op += SP + SP + SP + SP + "std::size_t const x2_idx =\n"; for (std::size_t d = 0; d < D; ++d) { - if (shapeX2_padded[d] == 1) + if (isOne(fDimShapeX2[d])) op += SP + SP + SP + SP + SP + "0u"; else op += SP + SP + SP + SP + SP + "out_" + std::to_string(d) - + " * " + std::to_string(stridesX2[d]) + "u"; + + " * (" + stridesX2[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -283,13 +450,13 @@ public: return SP + kname + " comparisonKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; - if (fShapeY.empty()) + if (fDimShapeY.empty()) throw std::runtime_error("SOFIE Comparision Op called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeY); + std::string totalElements = ConvertDimShapeToLength(fDimShapeY); std::string kname = "comparisonKernel_" + opName; std::stringstream out; @@ -301,10 +468,12 @@ public: << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNX1 << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNX2 << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; - + return out.str(); } diff --git a/core/inc/SOFIE/ROperator_Concat.hxx b/core/inc/SOFIE/ROperator_Concat.hxx index 36ede270..88d01f50 100644 --- a/core/inc/SOFIE/ROperator_Concat.hxx +++ b/core/inc/SOFIE/ROperator_Concat.hxx @@ -375,7 +375,7 @@ return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant || fIsOutputParamShape) return ""; opName = "op_" + opName; if (fOutputShape.empty()) @@ -386,10 +386,11 @@ auto outStrides = UTILITY::ComputeStrideFromShape(fOutputShape); - std::vector prefix(Nin); - prefix[0] = 0; + // cumulative offsets along the concat axis, as strings since a dim can be symbolic + std::vector prefix(Nin); + prefix[0] = "0"; for (std::size_t k = 1; k < Nin; ++k) - prefix[k] = prefix[k - 1] + std::stoul(fInputShapes[k - 1][fAxis].GetVal()); + prefix[k] = prefix[k - 1] + " + (" + fInputShapes[k - 1][fAxis].GetVal() + ")"; std::vector> inStrides(Nin); for (std::size_t k = 0; k < Nin; ++k) @@ -403,6 +404,8 @@ op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "std::array inputs,\n"; op += SP + SP + SP + "T* output,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -414,26 +417,26 @@ op += SP + SP + SP + SP + "remaining = elem_idx;\n"; for (std::size_t d = 0; d < D; ++d) { - std::string stride_val = outStrides[d].GetVal(); + std::string stride_val = "(" + outStrides[d].GetVal() + ")"; op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = remaining / " + stride_val + "u;\n"; + + " = remaining / " + stride_val + ";\n"; op += SP + SP + SP + SP + "remaining -= out_" + std::to_string(d) - + " * " + stride_val + "u;\n"; + + " * " + stride_val + ";\n"; } op += "\n"; op += SP + SP + SP + SP + "std::size_t chosen = 0;\n"; for (std::size_t k = 0; k < Nin; ++k) { - std::size_t end_k = prefix[k] + std::stoul(fInputShapes[k][fAxis].GetVal()); + std::string end_k = "(" + prefix[k] + " + (" + fInputShapes[k][fAxis].GetVal() + "))"; op += SP + SP + SP + SP + "chosen += static_cast(" - + std::to_string(end_k) + "u <= out_" + std::to_string(fAxis) + ");\n"; + + end_k + " <= out_" + std::to_string(fAxis) + ");\n"; } op += "\n"; op += SP + SP + SP + SP + "std::size_t const output_idx =\n"; for (std::size_t d = 0; d < D; ++d) { op += SP + SP + SP + SP + SP + "out_" + std::to_string(d) - + " * " + outStrides[d].GetVal() + "u"; + + " * (" + outStrides[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -442,10 +445,10 @@ op += SP + SP + SP + SP + SP + "(chosen == " + std::to_string(k) + "u) * (\n"; for (std::size_t d = 0; d < D; ++d) { std::string coord = (d == static_cast(fAxis)) - ? ("(out_" + std::to_string(d) + " - " + std::to_string(prefix[k]) + "u)") + ? ("(out_" + std::to_string(d) + " - (" + prefix[k] + "))") : ("out_" + std::to_string(d)); op += SP + SP + SP + SP + SP + SP + coord - + " * " + inStrides[k][d].GetVal() + "u"; + + " * (" + inStrides[k][d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : "\n"; } op += SP + SP + SP + SP + SP + ")"; @@ -466,7 +469,7 @@ return SP + "ConcatKernel_" + opName + " concatKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string OpName) override { + std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { if (fIsOutputConstant || fIsOutputParamShape) return ""; OpName = "op_" + OpName; if (fOutputShape.empty()) { @@ -493,7 +496,10 @@ out << SP << "auto const elementsPerGrid_"<(workDiv_" << OpName - << ", concatKernel_" << OpName << ", input_ptrs_" << OpName << ", alpaka::getPtrNative(deviceBuf_" << fOutput << "), static_cast(" << length << "));\n"; + << ", concatKernel_" << OpName << ", input_ptrs_" << OpName << ", alpaka::getPtrNative(deviceBuf_" << fOutput << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << length << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << OpName << ");\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_Constant.hxx b/core/inc/SOFIE/ROperator_Constant.hxx index 4fea387d..52554a1a 100644 --- a/core/inc/SOFIE/ROperator_Constant.hxx +++ b/core/inc/SOFIE/ROperator_Constant.hxx @@ -6,6 +6,8 @@ #include "SOFIE/RModel.hxx" #include +#include +#include namespace SOFIE{ @@ -23,6 +25,25 @@ private: std::string fAttrType; bool fIsConstantOfShape = false; + std::string ValueLiteral(T v) const { + std::ostringstream s; + if (fAttrType == "float" || fAttrType == "double") { + s << std::setprecision(fAttrType == "float" ? std::numeric_limits::max_digits10 + : std::numeric_limits::max_digits10) + << v; + std::string str = s.str(); + if (str.find_first_of(".eE") == std::string::npos) + str += ".0"; + if (fAttrType == "float") str += "f"; + return str; + } + if (fAttrType == "bool") + s << (v ? "true" : "false"); + else + s << v; + return s.str(); + } + public: ROperator_Constant(){} @@ -64,8 +85,8 @@ public: outShape.push_back(d); if (fValues.size() != 1) throw std::runtime_error("SOFIE ConstantOfShape Op value Tensor has invalid size " + std::to_string(fValues.size())); - // Register as a dynamic intermediate tensor — values will be filled at runtime - model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), outShape); + // Register as a dynamic intermediate tensor — values will be filled at runtime. + model.AddIntermediateTensor(fNY, GetTemplatedType(T()), outShape); // Store shape for code generation (use fShape for rank, values = 0 for symbolic dims) fShape.resize(outShape.size()); for (size_t i = 0; i < outShape.size(); i++) @@ -119,13 +140,59 @@ public: } std::string Generate(std::string /* OpName */) override { - // no code to generate here. Tensor are defined in Session constructor - return "//---------------------------------------\n"; + // For the static case, values are baked into the tensor at construction. + // But a dynamic ConstantOfShape has a runtime-sized output whose + // length isn't known until inference time, so it still needs an explicit fill. + if (!fIsConstantOfShape || fIsOutputConstant) + return "//---------------------------------------\n"; + std::stringstream out; + std::string length = ConvertDimShapeToLength(fDimShape); + out << "\n//------ CONSTANTOFSHAPE (dynamic)\n"; + out << SP << "for (size_t i = 0; i < " << length << "; i++)\n"; + out << SP << SP << "tensor_" << fNY << "[i] = " << ValueLiteral(fValues[0]) << ";\n"; + return out.str(); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + if (!fIsConstantOfShape || fIsOutputConstant) return ""; + opName = "op_" + opName; + std::string kname = "ConstantOfShapeFillKernel_" + opName; + std::string op; + op = "\n//------ CONSTANTOFSHAPE_FILL_KERNEL_ALPAKA (dynamic)\n"; + op += SP + "struct " + kname + " {\n"; + op += SP + SP + "template\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const& acc, T* __restrict__ out, T const value, std::size_t const totalElements) const {\n"; + op += SP + SP + SP + "auto const idx = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "auto const stride = alpaka::getWorkDiv(acc)[0];\n"; + op += SP + SP + SP + "for (std::size_t i = idx; i < totalElements; i += stride) out[i] = value;\n"; + op += SP + SP + "}\n"; + op += SP + "};\n"; + return op; + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override { + if (!fIsConstantOfShape || fIsOutputConstant) return ""; + opName = "op_" + opName; + return SP + "ConstantOfShapeFillKernel_" + opName + " constantOfShapeFillKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string /* OpName */) override { - // no code to generate here. Tensor are defined in Session constructor - return "//---------------------------------------\n"; + std::string Generate_GPU_ALPAKA(std::string opName) override { + if (!fIsConstantOfShape || fIsOutputConstant) + return "//---------------------------------------\n"; + opName = "op_" + opName; + std::string length = ConvertDimShapeToLength(fDimShape); + + std::stringstream out; + out << "\n//------ CONSTANTOFSHAPE_GPU_ALPAKA (dynamic)\n"; + out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{static_cast(" << length << ")});\n"; + out << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; + out << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName + << ", constantOfShapeFillKernel_" << opName + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", " << ValueLiteral(fValues[0]) + << ", static_cast(" << length << "));\n"; + out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; + return out.str(); } }; diff --git a/core/inc/SOFIE/ROperator_Conv.hxx b/core/inc/SOFIE/ROperator_Conv.hxx index 1a93b013..df4889c9 100644 --- a/core/inc/SOFIE/ROperator_Conv.hxx +++ b/core/inc/SOFIE/ROperator_Conv.hxx @@ -45,6 +45,27 @@ private: size_t fDim; // dimension of the convolution + //the GEMM this conv lowers to and the extents defining it, shared by the GPU launch and the cuBLASLt layout + struct GemmInfo { + std::string oDepth, oHeight, oWidth, iDepth, iHeight, iWidth; + std::string m; //output spatial size per channel + size_t n, k; //output channels per group, input channels per group * kernel volume + }; + + GemmInfo GetGemmInfo() const { + GemmInfo gemm; + gemm.oDepth = (fDim > 2) ? fShapeY[2].GetVal() : "1"; + gemm.oHeight = (fDim > 1) ? fShapeY[fDim].GetVal() : "1"; + gemm.oWidth = fShapeY[fDim + 1].GetVal(); + gemm.iDepth = (fDim > 2) ? fShapeX[2].GetVal() : "1"; + gemm.iHeight = (fDim > 1) ? fShapeX[fDim].GetVal() : "1"; + gemm.iWidth = fShapeX[fDim + 1].GetVal(); + gemm.m = UTILITY::ComputeStrideFromShape(fShapeY)[1].GetVal(); //channel stride = oD*oH*oW + gemm.n = fShapeW[0] / fAttrGroup; + gemm.k = fShapeW[1] * fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2]; + return gemm; + } + public: @@ -95,7 +116,8 @@ public: } // function returning output shape given input - std::vector DoShapeInference(const std::vector & input, const std::vector & weight) { + using ROperator::ShapeInference; + std::vector ShapeInference(const std::vector & input, const std::vector & weight) { // shape of convolution input has to be (according to ONNX): N x C x H x W // Where N : batch size, C : input channels, H : input height, W : input width @@ -202,7 +224,7 @@ public: } else { // general case (stride not 1) int64_t v = pad - kernel; std::string outStr = "((" + inputDim.param + "+" + std::to_string(v) + ")/" - + std::to_string(stride) + "1)"; + + std::to_string(stride) + "+1)"; return Dim{ outStr, static_cast(-1)}; } } @@ -258,7 +280,7 @@ public: std::cout << fNW << " : " << ConvertShapeToString(fShapeW) << std::endl; throw std::runtime_error("TMVA SOFIE Conv Op input weight tensor" + fNW + " is not of 3,4 or 5 dimensions"); } - fShapeY = DoShapeInference(fShapeX, fShapeW); + fShapeY = ShapeInference(fShapeX, fShapeW); model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); if (fNB != "") { if (!model.CheckIfTensorAlreadyExist(fNB)) { @@ -605,22 +627,11 @@ public: if (fShapeX.empty() || fShapeW.empty() || fShapeY.empty()) throw std::runtime_error("TMVA SOFIE Conv Op called to Generate without being initialized first"); - 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 iDepth = (fDim > 2) ? fShapeX[2].dim : 1; - size_t iHeight = (fDim > 1) ? fShapeX[fDim].dim : 1; - size_t iWidth = fShapeX[fDim + 1].dim; size_t kHeight = (fDim > 1) ? fShapeW[fDim] : 1; size_t kWidth = fShapeW[fDim + 1]; size_t kDepth = (fDim > 2) ? fShapeW[2] : 1; size_t kernelSize = fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2]; - size_t colRows = fShapeW[1] * kernelSize; - size_t colCols = oDepth * oHeight * oWidth; - size_t colElements = colRows * colCols; - size_t outChannels = fShapeW[0]; - size_t spatialSize = oDepth * oHeight * oWidth; // Strides for weight vectorisation size_t id = (fDim > 2) ? fDim - 3 : 2; @@ -631,11 +642,8 @@ public: size_t dstrideDil = fAttrDilations[id] * fAttrKernelShape[ih] * fAttrKernelShape[iw]; size_t icstrideDil = fAttrKernelShape[id] * fAttrKernelShape[ih] * fAttrKernelShape[iw]; size_t ocstrideDil = fShapeW[1] * icstrideDil; - size_t hstride = kWidth; - size_t dstride = kHeight * kWidth; size_t icstride = kHeight * kWidth * kDepth; size_t ocstride = fShapeW[1] * icstride; - size_t wTotalElements = ConvertShapeToLength(fShapeW); // effective (dilation-expanded) kernel extents, used for the dense im2col decode size_t kHeightEff = (fDim > 1) ? fAttrKernelShape[ih] : 1; @@ -702,6 +710,8 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ input,\n"; op += SP + SP + SP + "T* __restrict__ col,\n"; + op += SP + SP + SP + "std::size_t const oDepth, std::size_t const oHeight, std::size_t const oWidth,\n"; + op += SP + SP + SP + "std::size_t const iDepth, std::size_t const iHeight, std::size_t const iWidth,\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -710,8 +720,8 @@ public: op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - op += SP + SP + SP + SP + "std::size_t const col_row = elem_idx / " + std::to_string(colCols) + "u;\n"; - op += SP + SP + SP + SP + "std::size_t const col_col = elem_idx % " + std::to_string(colCols) + "u;\n\n"; + op += SP + SP + SP + SP + "std::size_t const col_row = elem_idx / (oDepth * oHeight * oWidth);\n"; + op += SP + SP + SP + SP + "std::size_t const col_col = elem_idx % (oDepth * oHeight * oWidth);\n\n"; op += SP + SP + SP + SP + "std::size_t const ic = col_row / " + std::to_string(kernelSize) + "u;\n"; op += SP + SP + SP + SP + "std::size_t const k_rem = col_row % " + std::to_string(kernelSize) + "u;\n"; @@ -730,13 +740,13 @@ public: } if (fDim > 2) { - op += SP + SP + SP + SP + "std::size_t const od = col_col / " + std::to_string(oHeight * oWidth) + "u;\n"; - op += SP + SP + SP + SP + "std::size_t const oh = (col_col / " + std::to_string(oWidth) + "u) % " + std::to_string(oHeight) + "u;\n"; - op += SP + SP + SP + SP + "std::size_t const ow = col_col % " + std::to_string(oWidth) + "u;\n\n"; + op += SP + SP + SP + SP + "std::size_t const od = col_col / (oHeight * oWidth);\n"; + op += SP + SP + SP + SP + "std::size_t const oh = (col_col / oWidth) % oHeight;\n"; + op += SP + SP + SP + SP + "std::size_t const ow = col_col % oWidth;\n\n"; } else if (fDim > 1) { op += SP + SP + SP + SP + "std::size_t const od = 0u;\n"; - op += SP + SP + SP + SP + "std::size_t const oh = col_col / " + std::to_string(oWidth) + "u;\n"; - op += SP + SP + SP + SP + "std::size_t const ow = col_col % " + std::to_string(oWidth) + "u;\n\n"; + op += SP + SP + SP + SP + "std::size_t const oh = col_col / oWidth;\n"; + op += SP + SP + SP + SP + "std::size_t const ow = col_col % oWidth;\n\n"; } else { op += SP + SP + SP + SP + "std::size_t const od = 0u;\n"; op += SP + SP + SP + SP + "std::size_t const oh = 0u;\n"; @@ -770,15 +780,15 @@ public: } op += SP + SP + SP + SP + "bool const in_bounds =\n"; - op += SP + SP + SP + SP + SP + "id_in >= 0 && id_in < " + std::to_string(iDepth) + " &&\n"; - op += SP + SP + SP + SP + SP + "ih_in >= 0 && ih_in < " + std::to_string(iHeight) + " &&\n"; - op += SP + SP + SP + SP + SP + "iw_in >= 0 && iw_in < " + std::to_string(iWidth) + ";\n\n"; + op += SP + SP + SP + SP + SP + "id_in >= 0 && id_in < static_cast(iDepth) &&\n"; + op += SP + SP + SP + SP + SP + "ih_in >= 0 && ih_in < static_cast(iHeight) &&\n"; + op += SP + SP + SP + SP + SP + "iw_in >= 0 && iw_in < static_cast(iWidth) ;\n\n"; op += SP + SP + SP + SP + "if (in_bounds) {\n"; op += SP + SP + SP + SP + SP + "std::size_t const in_idx =\n"; - op += SP + SP + SP + SP + SP + SP + "ic * " + std::to_string(iDepth * iHeight * iWidth) + "u +\n"; - op += SP + SP + SP + SP + SP + SP + "static_cast(id_in) * " + std::to_string(iHeight * iWidth) + "u +\n"; - op += SP + SP + SP + SP + SP + SP + "static_cast(ih_in) * " + std::to_string(iWidth) + "u +\n"; + op += SP + SP + SP + SP + SP + SP + "ic * (iDepth * iHeight * iWidth) +\n"; + op += SP + SP + SP + SP + SP + SP + "static_cast(id_in) * (iHeight * iWidth) +\n"; + op += SP + SP + SP + SP + SP + SP + "static_cast(ih_in) * iWidth +\n"; op += SP + SP + SP + SP + SP + SP + "static_cast(iw_in);\n"; op += SP + SP + SP + SP + SP + "col[elem_idx] = input[in_idx];\n"; op += SP + SP + SP + SP + "} else {\n"; @@ -798,6 +808,7 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ bias,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; + op += SP + SP + SP + "std::size_t const spatialSize,\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -805,7 +816,7 @@ public: op += SP + SP + SP + "auto const grid_thread_extent = alpaka::getWorkDiv(acc)[0];\n\n"; op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n"; - op += SP + SP + SP + SP + "std::size_t const channel = elem_idx / " + std::to_string(spatialSize) + "u;\n"; + op += SP + SP + SP + SP + "std::size_t const channel = elem_idx / spatialSize;\n"; op += SP + SP + SP + SP + "output[elem_idx] = bias[channel];\n"; op += SP + SP + SP + "}\n"; op += SP + SP + "}\n"; @@ -830,25 +841,16 @@ public: if (fShapeX.empty() || fShapeW.empty() || fShapeY.empty()) throw std::runtime_error("SOFIE Conv Op called to Generate without being initialized first"); - 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 iDepth = (fDim > 2) ? fShapeX[2].dim : 1; - size_t iHeight = (fDim > 1) ? fShapeX[fDim].dim : 1; - size_t iWidth = fShapeX[fDim + 1].dim; - size_t outChannels = fShapeW[0]; - size_t kernelSize = fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2]; - // gemm dimensions computed from shape members - size_t gemm_n = outChannels; // output channels - size_t gemm_k = fShapeW[1] * kernelSize; // input channels/group * kernel volume - size_t gemm_m = oDepth * oHeight * oWidth; // output spatial size per channel - size_t colElements = gemm_k * gemm_m; // colRows * colCols - size_t wTotal = ConvertShapeToLength(fShapeW); - - // For group conv: per-group output channels and _f offset - size_t gemm_n_group = gemm_n / fAttrGroup; // output channels produced by a single group - size_t groupFOffset = gemm_n_group * gemm_k; // elements of _f per group + auto gemm = GetGemmInfo(); + auto strideX = UTILITY::ComputeStrideFromShape(fShapeX); //[0] = one sample (C*D*H*W), [1] = one channel (D*H*W) + auto strideY = UTILITY::ComputeStrideFromShape(fShapeY); + std::string bsize = fShapeX[0].GetVal(); + std::string colElements = std::to_string(gemm.k) + " * " + gemm.m; // colRows * colCols + size_t wTotal = ConvertShapeToLength(fShapeW); + size_t groupFOffset = gemm.n * gemm.k; // elements of _f per group + std::string im2colExtents; //the six extent args of Im2ColKernel, same order as its signature + for (const std::string &e : {gemm.oDepth, gemm.oHeight, gemm.oWidth, gemm.iDepth, gemm.iHeight, gemm.iWidth}) + im2colExtents += ", static_cast(" + e + ")"; std::stringstream out; out << "\n//------ CONV_GPU_ALPAKA\n"; @@ -858,7 +860,6 @@ public: for (size_t d = 0; d < fDim; ++d) if (fAttrDilations[d] > 1) hasDilation = true; if (hasDilation) { out << SP << "alpaka::memset(queue, deviceBuf_" << convK << ", 0);\n"; - out << SP << "alpaka::wait(queue);\n"; } // ----------------------------------------------------------------------- @@ -873,17 +874,16 @@ public: << ", alpaka::getPtrNative(deviceBuf_" << fNW << ")" << ", alpaka::getPtrNative(deviceBuf_" << convK << ")" << ", static_cast(" << wTotal << "));\n"; - out << SP << SP << "alpaka::wait(queue);\n"; out << SP << "}\n\n"; // ----------------------------------------------------------------------- // Step 2: Batch loop // ----------------------------------------------------------------------- out << SP << "for (std::size_t n = 0; n < " << bsize << "; n++) {\n\n"; - out << SP << SP << "std::size_t const x_offset = n * " - << fShapeX[1].dim * iDepth * iHeight * iWidth << "u;\n"; - out << SP << SP << "std::size_t const out_offset = n * " - << fShapeY[1].dim * gemm_m << "u;\n\n"; + out << SP << SP << "std::size_t const x_offset = n * (" + << strideX[0].GetVal() << ");\n"; + out << SP << SP << "std::size_t const out_offset = n * (" + << strideY[0].GetVal() << ");\n\n"; // ----------------------------------------------------------------------- // Step 3 + 4: Im2Col then GEMM — structure differs for grouped vs non-grouped @@ -898,12 +898,12 @@ public: out << SP << SP << SP << "alpaka::exec(queue, workDiv_im2col, im2colKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNX << ") + x_offset" << ", alpaka::getPtrNative(deviceBuf_" << imcol << ")" + << im2colExtents << ", static_cast(" << colElements << "));\n"; - out << SP << SP << SP << "alpaka::wait(queue);\n"; out << SP << SP << "}\n\n"; if (!fNB.empty()) { - size_t biasElements = gemm_n * gemm_m; + std::string biasElements = std::to_string(gemm.n) + " * " + gemm.m; out << SP << SP << "// Step 4a: broadcast bias into output slice\n"; out << SP << SP << "{\n"; out << SP << SP << SP << "auto const elementsPerThread_bias = Vec::all(static_cast(1));\n"; @@ -912,34 +912,32 @@ public: out << SP << SP << SP << "alpaka::exec(queue, workDiv_bias, biasBroadcastKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNB << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNY << ") + out_offset" + << ", static_cast(" << gemm.m << ")" << ", static_cast(" << biasElements << "));\n"; - out << SP << SP << SP << "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 + << 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 + << 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"; } else { // Grouped convolution: im2col and GEMM per group with group-adjusted input pointer. // Each group processes fShapeW[1] input channels starting at g * fShapeW[1]. out << SP << SP << "for (std::size_t g = 0; g < " << fAttrGroup << "; g++) {\n\n"; - out << SP << SP << SP << "std::size_t const g_in_offset = x_offset + g * " - << fShapeW[1] * iDepth * iHeight * iWidth << "u;\n"; - out << SP << SP << SP << "std::size_t const g_out_offset = out_offset + g * " - << gemm_n_group * gemm_m << "u;\n"; + out << SP << SP << SP << "std::size_t const g_in_offset = x_offset + g * (" + << std::to_string(fShapeW[1]) + " * " + strideX[1].GetVal() << ");\n"; + out << SP << SP << SP << "std::size_t const g_out_offset = out_offset + g * (" + << std::to_string(gemm.n) + " * " + gemm.m << ");\n"; out << SP << SP << SP << "std::size_t const f_offset = g * " << groupFOffset << "u;\n\n"; out << SP << SP << SP << "// im2col for group g (reads only this group's input channels)\n"; @@ -950,37 +948,35 @@ public: out << SP << SP << SP << SP << "alpaka::exec(queue, workDiv_im2col, im2colKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNX << ") + g_in_offset" << ", alpaka::getPtrNative(deviceBuf_" << imcol << ")" + << im2colExtents << ", static_cast(" << colElements << "));\n"; - out << SP << SP << SP << SP << "alpaka::wait(queue);\n"; out << SP << SP << SP << "}\n\n"; if (!fNB.empty()) { - size_t groupBiasElements = gemm_n_group * gemm_m; + std::string groupBiasElements = std::to_string(gemm.n) + " * " + gemm.m; out << SP << SP << SP << "// Broadcast group bias\n"; out << SP << SP << SP << "{\n"; out << SP << SP << SP << SP << "auto const elementsPerThread_bias = Vec::all(static_cast(1));\n"; out << SP << SP << SP << SP << "auto const elementsPerGrid_bias = Vec::all(Idx{" << groupBiasElements << "});\n"; out << SP << SP << SP << SP << "auto const workDiv_bias = sofie_workdiv(elementsPerGrid_bias);\n"; out << SP << SP << SP << SP << "alpaka::exec(queue, workDiv_bias, biasBroadcastKernel_" << opName - << ", alpaka::getPtrNative(deviceBuf_" << fNB << ") + g * " << gemm_n_group + << ", alpaka::getPtrNative(deviceBuf_" << fNB << ") + g * " << gemm.n << ", alpaka::getPtrNative(deviceBuf_" << fNY << ") + g_out_offset" + << ", static_cast(" << gemm.m << ")" << ", static_cast(" << groupBiasElements << "));\n"; - out << SP << SP << SP << SP << "alpaka::wait(queue);\n"; out << SP << SP << SP << "}\n\n"; out << SP << SP << SP << "blas.matmul('n', 'n', " - << gemm_m << ", " << gemm_n_group << ", " << gemm_k + << gemm.m << ", " << gemm.n << ", " << gemm.k << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << imcol << ")" << ", alpaka::getPtrNative(deviceBuf_" << convK << ") + f_offset" << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << fNY << ") + g_out_offset);\n\n"; } else { out << SP << SP << SP << "blas.matmul('n', 'n', " - << gemm_m << ", " << gemm_n_group << ", " << gemm_k + << gemm.m << ", " << gemm.n << ", " << gemm.k << ", 1.0f, alpaka::getPtrNative(deviceBuf_" << imcol << ")" << ", alpaka::getPtrNative(deviceBuf_" << convK << ") + f_offset" << ", 0.0f, alpaka::getPtrNative(deviceBuf_" << fNY << ") + g_out_offset);\n\n"; } - // Wait for GEMM to finish before next group's im2col overwrites the shared _xcol buffer. - out << SP << SP << SP << "alpaka::wait(queue);\n\n"; out << SP << SP << "}\n"; // end group loop } @@ -994,17 +990,9 @@ public: std::string GetBlasConfig(){ - 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 kSize_ = fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2]; - size_t gemm_n_ = fShapeW[0] / fAttrGroup; - size_t gemm_k_ = fShapeW[1] * kSize_; - size_t gemm_m_ = oDepth_ * oHeight_ * oWidth_; - auto lda = std::to_string(gemm_m_); // ld for xcol^T (gemm_m×gemm_k col-major) - auto ldb = std::to_string(gemm_k_); // ld for xf^T (gemm_k×gemm_n col-major) - auto ldc = std::to_string(gemm_m_); // ld for y^T (gemm_m×gemm_n col-major) - return std::to_string(gemm_m_) + ", " + std::to_string(gemm_n_) + ", " + std::to_string(gemm_k_) + ", " + lda + ", " + ldb + ", " + ldc + ", 'n', 'n'"; + auto gemm = GetGemmInfo(); + std::string n = std::to_string(gemm.n), k = std::to_string(gemm.k); + return gemm.m + ", " + n + ", " + k + ", " + gemm.m + ", " + k + ", " + gemm.m + ", 'n', 'n', Epilogue::Default"; } }; diff --git a/core/inc/SOFIE/ROperator_CumSum.hxx b/core/inc/SOFIE/ROperator_CumSum.hxx index 6ec65d20..892e3e28 100644 --- a/core/inc/SOFIE/ROperator_CumSum.hxx +++ b/core/inc/SOFIE/ROperator_CumSum.hxx @@ -149,7 +149,7 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShape.empty()) throw std::runtime_error("SOFIE CumSum GPU kernel called without initialization"); @@ -160,6 +160,8 @@ public: std::vector dims(fRank); for (size_t i = 0; i < fRank; ++i) dims[i] = fShape[i].GetVal(); + auto sz = [](const std::string &e) { return "static_cast(" + e + ")"; }; + std::vector outerShape; for (size_t i = 0; i < fRank; ++i) if (i != ax) outerShape.emplace_back(fShape[i]); @@ -182,6 +184,8 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ X,\n"; op += SP + SP + SP + "T* __restrict__ Y,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const outerLen,\n"; op += SP + SP + SP + "std::size_t const axLen) const {\n\n"; @@ -192,7 +196,7 @@ public: for (size_t oi = 0; oi < outerDimIdx.size(); ++oi) { size_t di = outerDimIdx[oi]; op += SP + SP + SP + "std::size_t const d_" + std::to_string(di) - + " = (tid / " + outerStrides[oi].GetVal() + "u) % " + dims[di] + "u;\n"; + + " = (tid / " + sz(outerStrides[oi].GetVal()) + ") % " + sz(dims[di]) + ";\n"; } op += "\n"; } @@ -201,14 +205,14 @@ public: bool firstTerm = true; for (size_t di : outerDimIdx) { if (!firstTerm) op += " +"; - op += " d_" + std::to_string(di) + " * " + strides[di].GetVal() + "u"; + op += " d_" + std::to_string(di) + " * " + sz(strides[di].GetVal()); firstTerm = false; } if (firstTerm) op += " 0u"; op += ";\n\n"; // axis stride - op += SP + SP + SP + "std::size_t const ax_stride = " + strides[ax].GetVal() + "u;\n\n"; + op += SP + SP + SP + "std::size_t const ax_stride = " + sz(strides[ax].GetVal()) + ";\n\n"; // cumsum loop std::string excl = (fExclusive ? "true" : "false"); @@ -249,7 +253,7 @@ public: return SP + kname + " cumSumKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShape.empty()) throw std::runtime_error("SOFIE CumSum GPU dispatch called without initialization"); @@ -269,8 +273,10 @@ public: out << SP << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName << ", cumSumKernel_" << opName << ", " << "alpaka::getPtrNative(deviceBuf_" << fNX << "), " - << "alpaka::getPtrNative(deviceBuf_" << fNY << "), " - << "static_cast(" << outerLen << "), " + << "alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << outerLen << "), " << "static_cast(" << axLen << "));\n"; out << SP << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; out << SP << "}\n"; diff --git a/core/inc/SOFIE/ROperator_Expand.hxx b/core/inc/SOFIE/ROperator_Expand.hxx index 95955ed7..d016126f 100644 --- a/core/inc/SOFIE/ROperator_Expand.hxx +++ b/core/inc/SOFIE/ROperator_Expand.hxx @@ -164,7 +164,22 @@ public: return out.str(); } -std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + // fShapeX left-padded to fShapeY's rank with Dim{1} entries + std::vector PaddedShapeX() const { + std::vector padded = fShapeX; + if (padded.size() < fShapeY.size()) + padded.insert(padded.begin(), fShapeY.size() - padded.size(), Dim{1}); + return padded; + } + + bool NeedsBroadcast() const { + if (fShapeX.size() != fShapeY.size()) return true; + for (size_t i = 0; i < fShapeX.size(); ++i) + if (fShapeX[i] != fShapeY[i]) return true; + return false; + } + +std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; if (fInitialized) return ""; @@ -172,39 +187,13 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { if (fShapeY.empty()) throw std::runtime_error("SOFIE Expand Op called to Generate without being initialized first"); - // Can only generate a static kernel if all dimensions are concrete values - auto isStatic = [](const std::vector& shape) { - return std::all_of(shape.begin(), shape.end(), - [](const Dim& d){ return !d.isParam; }); - }; - if (!isStatic(fShapeX) || !isStatic(fShapeY)) return ""; - - // Check if broadcast is actually needed - bool needsBroadcast = (fShapeX.size() != fShapeY.size()); - if (!needsBroadcast) { - needsBroadcast = std::any_of(fShapeX.begin(), fShapeX.end(), - [&](const Dim& d) { - size_t i = &d - fShapeX.data(); - return fShapeX[i].dim != fShapeY[i].dim; - }); - } - if (!needsBroadcast) return ""; // same static shape — just a memcpy + if (!NeedsBroadcast()) return ""; const std::size_t D = fShapeY.size(); - - // Left-pad fShapeX with dim=1 entries to match rank of fShapeY - std::vector shapeX_padded(D, 1); - size_t offset = D - fShapeX.size(); - for (size_t i = 0; i < fShapeX.size(); ++i) - shapeX_padded[offset + i] = fShapeX[i].dim; - - std::vector shapeY_int(D); - for (size_t i = 0; i < D; ++i) - shapeY_int[i] = fShapeY[i].dim; + std::vector shapeX_padded = PaddedShapeX(); auto stridesX = UTILITY::ComputeStrideFromShape(shapeX_padded); - auto stridesY = UTILITY::ComputeStrideFromShape(shapeY_int); - std::size_t totalElements = ConvertShapeToLength(shapeY_int); + auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); std::string kname = "ExpandKernel_" + opName; @@ -215,8 +204,10 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { op += SP + SP + "ALPAKA_FN_ACC void operator()(\n"; op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ input,\n"; - op += SP + SP + SP + "T* __restrict__ output,\n"; - op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; + op += SP + SP + SP + "T* __restrict__ output"; + for (auto &p : dynParamNames) + op += ",\n" + SP + SP + SP + "std::size_t const " + p; + op += ",\n" + SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; op += SP + SP + SP + "if (global_thread_idx >= totalElements) return;\n"; @@ -224,24 +215,23 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - // Decompose output linear index using compile-time output strides + // Decompose output linear index using output strides (literal or symbolic) for (std::size_t d = 0; d < D; ++d) { op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(stridesY[d]) + "u) % " - + std::to_string(shapeY_int[d]) + "u;\n"; + + " = (elem_idx / static_cast(" + stridesY[d].GetVal() + ")) % " + + "static_cast(" + fShapeY[d].GetVal() + ");\n"; } op += "\n"; - // Input index: broadcast dims (shapeX_padded[d]==1) contribute 0 — - // compiler eliminates zero terms entirely, no runtime branch + // Input index: broadcast dims (shapeX_padded[d] == 1) contribute 0 op += SP + SP + SP + SP + "std::size_t const input_idx =\n"; for (std::size_t d = 0; d < D; ++d) { - if (shapeX_padded[d] == 1) { + if (shapeX_padded[d].GetVal() == "1") { op += SP + SP + SP + SP + SP + "0u"; } else { op += SP + SP + SP + SP + SP + "out_" + std::to_string(d) - + " * " + std::to_string(stridesX[d]) + "u"; + + " * static_cast(" + stridesX[d].GetVal() + ")"; } op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -257,27 +247,14 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override { if (fIsOutputConstant) return ""; if (fInitialized) return ""; - - auto isStatic = [](const std::vector& shape) { - return std::all_of(shape.begin(), shape.end(), - [](const Dim& d){ return !d.isParam; }); - }; - if (!isStatic(fShapeX) || !isStatic(fShapeY)) return ""; - - // Check if broadcast is actually needed - bool needsBroadcast = (fShapeX.size() != fShapeY.size()); - if (!needsBroadcast) { - for (size_t i = 0; i < fShapeX.size(); ++i) - if (fShapeX[i].dim != fShapeY[i].dim) { needsBroadcast = true; break; } - } - if (!needsBroadcast) return ""; + if (!NeedsBroadcast()) return ""; opName = "op_" + opName; std::string kname = "ExpandKernel_" + opName; return SP + kname + " expandKernel_" + opName + ";\n"; } -std::string Generate_GPU_ALPAKA(std::string opName) override { +std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; if (fShapeY.empty()) @@ -291,52 +268,27 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { return ""; } - auto isStatic = [](const std::vector& shape) { - return std::all_of(shape.begin(), shape.end(), - [](const Dim& d){ return !d.isParam; }); - }; - bool staticShapes = isStatic(fShapeX) && isStatic(fShapeY); - - // Check if broadcast is actually needed for static shapes - bool needsBroadcast = !staticShapes; // dynamic always needs runtime broadcast - if (staticShapes) { - needsBroadcast = (fShapeX.size() != fShapeY.size()); - if (!needsBroadcast) { - for (size_t i = 0; i < fShapeX.size(); ++i) - if (fShapeX[i].dim != fShapeY[i].dim) { needsBroadcast = true; break; } - } - } - - if (!needsBroadcast) { - // Same static shape — device-to-device copy + if (!NeedsBroadcast()) { + // Same shape — device-to-device copy out << SP << "alpaka::memcpy(queue, deviceBuf_" << fNY << ", deviceBuf_" << fNX << ");\n"; out << SP << "alpaka::wait(queue);\n"; return out.str(); } - if (!staticShapes) { - // Dynamic shapes — not yet supported on GPU, throw a clear error - throw std::runtime_error( - "SOFIE Expand GPU: dynamic shapes are not yet supported for GPU inference. " - "Tensor " + fNX + " has a dynamic shape."); - } - - // Static broadcast — launch the expand kernel - std::vector shapeY_int(fShapeY.size()); - for (size_t i = 0; i < fShapeY.size(); ++i) - shapeY_int[i] = fShapeY[i].dim; - std::size_t totalElements = ConvertShapeToLength(shapeY_int); + std::string totalElements = ConvertDimShapeToLength(fShapeY); std::string kname = "expandKernel_" + opName; out << SP << "auto const elementsPerThread_" << opName << " = Vec::all(static_cast(1));\n"; - out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{" << totalElements << "});\n"; + out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{static_cast(" << totalElements << ")});\n"; out << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; out << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; out << SP <<"alpaka::enqueue(queue, task_" << opName << ");\n"; return out.str(); diff --git a/core/inc/SOFIE/ROperator_Gather.hxx b/core/inc/SOFIE/ROperator_Gather.hxx index 3c16f180..2d2f955b 100644 --- a/core/inc/SOFIE/ROperator_Gather.hxx +++ b/core/inc/SOFIE/ROperator_Gather.hxx @@ -28,6 +28,8 @@ private: std::vector fIndices; // indices vector in case they are known at initialization + std::vector fOutputShapeData; // in case output is a shape tensor we store here the shape value data (can be parametric) + std::string fType; public: @@ -69,19 +71,27 @@ public: // case indices tensor is initialized if (model.IsInitializedTensor(fNIndices)) { // empty shape Indices is a scalar value for the indices + bool hasNegativeIndex = false; size_t indicesLength = ConvertShapeToLength(model.GetTensorShape(fNIndices)); - int64_t* indicesData = static_cast(model.GetInitializedTensorData(fNIndices).get()); + int64_t* data = static_cast(model.GetInitializedTensorData(fNIndices).get()); + // copy in a vector since we may need to update the values in case of negative indices + fIndices = std::vector(data, data + indicesLength); // update indices data in case of negative dim values for (size_t i = 0; i < indicesLength; i++) { // move this at generation time? if (!fShapeX[fAttrAxis].isParam) { - if (indicesData[i] < 0) { - indicesData[i] += fShapeX[fAttrAxis].dim; + if (fIndices[i] < 0) { + hasNegativeIndex = true; + fIndices[i] += fShapeX[fAttrAxis].dim; } } } - // Save in a vector gather Indices of size q - fIndices = std::vector(indicesData, indicesData + indicesLength); + // for negative indices we need to add an extra constant tensor + if (hasNegativeIndex) { + std::string nameIndicesUpdated = fNIndices + "_updated"; + model.AddConstantTensor(nameIndicesUpdated, model.GetTensorShape(fNIndices), fIndices.data()); + fNIndices = nameIndicesUpdated; + } } // Output shape if (model.Verbose()) @@ -120,17 +130,17 @@ public: else if (model.IsShapeTensor(fNX) && q <=1 && fIndices.size() > 0) { auto inputData = model.GetShapeTensorValues(fNX); // if r == 1 and q<=1 then output length is 1 (is a scalar or tensor of size1) - std::vector outputData(1); - outputData[0] = inputData[fIndices[0]]; - if (outputData[0].isParam) { - fIsOutputConstant = true; + fOutputShapeData.resize(1); + fOutputShapeData[0] = inputData[fIndices[0]]; + if (fOutputShapeData[0].isParam) { + fIsOutputParamShape = true; // shapeY can be scalar or vector of size1 - model.AddShapeTensor(fNY, outputData, fShapeY.size() == 0); + model.AddShapeTensor(fNY, fOutputShapeData, fShapeY.size() == 0); if (model.Verbose()) std::cout << "Gather: " << fNX << " " << ConvertDimShapeToString(fShapeX) << " -> " << fNY << " with shape " << ConvertDimShapeToString(fShapeY) - << " and values " << ConvertDimShapeToString(outputData) << " (shape) " << std::endl; + << " and values " << ConvertDimShapeToString(fOutputShapeData) << " (shape) " << std::endl; } else { - int64_t value = static_cast(outputData[0].dim); + int64_t value = static_cast(fOutputShapeData[0].dim); auto shapeY = ConvertShapeToInt(fShapeY); model.AddConstantTensor(fNY, shapeY, &value); fIsOutputConstant = true; @@ -139,7 +149,7 @@ public: << " and values {" << value << "} (constant) " << std::endl; } } - if (!fIsOutputConstant) { + if (!fIsOutputConstant && !fIsOutputParamShape) { // Add output tensor model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); fType = ConvertTypeToString(model.GetTensorType(fNX)); @@ -158,6 +168,13 @@ public: out << "//--------------------(constant)----------\n"; return out.str(); } + if (fIsOutputParamShape) { + out << "//--------------------(shape)----------\n"; + for (int i = 0; i < static_cast(fOutputShapeData.size()); i++) { + out << SP << "tensor_" << fNY << "[" << i << " ] = " << fOutputShapeData[i].GetVal() << ";\n"; + } + return out.str(); + } // The shape of the output is q + r - 1 size_t r = fShapeX.size(); // Indices of shape q @@ -276,13 +293,12 @@ public: return out.str(); } -std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { - if (fIsOutputConstant) return ""; +std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { + if (fIsOutputConstant || fIsOutputParamShape) return ""; opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE Gather Op called to Generate without being initialized first"); - const std::size_t D = fShapeY.size(); // output rank = q + r - 1 const std::size_t r = fShapeX.size(); const std::size_t q = fShapeIndices.size(); @@ -301,6 +317,8 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { op += SP + SP + SP + "T const* __restrict__ input,\n"; op += SP + SP + SP + "int64_t const* __restrict__ indices,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -309,11 +327,7 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - for (std::size_t d = 0; d < D; ++d) { - op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + stridesY[d].GetVal() + "u) % " - + fShapeY[d].GetVal() + "u;\n"; - } + EmitOutputCoordsFromThreadIdx(op, SP + SP + SP + SP, stridesY, fShapeY); op += "\n"; // Output dims [axis ... axis+q) correspond to the indices tensor dims [0 ... q) @@ -325,14 +339,14 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { for (std::size_t i = 0; i < q; ++i) { op += SP + SP + SP + SP + SP + "out_" + std::to_string(fAttrAxis + i) - + " * " + stridesIndices[i].GetVal() + "u"; + + " * (" + stridesIndices[i].GetVal() + ")"; op += (i + 1 < q) ? " +\n" : ";\n"; } } op += "\n"; op += SP + SP + SP + SP + "int64_t k = indices[i_index];\n"; - op += SP + SP + SP + SP + "if (k < 0) k += " + fShapeX[fAttrAxis].GetVal() + ";\n"; + op += SP + SP + SP + SP + "if (k < 0) k += (" + fShapeX[fAttrAxis].GetVal() + ");\n"; op += SP + SP + SP + SP + "if (k < 0) k = 0;\n"; op += SP + SP + SP + SP + "if (k >= static_cast(" + fShapeX[fAttrAxis].GetVal() + ")) " + "k = static_cast(" + fShapeX[fAttrAxis].GetVal() + ") - 1;\n\n"; @@ -342,15 +356,15 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { // + sum over j in [axis+1, r): out_{j-1+q} * stridesX[j] // (the dims after axis in Y are shifted by q-1 relative to X) op += SP + SP + SP + SP + "std::size_t const input_idx =\n"; - op += SP + SP + SP + SP + SP + "static_cast(k) * " + stridesX[fAttrAxis].GetVal() + "u"; + op += SP + SP + SP + SP + SP + "static_cast(k) * (" + stridesX[fAttrAxis].GetVal() + ")"; for (std::size_t j = 0; j < static_cast(fAttrAxis); ++j) { op += " +\n" + SP + SP + SP + SP + SP - + "out_" + std::to_string(j) + " * " + stridesX[j].GetVal() + "u"; + + "out_" + std::to_string(j) + " * (" + stridesX[j].GetVal() + ")"; } for (std::size_t j = fAttrAxis + 1; j < r; ++j) { // in Y, the coord for X's dim j lives at output dim q + j - 1 op += " +\n" + SP + SP + SP + SP + SP - + "out_" + std::to_string(q + j - 1) + " * " + stridesX[j].GetVal() + "u"; + + "out_" + std::to_string(q + j - 1) + " * (" + stridesX[j].GetVal() + ")"; } op += ";\n\n"; @@ -363,14 +377,25 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { } std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override { - if (fIsOutputConstant) return ""; + if (fIsOutputConstant || fIsOutputParamShape) return ""; opName = "op_" + opName; std::string kname = "GatherKernel_" + opName; return SP + kname + " gatherKernel_" + opName + ";\n"; } -std::string Generate_GPU_ALPAKA(std::string opName) override { +std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; + if (fIsOutputParamShape) { + std::stringstream out; + out << "\n//------ GATHER (shape) GPU\n"; + for (int i = 0; i < static_cast(fOutputShapeData.size()); i++) { + out << SP << "tensor_" << fNY << "[" << i << "] = " << fOutputShapeData[i].GetVal() << ";\n"; + } + out << SP << "auto hostBuf_" << fNY << " = alpaka::createView(hostAcc, tensor_" << fNY + << ", " << fOutputShapeData.size() << ");\n"; + out << SP << "alpaka::memcpy(queue, deviceBuf_" << fNY << ", hostBuf_" << fNY << ");\n"; + return out.str(); + } opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE Gather Op called to Generate without being initialized first"); @@ -387,8 +412,10 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNIndices << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; return out.str(); } @@ -397,4 +424,4 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { }//SOFIE -#endif //SOFIE_ROPERATOR_RELU +#endif //SOFIE_ROPERATOR_GATHER diff --git a/core/inc/SOFIE/ROperator_GatherND.hxx b/core/inc/SOFIE/ROperator_GatherND.hxx index ffcdab86..c58d4591 100644 --- a/core/inc/SOFIE/ROperator_GatherND.hxx +++ b/core/inc/SOFIE/ROperator_GatherND.hxx @@ -22,12 +22,16 @@ private: std::string fNIndices; std::string fNY; - std::vector fShapeData; - std::vector fShapeIndices; - std::vector fShapeY; + std::vector fShapeData; + std::vector fShapeIndices; + std::vector fShapeY; std::string fType; + // A size expression (a literal, or a symbolic dim name/expression from + // Dim::GetVal()) wrapped as an explicit std::size_t cast. + static std::string sz(const std::string &e) { return "static_cast(" + e + ")"; } + public: ROperator_GatherND() {} ROperator_GatherND(int64_t batchDims, @@ -57,13 +61,12 @@ public: if (!model.CheckIfTensorAlreadyExist(fNIndices)) throw std::runtime_error("SOFIE GatherND: indices tensor " + fNIndices + " not found in model"); - fShapeData = model.GetTensorShape(fNData); - fShapeIndices = model.GetTensorShape(fNIndices); + fShapeData = model.GetDimTensorShape(fNData); + fShapeIndices = model.GetDimTensorShape(fNIndices); size_t r = fShapeData.size(); size_t q = fShapeIndices.size(); size_t b = static_cast(fBatchDims); - size_t last_idx_dim = fShapeIndices.back(); if (r < 1) throw std::runtime_error("SOFIE GatherND: data rank must be >= 1"); @@ -71,6 +74,11 @@ public: throw std::runtime_error("SOFIE GatherND: indices rank must be >= 1"); if (b >= std::min(q, r)) throw std::runtime_error("SOFIE GatherND: batch_dims must be < min(q, r)"); + + if (fShapeIndices.back().isParam) + throw std::runtime_error("SOFIE GatherND: the last indices dimension (index tuple length) " + "must be static - a dynamic index-tuple length is not supported"); + size_t last_idx_dim = fShapeIndices.back().dim; if (last_idx_dim > r - b) throw std::runtime_error("SOFIE GatherND: indices_shape[-1] must be <= r - batch_dims"); @@ -94,10 +102,10 @@ public: fType = ConvertTypeToString(model.GetTensorType(fNData)); if (model.Verbose()) - std::cout << "GatherND: data " << ConvertShapeToString(fShapeData) - << " indices " << ConvertShapeToString(fShapeIndices) + std::cout << "GatherND: data " << ConvertDimShapeToString(fShapeData) + << " indices " << ConvertDimShapeToString(fShapeIndices) << " batch_dims=" << fBatchDims - << " -> " << fNY << " " << ConvertShapeToString(fShapeY) << std::endl; + << " -> " << fNY << " " << ConvertDimShapeToString(fShapeY) << std::endl; } std::string Generate(std::string opName) override { @@ -108,39 +116,39 @@ public: size_t r = fShapeData.size(); size_t q = fShapeIndices.size(); size_t b = static_cast(fBatchDims); - size_t last_idx_dim = fShapeIndices.back(); + size_t last_idx_dim = fShapeIndices.back().dim; auto stridesData = UTILITY::ComputeStrideFromShape(fShapeData); auto stridesIndices = UTILITY::ComputeStrideFromShape(fShapeIndices); auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); - size_t totalOutput = ConvertShapeToLength(fShapeY); + std::string totalOutput = ConvertDimShapeToLength(fShapeY); std::stringstream out; out << SP << "//--------- GatherND operator " << opName << "\n"; - out << SP << "for (size_t out_idx = 0; out_idx < " << totalOutput << "; out_idx++) {\n"; + out << SP << "for (size_t out_idx = 0; out_idx < static_cast(" << totalOutput << "); out_idx++) {\n"; out << SP << SP << "size_t rem = out_idx;\n"; size_t Dy = fShapeY.size(); for (size_t d = 0; d < Dy; ++d) { - out << SP << SP << "size_t oy_" << d << " = rem / " << stridesY[d] << ";\n"; - out << SP << SP << "rem %= " << stridesY[d] << ";\n"; + out << SP << SP << "size_t oy_" << d << " = rem / " << stridesY[d].GetVal() << ";\n"; + out << SP << SP << "rem %= " << stridesY[d].GetVal() << ";\n"; } out << SP << SP << "size_t idx_base = 0;\n"; for (size_t i = 0; i < b; ++i) - out << SP << SP << "idx_base += oy_" << i << " * " << stridesIndices[i] << ";\n"; + out << SP << SP << "idx_base += oy_" << i << " * " << stridesIndices[i].GetVal() << ";\n"; for (size_t i = b; i + 1 < q; ++i) - out << SP << SP << "idx_base += oy_" << i << " * " << stridesIndices[i] << ";\n"; + out << SP << SP << "idx_base += oy_" << i << " * " << stridesIndices[i].GetVal() << ";\n"; out << SP << SP << "size_t data_idx = 0;\n"; for (size_t i = 0; i < b; ++i) - out << SP << SP << "data_idx += oy_" << i << " * " << stridesData[i] << ";\n"; + out << SP << SP << "data_idx += oy_" << i << " * " << stridesData[i].GetVal() << ";\n"; out << SP << SP << "for (size_t k = 0; k < " << last_idx_dim << "; k++) {\n"; out << SP << SP << SP << "int64_t idx_val = tensor_" << fNIndices - << "[idx_base + k * " << stridesIndices[q - 1] << "];\n"; + << "[idx_base + k * " << stridesIndices[q - 1].GetVal() << "];\n"; out << SP << SP << SP << "if (idx_val < 0) idx_val += " << "static_cast(tensor_" << fNData << "_shape[" << b << " + k]);\n"; out << SP << SP << SP << "data_idx += static_cast(idx_val) * " << "data_stride_b_plus_k_" << opName << "[k];\n"; @@ -151,7 +159,7 @@ public: size_t y_trailing_start = b + (q - b - 1); for (size_t i = b + last_idx_dim; i < r; ++i) { size_t oy_dim = y_trailing_start + (i - (b + last_idx_dim)); - out << SP << SP << "data_idx += oy_" << oy_dim << " * " << stridesData[i] << ";\n"; + out << SP << SP << "data_idx += oy_" << oy_dim << " * " << stridesData[i].GetVal() << ";\n"; } out << SP << SP << "tensor_" << fNY << "[out_idx] = tensor_" << fNData << "[data_idx];\n"; @@ -160,7 +168,7 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE GatherND called to Generate without being initialized first"); @@ -168,14 +176,13 @@ public: size_t r = fShapeData.size(); size_t q = fShapeIndices.size(); size_t b = static_cast(fBatchDims); - size_t last_idx_dim = fShapeIndices.back(); + size_t last_idx_dim = fShapeIndices.back().dim; auto stridesData = UTILITY::ComputeStrideFromShape(fShapeData); auto stridesIndices = UTILITY::ComputeStrideFromShape(fShapeIndices); auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY); size_t Dy = fShapeY.size(); - size_t totalOutput = ConvertShapeToLength(fShapeY); std::string kname = "GatherNDKernel_" + opName; @@ -188,6 +195,8 @@ public: op += SP + SP + SP + "T const* __restrict__ data,\n"; op += SP + SP + SP + "int64_t const* __restrict__ indices,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -198,8 +207,8 @@ public: for (size_t d = 0; d < Dy; ++d) { op += SP + SP + SP + SP + "std::size_t const oy_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(stridesY[d]) + "u) % " - + std::to_string(fShapeY[d]) + "u;\n"; + + " = (elem_idx / " + sz(stridesY[d].GetVal()) + ") % " + + sz(fShapeY[d].GetVal()) + ";\n"; } op += "\n"; @@ -210,7 +219,7 @@ public: for (size_t i = 0; i < q - 1; ++i) { op += SP + SP + SP + SP + SP + (first ? "" : "+ ") - + "oy_" + std::to_string(i) + " * " + std::to_string(stridesIndices[i]) + "u\n"; + + "oy_" + std::to_string(i) + " * " + sz(stridesIndices[i].GetVal()) + "\n"; first = false; } if (first) op += SP + SP + SP + SP + SP + "0u\n"; // q==1: scalar index tuple @@ -221,7 +230,7 @@ public: for (size_t i = 0; i < b; ++i) { op += SP + SP + SP + SP + SP + (first ? "" : "+ ") - + "oy_" + std::to_string(i) + " * " + std::to_string(stridesData[i]) + "u\n"; + + "oy_" + std::to_string(i) + " * " + sz(stridesData[i].GetVal()) + "\n"; first = false; } if (first) op += SP + SP + SP + SP + SP + "0u\n"; @@ -237,10 +246,10 @@ public: + std::to_string(idx_offset) + "u];\n"; op += SP + SP + SP + SP + SP + "if (idx_val < 0) idx_val += " - + std::to_string(fShapeData[data_axis]) + ";\n"; + + fShapeData[data_axis].GetVal() + ";\n"; op += SP + SP + SP + SP + SP + "data_idx += static_cast(idx_val) * " - + std::to_string(stridesData[data_axis]) + "u;\n"; + + sz(stridesData[data_axis].GetVal()) + ";\n"; op += SP + SP + SP + SP + "}\n"; } op += "\n"; @@ -250,7 +259,7 @@ public: size_t oy_dim = y_trailing_start + (i - (b + last_idx_dim)); op += SP + SP + SP + SP + "data_idx += oy_" + std::to_string(oy_dim) - + " * " + std::to_string(stridesData[i]) + "u;\n"; + + " * " + sz(stridesData[i].GetVal()) + ";\n"; } op += "\n"; @@ -268,25 +277,27 @@ public: return SP + kname + " gatherNDKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE GatherND called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeY); + std::string totalElements = ConvertDimShapeToLength(fShapeY); std::string kname = "gatherNDKernel_" + opName; std::stringstream out; out << "\n//------ GATHERND_GPU_ALPAKA\n"; out << SP << "auto const elementsPerThread_" << opName << " = Vec::all(static_cast(1));\n"; - out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{" << totalElements << "});\n"; + out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{static_cast(" << totalElements << ")});\n"; out << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; out << SP << "alpaka::exec(queue, workDiv_" << opName << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNData << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNIndices << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; out << SP <<"alpaka::wait(queue);\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_Gemm.hxx b/core/inc/SOFIE/ROperator_Gemm.hxx index 6438d611..b90f005f 100644 --- a/core/inc/SOFIE/ROperator_Gemm.hxx +++ b/core/inc/SOFIE/ROperator_Gemm.hxx @@ -220,11 +220,10 @@ namespace SOFIE{ // try low rank factorization of weight matrix B, if enabled on the model. // restrict to the plain 2D case (no MatMul batch stacking); both transB == 0 - // and transB == 1 (the latter is how e.g. torch.nn.Linear layers are - // universally exported to ONNX Gemm) are supported: the generated low-rank + // and transB == 1 are supported: the generated low-rank // Bin/Bout tensors are always stored logically untransposed (k x rank) and - // (rank x n) regardless of how the original B was stored, so Generate() / - // Generate_GPU_ALPAKA() never need to special-case transB for the chained + // (rank x n) regardless of how the original B was stored, so Generate methods + // never need to special-case transB for the chained // low-rank calls -- when transB == 1 we just transpose the weight data once // here, before factorizing, to get it into that same logical orientation. if (model.LowRankFactorize() && !fIsDynamic && !appendOne && @@ -340,9 +339,6 @@ namespace SOFIE{ std::string Generate(std::string opName) override { opName = "op_" + opName; - // if (fShapeA.empty() || fShapeB.empty() || fShapeY.empty() || (fNC != "" && fShapeC.empty())) { - // throw std::runtime_error("SOFIE Gemm Op called to Generate without being initialized first"); - // } std::stringstream out; out << "\n//--------- Gemm " << opName << " " << ConvertDimShapeToString(fShapeA) << " * " << ConvertDimShapeToString(fShapeB) << " -> " << ConvertDimShapeToString(fShapeY) << "\n"; @@ -628,13 +624,6 @@ namespace SOFIE{ } std::stringstream out; out << "\n//--------- Gemm_GPU_ALPAKA\n"; - // Note: alpaka::wait(queue) intentionally removed here. - // Operations are enqueued asynchronously on the Alpaka queue's CUDA - // stream. Synchronisation only happens once per inference at the - // alpaka::wait(queue) call in _infer_impl's tail and at the - // cudaDeviceSynchronize in the benchmark harness. Adding a wait - // before every GEMM stalls the CPU<->GPU pipeline and is the primary - // cause of SOFIE being slower than ONNXRuntime. out << SP << "char " << opName << "_transA = " << (fAttrTransA ? "\'t\'" : "\'n\'") << ";\n"; out << SP << "char " << opName << "_transB = " << (fAttrTransB ? "\'t\'" : "\'n\'") << ";\n"; // need to consider case A and B have dim > 2 (for MatMul) @@ -693,29 +682,23 @@ namespace SOFIE{ // exclude case where we have only 1's in the additional dims bool doStackMul = dimY > 2 && ( fIsDynamic || std::stoi(lengthExtra) > 1); - // Compute per-iteration strides for each buffer when stacking. - // m/n/k are std::string from Dim::GetVal(); stoi() is safe for static shapes. + // B is a shared weight (broadcast over the stacked/batch dimension) when all its + // leading dims (beyond the 2 matrix dims) are statically known to be 1 + bool bLeadingDimsAllOne = true; + for (int64_t i = 0; i < dimB - 2; i++) { + if (fShapeB[i].isParam || fShapeB[i].dim != 1) { bLeadingDimsAllOne = false; break; } + } + + std::string strideAExpr = "(" + m + ") * (" + k + ")"; + std::string strideBExpr = bLeadingDimsAllOne ? "" : ("(" + n + ") * (" + k + ")"); + std::string strideYExpr = "(" + m + ") * (" + n + ")"; + std::string strideCExpr = (!fNC.empty() && !fBroadcastBias) ? lengthGemm : ""; + size_t strideA = 0, strideB = 0, strideY = 0, strideC = 0; - // GPU optimisation flags (static shapes only): - // batchCollapseB — strideB==0: B is the shared weight, so replace the N-iteration - // loop with a single cuBLASLt GEMM whose batch dimension is - // folded into the "n_sofie" parameter (n_sofie = m_onnx * N). - // This turns 30 per-token GEMM launches into one kernel call. - // useSBatched — both strides non-zero AND no bias: use cublasSgemmStridedBatched - // so the GPU driver schedules all N GEMMs in one call. - // (Bias epilogue is not available on the strided-batched path, so - // this only applies to pure MatMul ops such as softmax(QK^T)·V.) bool batchCollapseB = false; bool useSBatched = false; if (doStackMul && !fIsDynamic) { strideA = static_cast(std::stoi(m)) * static_cast(std::stoi(k)); - // B is a shared weight (broadcast over the stacked/batch dimension) when all its - // leading dims (beyond the 2 matrix dims) are 1. In that case strideB must be 0 - // so every iteration reads from the same B slice — not i * n*k (which goes OOB). - bool bLeadingDimsAllOne = true; - for (int64_t i = 0; i < dimB - 2; i++) { - if (fShapeB[i].dim != 1) { bLeadingDimsAllOne = false; break; } - } strideB = bLeadingDimsAllOne ? 0 : static_cast(std::stoi(n)) * static_cast(std::stoi(k)); strideY = static_cast(std::stoi(m)) * static_cast(std::stoi(n)); @@ -725,26 +708,20 @@ namespace SOFIE{ useSBatched = !batchCollapseB && fNC.empty(); } - // Emit the loop only for the serial fallback path (dynamic shapes, or static - // shapes where both A and B vary per iteration AND a bias epilogue is needed). bool useSerialLoop = doStackMul && !batchCollapseB && !useSBatched; - if (useSerialLoop || (doStackMul && fIsDynamic)) { + if (useSerialLoop) { out << SP << "size_t " << opName << "_yoffset = 0;\n"; out << SP << "for (int i = 0; i < " << lengthExtra << "; i++){\n"; } - // Use getPtrNative() for all args so the raw-pointer overload is selected - // regardless of whether each buffer is a BufXxx or ViewPlainPtr. - // For the loop path, add per-iteration offsets; for the collapsed/batched - // paths, use base pointers (the whole contiguous tensor is processed at once). - std::string pA = "alpaka::getPtrNative(deviceBuf_" + fNA + ")"; - std::string pB = "alpaka::getPtrNative(deviceBuf_" + fNB + ")"; + std::string pA = "static_cast(alpaka::getPtrNative(deviceBuf_" + fNA + "))"; + std::string pB = "static_cast(alpaka::getPtrNative(deviceBuf_" + fNB + "))"; std::string pY = "alpaka::getPtrNative(deviceBuf_" + fNY + ")"; - if (useSerialLoop && !fIsDynamic) { - pA += " + i * " + std::to_string(strideA); - if (strideB > 0) pB += " + i * " + std::to_string(strideB); - // strideB == 0: B is a shared weight, pointer stays at base - pY += " + i * " + std::to_string(strideY); + if (useSerialLoop) { + pA += " + i * (" + strideAExpr + ")"; + if (!strideBExpr.empty()) pB += " + i * (" + strideBExpr + ")"; + // strideB shared (empty): B is a shared weight, pointer stays at base + pY += " + i * (" + strideYExpr + ")"; } if (fLowRank) { @@ -770,7 +747,7 @@ namespace SOFIE{ out << SP << "blas.matmul(" << "'n', " << opName << "_transA, " << fLowRankRank << ", " << opName << "_m, " << opName << "_k, " - << opName << "_alpha, alpaka::getPtrNative(deviceBuf_" << fLowRankInName << "), " << pA + << opName << "_alpha, static_cast(alpaka::getPtrNative(deviceBuf_" << fLowRankInName << ")), " << pA << ", 0.f, tensor_" << opName << "_lrtmp);\n"; // step 2: Y = 1 * tmp * Bout (+ bias). tmp and Bout are both untransposed @@ -780,37 +757,24 @@ namespace SOFIE{ const char *callFn = (fActivation == EActivationType::RELU) ? "blas.gemmrelu(" : "blas.gemm("; out << SP << callFn << "'n', 'n', " << opName << "_n, " << opName << "_m, " << fLowRankRank << ", " - << "1.f, alpaka::getPtrNative(deviceBuf_" << fLowRankOutName << "), tensor_" << opName << "_lrtmp, " + << "1.f, static_cast(alpaka::getPtrNative(deviceBuf_" << fLowRankOutName << ")), static_cast(tensor_" << opName << "_lrtmp), " << opName << "_beta, " << pC << ", " << pY << ");\n"; } else { out << SP << "blas.matmul('n', 'n', " << opName << "_n, " << opName << "_m, " << fLowRankRank << ", " - << "1.f, alpaka::getPtrNative(deviceBuf_" << fLowRankOutName << "), tensor_" << opName << "_lrtmp, " + << "1.f, static_cast(alpaka::getPtrNative(deviceBuf_" << fLowRankOutName << ")), static_cast(tensor_" << opName << "_lrtmp), " << opName << "_beta, " << pY << ");\n"; } } else if (useSBatched) { - // ---------------------------------------------------------------- - // gemmStridedBatched: both A and B vary per batch (e.g. per attention - // head), and there is no bias. Uses cublasSgemmStridedBatched via - // the legacy cuBLAS handle so all N GEMMs are issued in one driver call. - // - // sofieBLAS convention (column-major transpose trick): - // transa_sofie = transB_onnx, transb_sofie = transA_onnx - // m_sofie = n_onnx, n_sofie = m_onnx - // A_sofie = fNB, B_sofie = fNA - // lda = m_sofie (leading dim of A when transA_sofie='n') - // ldb = k (leading dim of B when transB_sofie='n') - // ldc = m_sofie (leading dim of C) - // ---------------------------------------------------------------- - size_t m_sofie = static_cast(std::stoi(n)); // ONNX n - size_t n_sofie = static_cast(std::stoi(m)); // ONNX m + size_t m_sofie = static_cast(std::stoi(n)); + size_t n_sofie = static_cast(std::stoi(m)); size_t k_val = static_cast(std::stoi(k)); - size_t lda = m_sofie; // transA_sofie='n' - size_t ldb = k_val; // transB_sofie='n' + size_t lda = m_sofie; + size_t ldb = k_val; size_t ldc = m_sofie; size_t sA = m_sofie * k_val; // stride per batch for fNB - size_t sB = k_val * n_sofie; // stride per batch for fNA (= strideA_onnx) - size_t sC = m_sofie * n_sofie; // stride per batch for fNY (= strideY) + size_t sB = k_val * n_sofie; // stride per batch for fNA + size_t sC = m_sofie * n_sofie; // stride per batch for fNY size_t batchCount = static_cast(std::stoi(lengthExtra)); out << SP << "blas.gemmStridedBatched(" << opName << "_transB, " << opName << "_transA, " @@ -825,26 +789,13 @@ namespace SOFIE{ << ldc << ", " << sC << ", " << batchCount << ");\n"; } else if (!fNC.empty()) { - // ---------------------------------------------------------------- - // GEMM with bias: Y = alpha * op(A) * op(B) + bias - // cuBLAS is column-major so we swap A↔B and transA↔transB - // (row-major C=A*B ↔ col-major C^T = B^T * A^T). - // The epilogue fuses the bias-add (and optional ReLU/GELU) in the - // same kernel, avoiding a separate element-wise pass. - // - // For batch-collapse (batchCollapseB), use m*batchCount so that all - // tokens are processed in a single cuBLASLt kernel launch instead of N. - // The bias vector is broadcast across all columns by the epilogue. - // ---------------------------------------------------------------- std::string call_m = batchCollapseB ? std::to_string(static_cast(std::stoi(m)) * static_cast(std::stoi(lengthExtra))) : (opName + "_m"); std::string pC = "alpaka::getPtrNative(deviceBuf_" + fNC + ")"; - if (useSerialLoop && !fIsDynamic) { - if (!fBroadcastBias && strideC > 0) - pC += " + i * " + std::to_string(strideC); - } + if (useSerialLoop && !strideCExpr.empty()) + pC += " + i * (" + strideCExpr + ")"; if (fActivation == EActivationType::RELU) { out << SP << "blas.gemmrelu(" << opName << "_transB, " << opName << "_transA, " @@ -863,12 +814,7 @@ namespace SOFIE{ } else { // ---------------------------------------------------------------- // Pure MatMul (no bias): Y = alpha * op(A) * op(B) - // This covers: - // • Scaled Dot-Product Attention: softmax(QK^T/√d) @ V - // • Any other no-bias matrix multiplication - // Previously this branch emitted nothing (empty loop body), which - // caused the attention output to be silently uninitialized. - // For batch-collapse, use m*batchCount for the same reason as above. + // This covers Scaled Dot-Product Attention and other no-bias matrix multiplication // ---------------------------------------------------------------- std::string call_m = batchCollapseB ? std::to_string(static_cast(std::stoi(m)) * static_cast(std::stoi(lengthExtra))) @@ -882,14 +828,12 @@ namespace SOFIE{ << opName << "_beta, " << pY << ");\n"; } - if (useSerialLoop || (doStackMul && fIsDynamic)) { + if (useSerialLoop) { out << SP << "}\n"; // end of loop on the stacked multiplication } // GEMM+LeakyReLU fusion (GPU): cuBLASLt has no native LeakyReLU epilogue, // so we emit a cheap in-place ALPAKA kernel immediately after the GEMM. - // This avoids allocating a separate intermediate buffer and saves one - // GPU kernel launch compared to a standalone LeakyReLU operator. if (fActivation == EActivationType::LEAKYRELU) { std::string numElem = ConvertDimShapeToLength(fShapeY); out << SP << "//--- GEMM+LeakyReLU in-place fusion\n"; @@ -948,13 +892,8 @@ namespace SOFIE{ auto ldb = (fAttrTransB ? k : n); auto ldc = n; std::string transFlags = std::string(fAttrTransB ? "'t'" : "'n'") + ", " + (fAttrTransA ? "'t'" : "'n'"); - - // For stacked (batched) GEMMs on static shapes, return the layout that - // matches the actual call emitted by Generate_GPU_ALPAKA: - // - batch-collapse (strideB==0): single GEMM with n_sofie = m_onnx * batchCount - // → register the batched layout - // - gemmStridedBatched (both strides non-zero, no bias): uses legacy cuBLAS, - // no cuBLASLt layout needed → return "" + std::string epilogue = fNC.empty() ? "Epilogue::Default" + : (fActivation == EActivationType::RELU ? "Epilogue::ReluBias" : "Epilogue::Bias"); if (dimY > 2 && !fIsDynamic) { std::vector sExtra; for (int64_t i = 0; i < dimY - 2; i++) sExtra.push_back(fShapeY[i]); @@ -967,21 +906,16 @@ namespace SOFIE{ if (bLeadingDimsAllOne) { // batch-collapse: register layout for the full-batch GEMM auto m_batched = std::to_string(std::stoi(m) * std::stoi(lengthExtra)); - return n+", "+m_batched+", "+k+", "+ldb+", "+lda+", "+ldc+", "+transFlags; + return n+", "+m_batched+", "+k+", "+ldb+", "+lda+", "+ldc+", "+transFlags+", "+epilogue; } else if (fNC.empty()) { - // gemmStridedBatched: legacy cuBLAS, no cuBLASLt layout needed return ""; } - // else: serial loop with bias — fall through to per-iteration layout } } - return n+", "+m+", "+k+", "+ldb+", "+lda+", "+ldc+", "+transFlags; + return n+", "+m+", "+k+", "+ldb+", "+lda+", "+ldc+", "+transFlags+", "+epilogue; } - // low rank factorized Gemm issues two chained GEMM calls (see Generate_GPU_ALPAKA) - // of different shapes, so two cuBLASLt layouts need to be pre-registered instead - // of the single one GetBlasConfig() computes for the dense case. std::vector GetBlasConfigs() override { if (!fLowRank) return ROperator::GetBlasConfigs(); @@ -992,16 +926,17 @@ namespace SOFIE{ int64_t dimB = fShapeB.size(); auto n = (fAttrTransB ? fShapeB[dimB-2].GetVal() : fShapeB[dimB-1].GetVal()); std::string rankStr = std::to_string(fLowRankRank); - // Bin is always stored logically untransposed (see Initialize()), regardless - // of the original op's transB attribute, so the B-side flag here is always 'n'. std::string transFlags1 = std::string("'n', ") + (fAttrTransA ? "'t'" : "'n'"); auto lda1 = (fAttrTransA ? m : k); + std::string epilogue = fNC.empty() ? "Epilogue::Default" + : (fActivation == EActivationType::RELU ? "Epilogue::ReluBias" : "Epilogue::Bias"); + // step 1: tmp (m x rank) = op(A) * Bin - std::string cfg1 = rankStr+", "+m+", "+k+", "+rankStr+", "+lda1+", "+rankStr+", "+transFlags1; + std::string cfg1 = rankStr+", "+m+", "+k+", "+rankStr+", "+lda1+", "+rankStr+", "+transFlags1+", Epilogue::Default"; - // step 2: Y (m x n) = tmp * Bout (+ bias) -- both operands untransposed - std::string cfg2 = n+", "+m+", "+rankStr+", "+n+", "+rankStr+", "+n+", 'n', 'n'"; + // step 2: Y (m x n) = tmp * Bout (+ bias) + std::string cfg2 = n+", "+m+", "+rankStr+", "+n+", "+rankStr+", "+n+", 'n', 'n', "+epilogue; return {cfg1, cfg2}; } diff --git a/core/inc/SOFIE/ROperator_GroupNorm.hxx b/core/inc/SOFIE/ROperator_GroupNorm.hxx index 30844c9c..7e1abcb4 100644 --- a/core/inc/SOFIE/ROperator_GroupNorm.hxx +++ b/core/inc/SOFIE/ROperator_GroupNorm.hxx @@ -149,12 +149,18 @@ public: return out.str(); } + // Threads cooperating on one (n,g) group. A group's gn_gs*gn_spatial elements are + // contiguous in memory (c = g*gn_gs+ci, address = n*C*spatial + c*spatial + s = + // (row base) + ci*spatial + s, and ci*spatial+s enumerates [0, gn_gs*spatial) + // contiguously) + static constexpr size_t kBlockSize = 256; + std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { opName = "op_" + opName; if (fShapeX.empty()) throw std::runtime_error("SOFIE GroupNorm GPU kernel called without initialization"); - std::string G = std::to_string(fNumGroups); + std::string bs = std::to_string(kBlockSize); std::string eps = std::to_string(fAttrEpsilon); std::string kname = "GroupNormKernel_" + opName; @@ -169,50 +175,64 @@ public: if (!fNBias.empty()) op += SP + SP + SP + "T const* __restrict__ bias,\n"; op += SP + SP + SP + "T* __restrict__ Y,\n"; - op += SP + SP + SP + "std::size_t const gn_N,\n"; op += SP + SP + SP + "std::size_t const gn_C,\n"; op += SP + SP + SP + "std::size_t const gn_spatial,\n"; - op += SP + SP + SP + "std::size_t const gn_G) const {\n\n"; - - op += SP + SP + SP + "auto const tid = alpaka::getIdx(acc)[0];\n"; - op += SP + SP + SP + "std::size_t const total = gn_N * gn_G;\n"; - op += SP + SP + SP + "if (tid >= total) return;\n\n"; - op += SP + SP + SP + "std::size_t const n = tid / gn_G;\n"; - op += SP + SP + SP + "std::size_t const g = tid % gn_G;\n"; - op += SP + SP + SP + "std::size_t const gn_gs = gn_C / gn_G;\n"; + op += SP + SP + SP + "std::size_t const gn_G,\n"; + op += SP + SP + SP + "std::size_t const numRows) const {\n\n"; + + // one block per (n,g) row - a block-shared reduction buffer, reused across the + // mean and variance passes + op += SP + SP + SP + "auto& sred = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto const row = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "auto const tid = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (row >= numRows) return;\n\n"; + + op += SP + SP + SP + "std::size_t const gn_gs = gn_C / gn_G;\n"; + op += SP + SP + SP + "std::size_t const rowLen = gn_gs * gn_spatial;\n"; + op += SP + SP + SP + "std::size_t const g = row % gn_G;\n"; + op += SP + SP + SP + "std::size_t const base = row * rowLen;\n"; op += SP + SP + SP + "T const gn_eps = static_cast(" + eps + ");\n\n"; - op += SP + SP + SP + "// mean\n"; - op += SP + SP + SP + "T mean = static_cast(0);\n"; - op += SP + SP + SP + "for (std::size_t ci = 0; ci < gn_gs; ++ci) {\n"; - op += SP + SP + SP + SP + "std::size_t c = g * gn_gs + ci;\n"; - op += SP + SP + SP + SP + "for (std::size_t s = 0; s < gn_spatial; ++s)\n"; - op += SP + SP + SP + SP + SP + "mean += X[n * gn_C * gn_spatial + c * gn_spatial + s];\n"; - op += SP + SP + SP + "}\n"; - op += SP + SP + SP + "mean /= static_cast(gn_gs * gn_spatial);\n\n"; - - op += SP + SP + SP + "// variance\n"; - op += SP + SP + SP + "T var = static_cast(0);\n"; - op += SP + SP + SP + "for (std::size_t ci = 0; ci < gn_gs; ++ci) {\n"; - op += SP + SP + SP + SP + "std::size_t c = g * gn_gs + ci;\n"; - op += SP + SP + SP + SP + "for (std::size_t s = 0; s < gn_spatial; ++s) {\n"; - op += SP + SP + SP + SP + SP + "T d = X[n * gn_C * gn_spatial + c * gn_spatial + s] - mean;\n"; - op += SP + SP + SP + SP + SP + "var += d * d;\n"; - op += SP + SP + SP + SP + "}\n"; + auto treeReduce = [&](const std::string & ind) { + std::string s; + size_t half = kBlockSize / 2; + while (half > 0) { + s += ind + "if (tid < " + std::to_string(half) + "u) sred[tid] += sred[tid + " + std::to_string(half) + "u];\n"; + s += ind + "alpaka::syncBlockThreads(acc);\n"; + half >>= 1; + } + return s; + }; + + op += SP + SP + SP + "// mean: per-thread partial sum, then block-shared tree reduction\n"; + op += SP + SP + SP + "T sum = static_cast(0);\n"; + op += SP + SP + SP + "for (std::size_t l = tid; l < rowLen; l += " + bs + "u) sum += X[base + l];\n"; + op += SP + SP + SP + "sred[tid] = sum;\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; + op += treeReduce(SP + SP + SP); + op += SP + SP + SP + "T const mean = sred[0] / static_cast(rowLen);\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; // all reads of sred[0] done before sred is reused below + + op += SP + SP + SP + "// variance: same pattern, from deviations around mean\n"; + op += SP + SP + SP + "T devsq = static_cast(0);\n"; + op += SP + SP + SP + "for (std::size_t l = tid; l < rowLen; l += " + bs + "u) {\n"; + op += SP + SP + SP + SP + "T d = X[base + l] - mean;\n"; + op += SP + SP + SP + SP + "devsq += d * d;\n"; op += SP + SP + SP + "}\n"; - op += SP + SP + SP + "var /= static_cast(gn_gs * gn_spatial);\n"; + op += SP + SP + SP + "sred[tid] = devsq;\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; + op += treeReduce(SP + SP + SP); + op += SP + SP + SP + "T const var = sred[0] / static_cast(rowLen);\n"; op += SP + SP + SP + "T const inv_std = static_cast(1) / sqrt(acc, var + gn_eps);\n\n"; op += SP + SP + SP + "// normalize + scale + bias\n"; - op += SP + SP + SP + "for (std::size_t ci = 0; ci < gn_gs; ++ci) {\n"; - op += SP + SP + SP + SP + "std::size_t c = g * gn_gs + ci;\n"; - op += SP + SP + SP + SP + "for (std::size_t s = 0; s < gn_spatial; ++s) {\n"; - op += SP + SP + SP + SP + SP + "std::size_t idx = n * gn_C * gn_spatial + c * gn_spatial + s;\n"; - op += SP + SP + SP + SP + SP + "T v = (X[idx] - mean) * inv_std * scale[c];\n"; + op += SP + SP + SP + "for (std::size_t l = tid; l < rowLen; l += " + bs + "u) {\n"; + op += SP + SP + SP + SP + "std::size_t idx = base + l;\n"; + op += SP + SP + SP + SP + "std::size_t c = g * gn_gs + l / gn_spatial;\n"; + op += SP + SP + SP + SP + "T v = (X[idx] - mean) * inv_std * scale[c];\n"; if (!fNBias.empty()) - op += SP + SP + SP + SP + SP + "v += bias[c];\n"; - op += SP + SP + SP + SP + SP + "Y[idx] = v;\n"; - op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + "v += bias[c];\n"; + op += SP + SP + SP + SP + "Y[idx] = v;\n"; op += SP + SP + SP + "}\n"; op += SP + SP + "}\n"; @@ -232,23 +252,24 @@ public: throw std::runtime_error("SOFIE GroupNorm GPU dispatch called without initialization"); std::string G = std::to_string(fNumGroups); + std::string numRows = "static_cast((" + fN + ") * " + G + ")"; std::string args = "alpaka::getPtrNative(deviceBuf_" + fNX + "), " "alpaka::getPtrNative(deviceBuf_" + fNScale + "), "; if (!fNBias.empty()) args += "alpaka::getPtrNative(deviceBuf_" + fNBias + "), "; args += "alpaka::getPtrNative(deviceBuf_" + fNY + "), " - "static_cast(" + fN + "), " "static_cast(" + fC + "), " "static_cast(" + fSpatial + "), " - "static_cast(" + G + ")"; - - std::string totalWork = "static_cast((" + fN + ") * " + G + ")"; + "static_cast(" + G + "), " + + numRows; std::stringstream out; out << "\n//------ GROUPNORM_GPU_ALPAKA\n"; - out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{" << totalWork << "});\n"; - out << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; + out << SP << "alpaka::WorkDivMembers workDiv_" << opName << "(\n"; + out << SP << SP << "Vec::all(" << numRows << "),\n"; + out << SP << SP << "Vec::all(Idx{" << kBlockSize << "u}),\n"; + out << SP << SP << "Vec::all(Idx{1u}));\n"; out << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName << ", groupNormKernel_" << opName << ", " << args << ");\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; diff --git a/core/inc/SOFIE/ROperator_Identity.hxx b/core/inc/SOFIE/ROperator_Identity.hxx index 43688cf6..960c93b7 100644 --- a/core/inc/SOFIE/ROperator_Identity.hxx +++ b/core/inc/SOFIE/ROperator_Identity.hxx @@ -19,7 +19,7 @@ private: bool fIsInputInitialized = false; std::string fNX; std::string fNY; - std::vector fShape; + std::vector fShape; public: ROperator_Identity(){} @@ -43,14 +43,14 @@ public: if (model.CheckIfTensorAlreadyExist(fNX) == false){ throw std::runtime_error("SOFIE Identity Op Input Tensor is not found in model"); } - fShape = model.GetTensorShape(fNX); + fShape = model.GetDimTensorShape(fNX); if (model.IsInitializedTensor(fNX)) { // we need to check if is a weight (initialized) or a constant tensor // in the first case we need to create a constant tensor with the output, in teh second we // need to generate the identy code in the GenerateInitCode if (model.IsConstantTensor(fNX)) { auto inputData = static_cast(model.GetInitializedTensorData(fNX).get()); - model.AddConstantTensor(fNY, fShape, inputData); + model.AddConstantTensor(fNY, ConvertShapeToInt(fShape), inputData); fIsOutputConstant = true; } else { fIsInputInitialized = true; @@ -87,8 +87,6 @@ public: } std::string GenerateInitCode_GPU_ALPAKA() override { - // For initialized (weight) tensors: the device buffer for X is already populated by - // MoveInitializedTensorsToBuffers_ALPAKA(); copy it into the Y device buffer. if (!fIsInputInitialized) return ""; std::stringstream out; out << "\n//------ IDENTITY (init)\n"; @@ -105,9 +103,7 @@ public: } std::stringstream out; out << "\n//------ IDENTITY\n"; - // Device buffers cannot simply be aliased; perform an explicit device-to-device copy. out << SP << "alpaka::memcpy(queue, deviceBuf_" << fNY << ", deviceBuf_" << fNX << ");\n"; - out << SP << "alpaka::wait(queue);\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_LayerNormalization.hxx b/core/inc/SOFIE/ROperator_LayerNormalization.hxx index dd78bae6..da73ca92 100644 --- a/core/inc/SOFIE/ROperator_LayerNormalization.hxx +++ b/core/inc/SOFIE/ROperator_LayerNormalization.hxx @@ -322,14 +322,14 @@ public: throw std::runtime_error("TMVA::SOFIE LayerNormalization called to Generate without being initialized first"); // ----------------------------------------------------------------------- - // Parallel block-per-row strategy (for static normalizedLength ≤ 1024): + // Parallel block-per-row strategy (for any static normalizedLength): // • One block per row (axes element). - // • blockSize = next power-of-2 ≥ normalizedLength, capped at 1024. - // • Each thread loads one element, two shared-memory tree reductions - // compute mean then variance; final pass normalises in parallel. - // This replaces the previous single-thread-per-row serial scan. - // For dynamic shapes or normalizedLength > 1024, fall back to the original - // serial kernel (one thread per row, explicit loops). + // • blockSize = next power-of-2 >= normalizedLength, capped at 256 + // • Each thread handles itemsPerThread = ceil(normalizedLength/blockSize) + // elements, strided by blockSize, accumulating a local partial sum + // before the shared-memory tree reduction. + // Two shared-memory tree reductions compute mean then variance; a final + // pass normalises in parallel. // ----------------------------------------------------------------------- // Determine whether we can use the parallel path @@ -337,14 +337,21 @@ public: bool canParallel = false; try { normLenVal = std::stoul(fNormalizedLength); - canParallel = (normLenVal > 0 && normLenVal <= 1024); + canParallel = (normLenVal > 0); } catch (...) {} - // Compute blockSize = next power-of-2 >= normLenVal + // blockSize = next power-of-2 >= normLenVal, capped well under the 1024 + // hardware max: each thread carries a small vals[]/idxs[] array across + // passes, and a full 1024-thread block leaves very little register headroom + // once the kernel is compiled without heavy optimization, so a + // smaller, safer default is used and itemsPerThread absorbs the rest of a + // large normalizedLength. size_t blockSize = 1; if (canParallel) { - while (blockSize < normLenVal) blockSize <<= 1; + while (blockSize < normLenVal && blockSize < 256) blockSize <<= 1; } + // number of strided elements each thread accumulates within its row + size_t itemsPerThread = canParallel ? (normLenVal + blockSize - 1) / blockSize : 1; // Each thread handles one "row" — one element of the axes dims [0..axis) // and iterates over all normalized dims [axis..size) @@ -385,18 +392,77 @@ public: op += SP + SP + SP + "T* __restrict__ Y,\n"; op += SP + SP + SP + "std::size_t const axesLength) const {\n\n"; + // Sum of "axis_i * stride_i" over outer axes where the scale/bias tensor + // actually varies (dim != 1), joined with " +\n" (collecting terms first + // avoids mis-placing the separator once a base has more than one term). + // Shared by both the parallel and serial kernel bodies below. + auto axisBase = [&](const std::vector &shapeVec, const std::vector &strideVec, + const std::string &ind) -> std::string { + std::vector terms; + for (size_t i = 0; i < fAxis; ++i) { + if (shapeVec[i].dim != 1) + terms.push_back("axis_" + std::to_string(i) + " * " + strideVec[i].GetVal() + "u"); + } + if (terms.empty()) return ind + "0u"; + std::string s; + for (size_t t = 0; t < terms.size(); ++t) + s += ind + terms[t] + (t + 1 < terms.size() ? " +\n" : ""); + return s; + }; + if (canParallel) { // --------------------------------------------------------------- - // PARALLEL PATH: one block per row, blockSize threads per block. - // Each thread handles one element in the normalised dimension. + // PARALLEL PATH: one block per row, blockSize threads per block, + // each thread strided over itemsPerThread elements of the row. // Two shared-memory tree reductions compute mean then variance. // --------------------------------------------------------------- std::string bs = std::to_string(blockSize); std::string nl = fNormalizedLength; // e.g. "64" std::string eps = std::to_string(fAttrEpsilon); + std::string ipt = std::to_string(itemsPerThread); + + // Build the norm-dim strides (strides within the flattened normalised space) + auto normShape = fNormalizedShape; // dims [fAxis .. fSize-1] + auto normInner = UTILITY::ComputeStrideFromShape(normShape); + + // Emits norm_offset/s_norm_offset[/b_norm_offset] for a given linear index + // into the flattened normalised space (idxVar), reusing the same + // decomposition logic regardless of which loop calls it. + auto emitOffsets = [&](const std::string &idxVar, const std::string &ind) -> std::string { + std::string s; + if (fSize - fAxis == 1) { + s += ind + "std::size_t const norm_offset = " + idxVar + " * " + strides[fAxis].GetVal() + "u;\n"; + s += ind + "std::size_t const s_norm_offset = "; + s += (fShapeScale[fAxis].dim != 1) ? (idxVar + " * " + scaleStrides[fAxis].GetVal() + "u;\n") : "0u;\n"; + if (!fNB.empty()) { + s += ind + "std::size_t const b_norm_offset = "; + s += (fShapeB[fAxis].dim != 1) ? (idxVar + " * " + biasStrides[fAxis].GetVal() + "u;\n") : "0u;\n"; + } + } else { + s += ind + "std::size_t norm_offset = 0u, s_norm_offset = 0u"; + if (!fNB.empty()) s += ", b_norm_offset = 0u"; + s += ";\n"; + s += ind + "{\n"; + s += ind + SP + "std::size_t norm_rem = " + idxVar + ";\n"; + for (size_t j = fAxis; j < fSize; ++j) { + size_t ji = j - fAxis; + s += ind + SP + "{ std::size_t nj = norm_rem / " + normInner[ji].GetVal() + "u;" + + " norm_rem %= " + normInner[ji].GetVal() + "u;" + + " norm_offset += nj * " + strides[j].GetVal() + "u;"; + if (fShapeScale[j].dim != 1) + s += " s_norm_offset += nj * " + scaleStrides[j].GetVal() + "u;"; + if (!fNB.empty() && fShapeB[j].dim != 1) + s += " b_norm_offset += nj * " + biasStrides[j].GetVal() + "u;"; + s += " }\n"; + } + s += ind + "}\n"; + } + return s; + }; op += SP + SP + SP + "// Block-parallel LayerNorm: one block per row, " - + bs + " threads per block, " + nl + " active.\n"; + + bs + " threads per block, " + ipt + " item(s) per thread, " + + nl + " active elements.\n"; op += SP + SP + SP + "auto& shmem = alpaka::declareSharedVar(acc);\n"; op += SP + SP + SP + "auto const row = alpaka::getIdx(acc)[0];\n"; op += SP + SP + SP + "auto const tid = alpaka::getIdx(acc)[0];\n"; @@ -424,53 +490,27 @@ public: } } - // Map thread id → index within normalised dims. - // For each normalised dim j, the "within-norm" stride is the product of - // dimensions after it: normInnerStrides[j-fAxis] computed at code-gen time. - // Then: norm_offset = sum_j( (tid / normInnerStride[j]) % dim[j] * stride[j] ) - // For the common 1D normalised case this simplifies to: norm_offset = tid * stride[fAxis] - - // Build the norm-dim strides (strides within the flattened normalised space) - auto normShape = fNormalizedShape; // dims [fAxis .. fSize-1] - auto normInner = UTILITY::ComputeStrideFromShape(normShape); - - op += SP + SP + SP + "bool const in_range = (tid < " + nl + "u);\n"; - op += SP + SP + SP + "std::size_t norm_offset = 0u;\n"; - op += SP + SP + SP + "std::size_t s_norm_offset = 0u;\n"; - if (!fNB.empty()) - op += SP + SP + SP + "std::size_t b_norm_offset = 0u;\n"; - op += SP + SP + SP + "if (in_range) {\n"; - - if (fSize - fAxis == 1) { - // Single normalised dim — simplest case - op += SP + SP + SP + SP + "norm_offset = tid * " + strides[fAxis].GetVal() + "u;\n"; - if (fShapeScale[fAxis].dim != 1) - op += SP + SP + SP + SP + "s_norm_offset = tid * " + scaleStrides[fAxis].GetVal() + "u;\n"; - if (!fNB.empty() && fShapeB[fAxis].dim != 1) - op += SP + SP + SP + SP + "b_norm_offset = tid * " + biasStrides[fAxis].GetVal() + "u;\n"; - } else { - // Multi-dim normalised space - op += SP + SP + SP + SP + "std::size_t norm_rem = tid;\n"; - for (size_t j = fAxis; j < fSize; ++j) { - size_t ji = j - fAxis; - op += SP + SP + SP + SP + "{ std::size_t nj = norm_rem / " + normInner[ji].GetVal() + "u;" - + " norm_rem %= " + normInner[ji].GetVal() + "u;" - + " norm_offset += nj * " + strides[j].GetVal() + "u;"; - if (fShapeScale[j].dim != 1) - op += " s_norm_offset += nj * " + scaleStrides[j].GetVal() + "u;"; - if (!fNB.empty() && fShapeB[j].dim != 1) - op += " b_norm_offset += nj * " + biasStrides[j].GetVal() + "u;"; - op += " }\n"; - } - } + // --- Pass 1: each thread strided-accumulates its items, caching the + // loaded values and their flat indices for reuse in passes 2 and 3 --- + op += SP + SP + SP + "constexpr std::size_t itemsPerThread = " + ipt + "u;\n"; + op += SP + SP + SP + "T vals[itemsPerThread];\n"; + op += SP + SP + SP + "std::size_t idxs[itemsPerThread];\n"; + op += SP + SP + SP + "T local_sum = static_cast(0);\n"; + op += SP + SP + SP + "for (std::size_t k = 0; k < itemsPerThread; ++k) {\n"; + op += SP + SP + SP + SP + "std::size_t const j = tid + k * " + bs + "u;\n"; + op += SP + SP + SP + SP + "if (j < " + nl + "u) {\n"; + op += emitOffsets("j", SP + SP + SP + SP + SP); + op += SP + SP + SP + SP + SP + "idxs[k] = row_base + norm_offset;\n"; + op += SP + SP + SP + SP + SP + "vals[k] = X[idxs[k]];\n"; + op += SP + SP + SP + SP + SP + "local_sum += vals[k];\n"; + op += SP + SP + SP + SP + "} else {\n"; + op += SP + SP + SP + SP + SP + "idxs[k] = 0u;\n"; + op += SP + SP + SP + SP + SP + "vals[k] = static_cast(0);\n"; + op += SP + SP + SP + SP + "}\n"; op += SP + SP + SP + "}\n\n"; - op += SP + SP + SP + "std::size_t const norm_idx = row_base + norm_offset;\n"; - op += SP + SP + SP + "T const val = in_range ? X[norm_idx] : static_cast(0);\n\n"; - - // --- Pass 1: parallel mean --- op += SP + SP + SP + "// Pass 1: compute mean via shared-memory tree reduction\n"; - op += SP + SP + SP + "shmem[tid] = val;\n"; + op += SP + SP + SP + "shmem[tid] = local_sum;\n"; op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; size_t half = blockSize / 2; while (half > 0) { @@ -481,10 +521,14 @@ public: op += SP + SP + SP + "T const mean = shmem[0] / static_cast(" + nl + ");\n"; op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; - // --- Pass 2: parallel variance --- + // --- Pass 2: parallel variance, reusing the values cached in pass 1 --- op += SP + SP + SP + "// Pass 2: compute variance\n"; - op += SP + SP + SP + "T const diff = val - mean;\n"; - op += SP + SP + SP + "shmem[tid] = in_range ? diff * diff : static_cast(0);\n"; + op += SP + SP + SP + "T local_sq = static_cast(0);\n"; + op += SP + SP + SP + "for (std::size_t k = 0; k < itemsPerThread; ++k) {\n"; + op += SP + SP + SP + SP + "std::size_t const j = tid + k * " + bs + "u;\n"; + op += SP + SP + SP + SP + "if (j < " + nl + "u) { T const diff = vals[k] - mean; local_sq += diff * diff; }\n"; + op += SP + SP + SP + "}\n"; + op += SP + SP + SP + "shmem[tid] = local_sq;\n"; op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; half = blockSize / 2; while (half > 0) { @@ -502,49 +546,29 @@ public: op += SP + SP + SP + "if (tid == 0u) out_invstd[row] = invStdDev;\n"; op += "\n"; - // --- Pass 3: normalise, scale, bias --- + // --- Pass 3: normalise, scale, bias, using the cached values/indices --- op += SP + SP + SP + "// Pass 3: normalize + scale + bias\n"; - op += SP + SP + SP + "if (in_range) {\n"; - // scale base (axis contribution) - op += SP + SP + SP + SP + "std::size_t const scale_base =\n"; - { - bool any = false; - for (size_t i = 0; i < fAxis; ++i) { - if (fShapeScale[i].dim != 1) { - op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) - + " * " + scaleStrides[i].GetVal() + "u"; - if (any) op += " +\n"; - any = true; - } - } - if (!any) op += SP + SP + SP + SP + SP + "0u"; - op += ";\n"; - } - op += SP + SP + SP + SP + "T out_val = scale[scale_base + s_norm_offset] * invStdDev * (val - mean);\n"; - - if (!fNB.empty()) { - op += SP + SP + SP + SP + "std::size_t const bias_base =\n"; - bool any = false; - for (size_t i = 0; i < fAxis; ++i) { - if (fShapeB[i].dim != 1) { - op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) - + " * " + biasStrides[i].GetVal() + "u"; - if (any) op += " +\n"; - any = true; - } - } - if (!any) op += SP + SP + SP + SP + SP + "0u"; - op += ";\n"; - op += SP + SP + SP + SP + "out_val += bias[bias_base + b_norm_offset];\n"; - } + // scale/bias base (axis contribution) is invariant across a thread's + // items, so it is computed once, outside the per-item loop. + op += SP + SP + SP + "std::size_t const scale_base =\n" + axisBase(fShapeScale, scaleStrides, SP + SP + SP + SP) + ";\n"; + if (!fNB.empty()) + op += SP + SP + SP + "std::size_t const bias_base =\n" + axisBase(fShapeB, biasStrides, SP + SP + SP + SP) + ";\n"; - op += SP + SP + SP + SP + "Y[norm_idx] = out_val;\n"; - op += SP + SP + SP + "}\n"; // end in_range + op += SP + SP + SP + "for (std::size_t k = 0; k < itemsPerThread; ++k) {\n"; + op += SP + SP + SP + SP + "std::size_t const j = tid + k * " + bs + "u;\n"; + op += SP + SP + SP + SP + "if (j < " + nl + "u) {\n"; + op += emitOffsets("j", SP + SP + SP + SP + SP); + op += SP + SP + SP + SP + SP + "T out_val = scale[scale_base + s_norm_offset] * invStdDev * (vals[k] - mean);\n"; + if (!fNB.empty()) + op += SP + SP + SP + SP + SP + "out_val += bias[bias_base + b_norm_offset];\n"; + op += SP + SP + SP + SP + SP + "Y[idxs[k]] = out_val;\n"; + op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + "}\n"; } else { // --------------------------------------------------------------- - // SERIAL PATH (dynamic shapes or normalizedLength > 1024): + // SERIAL PATH (only for a dynamic/runtime normalizedLength with no known upper bound): // one thread per row, explicit loops over normalized dims. // --------------------------------------------------------------- op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -572,32 +596,10 @@ public: } } - op += SP + SP + SP + SP + "std::size_t const scale_base =\n"; - { bool any = false; - for (size_t i = 0; i < fAxis; ++i) { - if (fShapeScale[i].dim != 1) { - op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) - + " * " + scaleStrides[i].GetVal() + "u"; - if (any) op = " +\n" + op; any = true; - } - } - if (!any) op += SP + SP + SP + SP + SP + "0u"; - op += ";\n\n"; - } + op += SP + SP + SP + SP + "std::size_t const scale_base =\n" + axisBase(fShapeScale, scaleStrides, SP + SP + SP + SP + SP) + ";\n\n"; - if (!fNB.empty()) { - op += SP + SP + SP + SP + "std::size_t const bias_base =\n"; - bool any = false; - for (size_t i = 0; i < fAxis; ++i) { - if (fShapeB[i].dim != 1) { - op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) - + " * " + biasStrides[i].GetVal() + "u"; - if (any) op = " +\n" + op; any = true; - } - } - if (!any) op += SP + SP + SP + SP + SP + "0u"; - op += ";\n\n"; - } + if (!fNB.empty()) + op += SP + SP + SP + SP + "std::size_t const bias_base =\n" + axisBase(fShapeB, biasStrides, SP + SP + SP + SP + SP) + ";\n\n"; op += SP + SP + SP + SP + "T mean = static_cast(0);\n"; for (size_t j = fAxis; j < fSize; ++j) @@ -686,10 +688,10 @@ public: // Determine parallel vs serial (same logic as kernel generation) size_t normLenVal2 = 0; bool canParallel2 = false; - try { normLenVal2 = std::stoul(fNormalizedLength); canParallel2 = (normLenVal2 > 0 && normLenVal2 <= 1024); } + try { normLenVal2 = std::stoul(fNormalizedLength); canParallel2 = (normLenVal2 > 0); } catch (...) {} size_t blockSize2 = 1; - if (canParallel2) { while (blockSize2 < normLenVal2) blockSize2 <<= 1; } + if (canParallel2) { while (blockSize2 < normLenVal2 && blockSize2 < 256) blockSize2 <<= 1; } std::string args = "alpaka::getPtrNative(deviceBuf_" + fNX + "), " diff --git a/core/inc/SOFIE/ROperator_MambaScan.hxx b/core/inc/SOFIE/ROperator_MambaScan.hxx index 1a2c449b..4f595451 100644 --- a/core/inc/SOFIE/ROperator_MambaScan.hxx +++ b/core/inc/SOFIE/ROperator_MambaScan.hxx @@ -128,7 +128,7 @@ public: out += SP + SP + SP + SP + "T const dt = delta[b*D_dim*L + d*L + t];\n"; out += SP + SP + SP + SP + "T const u_val = u[b*D_dim*L + d*L + t];\n"; out += SP + SP + SP + SP + "for (std::size_t n = 0; n < N; ++n) {\n"; - out += SP + SP + SP + SP + SP + "T const dA = SOFIE_DEVICE_exp(acc, dt * A[d*N + n]);\n"; + out += SP + SP + SP + SP + SP + "T const dA = alpaka::math::exp(acc, dt * A[d*N + n]);\n"; out += SP + SP + SP + SP + SP + "T const dBu = dt * B[b*N*L + n*L + t] * u_val;\n"; out += SP + SP + SP + SP + SP + "state[b*D_dim*N + d*N + n] = dA * state[b*D_dim*N + d*N + n] + dBu;\n"; out += SP + SP + SP + SP + "}\n"; diff --git a/core/inc/SOFIE/ROperator_Pool.hxx b/core/inc/SOFIE/ROperator_Pool.hxx index eeeb5908..9fdf266c 100644 --- a/core/inc/SOFIE/ROperator_Pool.hxx +++ b/core/inc/SOFIE/ROperator_Pool.hxx @@ -51,6 +51,9 @@ private: std::vector fShapeX; std::vector fShapeY; + std::string fBatchExpr; + std::vector fShapeYDim; + std::string fType; size_t fDim; // dimension of the MaxPool @@ -207,12 +210,24 @@ public: throw std::runtime_error("SOFIE Pool op Input Tensor " + fNX + " is not found in model"); } - fShapeX = model.GetTensorShape(fNX); - if (fShapeX.size() < 3 || fShapeX.size() > 5) { - std::cout << fNX << " : " << ConvertShapeToString(fShapeX) << std::endl; + auto dimShapeX = model.GetDimTensorShape(fNX); + if (dimShapeX.size() < 3 || dimShapeX.size() > 5) { + std::cout << fNX << " : " << ConvertDimShapeToString(dimShapeX) << std::endl; throw std::runtime_error("SOFIE Pool Op input data tensor" + fNX + " is not of 3,4 or 5 dimensions"); } + // channel and spatial dims must be static (see fBatchExpr's comment); only the + // batch dim may be dynamic. + for (size_t i = 1; i < dimShapeX.size(); i++) { + if (dimShapeX[i].isParam) + throw std::runtime_error("SOFIE Pool Op: only the batch (first) dimension of " + + fNX + " may be dynamic - channel and spatial dimensions must be static"); + } + fBatchExpr = dimShapeX[0].GetVal(); + fShapeX.resize(dimShapeX.size()); + fShapeX[0] = dimShapeX[0].isParam ? size_t(1) : dimShapeX[0].dim; + for (size_t i = 1; i < dimShapeX.size(); i++) + fShapeX[i] = dimShapeX[i].dim; fDim = fShapeX.size() - 2; // case of GlobalAveragePool. It is a pool case with kernel shape == image shape if (fPoolMode == GlobalAveragePool) { @@ -229,7 +244,9 @@ public: } // find shape of Y and add it in the list of intermediate tensors fShapeY = ShapeInference({fShapeX})[0]; - model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); + fShapeYDim.assign(fShapeY.begin(), fShapeY.end()); + fShapeYDim[0] = dimShapeX[0]; // restore the true (possibly symbolic) batch dim + model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeYDim); // need cmath for INFINITY when using MaxPool if (fPoolMode == MaxPool) model.AddNeededStdLib("cmath"); @@ -330,7 +347,7 @@ public: if(fDim==1){ // loop on batches and channels out << SP << "size_t outIndex = 0;\n"; - out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n"; + out << SP << "for (size_t n = 0; n < static_cast(" << fBatchExpr << ") * " << fShapeX[1] << "; n++) {\n"; out << SP << SP << "size_t inputOffset = n*" << fShapeX[2] << ";\n"; out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n"; // loop on elements of filter region to compute maximum @@ -372,7 +389,7 @@ public: else if(fDim==2){ // loop on batches and channels out << SP << "size_t outIndex = 0;\n"; - out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n"; + out << SP << "for (size_t n = 0; n < static_cast(" << fBatchExpr << ") * " << fShapeX[1] << "; n++) {\n"; out << SP << SP << "size_t inputOffset = n*" << fShapeX[2]*fShapeX[3] << ";\n"; out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n"; out << SP << SP << SP << "for (int j = wmin; j < wmax; j+=" << fAttrStrides[1] << ") {\n"; @@ -418,7 +435,7 @@ public: else if(fDim==3){ // loop on batches and channels out << SP << "size_t outIndex = 0;\n"; - out << SP << "for (size_t n = 0; n < " << fShapeX[0]*fShapeX[1] << "; n++) {\n"; + out << SP << "for (size_t n = 0; n < static_cast(" << fBatchExpr << ") * " << fShapeX[1] << "; n++) {\n"; out << SP << SP << "size_t inputOffset = n*" << fShapeX[2]*fShapeX[3]*fShapeX[4] << ";\n"; out << SP << SP << "for (int i = hmin; i < hmax; i+=" << fAttrStrides[0] << ") {\n"; out << SP << SP << SP << "for (int j = wmin; j < wmax; j+=" << fAttrStrides[1] << ") {\n"; @@ -487,7 +504,7 @@ public: bool doPadding = false; for (auto & e : fAttrPads) doPadding |= (e > 0); // count_include_pad == 0 with padding: divide by the in-bounds cells counted - // at run time; otherwise by the constant kernel area (CPU Generate above). + // at run time; otherwise by the constant kernel area const bool runtimeCount = isAvg && fAttrCountIncludePad == 0 && doPadding; const std::string kname = (isAvg ? "AvgPoolKernel_" : "MaxPoolKernel_") + opName; @@ -660,13 +677,15 @@ public: return ""; const bool isAvg = (fPoolMode == AveragePool); - std::size_t totalOut = ConvertShapeToLength(fShapeY); + // Use fShapeYDim (the true, possibly-symbolic batch dim), not fShapeY (which + // holds a dummy batch value of 1 when the real batch dim is dynamic). + std::string totalOut = ConvertDimShapeToLength(fShapeYDim); std::string kname = (isAvg ? "avgPoolKernel_" : "maxPoolKernel_") + opName; std::stringstream out; out << "\n//------ " << (isAvg ? "AVGPOOL" : "MAXPOOL") << "_GPU_ALPAKA\n"; out << SP << "auto const elementsPerThread_" << fNY << " = Vec::all(static_cast(1));\n"; - out << SP << "auto const elementsPerGrid_" << fNY << " = Vec::all(Idx{" << totalOut << "});\n"; + out << SP << "auto const elementsPerGrid_" << fNY << " = Vec::all(Idx{static_cast(" << totalOut << ")});\n"; out << SP << "auto const workDiv_" << fNY << " = sofie_workdiv(elementsPerGrid_" << fNY << ");\n"; out << SP << "auto task_" << fNY << " = alpaka::createTaskKernel(workDiv_" << fNY diff --git a/core/inc/SOFIE/ROperator_RMSNorm.hxx b/core/inc/SOFIE/ROperator_RMSNorm.hxx index 1a16c65d..b079144e 100644 --- a/core/inc/SOFIE/ROperator_RMSNorm.hxx +++ b/core/inc/SOFIE/ROperator_RMSNorm.hxx @@ -143,7 +143,7 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeX.empty()) throw std::runtime_error("SOFIE RMSNorm GPU kernel called without initialization"); @@ -159,6 +159,10 @@ public: std::vector shape(fRank); for (size_t i = 0; i < fRank; ++i) shape[i] = fShapeX[i].GetVal(); + // Wraps a size expression (a literal, or a symbolic dim name/expression from + // Dim::GetVal()) as an explicit std::size_t cast. + auto sz = [](const std::string &e) { return "static_cast(" + e + ")"; }; + auto strides = UTILITY::ComputeStrideFromShape(fShapeX); auto axesStrides = UTILITY::ComputeStrideFromShape(fAxesShape); auto normStrides = UTILITY::ComputeStrideFromShape(fNormShape); @@ -172,8 +176,10 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ X,\n"; op += SP + SP + SP + "T const* __restrict__ scale,\n"; - op += SP + SP + SP + "T* __restrict__ Y,\n"; - op += SP + SP + SP + "std::size_t const axesLength) const {\n\n"; + op += SP + SP + SP + "T* __restrict__ Y"; + for (auto &p : dynParamNames) + op += ",\n" + SP + SP + SP + "std::size_t const " + p; + op += ",\n" + SP + SP + SP + "std::size_t const axesLength) const {\n\n"; std::string eps = std::to_string(fAttrEpsilon); std::string nl = fNormLength; @@ -191,7 +197,7 @@ public: if (fAxis > 0) { for (size_t i = 0; i < fAxis; ++i) { op += SP + SP + SP + "std::size_t const axis_" + std::to_string(i) - + " = (row / " + axesStrides[i].GetVal() + "u) % " + shape[i] + "u;\n"; + + " = (row / " + sz(axesStrides[i].GetVal()) + ") % " + sz(shape[i]) + ";\n"; } op += "\n"; } @@ -202,26 +208,26 @@ public: op += SP + SP + SP + SP + "0u;\n\n"; } else { for (size_t i = 0; i < fAxis; ++i) { - op += SP + SP + SP + SP + "axis_" + std::to_string(i) + " * " + strides[i].GetVal() + "u"; + op += SP + SP + SP + SP + "axis_" + std::to_string(i) + " * " + sz(strides[i].GetVal()); op += (i + 1 < fAxis) ? " +\n" : ";\n\n"; } } - op += SP + SP + SP + "bool const in_range = (tid < " + nl + "u);\n"; + op += SP + SP + SP + "bool const in_range = (tid < " + sz(nl) + ");\n"; op += SP + SP + SP + "std::size_t norm_offset = 0u;\n"; op += SP + SP + SP + "std::size_t s_offset = 0u;\n"; op += SP + SP + SP + "if (in_range) {\n"; if (fRank - fAxis == 1) { - op += SP + SP + SP + SP + "norm_offset = tid * " + strides[fAxis].GetVal() + "u;\n"; - op += SP + SP + SP + SP + "s_offset = tid * " + normStrides[0].GetVal() + "u;\n"; + op += SP + SP + SP + SP + "norm_offset = tid * " + sz(strides[fAxis].GetVal()) + ";\n"; + op += SP + SP + SP + SP + "s_offset = tid * " + sz(normStrides[0].GetVal()) + ";\n"; } else { op += SP + SP + SP + SP + "std::size_t rem = tid;\n"; for (size_t j = fAxis; j < fRank; ++j) { size_t ji = j - fAxis; - op += SP + SP + SP + SP + "{ std::size_t nj = rem / " + normStrides[ji].GetVal() - + "u; rem %= " + normStrides[ji].GetVal() - + "u; norm_offset += nj * " + strides[j].GetVal() - + "u; s_offset += nj * " + normStrides[ji].GetVal() + "u; }\n"; + op += SP + SP + SP + SP + "{ std::size_t nj = rem / " + sz(normStrides[ji].GetVal()) + + "; rem %= " + sz(normStrides[ji].GetVal()) + + "; norm_offset += nj * " + sz(strides[j].GetVal()) + + "; s_offset += nj * " + sz(normStrides[ji].GetVal()) + "; }\n"; } } op += SP + SP + SP + "}\n\n"; @@ -256,7 +262,7 @@ public: if (fAxis > 0) { for (size_t i = 0; i < fAxis; ++i) { op += SP + SP + SP + SP + "std::size_t const axis_" + std::to_string(i) - + " = (row / " + axesStrides[i].GetVal() + "u) % " + shape[i] + "u;\n"; + + " = (row / " + sz(axesStrides[i].GetVal()) + ") % " + sz(shape[i]) + ";\n"; } op += "\n"; } @@ -266,7 +272,7 @@ public: op += SP + SP + SP + SP + SP + "0u;\n\n"; } else { for (size_t i = 0; i < fAxis; ++i) { - op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) + " * " + strides[i].GetVal() + "u"; + op += SP + SP + SP + SP + SP + "axis_" + std::to_string(i) + " * " + sz(strides[i].GetVal()); op += (i + 1 < fAxis) ? " +\n" : ";\n\n"; } } @@ -274,9 +280,9 @@ public: op += SP + SP + SP + SP + "T rms_sum = static_cast(0);\n"; for (size_t j = fAxis; j < fRank; ++j) op += SP + SP + SP + SP + "for (std::size_t n_" + std::to_string(j) - + " = 0; n_" + std::to_string(j) + " < " + shape[j] + "u; ++n_" + std::to_string(j) + ") {\n"; + + " = 0; n_" + std::to_string(j) + " < " + sz(shape[j]) + "; ++n_" + std::to_string(j) + ") {\n"; op += SP + SP + SP + SP + SP + "std::size_t const idx = row_base"; - for (size_t j = fAxis; j < fRank; ++j) op += " + n_" + std::to_string(j) + " * " + strides[j].GetVal() + "u"; + for (size_t j = fAxis; j < fRank; ++j) op += " + n_" + std::to_string(j) + " * " + sz(strides[j].GetVal()); op += ";\n"; op += SP + SP + SP + SP + SP + "T v = X[idx]; rms_sum += v * v;\n"; for (size_t j = fAxis; j < fRank; ++j) op += SP + SP + SP + SP + "}\n"; @@ -285,16 +291,16 @@ public: for (size_t j = fAxis; j < fRank; ++j) op += SP + SP + SP + SP + "for (std::size_t n_" + std::to_string(j) - + " = 0; n_" + std::to_string(j) + " < " + shape[j] + "u; ++n_" + std::to_string(j) + ") {\n"; + + " = 0; n_" + std::to_string(j) + " < " + sz(shape[j]) + "; ++n_" + std::to_string(j) + ") {\n"; op += SP + SP + SP + SP + SP + "std::size_t const idx = row_base"; - for (size_t j = fAxis; j < fRank; ++j) op += " + n_" + std::to_string(j) + " * " + strides[j].GetVal() + "u"; + for (size_t j = fAxis; j < fRank; ++j) op += " + n_" + std::to_string(j) + " * " + sz(strides[j].GetVal()); op += ";\n"; op += SP + SP + SP + SP + SP + "std::size_t s_idx = "; { bool first = true; for (size_t j = fAxis; j < fRank; ++j) { if (!first) op += " + "; - op += "n_" + std::to_string(j) + " * " + normStrides[j - fAxis].GetVal() + "u"; + op += "n_" + std::to_string(j) + " * " + sz(normStrides[j - fAxis].GetVal()); first = false; } if (fRank == fAxis) op += "0u"; @@ -317,7 +323,7 @@ public: return SP + kname + " rmsNormKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeX.empty()) throw std::runtime_error("SOFIE RMSNorm GPU dispatch called without initialization"); @@ -332,8 +338,10 @@ public: std::string args = "alpaka::getPtrNative(deviceBuf_" + fNX + "), " "alpaka::getPtrNative(deviceBuf_" + fNScale + "), " - "alpaka::getPtrNative(deviceBuf_" + fNY + "), " - "static_cast(" + fAxesLength + ")"; + "alpaka::getPtrNative(deviceBuf_" + fNY + ")"; + for (auto &p : dynParamNames) + args += ", static_cast(" + p + ")"; + args += ", static_cast(" + fAxesLength + ")"; std::stringstream out; out << "\n//------ RMSNORM_GPU_ALPAKA\n"; diff --git a/core/inc/SOFIE/ROperator_RWKV_WKV6.hxx b/core/inc/SOFIE/ROperator_RWKV_WKV6.hxx index d5196bf4..92dc594b 100644 --- a/core/inc/SOFIE/ROperator_RWKV_WKV6.hxx +++ b/core/inc/SOFIE/ROperator_RWKV_WKV6.hxx @@ -100,7 +100,7 @@ public: out << SP << SP << SP << SP << "}\n"; // State update out << SP << SP << SP << SP << "for (size_t i = 0; i < " << Dh << "; ++i) {\n"; - out << SP << SP << SP << SP << SP << fType << " decay_i = std::exp(" << W << "[rBase+i]);\n"; + out << SP << SP << SP << SP << SP << fType << " decay_i = std::exp(-std::exp(" << W << "[rBase+i]));\n"; out << SP << SP << SP << SP << SP << "for (size_t j = 0; j < " << Dh << "; ++j)\n"; out << SP << SP << SP << SP << SP << SP << S << "[sBase+i*" << Dh << "+j] = decay_i * " << S << "[sBase+i*" << Dh << "+j] + " << K << "[rBase+i] * " << V << "[rBase+j];\n"; out << SP << SP << SP << SP << "}\n"; @@ -141,13 +141,13 @@ public: out += SP + SP + SP + SP + SP + "T y_j = static_cast(0);\n"; out += SP + SP + SP + SP + SP + "for (std::size_t i = 0; i < Dh; ++i) {\n"; out += SP + SP + SP + SP + SP + SP + "T const kv_ij = k[rBase+i] * v[rBase+j];\n"; - out += SP + SP + SP + SP + SP + SP + "y_j += r[rBase+i] * (state[sBase+i*Dh+j] + SOFIE_DEVICE_exp(acc, u[uBase+i]) * kv_ij);\n"; + out += SP + SP + SP + SP + SP + SP + "y_j += r[rBase+i] * (state[sBase+i*Dh+j] + alpaka::math::exp(acc, u[uBase+i]) * kv_ij);\n"; out += SP + SP + SP + SP + SP + "}\n"; out += SP + SP + SP + SP + SP + "y[rBase+j] = y_j;\n"; out += SP + SP + SP + SP + "}\n"; // Update state out += SP + SP + SP + SP + "for (std::size_t i = 0; i < Dh; ++i) {\n"; - out += SP + SP + SP + SP + SP + "T const decay_i = SOFIE_DEVICE_exp(acc, w[rBase+i]);\n"; + out += SP + SP + SP + SP + SP + "T const decay_i = alpaka::math::exp(acc, -alpaka::math::exp(acc, w[rBase+i]));\n"; out += SP + SP + SP + SP + SP + "for (std::size_t j = 0; j < Dh; ++j)\n"; out += SP + SP + SP + SP + SP + SP + "state[sBase+i*Dh+j] = decay_i * state[sBase+i*Dh+j] + k[rBase+i] * v[rBase+j];\n"; out += SP + SP + SP + SP + "}\n"; diff --git a/core/inc/SOFIE/ROperator_Range.hxx b/core/inc/SOFIE/ROperator_Range.hxx index 8ea17d9c..dca91a62 100644 --- a/core/inc/SOFIE/ROperator_Range.hxx +++ b/core/inc/SOFIE/ROperator_Range.hxx @@ -7,6 +7,7 @@ #include #include +#include namespace SOFIE{ @@ -22,6 +23,38 @@ private: std::vector fShape; std::string fType; + // set in Initialize() whenever the corresponding scalar input resolves to a value known + // at codegen time (an initializer, or a non-parametric shape-tensor entry); used to bake + // the value into the generated code as a literal instead of reading it through a pointer + bool fStartIsConst = false; + T fStartConst{}; + bool fLimitIsConst = false; + T fLimitConst{}; + bool fDeltaIsConst = false; + T fDeltaConst{}; + + std::string ValueLiteral(T v) const { + std::ostringstream s; + if (fType == "float") + s << std::setprecision(9) << v << "f"; + else + s << v; + return s.str(); + } + + // expression used to read start/limit/delta in generated code: a literal when the + // value was already known at codegen time, otherwise a dereference of the host pointer + std::string StartExpr() const { return fStartIsConst ? ValueLiteral(fStartConst) : ("*tensor_" + fNStart); } + std::string LimitExpr() const { return fLimitIsConst ? ValueLiteral(fLimitConst) : ("*tensor_" + fNLimit); } + std::string DeltaExpr() const { return fDeltaIsConst ? ValueLiteral(fDeltaConst) : ("*tensor_" + fNDelta); } + + // element count computed at run time from the three scalar inputs; shared by the CPU loop + // and the GPU launch. Each operand is either a literal or a host-pointer read, per *Expr(). + std::string RuntimeSizeExpr() const { + return "static_cast(std::max(std::ceil((static_cast(" + LimitExpr() + + ") - static_cast(" + StartExpr() + ")) / static_cast(" + DeltaExpr() + ")), 0.0f))"; + } + public: ROperator_Range(){} @@ -35,15 +68,10 @@ public: } static_assert( (std::is_same_v || std::is_same_v), "TMVA::SOFIE - Unsupported type by Range operator"); - } - - std::vector TypeInference(std::vector input) override { - return input; - } - - std::vector> ShapeInference(std::vector> input) override { - auto ret = input; //suggest copy to compiler - return ret; + { + fInputTensorNames = { fNStart, fNLimit, fNDelta }; + fOutputTensorNames = { fNOutput }; + } } void Initialize(RModel& model) override { @@ -61,32 +89,92 @@ public: std::runtime_error("SOFIE Range Op Input Tensor " + fNDelta + "is not found in model"); } ETensorType type = ConvertStringToType(fType); - if (model.IsInitializedTensor(fNStart) && model.IsInitializedTensor(fNDelta) && model.IsInitializedTensor(fNLimit)) { - T * start = static_cast(model.GetInitializedTensorData(fNStart).get()); - T * limit = static_cast(model.GetInitializedTensorData(fNLimit).get()); - T * delta = static_cast(model.GetInitializedTensorData(fNDelta).get()); - if (!start || !delta || !limit) - std::runtime_error("SOFIE Range Op Input Tensor has invalid input data"); - T a = *start; - T b = *limit; - T d = *delta; - int number_of_elements = std::max( static_cast(std::ceil( (b - a) / d )) , 0. ); + + + + auto analyzeInput = [&](const std::string & tName, T & value, Dim & dim) { + int ftype = 0; // type of input (0 intermediate, 1 constant , 2 shape) + if (model.IsInitializedTensor(tName)) { + T * data = static_cast(model.GetInitializedTensorData(tName).get()); + if (!data) + throw std::runtime_error("SOFIE Range Op Input Tensor has invalid input data"); + value = *data; + ftype = 1; + } else if (model.IsShapeTensor(tName)) { + auto data = model.GetShapeTensorValues(tName); + dim = data[0]; + if (!dim.isParam) { + value = static_cast(dim.dim); + ftype = 1; + } else + ftype = 2; + } + return ftype; + }; + + T start_value{}; + T limit_value{}; + T delta_value{}; + Dim start_dim{}; + Dim limit_dim{}; + Dim delta_dim{}; + int res1 = analyzeInput(fNStart, start_value, start_dim); + int res2 = analyzeInput(fNLimit, limit_value, limit_dim); + int res3 = analyzeInput(fNDelta, delta_value, delta_dim); + if (res1 == 1) { fStartIsConst = true; fStartConst = start_value; } + if (res2 == 1) { fLimitIsConst = true; fLimitConst = limit_value; } + if (res3 == 1) { fDeltaIsConst = true; fDeltaConst = delta_value; } + if (res1 == 0 || res2 == 0 || res3 == 0) { + // cannot know at compile time- need to do fully at run time + // + fShape = {Dim{"range_size_" + fNStart + "_" + fNLimit}}; + model.AddDynamicTensor(fNOutput, type, fShape); + } else if (res1 == 1 && res2 == 1 && res3 == 1) { + size_t number_of_elements = std::max(static_cast(std::ceil((limit_value - start_value) / delta_value )) , 0 ); + fIsOutputConstant = true; + + // compute output std::vector output(number_of_elements); - for (int i=0; i shape = {static_cast(number_of_elements)}; + std::vector shape = {number_of_elements}; model.AddConstantTensor(fNOutput,shape, output.data()); - fIsOutputConstant = true; - // set the input tensor not writable - model.SetNotWritableInitializedTensor(fNStart); - model.SetNotWritableInitializedTensor(fNDelta); - model.SetNotWritableInitializedTensor(fNLimit); - } - else { - fShape = {Dim{"range_size"}}; - model.AddDynamicTensor(fNOutput, type, fShape); + fShape = ConvertShapeToDim(shape); + + } else { // case of a shape tensor + std::string start = (res1 == 1) ? std::to_string(start_value) : start_dim.GetVal(); + std::string limit = (res2 == 1) ? std::to_string(limit_value) : limit_dim.GetVal(); + std::string delta = (res3 == 1) ? std::to_string(delta_value) : delta_dim.GetVal(); + std::stringstream s; + if (type == ETensorType::FLOAT ) { + if (delta_value == 1) + s << "std::max(std::ceil("<< limit << " - " << start << "),0.0f)"; + else + s << "std::max(std::ceil(("<< limit << " - " << start << ")/" << delta << "),0.0f)"; + } else if (type == ETensorType::INT64 ) { + if (delta == "1") { + if (start == "0") + s << limit; + else + s << "std::max((" << limit << " - " << start << "),0L)"; + } else { + if (start == "0") + s << "((" << limit << ")/" << delta << ")"; + else + s << "std::max((" << limit << " - " << start << ")/"<< delta << "),0L)"; + } + } else { + throw + std::runtime_error("SOFIE Range Op Input Tensor " + ConvertTypeToString(type) + "is not supported"); + } + + + fShape = { Dim {s.str(), static_cast(-1)} }; + model.AddDynamicTensor(fNOutput,type, fShape); } + + if (model.Verbose()) { std::cout << "Range -> output is " << fNOutput << " : " << ConvertDimShapeToString(fShape); if (fIsOutputConstant) std::cout << " : " << ConvertValuesToString(model.GetTensorData(fNOutput)); @@ -94,31 +182,134 @@ public: } } - std::string Generate(std::string OpName) override { + std::string Generate(std::string opName) override { std::stringstream out; - out << "\n//------ Range\n"; + out << "\n//------ Range " << opName << "---> " << ConvertDimShapeToString(fShape) << "\n"; if (fIsOutputConstant) return out.str(); - OpName = "op_" + OpName; + opName = "op_" + opName; if (fShape.empty()) { throw std::runtime_error("SOFIE Range operator called to Generate without being initialized first"); } - std::string sizeName = fShape[0].param; - out << SP << "size_t " << sizeName << " = static_cast(std::max(std::ceil((static_cast(*tensor_" << fNLimit << ") - static_cast(*tensor_" << fNStart << ")) / static_cast(*tensor_" << fNDelta << ")), 0.0f));\n"; - out << SP << "if (" << sizeName << " > " << "fTensor_" << fNOutput << ".size() ){\n"; - out << SP << SP << "fTensor_" << fNOutput << ".resize(" << sizeName << ");\n"; - // need to re-initialized pointer to tensor data - out << SP << SP << "tensor_" << fNOutput << " = fTensor_" << fNOutput << ".data();\n"; + std::string outputSizeVar; + std::string outputSize = fShape[0].param; + if (outputSize.find("range_size") != std::string::npos) { + outputSizeVar = outputSize; + outputSize = RuntimeSizeExpr(); + } else { + outputSizeVar = "range_" + opName; + } + out << SP << "size_t " << outputSizeVar << " = " << outputSize << ";\n"; + out << SP << "for (size_t i = 0; i < " << outputSizeVar << "; i++) {\n"; + out << SP << SP << "tensor_" << fNOutput << "[i] = " << StartExpr() << " + i * (" << DeltaExpr() << ");\n"; out << SP << "}\n"; - out << SP << "for (size_t i = 0; i < " << sizeName << "; i++) {\n"; - out << SP << SP << "fTensor_" << fNOutput << "[i] = *tensor_" << fNStart << " + i * (*tensor_" << fNDelta << ");\n"; + + return out.str(); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + if (fIsOutputConstant) return ""; + opName = "op_" + opName; + std::string op; + op += "\n//------ RANGE_KERNEL_ALPAKA\n"; + op += SP + "struct RangeKernel_" + opName + " {\n"; + // TS/TD are deduced independently as either T (a compile-time-constant start/delta, + // baked in and passed by value - no device read) or T const* (a runtime value read + // from device memory), so the same kernel covers both without duplicating the body. + op += SP + SP + "template\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const& acc, TS start, TD delta, T* output, std::size_t n) const {\n"; + op += SP + SP + SP + "auto const idx = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (idx >= n) return;\n"; + op += SP + SP + SP + "T startVal; if constexpr (std::is_pointer_v) startVal = start[0]; else startVal = static_cast(start);\n"; + op += SP + SP + SP + "T deltaVal; if constexpr (std::is_pointer_v) deltaVal = delta[0]; else deltaVal = static_cast(delta);\n"; + op += SP + SP + SP + "output[idx] = startVal + static_cast(idx) * deltaVal;\n"; + op += SP + SP + "}\n"; + op += SP + "};\n"; + return op; + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override { + if (fIsOutputConstant) return ""; + opName = "op_" + opName; + return SP + "RangeKernel_" + opName + " rangeKernel_" + opName + ";\n"; + } + + std::string Generate_GPU_ALPAKA(std::string opName) override { + std::stringstream out; + out << "\n//------ Range (GPU) " << opName << " ---> " << ConvertDimShapeToString(fShape) << "\n"; + if (fIsOutputConstant) return out.str(); + if (fShape.empty()) { + throw std::runtime_error("SOFIE Range operator called to Generate without being initialized first"); + } + + opName = "op_" + opName; + std::string outputSize = fShape[0].param; + if (outputSize.find("range_size") != std::string::npos) { + /* + * Run-time size: at least one of start/limit/delta is produced by another operator, so + * its value exists only on the device during inference. The generated code copies just + * the operands that are NOT already known at codegen time to the host (a constant + * start/delta needs no round trip, computes the element count with the same expression + * the CPU code uses, and declares it under the + * name the downstream launches already reference. The output buffer was allocated in + * the constructor from the value passed there under that same name, so the count is + * checked against it before the kernel runs. + */ + std::vector inputs; + for (auto &in : {fNStart, fNLimit, fNDelta}) { + bool isConst = (in == fNStart && fStartIsConst) || (in == fNLimit && fLimitIsConst) || + (in == fNDelta && fDeltaIsConst); + if (!isConst && std::find(inputs.begin(), inputs.end(), in) == inputs.end()) + inputs.push_back(in); + } + std::string sizeMember = GetMemberNameForDimShape(outputSize); + + out << SP << "size_t " << outputSize << ";\n"; + if (!inputs.empty()) { + out << SP << "{\n"; + for (auto &in : inputs) { + out << SP << SP << "auto host_" << in << " = alpaka::allocBuf<" << fType << ", Idx>(host, Ext1D::all(Idx{1}));\n"; + out << SP << SP << "alpaka::memcpy(queue, host_" << in << ", deviceBuf_" << in << ");\n"; + } + out << SP << SP << "alpaka::wait(queue);\n"; + for (auto &in : inputs) { + out << SP << SP << "const " << fType << "* tensor_" << in << " = alpaka::getPtrNative(host_" << in << ");\n"; + } + out << SP << SP << outputSize << " = " << RuntimeSizeExpr() << ";\n"; + out << SP << "}\n"; + } else { + out << SP << outputSize << " = " << RuntimeSizeExpr() << ";\n"; + } + out << SP << "if (" << outputSize << " > " << sizeMember << ") {\n"; + out << SP << SP << "throw std::runtime_error(\"SOFIE Range " << opName + << ": run-time size exceeds the size given at construction (" << outputSize << ")\");\n"; + out << SP << "}\n"; + } + + // start/delta are passed by value (baked-in literal) when known at codegen time, so no + // device read is needed for them at all; otherwise fall back to the device pointer. + std::string startArg = fStartIsConst ? ValueLiteral(fStartConst) + : ("alpaka::getPtrNative(deviceBuf_" + fNStart + ")"); + std::string deltaArg = fDeltaIsConst ? ValueLiteral(fDeltaConst) + : ("alpaka::getPtrNative(deviceBuf_" + fNDelta + ")"); + + out << SP << "{\n"; + out << SP << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{" << outputSize << "});\n"; + out << SP << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; + out << SP << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName + << ", rangeKernel_" << opName + << ", " << startArg + << ", " << deltaArg + << ", alpaka::getPtrNative(deviceBuf_" << fNOutput << ")" + << ", static_cast(" << outputSize << "));\n"; + out << SP << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; out << SP << "}\n"; return out.str(); } }; }//SOFIE - + #endif //SOFIE_ROPERATOR_RANGE diff --git a/core/inc/SOFIE/ROperator_Reduce.hxx b/core/inc/SOFIE/ROperator_Reduce.hxx index f3e7170f..ae321f1d 100644 --- a/core/inc/SOFIE/ROperator_Reduce.hxx +++ b/core/inc/SOFIE/ROperator_Reduce.hxx @@ -12,7 +12,6 @@ #include #include - namespace SOFIE{ enum EReduceOpMode { ReduceMean, ReduceSum, ReduceSumSquare, ReduceProd, ReduceL2, ReduceMax, InvalidReduceOp }; @@ -22,15 +21,43 @@ class ROperator_Reduce final : public ROperator { private: /* Attributes*/ + bool fInputDimShape = false; int fkeepdims = 1; //default value std::vector fAttrAxes; EReduceOpMode fReduceOpMode; std::string fNX; std::string fNAxes; std::string fNY; - std::vector fShapeX; - std::vector fShapeY; - std::vector fShapeYNotPruned; // needed for fKeepdims=0 + std::vector fShapeX; + std::vector fShapeY; + std::vector fShapeYNotPruned; // needed for fKeepdims=0 + + // GPU kernel launch tuning parameters + struct GpuLaunchParams { std::size_t blockSize = 256; std::size_t groups = 1; }; + + + GpuLaunchParams ComputeGpuLaunchParams() const { + GpuLaunchParams p; + if (fInputDimShape || fShapeX.empty() || fShapeY.empty()) + return p; + std::size_t inLen = std::stoul(ConvertDimShapeToLength(fShapeX)); + std::size_t outLen = std::stoul(ConvertDimShapeToLength(fShapeY)); + if (outLen == 0) + return p; + std::size_t redLen = inLen / outLen; + + std::size_t bs = 32; + while (bs < redLen && bs < 256) bs *= 2; + p.blockSize = bs; + + constexpr std::size_t kTargetBlocks = 132; + if (outLen < kTargetBlocks) { + std::size_t maxGroups = kTargetBlocks / outLen; + std::size_t chunksAvailable = (redLen + bs - 1) / bs; + p.groups = std::max(1, std::min(maxGroups, chunksAvailable)); + } + return p; + } public: @@ -46,11 +73,12 @@ public: } std::vector GetStdLibs() override { + std::vector libs = { std::string("memory") }; if (fReduceOpMode == ReduceL2) - return { std::string("cmath") }; + libs.push_back("cmath"); if (fReduceOpMode == ReduceMax) - return { std::string("limits") }; - return {}; + libs.push_back("limits"); + return libs; } ROperator_Reduce(){} @@ -72,18 +100,17 @@ public: } // shape of output tensors given input tensors - std::vector> ShapeInference(std::vector> input) override { - auto ret = input; //suggest copy to compiler - auto & outputShape = ret[0]; + using ROperator::ShapeInference; + std::vector ShapeInference(const std::vector & input) { + auto ret = input; + auto & outputShape = ret; for (size_t j = 0; j < fAttrAxes.size(); j++) { if (fAttrAxes[j] < 0) fAttrAxes[j] += outputShape.size(); - if (fAttrAxes[j] < 0 || (size_t) fAttrAxes[j] >= outputShape.size() ) + if (fAttrAxes[j] < 0 || (size_t) fAttrAxes[j] >= outputShape.size()) throw std::runtime_error("SOFIE Reduce Op - invalid axes values " + std::to_string(fAttrAxes[j])); - // set to 1 the reduced dims - outputShape[fAttrAxes[j]] = 1; + outputShape[fAttrAxes[j]] = Dim{1}; } fShapeYNotPruned = outputShape; - // in case of pruning dimension we need to sort axes attributes if (fkeepdims == 0) { auto ax = fAttrAxes; std::sort(ax.begin(), ax.end()); @@ -92,12 +119,13 @@ public: if (outputShape.size() > 1) { outputShape.erase(outputShape.begin() + ax[j]); for (size_t k = j+1; k < ax.size(); k++) - ax[k] -= 1; // decrease by one since we have removed a value + ax[k] -= 1; } } } return ret; } + void Initialize(RModel& model) override { fUseSession = model.UseSession(); @@ -106,7 +134,9 @@ public: // input must be a graph input, or already initialized intermediate tensor throw std::runtime_error("SOFIE Reduce Op Input Tensor " + fNX + " is not found in model"); } - fShapeX = model.GetTensorShape(fNX); + fShapeX = model.GetDimTensorShape(fNX); + if (model.IsDynamicTensor(fNX)) + fInputDimShape = true; // check if tensor with axes is provided if (!fNAxes.empty()) { auto ax_shptr = model.GetInitializedTensorData(fNAxes); @@ -121,22 +151,32 @@ public: fAttrAxes[i] = i; } // find shape of Y and add it in the list of intermediate tensors - fShapeY = ShapeInference({fShapeX})[0]; + fShapeY = ShapeInference(fShapeX); model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY); if (model.Verbose()){ - std::cout << Name() << " : " << fNX << " -> " << fNY << " shape " << ConvertShapeToString(fShapeY) << std::endl; + std::cout << Name() << " : " << fNX << " -> " << fNY << " shape " << ConvertDimShapeToString(fShapeY) << std::endl; } model.AddNeededStdLib("algorithm"); } + bool IsReducedAxis(size_t d) const { + return std::find(fAttrAxes.begin(), fAttrAxes.end(), (int64_t)d) != fAttrAxes.end(); + } + + // number of input elements reduced into each output element, as generated-code text; + // used by the CPU loops and passed to the GPU kernel by the launch + std::string ReducedLengthExpr() const { + return "(" + ConvertDimShapeToLength(fShapeX) + ") / (" + ConvertDimShapeToLength(fShapeY) + ")"; + } + std::string Generate(std::string opName) override { opName = "op_" + opName; if (fShapeX.empty() || fShapeY.empty()) { throw std::runtime_error("SOFIE Reduce Op called to Generate without being initialized first"); } - size_t inputLength = SOFIE::ConvertShapeToLength(fShapeX); - size_t outputLength = SOFIE::ConvertShapeToLength(fShapeY); + std::string inputLength = SOFIE::ConvertDimShapeToLength(fShapeX); + std::string outputLength = SOFIE::ConvertDimShapeToLength(fShapeY); auto inputStrides = SOFIE::UTILITY::ComputeStrideFromShape(fShapeX); // output stride (or not pruned vector) @@ -158,7 +198,7 @@ public: int kmin = fShapeX.size()-fAttrAxes.size(); for (int k = fShapeX.size()-1; k >= kmin; k--) { // if k is not a reduced axis is not last ones - if (std::find(fAttrAxes.begin(), fAttrAxes.end(), k) == fAttrAxes.end()) { + if (!IsReducedAxis(k)) { reduceDims = kMiddle; break; } @@ -168,18 +208,21 @@ public: // check if at the beginning for (size_t k = 0; k < fAttrAxes.size(); k++) { // if k is not a reduced axis is not first ones - if (std::find(fAttrAxes.begin(), fAttrAxes.end(), k) == fAttrAxes.end()) { + if (!IsReducedAxis(k)) { reduceDims = kMiddle; break; } } } - size_t reducedLength = inputLength / outputLength; + std::string reducedLength; + if (fInputDimShape) { + reducedLength = "reducedLength_" + opName; + out << SP << "size_t " << reducedLength << " = " << ReducedLengthExpr() << ";\n"; + } else { + int rLength = std::stoi(inputLength) / std::stoi(outputLength); + reducedLength = std::to_string(rLength); + } if (reduceDims == kLast) { - //std::cout << "reduction for operator " << opName << " is last" << std::endl; - // new faster implementation using a single loop - // faster to loop first on reduced dimension and then output - // reset output tensors // loop on output dimensions out << SP << "for (size_t i = 0; i < " << outputLength << "; i++) {\n"; @@ -265,10 +308,10 @@ public: // here we find output index out << SP << SP << "size_t outputIndex = 0;\n"; for (size_t k = 0; k < dim; k++) { - if (std::find(fAttrAxes.begin(), fAttrAxes.end(), k) == fAttrAxes.end()) { + if (!IsReducedAxis(k)) { // do for not reducing axes - out << SP << SP << "size_t i_" << k << " = i / " << inputStrides[k] << " % " << fShapeX[k] << ";\n"; - out << SP << SP << "outputIndex += i_" << k << " * " << outputStrides[k] << ";\n"; + out << SP << SP << "size_t i_" << k << " = i / (" << inputStrides[k].GetVal() << ") % (" << fShapeX[k].GetVal() << ");\n"; + out << SP << SP << "outputIndex += i_" << k << " * (" << outputStrides[k].GetVal() << ");\n"; } } // now compute reduction @@ -300,38 +343,53 @@ public: return out.str(); } - // --------------------------------------------------------------------------- - // GPU kernel: one block per output element, 256 threads cooperatively reduce - // the slice via shared-memory tree reduction. - // This replaces the previous naive "one thread per output element" approach - // which serialised the entire reduction loop inside a single thread. - // --------------------------------------------------------------------------- - std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/, const std::vector &dynParamNames) override { if (fShapeX.empty() || fShapeY.empty()) throw std::runtime_error("SOFIE Reduce Op called to Generate without being initialized first"); const std::size_t Dx = fShapeX.size(); auto inputStrides = UTILITY::ComputeStrideFromShape(fShapeX); auto outputStrides = UTILITY::ComputeStrideFromShape(fShapeYNotPruned); - std::size_t inputLength = ConvertShapeToLength(fShapeX); - std::size_t outputLength = ConvertShapeToLength(fShapeY); - std::size_t reducedLength = inputLength / outputLength; // Partition axes into keep (non-reduced) and reduce sets. std::vector redAxes, keepAxes; for (std::size_t d = 0; d < Dx; ++d) { - if (std::find(fAttrAxes.begin(), fAttrAxes.end(), (int64_t)d) != fAttrAxes.end()) + if (IsReducedAxis(d)) redAxes.push_back(d); else keepAxes.push_back(d); } - // Row-major strides for decomposing the flat reduction index r into - // per-axis coordinates. + + enum EReduceDim { kFirst, kLast, kMiddle }; + EReduceDim reduceDims = kLast; + { + int kmin = (int)Dx - (int)redAxes.size(); + for (int k = (int)Dx - 1; k >= kmin; k--) { + if (!IsReducedAxis(k)) { reduceDims = kMiddle; break; } + } + } + if (reduceDims == kMiddle) { + reduceDims = kFirst; + for (std::size_t k = 0; k < redAxes.size(); k++) { + if (!IsReducedAxis(k)) { reduceDims = kMiddle; break; } + } + } + + // row-major strides for decomposing the flat reduction index into coordinates // redStrides[i] = product of fShapeX[redAxes[j]] for j > i - std::vector redStrides(redAxes.size(), 1); + std::vector redStrides(redAxes.size(), "1"); for (int ri = (int)redAxes.size() - 2; ri >= 0; --ri) - redStrides[ri] = redStrides[ri + 1] * fShapeX[redAxes[ri + 1]]; + redStrides[ri] = "(" + redStrides[ri + 1] + " * " + fShapeX[redAxes[ri + 1]].GetVal() + ")"; + + bool dyn = fInputDimShape; + GpuLaunchParams lp = ComputeGpuLaunchParams(); + // Dynamic shapes always take the groups-capable, symbolic-block-size + // form since the split decision can only be made at runtime. + bool useGroupsDecomp = dyn || lp.groups > 1; + std::string bsExpr = dyn ? "blockDim" : (std::to_string(lp.blockSize) + "u"); + std::string halfBsExpr = dyn ? "(blockDim / 2)" : (std::to_string(lp.blockSize / 2) + "u"); + std::string shmemCap = dyn ? "256" : std::to_string(lp.blockSize); std::string kname = "ReduceKernel_" + Name() + "_" + fNY; @@ -344,26 +402,47 @@ public: op += SP + SP + SP + "T const* __restrict__ input,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; op += SP + SP + SP + "std::size_t const reducedLength,\n"; - op += SP + SP + SP + "std::size_t const outputLength) const {\n\n"; + op += SP + SP + SP + "std::size_t const outputLength"; + if (useGroupsDecomp) + op += ",\n" + SP + SP + SP + "std::size_t const groups"; + for (auto &p : dynParamNames) + op += ",\n" + SP + SP + SP + "std::size_t const " + p; + op += ") const {\n\n"; - // ---- shared memory (fixed 256 slots, matches block size) ---- - op += SP + SP + SP + "auto& shmem = alpaka::declareSharedVar(acc);\n\n"; + // ---- shared memory ---- + op += SP + SP + SP + "auto& shmem = alpaka::declareSharedVar(acc);\n\n"; // ---- block/thread addressing ---- - // One block per output element; threads cooperate within the block. - op += SP + SP + SP + "auto const out_idx = alpaka::getIdx(acc)[0];\n"; + // groups == 1: one block per output element; output is the final Y buffer. + // groups > 1: `groups` blocks cooperate per output element, each over a + // slice of the reduction axis; output is an unfinalized scratch buffer of + // size outputLength*groups, combined by ReduceFinalizeKernel below. op += SP + SP + SP + "auto const thread_id = alpaka::getIdx(acc)[0];\n"; - op += SP + SP + SP + "if (out_idx >= outputLength) return;\n\n"; + if (dyn) + op += SP + SP + SP + "auto const blockDim = alpaka::getWorkDiv(acc)[0];\n"; + if (useGroupsDecomp) { + op += SP + SP + SP + "auto const flat_block = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "auto const out_idx = flat_block % outputLength;\n"; + op += SP + SP + SP + "auto const group_id = flat_block / outputLength;\n"; + op += SP + SP + SP + "if (flat_block >= outputLength * groups) return;\n\n"; + } else { + op += SP + SP + SP + "auto const out_idx = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (out_idx >= outputLength) return;\n\n"; + } // ---- decode output (keep-axis) coordinates from out_idx ---- - for (std::size_t d = 0; d < Dx; ++d) { - if (std::find(redAxes.begin(), redAxes.end(), d) == redAxes.end()) { - op += SP + SP + SP + "std::size_t const oy_" + std::to_string(d) - + " = (out_idx / " + std::to_string(outputStrides[d]) + "u) % " - + std::to_string(fShapeYNotPruned[d]) + "u;\n"; + // (only needed for the generic kMiddle path; kFirst/kLast compute + // in_idx directly from out_idx without per-axis coordinates) + if (reduceDims == kMiddle) { + for (std::size_t d = 0; d < Dx; ++d) { + if (!IsReducedAxis(d)) { + op += SP + SP + SP + "std::size_t const oy_" + std::to_string(d) + + " = (out_idx / (" + outputStrides[d].GetVal() + ")) % (" + + fShapeYNotPruned[d].GetVal() + ");\n"; + } } + op += "\n"; } - op += "\n"; // ---- thread-stride partial accumulation over reduction axis ---- std::string startVal; @@ -371,23 +450,36 @@ public: else if (Op == ReduceMax) startVal = "std::numeric_limits::lowest()"; else startVal = "static_cast(0)"; op += SP + SP + SP + "T partial = " + startVal + ";\n"; - op += SP + SP + SP + "for (std::size_t r = thread_id; r < reducedLength; r += 256u) {\n"; - - // Decode flat reduction index r into per-axis coordinates. - for (std::size_t ri = 0; ri < redAxes.size(); ++ri) { - std::size_t rd = redAxes[ri]; - op += SP + SP + SP + SP + "std::size_t const r_" + std::to_string(rd) - + " = (r / " + std::to_string(redStrides[ri]) + "u) % " - + std::to_string(fShapeX[rd]) + "u;\n"; - } + if (useGroupsDecomp) + op += SP + SP + SP + "for (std::size_t r = thread_id + group_id * " + bsExpr + "; r < reducedLength; r += " + + bsExpr + " * groups) {\n"; + else + op += SP + SP + SP + "for (std::size_t r = thread_id; r < reducedLength; r += " + bsExpr + ") {\n"; - // Compute flat input index. - op += SP + SP + SP + SP + "std::size_t const in_idx =\n"; - for (std::size_t d = 0; d < Dx; ++d) { - bool isReduced = std::find(redAxes.begin(), redAxes.end(), d) != redAxes.end(); - std::string coord = isReduced ? "r_" + std::to_string(d) : "oy_" + std::to_string(d); - op += SP + SP + SP + SP + SP + coord + " * " + std::to_string(inputStrides[d]) + "u"; - op += (d + 1 < Dx) ? " +\n" : ";\n"; + if (reduceDims == kLast) { + // reduced axes are the trailing contiguous axes: input is laid out + // as [outputLength, reducedLength] row-major. + op += SP + SP + SP + SP + "std::size_t const in_idx = out_idx * reducedLength + r;\n"; + } else if (reduceDims == kFirst) { + // reduced axes are the leading contiguous axes: input is laid out + // as [reducedLength, outputLength] row-major. + op += SP + SP + SP + SP + "std::size_t const in_idx = r * outputLength + out_idx;\n"; + } else { + // Decode flat reduction index r into per-axis coordinates. + for (std::size_t ri = 0; ri < redAxes.size(); ++ri) { + std::size_t rd = redAxes[ri]; + op += SP + SP + SP + SP + "std::size_t const r_" + std::to_string(rd) + + " = (r / (" + redStrides[ri] + ")) % (" + + fShapeX[rd].GetVal() + ");\n"; + } + + // Compute flat input index. + op += SP + SP + SP + SP + "std::size_t const in_idx =\n"; + for (std::size_t d = 0; d < Dx; ++d) { + std::string coord = IsReducedAxis(d) ? "r_" + std::to_string(d) : "oy_" + std::to_string(d); + op += SP + SP + SP + SP + SP + coord + " * (" + inputStrides[d].GetVal() + ")"; + op += (d + 1 < Dx) ? " +\n" : ";\n"; + } } // Partial accumulation step. @@ -406,8 +498,12 @@ public: op += SP + SP + SP + "shmem[thread_id] = partial;\n"; op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; - // ---- binary tree reduction within the block ---- - op += SP + SP + SP + "for (std::size_t s = 128u; s > 0u; s >>= 1u) {\n"; + // ---- shared-memory tree reduction down to one warp's worth of partials ---- + // Stops once the active thread count fits in a single warp; the final + // warp-width combine below uses shuffles instead of shared memory. + op += SP + SP + SP + "std::int32_t const warpSize = alpaka::warp::getSize(acc);\n"; + op += SP + SP + SP + "for (std::size_t s = " + halfBsExpr + + "; s >= static_cast(warpSize); s >>= 1u) {\n"; op += SP + SP + SP + SP + "if (thread_id < s) {\n"; if (Op == ReduceProd) op += SP + SP + SP + SP + SP + "shmem[thread_id] *= shmem[thread_id + s];\n"; @@ -419,49 +515,239 @@ public: op += SP + SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; op += SP + SP + SP + "}\n\n"; - // ---- thread 0 writes the final result ---- - op += SP + SP + SP + "if (thread_id == 0u) {\n"; - op += SP + SP + SP + SP + "T result = shmem[0];\n"; - if (Op == ReduceMean) - op += SP + SP + SP + SP + "result /= static_cast(" + std::to_string(reducedLength) + "u);\n"; - else if (Op == ReduceL2) - op += SP + SP + SP + SP + "result = std::sqrt(result);\n"; - op += SP + SP + SP + SP + "output[out_idx] = result;\n"; - op += SP + SP + SP + "}\n"; + op += SP + SP + SP + "if (thread_id < static_cast(warpSize)) {\n"; + op += SP + SP + SP + SP + "T val = shmem[thread_id];\n"; + op += SP + SP + SP + SP + "for (std::int32_t offset = warpSize / 2; offset > 0; offset >>= 1) {\n"; + op += SP + SP + SP + SP + SP + + "T const other = alpaka::warp::shfl_down(acc, val, static_cast(offset));\n"; + if (Op == ReduceProd) + op += SP + SP + SP + SP + SP + "val *= other;\n"; + else if (Op == ReduceMax) + op += SP + SP + SP + SP + SP + "if (other > val) val = other;\n"; + else + op += SP + SP + SP + SP + SP + "val += other;\n"; + op += SP + SP + SP + SP + "}\n\n"; + + op += SP + SP + SP + SP + "if (thread_id == 0u) {\n"; + if (dyn) { + op += SP + SP + SP + SP + SP + "T result = val;\n"; + op += SP + SP + SP + SP + SP + "if (groups == 1u) {\n"; + if (Op == ReduceMean) + op += SP + SP + SP + SP + SP + SP + "result /= static_cast(reducedLength);\n"; + else if (Op == ReduceL2) + op += SP + SP + SP + SP + SP + SP + "result = std::sqrt(result);\n"; + op += SP + SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + SP + "output[flat_block] = result;\n"; + } else if (lp.groups > 1) { + op += SP + SP + SP + SP + SP + "output[flat_block] = val;\n"; + } else { + op += SP + SP + SP + SP + SP + "T result = val;\n"; + if (Op == ReduceMean) + op += SP + SP + SP + SP + SP + "result /= static_cast(reducedLength);\n"; + else if (Op == ReduceL2) + op += SP + SP + SP + SP + SP + "result = std::sqrt(result);\n"; + op += SP + SP + SP + SP + SP + "output[out_idx] = result;\n"; + } + op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + "}\n"; // end if (thread_id < warpSize) op += SP + SP + "}\n"; // end operator() op += SP + "};\n"; // end struct + + if (useGroupsDecomp) { + std::string kname2 = "ReduceFinalizeKernel_" + Name() + "_" + fNY; + std::string groupsDesc = dyn ? "a runtime-determined number of" : std::to_string(lp.groups); + op += "\n//------ " + Name() + "_FINALIZE_KERNEL_ALPAKA (combine " + groupsDesc + + " partial blocks per output)\n"; + op += SP + "struct " + kname2 + " {\n"; + op += SP + SP + "template\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(\n"; + op += SP + SP + SP + "TAcc const& acc,\n"; + op += SP + SP + SP + "T const* __restrict__ partial,\n"; + op += SP + SP + SP + "T* __restrict__ output,\n"; + op += SP + SP + SP + "std::size_t const reducedLength,\n"; + op += SP + SP + SP + "std::size_t const outputLength,\n"; + op += SP + SP + SP + "std::size_t const groups) const {\n\n"; + op += SP + SP + SP + "auto const out_idx = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (out_idx >= outputLength) return;\n\n"; + op += SP + SP + SP + "T result = " + startVal + ";\n"; + op += SP + SP + SP + "for (std::size_t g = 0; g < groups; ++g) {\n"; + op += SP + SP + SP + SP + "T const v = partial[g * outputLength + out_idx];\n"; + if (Op == ReduceProd) + op += SP + SP + SP + SP + "result *= v;\n"; + else if (Op == ReduceMax) + op += SP + SP + SP + SP + "if (v > result) result = v;\n"; + else + op += SP + SP + SP + SP + "result += v;\n"; + op += SP + SP + SP + "}\n\n"; + if (Op == ReduceMean) + op += SP + SP + SP + "result /= static_cast(reducedLength);\n"; + else if (Op == ReduceL2) + op += SP + SP + SP + "result = std::sqrt(result);\n"; + op += SP + SP + SP + "output[out_idx] = result;\n"; + op += SP + SP + "}\n"; // end operator() + op += SP + "};\n"; // end struct + } + return op; } std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string /*opName*/) override { + bool dyn = fInputDimShape; + GpuLaunchParams lp = ComputeGpuLaunchParams(); + bool needsScratch = dyn || lp.groups > 1; std::string kname = "ReduceKernel_" + Name() + "_" + fNY; - return SP + kname + " reduceKernel_" + Name() + "_" + fNY + ";\n"; + std::string op = SP + kname + " reduceKernel_" + Name() + "_" + fNY + ";\n"; + if (needsScratch) { + std::string kname2 = "ReduceFinalizeKernel_" + Name() + "_" + fNY; + op += SP + kname2 + " reduceFinalizeKernel_" + Name() + "_" + fNY + ";\n"; + op += SP + "std::unique_ptr reduceScratch_" + fNY + ";\n"; + op += SP + "std::size_t reduceScratchCapacity_" + fNY + " = 0;\n"; + } + return op; } - std::string Generate_GPU_ALPAKA(std::string /*opName*/) override { + std::string Generate_GPU_ALPAKA(std::string /*opName*/, const std::vector &dynParamNames) override { if (fShapeX.empty() || fShapeY.empty()) throw std::runtime_error("SOFIE Reduce Op called to Generate without being initialized first"); - std::size_t inputLength = ConvertShapeToLength(fShapeX); - std::size_t outputLength = ConvertShapeToLength(fShapeY); - std::size_t reducedLength = inputLength / outputLength; + std::string outputLength = ConvertDimShapeToLength(fShapeY); + std::string reducedLength = ReducedLengthExpr(); std::string kname = "reduceKernel_" + Name() + "_" + fNY; + bool dyn = fInputDimShape; + GpuLaunchParams lp = ComputeGpuLaunchParams(); + std::string bs = std::to_string(lp.blockSize); + + std::string dynArgs; + for (auto &p : dynParamNames) dynArgs += ", static_cast(" + p + ")"; + std::stringstream out; out << "\n//------ " << Name() << "_GPU_ALPAKA\n"; - // Grid: one block per output element; Block: 256 threads cooperate to - // reduce the corresponding slice. - out << SP << "alpaka::WorkDivMembers workDiv_" << fNY << "(\n"; - out << SP << SP << "Vec::all(Idx{" << outputLength << "u}),\n"; - out << SP << SP << "Vec::all(Idx{256u}),\n"; - out << SP << SP << "Vec::all(Idx{1u}));\n"; - out << SP << "alpaka::exec(queue, workDiv_" << fNY - << ", " << kname - << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << reducedLength << "u)" - << ", static_cast(" << outputLength << "u));\n"; + + if (dyn) { + std::string rl = "reduceLaunch_" + fNY; + out << SP << "std::size_t const " << rl << "_reducedLength = " << reducedLength << ";\n"; + out << SP << "std::size_t const " << rl << "_outputLength = " << outputLength << ";\n"; + out << SP << "std::size_t " << rl << "_blockSize = 32u;\n"; + out << SP << "while (" << rl << "_blockSize < " << rl << "_reducedLength && " << rl + << "_blockSize < 256u) " << rl << "_blockSize *= 2u;\n"; + out << SP << "std::size_t " << rl << "_groups = 1u;\n"; + out << SP << "if (" << rl << "_outputLength > 0 && " << rl << "_outputLength < 132u) {\n"; + out << SP << SP << "std::size_t const " << rl << "_maxGroups = 132u / " << rl << "_outputLength;\n"; + out << SP << SP << "std::size_t const " << rl << "_chunks = (" << rl << "_reducedLength + " << rl + << "_blockSize - 1u) / " << rl << "_blockSize;\n"; + out << SP << SP << rl << "_groups = std::max(1u, std::min(" << rl << "_maxGroups, " << rl + << "_chunks));\n"; + out << SP << "}\n\n"; + + out << SP << "if (" << rl << "_groups > 1u) {\n"; + out << SP << SP << "std::size_t const " << rl << "_capacity = " << rl << "_outputLength * " << rl + << "_groups;\n"; + out << SP << SP << "if (!reduceScratch_" << fNY << " || reduceScratchCapacity_" << fNY << " < " << rl + << "_capacity) {\n"; + out << SP << SP << SP << "reduceScratch_" << fNY + << " = std::make_unique(alpaka::allocBuf(devAcc, Ext1D::all(Idx{" << rl + << "_capacity})));\n"; + out << SP << SP << SP << "reduceScratchCapacity_" << fNY << " = " << rl << "_capacity;\n"; + out << SP << SP << "}\n"; + out << SP << SP << "alpaka::WorkDivMembers workDiv_" << fNY << "(\n"; + out << SP << SP << SP << "Vec::all(Idx{" << rl << "_capacity}),\n"; + out << SP << SP << SP << "Vec::all(Idx{" << rl << "_blockSize}),\n"; + out << SP << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << SP << "alpaka::exec(queue, workDiv_" << fNY + << ", " << kname + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(*reduceScratch_" << fNY << ")" + << ", " << rl << "_reducedLength" + << ", " << rl << "_outputLength" + << ", " << rl << "_groups" + << dynArgs << ");\n\n"; + + std::string kname2 = "reduceFinalizeKernel_" + Name() + "_" + fNY; + out << SP << SP << "auto const elementsPerThread_" << fNY << "_finalize = Vec::all(static_cast(1));\n"; + out << SP << SP << "auto const elementsPerGrid_" << fNY << "_finalize = Vec::all(Idx{" << rl + << "_outputLength});\n"; + out << SP << SP << "auto const workDiv_" << fNY << "_finalize = sofie_workdiv(elementsPerGrid_" << fNY + << "_finalize);\n"; + out << SP << SP << "alpaka::exec(queue, workDiv_" << fNY << "_finalize" + << ", " << kname2 + << ", alpaka::getPtrNative(*reduceScratch_" << fNY << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", " << rl << "_reducedLength" + << ", " << rl << "_outputLength" + << ", " << rl << "_groups);\n"; + out << SP << "} else {\n"; + out << SP << SP << "alpaka::WorkDivMembers workDiv_" << fNY << "(\n"; + out << SP << SP << SP << "Vec::all(Idx{" << rl << "_outputLength}),\n"; + out << SP << SP << SP << "Vec::all(Idx{" << rl << "_blockSize}),\n"; + out << SP << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << SP << "alpaka::exec(queue, workDiv_" << fNY + << ", " << kname + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", " << rl << "_reducedLength" + << ", " << rl << "_outputLength" + << ", " << rl << "_groups" + << dynArgs << ");\n"; + out << SP << "}\n"; + } else if (lp.groups > 1) { + // Two-pass reduction: split the reduction axis across `groups` blocks + // per output element so a small number of outputs (which would + // otherwise launch just outputLength blocks total, badly + // under-occupying the GPU) still gets enough parallel work. Pass 1 + // writes unfinalized partials to a scratch buffer; pass 2 combines the + // `groups` partials per output and applies the finalization. The + // buffer size is a fixed literal for static shapes, so this cache + // check allocates exactly once (on the first infer() call) and + // every later call just reuses it. + std::string totalBlocks = std::to_string(std::stoul(outputLength) * lp.groups); + out << SP << "if (!reduceScratch_" << fNY << ") {\n"; + out << SP << SP << "reduceScratch_" << fNY + << " = std::make_unique(alpaka::allocBuf(devAcc, Ext1D::all(Idx{" << totalBlocks + << "})));\n"; + out << SP << SP << "reduceScratchCapacity_" << fNY << " = " << totalBlocks << ";\n"; + out << SP << "}\n"; + + out << SP << "alpaka::WorkDivMembers workDiv_" << fNY << "(\n"; + out << SP << SP << "Vec::all(Idx{" << totalBlocks << "u}),\n"; + out << SP << SP << "Vec::all(Idx{" << bs << "u}),\n"; + out << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << "alpaka::exec(queue, workDiv_" << fNY + << ", " << kname + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(*reduceScratch_" << fNY << ")" + << ", static_cast(" << reducedLength << ")" + << ", static_cast(" << outputLength << ")" + << ", static_cast(" << lp.groups << "u)" + << dynArgs << ");\n"; + + std::string kname2 = "reduceFinalizeKernel_" + Name() + "_" + fNY; + out << SP << "auto const elementsPerThread_" << fNY << "_finalize = Vec::all(static_cast(1));\n"; + out << SP << "auto const elementsPerGrid_" << fNY << "_finalize = Vec::all(Idx{" << outputLength << "});\n"; + out << SP << "auto const workDiv_" << fNY << "_finalize = sofie_workdiv(elementsPerGrid_" << fNY << "_finalize);\n"; + out << SP << "alpaka::exec(queue, workDiv_" << fNY << "_finalize" + << ", " << kname2 + << ", alpaka::getPtrNative(*reduceScratch_" << fNY << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", static_cast(" << reducedLength << ")" + << ", static_cast(" << outputLength << ")" + << ", static_cast(" << lp.groups << "u));\n"; + } else { + // Grid: one block per output element; Block: threads cooperate to + // reduce the corresponding slice. + out << SP << "alpaka::WorkDivMembers workDiv_" << fNY << "(\n"; + out << SP << SP << "Vec::all(Idx{" << outputLength << "}),\n"; + out << SP << SP << "Vec::all(Idx{" << bs << "u}),\n"; + out << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << "alpaka::exec(queue, workDiv_" << fNY + << ", " << kname + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << ", static_cast(" << reducedLength << ")" + << ", static_cast(" << outputLength << ")" + << dynArgs << ");\n"; + } return out.str(); } diff --git a/core/inc/SOFIE/ROperator_SDPA.hxx b/core/inc/SOFIE/ROperator_SDPA.hxx index c02cf508..23554aa5 100644 --- a/core/inc/SOFIE/ROperator_SDPA.hxx +++ b/core/inc/SOFIE/ROperator_SDPA.hxx @@ -158,6 +158,40 @@ public: out += SP + SP + SP + "std::size_t const s = global_idx % S;\n\n"; out += SP + SP + SP + "std::size_t const qBase = b*H*S*D + h*S*D + s*D;\n\n"; + // Dv static -> fuse into a single online-softmax pass (the flash-attention + // recurrence): the QK dot product is then computed exactly once per (query, + // key) pair instead of twice (once to find the row max, once to weight V), + // by rescaling the running weighted-sum accumulator whenever a new max is + // found. Needs Dv known at codegen time to size the register accumulator; + // falls back to the two-pass version + bool canFuse = IsInteger(fDv); + if (canFuse) { + std::string dv = fDv; + out += SP + SP + SP + "// Fused online-softmax pass: max, sum and weighted-V accumulator\n"; + out += SP + SP + SP + "// updated together, each rescaled when a new row max is found.\n"; + out += SP + SP + SP + "T acc_v[" + dv + "] = {};\n"; + out += SP + SP + SP + "T max_score = -std::numeric_limits::max();\n"; + out += SP + SP + SP + "T sum_exp = static_cast(0);\n\n"; + + out += SP + SP + SP + "for (std::size_t j = 0; j < S; ++j) {\n"; + out += SP + SP + SP + SP + "std::size_t const kBase = b*H*S*D + h*S*D + j*D;\n"; + out += SP + SP + SP + SP + "T dot = static_cast(0);\n"; + out += SP + SP + SP + SP + "for (std::size_t d = 0; d < D; ++d) dot += Q[qBase+d] * K[kBase+d];\n"; + out += SP + SP + SP + SP + "T sc = dot * scale;\n"; + if (fHasMask) + out += SP + SP + SP + SP + "sc += mask[b*H*S*S + h*S*S + s*S + j];\n"; + out += SP + SP + SP + SP + "T const new_max = (sc > max_score) ? sc : max_score;\n"; + out += SP + SP + SP + SP + "T const correction = alpaka::math::exp(acc, max_score - new_max);\n"; + out += SP + SP + SP + SP + "T const p = alpaka::math::exp(acc, sc - new_max);\n"; + out += SP + SP + SP + SP + "sum_exp = sum_exp * correction + p;\n"; + out += SP + SP + SP + SP + "std::size_t const vBase = b*H*S*Dv + h*S*Dv + j*Dv;\n"; + out += SP + SP + SP + SP + "for (std::size_t d = 0; d < Dv; ++d)\n"; + out += SP + SP + SP + SP + SP + "acc_v[d] = acc_v[d] * correction + p * V[vBase+d];\n"; + out += SP + SP + SP + SP + "max_score = new_max;\n"; + out += SP + SP + SP + "}\n\n"; + out += SP + SP + SP + "for (std::size_t d = 0; d < Dv; ++d)\n"; + out += SP + SP + SP + SP + "Y[b*H*S*Dv + h*S*Dv + s*Dv + d] = acc_v[d] / sum_exp;\n"; + } else { out += SP + SP + SP + "// scores computed in thread-local storage\n"; out += SP + SP + SP + "T max_score = -std::numeric_limits::max();\n"; out += SP + SP + SP + "T sum_exp = static_cast(0);\n\n"; @@ -181,7 +215,7 @@ public: out += SP + SP + SP + SP + "T sc = dot * scale;\n"; if (fHasMask) out += SP + SP + SP + SP + "sc += mask[b*H*S*S + h*S*S + s*S + j];\n"; - out += SP + SP + SP + SP + "T const a_j = SOFIE_DEVICE_exp(acc, sc - max_score);\n"; + out += SP + SP + SP + SP + "T const a_j = alpaka::math::exp(acc, sc - max_score);\n"; out += SP + SP + SP + SP + "sum_exp += a_j;\n"; out += SP + SP + SP + SP + "std::size_t const vBase = b*H*S*Dv + h*S*Dv + j*Dv;\n"; out += SP + SP + SP + SP + "for (std::size_t d = 0; d < Dv; ++d)\n"; @@ -190,6 +224,7 @@ public: out += SP + SP + SP + "// Normalize\n"; out += SP + SP + SP + "for (std::size_t d = 0; d < Dv; ++d)\n"; out += SP + SP + SP + SP + "Y[b*H*S*Dv + h*S*Dv + s*Dv + d] /= sum_exp;\n"; + } out += SP + SP + "}\n" + SP + "};\n"; return out; } diff --git a/core/inc/SOFIE/ROperator_ScatterElements.hxx b/core/inc/SOFIE/ROperator_ScatterElements.hxx index 3cedaa7e..9a6e936f 100644 --- a/core/inc/SOFIE/ROperator_ScatterElements.hxx +++ b/core/inc/SOFIE/ROperator_ScatterElements.hxx @@ -285,7 +285,7 @@ public: // // Non-"add" reductions retain the original atomicXxx kernel. // ----------------------------------------------------------------------- - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) { throw std::runtime_error("SOFIE ScatterElements Op called to Generate without being initialized first"); @@ -318,6 +318,8 @@ public: op += SP + SP + SP + "int64_t const* I_sorted,\n"; // axis index, sorted op += SP + SP + SP + "T const* U,\n"; op += SP + SP + SP + "int32_t const* sortPerm,\n"; // argsort of I + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalUpdates,\n"; op += SP + SP + SP + "std::size_t const numFeatures) const {\n\n"; @@ -359,6 +361,8 @@ public: op += SP + SP + SP + "T* Y,\n"; op += SP + SP + SP + "int64_t const* I,\n"; op += SP + SP + SP + "T const* U,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -410,7 +414,7 @@ std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override return SP + "ScatterElementsKernel_" + opName + " scatterElementsKernel_" + opName + ";\n"; } -std::string Generate_GPU_ALPAKA(std::string opName) override { +std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) { throw std::runtime_error("SOFIE ScatterElements Op called to Generate without being initialized first"); @@ -442,8 +446,10 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNI << "_sortedI)" << ", alpaka::getPtrNative(deviceBuf_" << fNU << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNI << "_sortPerm)" - << ", static_cast(" << numRows << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNI << "_sortPerm)"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << numRows << ")" << ", static_cast(" << numFeatures << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; } else { @@ -455,8 +461,10 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { << ", scatterElementsKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNI << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNU << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNU << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; } return out.str(); diff --git a/core/inc/SOFIE/ROperator_ScatterND.hxx b/core/inc/SOFIE/ROperator_ScatterND.hxx index a2bb02da..e254dc5a 100644 --- a/core/inc/SOFIE/ROperator_ScatterND.hxx +++ b/core/inc/SOFIE/ROperator_ScatterND.hxx @@ -20,16 +20,18 @@ private: std::string fNY; std::string fReduction; - std::vector fShapeData; - std::vector fShapeIndices; - std::vector fShapeY; + std::vector fShapeData; + std::vector fShapeIndices; + std::vector fShapeY; - size_t fK = 0; - size_t fSliceSize = 1; - size_t fNumOuter = 1; + size_t fK = 0; + std::string fSliceSize; // host size expression + std::string fNumOuter; // host size expression std::string fType; + static std::string sz(const std::string &e) { return "static_cast(" + e + ")"; } + public: ROperator_ScatterND() {} @@ -64,8 +66,8 @@ public: if (!model.CheckIfTensorAlreadyExist(fNUpdates)) throw std::runtime_error("SOFIE ScatterND: updates tensor " + fNUpdates + " not found"); - fShapeData = model.GetTensorShape(fNData); - fShapeIndices = model.GetTensorShape(fNIndices); + fShapeData = model.GetDimTensorShape(fNData); + fShapeIndices = model.GetDimTensorShape(fNIndices); size_t r = fShapeData.size(); size_t q = fShapeIndices.size(); @@ -75,29 +77,30 @@ public: if (q < 1) throw std::runtime_error("SOFIE ScatterND: indices rank must be >= 1"); - fK = fShapeIndices.back(); + if (fShapeIndices.back().isParam) + throw std::runtime_error("SOFIE ScatterND: the last indices dimension (index tuple length) " + "must be static - a dynamic index-tuple length is not supported"); + fK = fShapeIndices.back().dim; if (fK > r) throw std::runtime_error("SOFIE ScatterND: indices.shape[-1] must be <= data rank"); - fNumOuter = 1; - for (size_t i = 0; i + 1 < q; i++) - fNumOuter *= fShapeIndices[i]; + std::vector outerShape(fShapeIndices.begin(), fShapeIndices.end() - 1); + fNumOuter = ConvertDimShapeToLength(outerShape); - fSliceSize = 1; - for (size_t i = fK; i < r; i++) - fSliceSize *= fShapeData[i]; + std::vector sliceShape(fShapeData.begin() + fK, fShapeData.end()); + fSliceSize = ConvertDimShapeToLength(sliceShape); fShapeY = fShapeData; model.AddIntermediateTensor(fNY, model.GetTensorType(fNData), fShapeY); fType = ConvertTypeToString(model.GetTensorType(fNData)); if (model.Verbose()) - std::cout << "ScatterND: data " << ConvertShapeToString(fShapeData) - << " indices " << ConvertShapeToString(fShapeIndices) + std::cout << "ScatterND: data " << ConvertDimShapeToString(fShapeData) + << " indices " << ConvertDimShapeToString(fShapeIndices) << " k=" << fK << " numOuter=" << fNumOuter << " sliceSize=" << fSliceSize << " reduction=" << (fReduction.empty() ? "none" : fReduction) - << " -> " << fNY << " " << ConvertShapeToString(fShapeY) << "\n"; + << " -> " << fNY << " " << ConvertDimShapeToString(fShapeY) << "\n"; } std::string GenerateInitCode() override { return ""; } @@ -107,17 +110,17 @@ public: if (fShapeY.empty()) throw std::runtime_error("SOFIE ScatterND: Generate called before Initialize"); - size_t dataSize = ConvertShapeToLength(fShapeY); + std::string dataSize = ConvertDimShapeToLength(fShapeY); auto stridesData = UTILITY::ComputeStrideFromShape(fShapeData); std::stringstream out; out << SP << "//------- ScatterND " << opName << "\n"; out << SP << "std::copy(tensor_" << fNData - << ", tensor_" << fNData << " + " << dataSize + << ", tensor_" << fNData << " + static_cast(" << dataSize << ")" << ", tensor_" << fNY << ");\n"; - out << SP << "for (std::size_t _i = 0; _i < " << fNumOuter << "; ++_i) {\n"; + out << SP << "for (std::size_t _i = 0; _i < static_cast(" << fNumOuter << "); ++_i) {\n"; out << SP << SP << "std::size_t _out_base = 0;\n"; for (size_t j = 0; j < fK; ++j) { @@ -125,15 +128,15 @@ public: out << SP << SP << SP << "int64_t _idx = tensor_" << fNIndices << "[_i * " << fK << " + " << j << "];\n"; out << SP << SP << SP << "if (_idx < 0) _idx += " - << static_cast(fShapeData[j]) << ";\n"; + << "static_cast(" << fShapeData[j].GetVal() << ");\n"; out << SP << SP << SP << "_out_base += static_cast(_idx) * " - << stridesData[j] << ";\n"; + << stridesData[j].GetVal() << ";\n"; out << SP << SP << "}\n"; } - out << SP << SP << "for (std::size_t _s = 0; _s < " << fSliceSize << "; ++_s) {\n"; + out << SP << SP << "for (std::size_t _s = 0; _s < static_cast(" << fSliceSize << "); ++_s) {\n"; out << SP << SP << SP << "std::size_t const _out_idx = _out_base + _s;\n"; - out << SP << SP << SP << "std::size_t const _upd_idx = _i * " << fSliceSize << " + _s;\n"; + out << SP << SP << SP << "std::size_t const _upd_idx = _i * static_cast(" << fSliceSize << ") + _s;\n"; if (fReduction.empty() || fReduction == "none") { out << SP << SP << SP << "tensor_" << fNY << "[_out_idx] = tensor_" << fNUpdates << "[_upd_idx];\n"; @@ -156,7 +159,7 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE ScatterND: Generate_GPU_Kernel_ALPAKA called before Initialize"); @@ -173,6 +176,8 @@ public: op += SP + SP + SP + "T* Y,\n"; op += SP + SP + SP + "int64_t const* indices,\n"; op += SP + SP + SP + "T const* updates,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const numOuter,\n"; op += SP + SP + SP + "std::size_t const sliceSize) const {\n\n"; @@ -188,9 +193,9 @@ public: op += SP + SP + SP + SP + SP + "int64_t idx = indices[i * " + std::to_string(fK) + "u + " + std::to_string(j) + "u];\n"; op += SP + SP + SP + SP + SP - + "if (idx < 0) idx += " + std::to_string(static_cast(fShapeData[j])) + ";\n"; + + "if (idx < 0) idx += " + fShapeData[j].GetVal() + ";\n"; op += SP + SP + SP + SP + SP - + "out_base += static_cast(idx) * " + std::to_string(stridesData[j]) + "u;\n"; + + "out_base += static_cast(idx) * " + sz(stridesData[j].GetVal()) + ";\n"; op += SP + SP + SP + SP + "}\n"; } @@ -224,7 +229,7 @@ public: return SP + kname + " scatterNDKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeY.empty()) throw std::runtime_error("SOFIE ScatterND: Generate_GPU_ALPAKA called before Initialize"); @@ -241,8 +246,10 @@ public: << ", scatterNDKernel_" << opName << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNIndices << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNUpdates << ")" - << ", static_cast(" << fNumOuter << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNUpdates << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << fNumOuter << ")" << ", static_cast(" << fSliceSize << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; diff --git a/core/inc/SOFIE/ROperator_Slice.hxx b/core/inc/SOFIE/ROperator_Slice.hxx index fb738cf1..652d6cd9 100644 --- a/core/inc/SOFIE/ROperator_Slice.hxx +++ b/core/inc/SOFIE/ROperator_Slice.hxx @@ -32,8 +32,6 @@ private: std::vector fShapeOutput; // output shape std::vector fOutputShapeData; // output shape data in case output is a shape param tensor - // saved Start/End.Steps are corrected from initial ONNX for negative/default values - // and are available for each axis std::vector fStart; // starting values of slices for all axes std::vector fEnd; // End values of slices for all axes std::vector fSteps; // step values of slices for all axes @@ -494,18 +492,22 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; if (fShapeInput.empty() || fShapeOutput.empty()) throw std::runtime_error("SOFIE Slice Op called to Generate without being initialized first"); + if (fIsStartUndef || fIsEndUndef) + throw std::runtime_error("SOFIE Slice Op " + opName + + ": GPU codegen does not support start/end supplied by a runtime tensor's data " + "(as opposed to a named dynamic shape dimension) - only CPU codegen supports this case"); + const std::size_t D = fShapeInput.size(); auto inputStrides = UTILITY::ComputeStrideFromShape(fShapeInput); auto outputStrides = UTILITY::ComputeStrideFromShape(fShapeOutput); - std::size_t totalElements = ConvertShapeToLength(fShapeOutput); std::string kname = "SliceKernel_" + opName; std::string op; @@ -516,6 +518,8 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ input,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -524,26 +528,18 @@ public: op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - for (std::size_t d = 0; d < D; ++d) { - op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + outputStrides[d].GetVal() + "u) % " - + fShapeOutput[d].GetVal() + "u;\n"; - } + EmitOutputCoordsFromThreadIdx(op, SP + SP + SP + SP, outputStrides, fShapeOutput); op += "\n"; - // Map each output coord back to input coord: - // input_coord[d] = fStart[d] + out_d * fSteps[d] - // Negative steps are supported naturally since fStart/fEnd/fSteps are - // already corrected for negative/default values during Initialize(). op += SP + SP + SP + SP + "std::size_t const input_idx =\n"; for (std::size_t d = 0; d < D; ++d) { // input coordinate for this dim: start + out_d * step - std::string input_coord = "(" + fStart[d].GetVal() - + " + out_" + std::to_string(d) - + " * " + fSteps[d].GetVal() + ")"; + std::string input_coord = "((" + fStart[d].GetVal() + + ") + out_" + std::to_string(d) + + " * (" + fSteps[d].GetVal() + "))"; op += SP + SP + SP + SP + SP + "static_cast(" + input_coord + ")" - + " * " + inputStrides[d].GetVal() + "u"; + + " * (" + inputStrides[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -561,13 +557,13 @@ public: return SP + kname + " sliceKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; if (fShapeInput.empty() || fShapeOutput.empty()) throw std::runtime_error("SOFIE Slice Op called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeOutput); + auto totalElements = ConvertDimShapeToLength(fShapeOutput); std::string kname = "sliceKernel_" + opName; std::stringstream out; @@ -578,8 +574,10 @@ public: out << SP << "alpaka::exec(queue, workDiv_" << opName << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNData << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNOutput << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNOutput << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_Softmax.hxx b/core/inc/SOFIE/ROperator_Softmax.hxx index 5626c0f5..5523d5a6 100644 --- a/core/inc/SOFIE/ROperator_Softmax.hxx +++ b/core/inc/SOFIE/ROperator_Softmax.hxx @@ -185,6 +185,140 @@ public: } return out.str(); } + + // Threads per row: pow2 (the tree reduction halves its stride), clamped to a warp and + // the 1024/block limit, 256 if dynamic. + static size_t BlockSize(const std::string &rowLength) { + if (!IsInteger(rowLength)) + return 256; + size_t n = std::stoul(rowLength); + size_t p = 1; + while (p < n) p <<= 1; + if (p < 32) p = 32; + if (p > 1024) p = 1024; + return p; + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { + if (fShape.empty()) + throw std::runtime_error("SOFIE Softmax called to Generate_GPU_Kernel_ALPAKA without being initialized first"); + + opName = "op_" + opName; + std::string kname = "SoftmaxKernel_" + opName; + + size_t axis = fAttrAxis < 0 ? fShape.size() + fAttrAxis : fAttrAxis; + //a kernel row = one slice along the axis: nElements is the row length, strideAxis the step between its elements + auto s = UTILITY::ComputeSliceInfo(fShape, axis); + std::string bs = std::to_string(BlockSize(s.nElements)); + + // block-per-row online softmax + + std::string op; + op = "\n//------ SOFTMAX_KERNEL_ALPAKA\n"; + op += SP + "struct " + kname + " {\n"; + op += SP + SP + "template\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(\n"; + op += SP + SP + SP + "TAcc const& acc,\n"; + op += SP + SP + SP + "T const* __restrict__ X,\n"; + op += SP + SP + SP + "T* __restrict__ Y,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; + op += SP + SP + SP + "std::size_t const numRows) const {\n\n"; + + // declared before the early return so every thread reaches the collective declaration + op += SP + SP + SP + "auto& smax = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto& ssum = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto const row = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "auto const tid = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (row >= numRows) return;\n\n"; + + op += SP + SP + SP + "std::size_t const axis_size = " + s.nElements + ";\n"; + op += SP + SP + SP + "std::size_t const inner_stride = " + s.strideAxis + ";\n"; + op += SP + SP + SP + "std::size_t const row_block = axis_size * inner_stride;\n"; + op += SP + SP + SP + "std::size_t const row_base = (row / inner_stride) * row_block + (row % inner_stride);\n\n"; + + op += SP + SP + SP + "// fused pass: running (max, sum) per thread\n"; + op += SP + SP + SP + "T m = X[row_base];\n"; + op += SP + SP + SP + "T d = static_cast(0);\n"; + op += SP + SP + SP + "for (std::size_t l = tid; l < axis_size; l += " + bs + "u) {\n"; + op += SP + SP + SP + SP + "T x = X[row_base + l * inner_stride];\n"; + op += SP + SP + SP + SP + "T m_new = (x > m) ? x : m;\n"; + op += SP + SP + SP + SP + "d = d * exp(acc, m - m_new) + exp(acc, x - m_new);\n"; + op += SP + SP + SP + SP + "m = m_new;\n"; + op += SP + SP + SP + "}\n"; + op += SP + SP + SP + "smax[tid] = m;\n"; + op += SP + SP + SP + "ssum[tid] = d;\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; + + // (m_a, d_a) + (m_b, d_b) = (max, d_a*exp(m_a-max) + d_b*exp(m_b-max)) + op += SP + SP + SP + "// combined (max, sum) tree reduction\n"; + op += SP + SP + SP + "for (std::size_t s = " + bs + "u / 2u; s >= 1u; s /= 2u) {\n"; + op += SP + SP + SP + SP + "if (tid < s) {\n"; + op += SP + SP + SP + SP + SP + "T m_a = smax[tid];\n"; + op += SP + SP + SP + SP + SP + "T m_b = smax[tid + s];\n"; + op += SP + SP + SP + SP + SP + "T m_r = (m_b > m_a) ? m_b : m_a;\n"; + op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * exp(acc, m_a - m_r) + ssum[tid + s] * exp(acc, m_b - m_r);\n"; + op += SP + SP + SP + SP + SP + "smax[tid] = m_r;\n"; + op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + "alpaka::syncBlockThreads(acc);\n"; + op += SP + SP + SP + "}\n"; + op += SP + SP + SP + "T const vmax = smax[0];\n"; + op += SP + SP + SP + "T const sum = ssum[0];\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; + + op += SP + SP + SP + "// normalize pass\n"; + op += SP + SP + SP + "T const inv = static_cast(1) / sum;\n"; + op += SP + SP + SP + "for (std::size_t l = tid; l < axis_size; l += " + bs + "u) {\n"; + op += SP + SP + SP + SP + "std::size_t const idx = row_base + l * inner_stride;\n"; + op += SP + SP + SP + SP + "T e = exp(acc, X[idx] - vmax) * inv;\n"; + op += SP + SP + SP + SP + "Y[idx] = e;\n"; + if (fLogSoftmax) + op += SP + SP + SP + SP + "Y[idx] = log(acc, e);\n"; + op += SP + SP + SP + "}\n"; + + op += SP + SP + "}\n"; + op += SP + "};\n"; + return op; + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override { + opName = "op_" + opName; + std::string kname = "SoftmaxKernel_" + opName; + return SP + kname + " softmaxKernel_" + opName + ";\n"; + } + + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { + if (fShape.empty()) + throw std::runtime_error("SOFIE Softmax called to Generate_GPU_ALPAKA without being initialized first"); + + opName = "op_" + opName; + std::string kname = "softmaxKernel_" + opName; + + size_t axis = fAttrAxis < 0 ? fShape.size() + fAttrAxis : fAttrAxis; + auto s = UTILITY::ComputeSliceInfo(fShape, axis); //nSlices = number of rows, nElements = row length + const size_t kBlock = BlockSize(s.nElements); //threads per row + std::string dynArgs; //shape params the kernel body may name (dynamic axis or stride) + for (auto &p : dynParamNames) dynArgs += ", static_cast(" + p + ")"; + + std::stringstream out; + out << "\n//------ SOFTMAX_GPU_ALPAKA\n"; + out << SP << "alpaka::WorkDivMembers workDiv_" << opName << "(\n"; + out << SP << SP << "Vec::all(static_cast(" << s.nSlices << ")),\n"; //blocks: one per row + out << SP << SP << "Vec::all(Idx{" << kBlock << "u}),\n"; //threads per block + out << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << "auto task_" << opName << " = alpaka::createTaskKernel(workDiv_" << opName + << ", " << kname + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" + << dynArgs + << ", static_cast(" << s.nSlices << "));\n"; + out << SP << "alpaka::enqueue(queue, task_" << opName << ");\n"; + return out.str(); + } + + std::vector GetStdLibs() override { + return { std::string("cmath") }; + } }; } // namespace SOFIE diff --git a/core/inc/SOFIE/ROperator_Split.hxx b/core/inc/SOFIE/ROperator_Split.hxx index 9604ca8c..1be1e4e9 100644 --- a/core/inc/SOFIE/ROperator_Split.hxx +++ b/core/inc/SOFIE/ROperator_Split.hxx @@ -20,9 +20,9 @@ private: std::string fNX; std::string fNSplit; std::vector fNYs; - std::vector fInputShape; + std::vector fInputShape; std::vector fSplit; - std::vector> fOutputShapes; + std::vector> fOutputShapes; @@ -53,7 +53,7 @@ public: if (model.CheckIfTensorAlreadyExist(fNX) == false){ //input must be a graph input, or already initialized intermediate tensor throw std::runtime_error("SOFIE Split Op Input Tensor is not found in model"); } - fInputShape = model.GetTensorShape(fNX); + fInputShape = model.GetDimTensorShape(fNX); // correct for negative axis if (fAxis < 0) fAxis += fInputShape.size(); @@ -64,15 +64,19 @@ public: size_t nsplit = fNYs.size(); // case split tensor is empty if (fNSplit.empty()) { + if (fInputShape[fAxis].isParam) + throw std::runtime_error("SOFIE Split - cannot compute an equal split of a dynamic " + "(runtime-sized) axis " + std::to_string(fAxis) + " - provide explicit split sizes instead"); + int64_t axisDim = static_cast(fInputShape[fAxis].dim); int64_t splitValue = 0; - if (fInputShape[fAxis] % nsplit == 0) { - splitValue = fInputShape[fAxis]/nsplit; + if (axisDim % nsplit == 0) { + splitValue = axisDim/nsplit; fSplit = std::vector(nsplit, splitValue); } else { // case of not equal splitting - splitValue = std::ceil(double(fInputShape[fAxis])/nsplit); + splitValue = std::ceil(double(axisDim)/nsplit); fSplit = std::vector(nsplit-1, splitValue); - fSplit.push_back(fInputShape[fAxis] % splitValue); + fSplit.push_back(axisDim % splitValue); } } else { // get split tensor values @@ -85,22 +89,24 @@ public: fSplit = std::vector(split_data, split_data + nsplit); } // compute now the output shapes - size_t tot_split = 0; + int64_t tot_split = 0; for (size_t i = 0; i < fNYs.size(); i++) { - std::vector outputShape = fInputShape; - outputShape[fAxis] = fSplit[i]; + std::vector outputShape = fInputShape; + outputShape[fAxis] = Dim{static_cast(fSplit[i])}; tot_split += fSplit[i]; model.AddIntermediateTensor(fNYs[i], model.GetTensorType(fNX), outputShape); fOutputShapes.push_back(outputShape); } - if (tot_split != fInputShape[fAxis]) + // the total can only be validated at graph-construction time when the split + // axis itself is static,for a dynamic axis this must hold at runtime instead + if (!fInputShape[fAxis].isParam && tot_split != static_cast(fInputShape[fAxis].dim)) throw std::runtime_error("SOFIE Split - Sum of split sizes must match the input dimension along the axis"); if (model.Verbose()) { - std::cout << "Split - input shape " << ConvertShapeToString(fInputShape) << " --> "; + std::cout << "Split - input shape " << ConvertDimShapeToString(fInputShape) << " --> "; for (auto & s : fOutputShapes) - std::cout << ConvertShapeToString(s) << " "; + std::cout << ConvertDimShapeToString(s) << " "; std::cout << std::endl; } } @@ -120,10 +126,10 @@ public: out << SP << "size_t " << OpName << "_axis_offset = 0;\n"; // unroll the loop on split outputs for (size_t i = 0; i < fNYs.size(); i++) { - size_t length = ConvertShapeToLength(fOutputShapes[i]); + std::string length = ConvertDimShapeToLength(fOutputShapes[i]); auto output_strides = UTILITY::ComputeStrideFromShape(fOutputShapes[i]); - out << SP << "for (int id = 0; id < " << length << " ; id++){\n"; + out << SP << "for (int id = 0; id < static_cast(" << length << ") ; id++){\n"; // convert output index to input index out << SP << SP << "int input_index = 0;\n"; out << SP << SP << "int remaining = id;\n"; @@ -131,12 +137,12 @@ public: for (size_t k = 0; k < fOutputShapes[i].size(); ++k) { out << SP << SP << "// dim " << k << "\n"; if (k < fOutputShapes[i].size()-1) { - out << SP << SP << "input_index += (int(remaining / " << output_strides[k] << ")"; + out << SP << SP << "input_index += (int(remaining / " << output_strides[k].GetVal() << ")"; // for the split axis we need to consider the offset in the splits when converting to input coordinates if (k == static_cast(fAxis) && i > 0) out << " + " << OpName << "_axis_offset"; - out << ") * " << input_strides[k] << ";\n"; - out << SP << SP << "remaining %= " << output_strides[k] << ";\n"; + out << ") * " << input_strides[k].GetVal() << ";\n"; + out << SP << SP << "remaining %= " << output_strides[k].GetVal() << ";\n"; } else { // for last dims all strides are one out << SP << SP << "input_index += remaining"; @@ -153,7 +159,7 @@ public: return out.str(); } -std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { +std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fOutputShapes.empty()) throw std::runtime_error("SOFIE Operator Split called to Generate without being initialized first"); @@ -165,11 +171,12 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { std::string op; op = "\n//------ SPLIT_KERNEL_ALPAKA\n"; - std::cout<<"Generating GPU kernel for Split operator with input shape "<< ConvertShapeToString(fInputShape) << " and output shapes : "; for (std::size_t i = 0; i < Nin; ++i) { - std::cout<<"Loop running for output "<(acc)[0];\n"; @@ -192,8 +201,8 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { for (std::size_t d = 0; d < D; ++d) { op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(outputStrides[d]) + "u) % " - + std::to_string(fOutputShapes[i][d]) + "u;\n"; + + " = (elem_idx / static_cast(" + outputStrides[d].GetVal() + ")) % " + + "static_cast(" + fOutputShapes[i][d].GetVal() + ");\n"; } op += "\n"; @@ -202,7 +211,7 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { std::string coord = (d == static_cast(fAxis)) ? ("(out_" + std::to_string(d) + " + " + std::to_string(axis_offset) + "u)") : ("out_" + std::to_string(d)); - op += SP + SP + SP + SP + SP + coord + " * " + std::to_string(inputStrides[d]) + "u"; + op += SP + SP + SP + SP + SP + coord + " * static_cast(" + inputStrides[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -211,7 +220,6 @@ std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { op += SP + SP + "}\n"; op += SP + "};\n\n"; } - std::cout<<"Finished generating GPU kernel for Split operator "< &dynParamNames) override { opName = "op_" + opName; if (fOutputShapes.empty()) throw std::runtime_error("SOFIE Operator Split called to Generate without being initialized first"); @@ -234,18 +242,20 @@ std::string Generate_GPU_ALPAKA(std::string opName) override { out << "\n//------ SPLIT_GPU_ALPAKA\n"; for (std::size_t i = 0; i < fNYs.size(); ++i) { - std::size_t length = ConvertShapeToLength(fOutputShapes[i]); + std::string length = ConvertDimShapeToLength(fOutputShapes[i]); std::string kname = "splitKernel_" + opName + "_" + std::to_string(i); out << SP << "{\n"; out << SP << SP << "auto const elementsPerThread_" << i << " = Vec::all(static_cast(1));\n"; - out << SP << SP << "auto const elementsPerGrid_" << i << " = Vec::all(Idx{" << length << "});\n"; + out << SP << SP << "auto const elementsPerGrid_" << i << " = Vec::all(Idx{static_cast(" << length << ")});\n"; out << SP << SP << "auto const workDiv_" << i << " = sofie_workdiv(elementsPerGrid_" << i << ");\n"; out << SP << SP << "auto task_" << opName << "_" << i << " = alpaka::createTaskKernel(workDiv_" << i << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNYs[i] << ")" - << ", static_cast(" << length << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNYs[i] << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << length << "));\n"; out << SP << "alpaka::enqueue(queue, task_" << opName << "_" << i << ");\n"; out << SP << "}\n"; } diff --git a/core/inc/SOFIE/ROperator_Tile.hxx b/core/inc/SOFIE/ROperator_Tile.hxx index 5a3921e2..6fb000d9 100644 --- a/core/inc/SOFIE/ROperator_Tile.hxx +++ b/core/inc/SOFIE/ROperator_Tile.hxx @@ -19,9 +19,8 @@ private: std::string fNRepeats; std::string fNInput; std::string fNY; - std::vector fShapeInput; - std::vector fShapeY; - std::vector fRepeats; + std::vector fShapeInput; + std::vector fShapeY; public: ROperator_Tile(){} @@ -37,11 +36,19 @@ public: return input; } - std::vector> ShapeInference(std::vector> input) override { - std::vector ret = input[0]; - for (size_t i = 0; i < input[1].size(); i++) - ret[i] = ret[i] * input[1][i]; - return {ret}; + using ROperator::ShapeInference; + std::vector ShapeInference(const std::vector & input, const std::vector repeat) { + std::vector ret = input; + for(size_t i=0; i < repeat.size(); i++) { + if (repeat[i] != 1) { + if (ret[i].isParam) { + ret[i] = Dim{ std::string(ret[i].GetVal() + "*" + std::to_string(repeat[i])), static_cast(-1) }; + } else { + ret[i]=Dim { ret[i].dim *repeat[i] }; + } + } + } + return ret; } void Initialize(RModel& model) override { @@ -50,8 +57,10 @@ public: if (model.CheckIfTensorAlreadyExist(fNRepeats) == false) throw std::runtime_error("SOFIE Tile Op Repeats Tensor is not found in model"); - fShapeInput = model.GetTensorShape(fNInput); + fShapeInput = model.GetDimTensorShape(fNInput); + // if repeats vector is not initialized we cannot deduce shape of output + // not support for time being this case if (!model.IsInitializedTensor(fNRepeats)) throw std::runtime_error("SOFIE Tile Op: non-initialized repeats input is not supported"); @@ -65,24 +74,17 @@ public: throw std::runtime_error("SOFIE Tile Op: repeats tensor must be 1D"); size_t num_elements = repeats_shape[0]; + std::vector repeats_vector(num_elements); + std::copy(repeats_data, repeats_data + num_elements, repeats_vector.begin()); - // Save repeats if known at generation time so the GPU kernel can bake - // fShapeInput[d] directly without needing a runtime repeats pointer. - // fRepeats is left empty if repeats are not initialized (future case), - // which will cause the kernel to use the runtime repeats pointer path. - fRepeats.resize(num_elements); - std::copy(repeats_data, repeats_data + num_elements, fRepeats.begin()); - if (fRepeats.size()){ - model.RemoveInitializedTensor(fNRepeats); - } - fShapeY = ShapeInference({fShapeInput, fRepeats})[0]; + fShapeY = ShapeInference(fShapeInput, repeats_vector); model.AddIntermediateTensor(fNY, model.GetTensorType(fNInput), fShapeY); if (model.Verbose()) - std::cout << "Tile: " << fNInput << " " << ConvertShapeToString(fShapeInput) - << " -> " << fNY << " with shape " << ConvertShapeToString(fShapeY) - << " given repeats " << ConvertShapeToString(fRepeats) << std::endl; + std::cout << "Tile: " << fNInput << " " << ConvertDimShapeToString(fShapeInput) + << " -> " << fNY << " with shape " << ConvertDimShapeToString(fShapeY) + << " given repeats " << ConvertShapeToString(repeats_vector) << std::endl; } std::string Generate(std::string OpName) override { @@ -91,59 +93,49 @@ public: throw std::runtime_error("SOFIE Tile Op called to Generate without being initialized first"); std::stringstream out; - std::string input = "tensor_" + fNInput; - std::string output = "tensor_" + fNY; - std::string repeats = "tensor_" + fNRepeats; - + std::string input = "tensor_" + fNInput; + std::string output = "tensor_" + fNY; out << "///-------- Tile operator\n"; - out << "{\n"; - - out << SP << "const int input_shape[" << fShapeInput.size() << "] = {"; - for (size_t i = 0; i < fShapeInput.size(); ++i) { - if (i > 0) out << ", "; - out << fShapeInput[i]; - } - out << "};\n"; - - out << SP << "int inputLength = " << ConvertShapeToLength(fShapeInput) << ";\n"; - out << SP << "int s = 1;\n"; - - // Read repeats from the tensor at runtime so the generated code remains - // correct even if repeats become a runtime input/intermediate in the future - out << SP << "for (int i = " << fShapeInput.size() - 1 << "; i >= 0; i--) {\n"; - out << SP << SP << "int r = " << repeats << "[i];\n"; - out << SP << SP << "int i_offset = 0, o_offset = 0;\n"; - out << SP << SP << "s = s * input_shape[i];\n"; - out << SP << SP << "if (i == " << fShapeInput.size() - 1 << ") {\n"; - out << SP << SP << SP << "for (int j = 0; j < inputLength / s; j++) {\n"; - out << SP << SP << SP << SP << "for (int k = 0; k < r; k++) {\n"; - out << SP << SP << SP << SP << SP << "std::copy(" << input << " + i_offset, " - << input << " + i_offset + s, " - << output << " + o_offset);\n"; - out << SP << SP << SP << SP << SP << "o_offset += s;\n"; - out << SP << SP << SP << SP << "}\n"; - out << SP << SP << SP << SP << "i_offset += s;\n"; - out << SP << SP << SP << "}\n"; - out << SP << SP << "} else {\n"; - out << SP << SP << SP << "for (int j = inputLength / s - 1; j >= 0; j--) {\n"; - out << SP << SP << SP << SP << "o_offset = j * s * r;\n"; - out << SP << SP << SP << SP << "i_offset = j * s;\n"; - out << SP << SP << SP << SP << "for (int k = 0; k < r; k++) {\n"; - out << SP << SP << SP << SP << SP << "std::copy(" << output << " + i_offset, " - << output << " + i_offset + s, " - << output << " + o_offset);\n"; - out << SP << SP << SP << SP << SP << "o_offset += s;\n"; - out << SP << SP << SP << SP << "}\n"; - out << SP << SP << SP << "}\n"; - out << SP << SP << "}\n"; - out << SP << SP << "s *= r;\n"; - out << SP << SP << "inputLength *= r;\n"; - out << SP << "}\n"; - out << "}\n"; + out << "{\n"; // add scope to re-use same names + out << "const size_t input_shape[" << fShapeInput.size() << "] = " << ConvertDimShapeToString(fShapeInput) << ";\n"; + + out << "int inputLength = " << ConvertDimShapeToLength(fShapeInput) << ";\n"; + out << "int s = 1;\n"; + // loop from inverse dim order + out << "for (int i = " << fShapeInput.size()-1 << "; i >=0; i--) {\n"; + out << SP << "int r = tensor_" << fNRepeats << "[i];\n"; + out << SP << "int i_offset = 0, o_offset = 0;\n"; + out << SP << "s = s * input_shape[i];\n"; + // case we have first copy + out << SP << "if (i == " << fShapeInput.size()-1 << ") {\n"; + out << SP << SP << "for (int j = 0; j < inputLength/s ; j++) {\n"; + out << SP << SP << SP << "for (int k = 0; k < r ; k++) {\n"; + out << SP << SP << SP << SP << "std::copy(" << input << "+ i_offset, " + << input << "+ i_offset + s, " << output << "+ o_offset);\n"; + out << SP << SP << SP << SP << "o_offset += s;\n"; + out << SP << SP << SP << "}\n"; // end k loop + out << SP << SP << SP << "i_offset += s;\n"; + out << SP << SP << "}\n"; // end j loop + out << SP << "} else {\n"; // second copy we do from output to output + // and we need to loop on j from reverse order to avoid re-writing in output tensor + out << SP << SP << "for (int j = inputLength/s - 1 ; j>=0; j--) {\n"; + out << SP << SP << SP << "o_offset = j*s*r;\n"; + out << SP << SP << SP << "i_offset = j*s;\n"; + out << SP << SP << SP << "for (int k = 0; k < r ; k++) {\n"; + out << SP << SP << SP << SP << "std::copy(" << output << "+ i_offset, " + << output << "+ i_offset + s, " << output << "+ o_offset);\n"; + out << SP << SP << SP << SP << "o_offset += s;\n"; + out << SP << SP << SP << "}\n"; // end k loop + out << SP << SP << "}\n"; // end j loop + out << SP << "}\n"; // end if + out << SP << "s *= r;\n"; + out << SP << "inputLength *= r;\n"; + out << "}\n"; // end i loop + out << "}\n"; // end of scope return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeInput.empty() || fShapeY.empty()) throw std::runtime_error("SOFIE Operator Tile called to Generate without being initialized first"); @@ -152,12 +144,6 @@ public: auto inputStrides = UTILITY::ComputeStrideFromShape(fShapeInput); auto outputStrides = UTILITY::ComputeStrideFromShape(fShapeY); - std::size_t totalElements = ConvertShapeToLength(fShapeY); - - // If fRepeats is populated, repeats were known at generation time and - // we can bake fShapeInput[d] as literals — no runtime repeats pointer needed. - // If fRepeats is empty (future: runtime repeats), pass repeats as a kernel arg. - bool repeatsKnown = !fRepeats.empty(); std::string kname = "TileKernel_" + opName; @@ -169,8 +155,8 @@ public: op += SP + SP + SP + "TAcc const& acc,\n"; op += SP + SP + SP + "T const* __restrict__ input,\n"; op += SP + SP + SP + "T* __restrict__ output,\n"; - if (!repeatsKnown) - op += SP + SP + SP + "int64_t const* __restrict__ repeats,\n"; + for (auto &p : dynParamNames) + op += SP + SP + SP + "std::size_t const " + p + ",\n"; op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; @@ -179,24 +165,15 @@ public: op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - // Decompose output linear index — output strides always compile-time - for (std::size_t d = 0; d < D; ++d) { - op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(outputStrides[d]) + "u) % " - + std::to_string(fShapeY[d]) + "u;\n"; - } + EmitOutputCoordsFromThreadIdx(op, SP + SP + SP + SP, outputStrides, fShapeY); op += "\n"; - // Input index: fShapeInput[d] is always a compile-time constant since - // it is the input tensor shape, never runtime-variable. - // When repeatsKnown, we bake it directly as a literal. - // When not repeatsKnown (future), we still use fShapeInput[d] as a - // literal for the % — repeats pointer is only needed if fShapeY is dynamic. + // Input index: tiling wraps each output coordinate back into the input shape op += SP + SP + SP + SP + "std::size_t const input_idx =\n"; for (std::size_t d = 0; d < D; ++d) { op += SP + SP + SP + SP + SP - + "(out_" + std::to_string(d) + " % " + std::to_string(fShapeInput[d]) + "u)" - + " * " + std::to_string(inputStrides[d]) + "u"; + + "(out_" + std::to_string(d) + " % (" + fShapeInput[d].GetVal() + "))" + + " * (" + inputStrides[d].GetVal() + ")"; op += (d + 1 < D) ? " +\n" : ";\n\n"; } @@ -214,22 +191,20 @@ public: return SP + kname + " tileKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { opName = "op_" + opName; if (fShapeInput.empty() || fShapeY.empty()) throw std::runtime_error("SOFIE Operator Tile called to Generate without being initialized first"); - bool repeatsKnown = !fRepeats.empty(); - std::size_t totalElements = ConvertShapeToLength(fShapeY); + std::string totalElements = ConvertDimShapeToLength(fShapeY); std::string kname = "tileKernel_" + opName; - // Build argument list once, reused for both getValidWorkDiv and exec std::string args = "alpaka::getPtrNative(deviceBuf_" + fNInput + "), " + "alpaka::getPtrNative(deviceBuf_" + fNY + ")"; - if (!repeatsKnown) - args += ", alpaka::getPtrNative(deviceBuf_" + fNRepeats + ")"; - args += ", static_cast(" + std::to_string(totalElements) + ")"; + for (auto &p : dynParamNames) + args += ", static_cast(" + p + ")"; + args += ", static_cast(" + totalElements + ")"; std::stringstream out; out << "\n//------ TILE_GPU_ALPAKA\n"; diff --git a/core/inc/SOFIE/ROperator_TopK.hxx b/core/inc/SOFIE/ROperator_TopK.hxx index 7db17682..6058fe92 100644 --- a/core/inc/SOFIE/ROperator_TopK.hxx +++ b/core/inc/SOFIE/ROperator_TopK.hxx @@ -5,6 +5,7 @@ #include "SOFIE/ROperator.hxx" #include "SOFIE/RModel.hxx" +#include #include @@ -18,15 +19,24 @@ private: int fAttrLargest; int fAttrSorted; - size_t fK; + Dim fTopKCount; // k clamped to the axis dimension: a number for a static axis, a + // std::min(k, axis) expression for a dynamic one + size_t fRequestedK; // the model's requested k, unclamped; sizes the GPU register buffers std::string fNK; std::string fNX; std::string fNVal; std::string fNInd; - std::vector fShapeX; - std::vector fShapeY; + std::vector fShapeX; + std::vector fShapeY; std::string fType; + // GPU tuning knobs for the hierarchical kernel below. fStagingB is the size of each + // thread's private register buffer. fBlockSize is the number of threads cooperating on + // one slice. + static constexpr size_t fStagingBCap = 64; + size_t fStagingB = fStagingBCap; + size_t fBlockSize = 256; + public: ROperator_TopK() {} ROperator_TopK(int attr_axis, int attr_largest, int attr_sorted, std::string nameK, std::string nameX, std::string nameVal, std::string nameInd) @@ -42,23 +52,10 @@ public: } std::vector TypeInference(std::vector input) override { - ETensorType ret = input[0]; - return {ret, ret}; - } - - std::vector> ShapeInference(std::vector> input) override { - if (input.size() != 2) { - throw std::runtime_error("SOFIE TopK Op Shape Inference needs exactly 2 input tensors"); - } - - auto shape = input[0]; // Shape format: [ m x n x o x p ... ] - - // set the dimension at the specified axis to k (fAttrAxis is checked before that is in the correct range - shape[fAttrAxis] = fK; // Modified shape: [ m x n x k x p ... ] - return {shape, shape}; + ETensorType ret = input[0]; + return {ret, ret}; } - void Initialize(RModel& model) override { if (model.CheckIfTensorAlreadyExist(fNX) == false) { // input must be a graph input, or already initialized intermediate tensor @@ -69,43 +66,44 @@ public: throw std::runtime_error("SOFIE TopK Op Input Tensor i.e. K is not found in model"); } - fShapeX = model.GetTensorShape(fNX); - auto fShapeK = model.GetTensorShape(fNK); + fShapeX = model.GetDimTensorShape(fNX); auto kptr = static_cast(model.GetInitializedTensorData(fNK).get()); - fK = *kptr; + size_t kval = *kptr; + fRequestedK = kval; model.SetNotWritableInitializedTensor(fNK); fAttrAxis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis; - if(static_cast(fAttrAxis) >= fShapeX.size()){ - throw - std::runtime_error("TMVA::SOFIE ONNX TopK op axis = "+ std::to_string(fAttrAxis) +" value exeeds size of tensor " +fNX+" of size "+std::to_string(fShapeX.size())+" ."); + if (static_cast(fAttrAxis) >= fShapeX.size()) { + throw std::runtime_error("SOFIE TopK op axis = " + std::to_string(fAttrAxis) + + " value exceeds size of tensor " + fNX + " of size " + std::to_string(fShapeX.size()) + " ."); + } + // fTopKCount cannot be larger than the axis dimension + if (fShapeX[fAttrAxis].isParam) { + fTopKCount = Dim{std::string("std::min(size_t(" + std::to_string(kval) + "), " + fShapeX[fAttrAxis].GetVal() + ")" ), static_cast(-1) }; + // axis size unknown at codegen time - rsV/rsI must be a fixed-size array, so + // fall back to the empirically-safe cap of 64. + fStagingB = fStagingBCap; + } else { + fTopKCount = Dim { std::min(kval, fShapeX[fAttrAxis].dim) }; + // a thread never buffers more than its strided share of the axis, so B need + // never exceed that share. + size_t itemsPerThread = (fShapeX[fAttrAxis].dim + fBlockSize - 1) / fBlockSize; + fStagingB = std::min(std::max(itemsPerThread, 1), fStagingBCap); } - // fK cannot be larger that axis dimension - fK = std::min(fK, fShapeX[fAttrAxis]); - // if(fK>fShapeX[fAttrAxis]){ - // throw - // std::runtime_error("TMVA::SOFIE ONNX TopK op k = "+ std::to_string(fK) +" value exeeds value of tensor " +fNX+" of size "+fShapeX.size()+" at axis= "+std::to_string(fAttrAxis)+"."); - // } - // fShapeX = model.GetTensorShape(fNX); // [ m x n x o x p ... ] - // if(k[0]>=fShapeX.size()){ - // throw - // std::runtime_error("TMVA::SOFIE ONNX TopK op k = "+ std::to_string(k[0]) +"value exeeds size of tensor " +fNX+" of size "+fShapeX.size()+" ."); - // } - // fShapeY.push_back(2); - // for (auto i : fShapeX) - // fShapeY.push_back(i); // [ 2 x m x n x o x p ... ] - // size_t axis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis; - // fShapeY[axis] = k[0]; // [ 2 x m x n x K x p ... ] - fShapeY=ShapeInference({fShapeX,fShapeK})[0]; - - // for(int i=0;i "< " << fNVal << " " << ConvertDimShapeToString(fShapeY) << std::endl; + } } std::string Generate(std::string OpName) override { @@ -114,66 +112,247 @@ public: throw std::runtime_error("SOFIE Operator TopK called to Generate without being initialized first"); } std::stringstream out; - size_t size = fShapeX.size(); - size_t axis = fAttrAxis < 0 ? size + fAttrAxis : fAttrAxis; out << "\n" << SP << "//------ TopK\n"; - size_t length=ConvertShapeToLength(fShapeX); - auto strideX = UTILITY::ComputeStrideFromShape(fShapeX); - auto strideY = UTILITY::ComputeStrideFromShape(fShapeX); // we perform loop on dimension before sorted axis and after sorted axis - size_t n_before = (axis>0) ? length/strideX[axis-1] : 1; - size_t n_after = strideX[axis]; - size_t n_elements = fShapeX[axis]; // number of elements to be sorted + size_t axis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis; + auto sx = UTILITY::ComputeSliceInfo(fShapeX, axis), sy = UTILITY::ComputeSliceInfo(fShapeY, axis); //X has n elements along the axis, Y has k: two layouts - // } out << SP << "{\n"; // to define a separate scope for the operator code - out << SP << "std::vector> elements(" << n_elements << ");\n"; + out << SP << "std::vector> elements(" << sx.nElements << ");\n"; // loop on elements before - if (n_before > 1) { - out << SP << "for (size_t i = 0; i < " << n_before << "; i++) {\n"; - out << SP << SP << "size_t xoffset = i*" << strideX[axis-1] << ";\n"; - out << SP << SP << "size_t yoffset = i*" << strideY[axis-1] << ";\n"; + if (sx.nBefore != "1") { + out << SP << "for (size_t i = 0; i < " << sx.nBefore << "; i++) {\n"; + out << SP << SP << "size_t xoffset = i*" << sx.strideBefore << ";\n"; + out << SP << SP << "size_t yoffset = i*" << sy.strideBefore << ";\n"; out << SP; } else { out << SP << "size_t xoffset = 0;\n"; out << SP << "size_t yoffset = 0;\n"; } - if (n_after > 1) - out << SP << "for (size_t j = 0; j < " << n_after << "; j++) {\n"; + if (sx.nAfter != "1") + out << SP << "for (size_t j = 0; j < " << sx.nAfter << "; j++) {\n"; else out << SP << "const size_t j = 0;\n"; // copy elements to be sorted in vector of pair - out << SP << SP << "for (size_t l = 0; l < " << n_elements << "; l++) {\n"; - out << SP << SP << SP << "elements[l] = std::make_pair(tensor_" << fNX << "[xoffset + " << strideX[axis] << "*l + j], l);\n"; + out << SP << SP << "for (size_t l = 0; l < " << sx.nElements << "; l++) {\n"; + out << SP << SP << SP << "elements[l] = std::make_pair(tensor_" << fNX << "[xoffset + " << sx.strideAxis << "*l + j], l);\n"; out << SP << SP << "}\n"; if (fAttrSorted) { - if (fAttrLargest) { - out<a,std::pairb){return (a.first!=b.first) ? (a.first>b.first) : a.second < b.second;});\n"; - - } else - out<a,std::pairb){return (a.first!=b.first) ? (a.firsta,std::pairb){return (a.first!=b.first) ? (a.first>b.first) : a.second < b.second;});\n"; + else + out << SP << SP << "std::partial_sort(elements.begin(),elements.begin()+" << fTopKCount << ",elements.end()," + << "[](std::paira,std::pairb){return (a.first!=b.first) ? (a.first 1) out << SP << SP << "}\n"; - if (n_before> 1) out << SP << "}\n"; + if (sx.nAfter != "1") out << SP << SP << "}\n"; + if (sx.nBefore != "1") out << SP << "}\n"; out << SP << "}\n"; // end operator scope return out.str(); } + + // Hierarchical GPU path: one block per slice, fBlockSize threads cooperating on that + // slice's whole axis via a block-shared result (sVals/sInds, capacity K, K = the + // model's requested k unclamped - the only value that needs to be a compile-time + // constant here) guarded by a spinlock, plus a fill counter (sCount). Each thread + // scans a strided subset of the axis into a tiny private register buffer (rsV/rsI, + // size fStagingB - independent of K). This mirrors the shape of + // real GPU top-k selection structures, simplified to use a + // block-shared array + lock for the merge instead of warp-shuffle-only cooperation: + // a shared-memory merge is both simpler and the more portable + // choice here. The spinlock is safe specifically because this is a *single-block* + // design, unlike the cross-block case. + + // Admission before touching the shared state: a candidate only gets buffered locally + // once it beats the current shared threshold (sVals[topKCount-1], once sCount reaches + // topKCount) - anything worse is dropped immediately, for free, without acquiring the + // lock. That read is deliberately lock-free: sCount and the shared array only ever get + // strictly tighter/more-populated over the scan, so an out-of-date value can make a + // thread buffer something that turns out unnecessary. It still has to be read + // through volatile-qualified pointers. + // + // A thread's local buffer is flushed into the shared result whenever it fills, and + // once more after the thread's own portion of the axis is exhausted - + // required for correctness in general (a short axis, or many threads, means a + // thread's last few buffered candidates may never trigger a fill-driven flush). + // + // fStagingB and fBlockSize are compile-time literals with no closed-form optimum + std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) override { + if (fShapeX.empty()) + throw std::runtime_error("SOFIE Operator TopK called to Generate without being initialized first"); + + std::string CMP = fAttrLargest ? ">" : "<"; + std::string K = std::to_string(fRequestedK); + std::string B = std::to_string(fStagingB); + std::string P = std::to_string(fBlockSize); + std::string better = fAttrLargest + ? "(v > sVals[p-1] || (v == sVals[p-1] && idx < sInds[p-1]))" + : "(v < sVals[p-1] || (v == sVals[p-1] && idx < sInds[p-1]))"; + std::string beatsWorst = fAttrLargest + ? "(v > sVals[topKCount-1] || (v == sVals[topKCount-1] && idx < sInds[topKCount-1]))" + : "(v < sVals[topKCount-1] || (v == sVals[topKCount-1] && idx < sInds[topKCount-1]))"; + // Admission-check variant for the one read site that happens outside the lock + std::string beatsWorstVolatile = fAttrLargest + ? "(v >= vSVals[topKCount-1])" + : "(v <= vSVals[topKCount-1])"; + std::string kname = "TopKKernel_" + fNVal; + + auto emitInsert = [&](const std::string &ind) { + std::string s; + s += ind + "if ((std::size_t)sCount < topKCount) {\n"; + s += ind + SP + "int p = sCount;\n"; + s += ind + SP + "while (p > 0 && " + better + ") { sVals[p] = sVals[p-1]; sInds[p] = sInds[p-1]; --p; }\n"; + s += ind + SP + "sVals[p] = v; sInds[p] = idx; sCount++;\n"; + s += ind + "} else if (" + beatsWorst + ") {\n"; + s += ind + SP + "int p = (int)topKCount - 1;\n"; + s += ind + SP + "while (p > 0 && " + better + ") { sVals[p] = sVals[p-1]; sInds[p] = sInds[p-1]; --p; }\n"; + s += ind + SP + "sVals[p] = v; sInds[p] = idx;\n"; + s += ind + "} else break;\n"; + return s; + }; + + std::string op = "\n//------ TopK_KERNEL_ALPAKA (hierarchical: per-thread staging + block-shared merge)\n"; + op += SP + "struct " + kname + " {\n"; + op += SP + SP + "template\n"; + op += SP + SP + "ALPAKA_FN_ACC void operator()(\n"; + op += SP + SP + SP + "TAcc const& acc,\n"; + op += SP + SP + SP + "T const* __restrict__ x,\n"; + op += SP + SP + SP + "T* __restrict__ vals,\n"; + op += SP + SP + SP + "int64_t* __restrict__ inds,\n"; + op += SP + SP + SP + "std::size_t const numSlices,\n"; + op += SP + SP + SP + "std::size_t const nAfter,\n"; + op += SP + SP + SP + "std::size_t const nElAxis,\n"; + op += SP + SP + SP + "std::size_t const topKCount,\n"; + op += SP + SP + SP + "std::size_t const strideXAxis,\n"; + op += SP + SP + SP + "std::size_t const strideXBefore,\n"; + op += SP + SP + SP + "std::size_t const strideYAxis,\n"; + op += SP + SP + SP + "std::size_t const strideYBefore) const {\n\n"; + + op += SP + SP + SP + "auto const slice = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "auto const tid = alpaka::getIdx(acc)[0];\n"; + op += SP + SP + SP + "if (slice >= numSlices) return;\n\n"; + + op += SP + SP + SP + "std::size_t const i = slice / nAfter;\n"; + op += SP + SP + SP + "std::size_t const j = slice % nAfter;\n"; + op += SP + SP + SP + "std::size_t const xbase = i * strideXBefore + j;\n"; + op += SP + SP + SP + "std::size_t const ybase = i * strideYBefore + j;\n\n"; + + op += SP + SP + SP + "auto& sVals = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto& sInds = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto& sCount = alpaka::declareSharedVar(acc);\n"; + op += SP + SP + SP + "auto& sLock = alpaka::declareSharedVar(acc);\n\n"; + + op += SP + SP + SP + "if (tid == 0) { sCount = 0; sLock = 0; }\n"; + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; + + op += SP + SP + SP + "T const volatile* vSVals = sVals;\n"; + op += SP + SP + SP + "int const volatile* vSCount = &sCount;\n\n"; + + op += SP + SP + SP + "T rsV[" + B + "]; int64_t rsI[" + B + "]; int rsN = 0;\n\n"; + + op += SP + SP + SP + "for (std::size_t idx = tid; idx < nElAxis; idx += " + P + "u) {\n"; + op += SP + SP + SP + SP + "T v = x[xbase + strideXAxis * idx];\n"; + op += SP + SP + SP + SP + "bool admit;\n"; + op += SP + SP + SP + SP + "if ((std::size_t)*vSCount < topKCount) admit = true;\n"; + op += SP + SP + SP + SP + "else admit = " + beatsWorstVolatile + ";\n"; + op += SP + SP + SP + SP + "if (admit) {\n"; + op += SP + SP + SP + SP + SP + "int p = rsN;\n"; + op += SP + SP + SP + SP + SP + "while (p > 0 && v " + CMP + " rsV[p-1]) { rsV[p] = rsV[p-1]; rsI[p] = rsI[p-1]; --p; }\n"; + op += SP + SP + SP + SP + SP + "rsV[p] = v; rsI[p] = (int64_t)idx; rsN++;\n"; + op += SP + SP + SP + SP + SP + "if (rsN == " + B + ") {\n"; + op += SP + SP + SP + SP + SP + SP + "while (alpaka::atomicCas(acc, &sLock, 0, 1, alpaka::hierarchy::Grids{}) != 0) {}\n"; + op += SP + SP + SP + SP + SP + SP + "alpaka::mem_fence(acc, alpaka::memory_scope::Device{});\n"; + op += SP + SP + SP + SP + SP + SP + "for (int r = 0; r < " + B + "; ++r) {\n"; + op += SP + SP + SP + SP + SP + SP + SP + "T v = rsV[r]; int64_t idx = rsI[r];\n"; + op += emitInsert(SP + SP + SP + SP + SP + SP + SP); + op += SP + SP + SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + SP + SP + "alpaka::mem_fence(acc, alpaka::memory_scope::Device{});\n"; + op += SP + SP + SP + SP + SP + SP + "alpaka::atomicExch(acc, &sLock, 0, alpaka::hierarchy::Grids{});\n"; + op += SP + SP + SP + SP + SP + SP + "rsN = 0;\n"; + op += SP + SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + "}\n\n"; + + op += SP + SP + SP + "// drain: flush whatever's left, required for correctness in general\n"; + op += SP + SP + SP + "// (e.g. a short axis may never trigger a fill-driven flush above)\n"; + op += SP + SP + SP + "if (rsN > 0) {\n"; + op += SP + SP + SP + SP + "while (alpaka::atomicCas(acc, &sLock, 0, 1, alpaka::hierarchy::Grids{}) != 0) {}\n"; + op += SP + SP + SP + SP + "alpaka::mem_fence(acc, alpaka::memory_scope::Device{});\n"; + op += SP + SP + SP + SP + "for (int r = 0; r < rsN; ++r) {\n"; + op += SP + SP + SP + SP + SP + "T v = rsV[r]; int64_t idx = rsI[r];\n"; + op += emitInsert(SP + SP + SP + SP + SP + SP); + op += SP + SP + SP + SP + "}\n"; + op += SP + SP + SP + SP + "alpaka::mem_fence(acc, alpaka::memory_scope::Device{});\n"; + op += SP + SP + SP + SP + "alpaka::atomicExch(acc, &sLock, 0, alpaka::hierarchy::Grids{});\n"; + op += SP + SP + SP + "}\n\n"; + + op += SP + SP + SP + "alpaka::syncBlockThreads(acc);\n\n"; + op += SP + SP + SP + "for (std::size_t s = tid; s < topKCount; s += " + P + "u) {\n"; + op += SP + SP + SP + SP + "vals[ybase + strideYAxis * s] = sVals[s];\n"; + op += SP + SP + SP + SP + "inds[ybase + strideYAxis * s] = sInds[s];\n"; + op += SP + SP + SP + "}\n"; + op += SP + SP + "}\n"; + op += SP + "};\n"; + + return op; + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string /*opName*/) override { + return SP + "TopKKernel_" + fNVal + " topKernel_" + fNVal + ";\n"; + } + + std::string Generate_GPU_ALPAKA(std::string opName) override { + opName = "op_" + opName; + if (fShapeX.empty()) + throw std::runtime_error("SOFIE Operator TopK called to Generate without being initialized first"); + + size_t axis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis; + auto sx = UTILITY::ComputeSliceInfo(fShapeX, axis), sy = UTILITY::ComputeSliceInfo(fShapeY, axis); //X has n elements along the axis, Y has k: two layouts + std::string numSlices = "((" + sx.nBefore + ")*(" + sx.nAfter + "))"; + std::string P = std::to_string(fBlockSize); + std::string n = fNVal; + + std::stringstream out; + out << "\n//-- TopK_GPU_ALPAKA (hierarchical)\n"; + out << SP << "{\n"; + out << SP << SP << "std::size_t const topkNumSlices_" << n << " = static_cast(" << numSlices << ");\n"; + out << SP << SP << "std::size_t const topkNElAxis_" << n << " = static_cast(" << sx.nElements << ");\n"; + out << SP << SP << "std::size_t const topkTopKCount_" << n << " = static_cast(" << fTopKCount.GetVal() << ");\n\n"; + + out << SP << SP << "alpaka::WorkDivMembers workDiv_" << n << "(\n"; + out << SP << SP << SP << "Vec::all(Idx{topkNumSlices_" << n << "}),\n"; + out << SP << SP << SP << "Vec::all(Idx{" << P << "u}),\n"; + out << SP << SP << SP << "Vec::all(Idx{1u}));\n"; + out << SP << SP << "alpaka::exec(queue, workDiv_" << n << ", topKernel_" << n + << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNVal << ")" + << ", alpaka::getPtrNative(deviceBuf_" << fNInd << ")" + << ", topkNumSlices_" << n + << ", static_cast(" << sx.nAfter << ")" + << ", topkNElAxis_" << n + << ", topkTopKCount_" << n + << ", static_cast(" << sx.strideAxis << ")" + << ", static_cast(" << sx.strideBefore << ")" + << ", static_cast(" << sy.strideAxis << ")" + << ", static_cast(" << sy.strideBefore << "));\n"; + out << SP << "}\n"; + return out.str(); + } + }; -} // nameSPace SOFIE +} // namespace SOFIE #endif // SOFIE_ROPERATOR_TOPK diff --git a/core/inc/SOFIE/ROperator_Transpose.hxx b/core/inc/SOFIE/ROperator_Transpose.hxx index 03dad41f..5c323198 100644 --- a/core/inc/SOFIE/ROperator_Transpose.hxx +++ b/core/inc/SOFIE/ROperator_Transpose.hxx @@ -23,10 +23,8 @@ private: std::string fNData; std::string fNOutput; - std::vector fShapeData; // used for initialized (constant) tensor case - std::vector fShapeOutput; // used for initialized (constant) tensor case - std::vector fDimShapeData; // used for dynamic/runtime tensor case - std::vector fDimShapeOutput; // used for dynamic/runtime tensor case + std::vector fDimShapeData; + std::vector fDimShapeOutput; public: @@ -70,21 +68,21 @@ public: } if (model.IsInitializedTensor(fNData)) { // Constant/initialized tensor: use concrete shapes and perform transpose at init time - fShapeData = model.GetTensorShape(fNData); + std::vector shapeData = model.GetTensorShape(fNData); if (fAttrPerm.empty()){ - fAttrPerm.reserve(fShapeData.size()); - for (int i = fShapeData.size() - 1; i >= 0; i--){ + fAttrPerm.reserve(shapeData.size()); + for (int i = shapeData.size() - 1; i >= 0; i--){ fAttrPerm.push_back(i); } } - std::vector> inputs = { fShapeData }; - fShapeOutput = ShapeInference(inputs).front(); + std::vector> inputs = { shapeData }; + std::vector shapeOutput = ShapeInference(inputs).front(); fIsOutputConstant = true; - auto inStrides = UTILITY::ComputeStrideFromShape(fShapeData); - auto outStrides = UTILITY::ComputeStrideFromShape(fShapeOutput); - size_t length = ConvertShapeToLength(fShapeOutput); + auto inStrides = UTILITY::ComputeStrideFromShape(shapeData); + auto outStrides = UTILITY::ComputeStrideFromShape(shapeOutput); + size_t length = ConvertShapeToLength(shapeOutput); auto inputData = static_cast(model.GetInitializedTensorData(fNData).get()); - size_t dim = fShapeData.size(); + size_t dim = shapeData.size(); std::vector outputIdx(dim); std::vector outputData(length); for (size_t i = 0; i < length; i++) { @@ -101,9 +99,12 @@ public: } outputData[i] = inputData[inputIndex]; } - model.AddConstantTensor(fNOutput, fShapeOutput, outputData.data()); + model.AddConstantTensor(fNOutput, shapeOutput, outputData.data()); + //keep the Dim members valid in every path + fDimShapeData = ConvertShapeToDim(shapeData); + fDimShapeOutput = ConvertShapeToDim(shapeOutput); if (model.Verbose()) { - std::cout << "Transpose: output is a constant tensor " << ConvertShapeToString(fShapeOutput) << " : " + std::cout << "Transpose: output is a constant tensor " << ConvertShapeToString(shapeOutput) << " : " << ConvertValuesToString(outputData) << std::endl; } } else { @@ -130,16 +131,13 @@ public: std::string Generate(std::string OpName) override { if (fIsOutputConstant) return ""; //no op for constant tensors OpName = "op_" + OpName; - // Use Dim shapes when available (dynamic case), else convert from concrete shapes - auto dimShapeData = fDimShapeData.empty() ? ConvertShapeToDim(fShapeData) : fDimShapeData; - auto dimShapeOutput = fDimShapeOutput.empty() ? ConvertShapeToDim(fShapeOutput) : fDimShapeOutput; - if (dimShapeData.empty() || dimShapeOutput.empty()){ + if (fDimShapeData.empty() || fDimShapeOutput.empty()){ throw std::runtime_error("SOFIE Transpose Op called to Generate without being initialized first"); } - int dim = dimShapeData.size(); - auto inStrides = UTILITY::ComputeStrideFromShape(dimShapeData); - auto outStrides = UTILITY::ComputeStrideFromShape(dimShapeOutput); - std::string length = ConvertDimShapeToLength(dimShapeOutput); + int dim = fDimShapeData.size(); + auto inStrides = UTILITY::ComputeStrideFromShape(fDimShapeData); + auto outStrides = UTILITY::ComputeStrideFromShape(fDimShapeOutput); + std::string length = ConvertDimShapeToLength(fDimShapeOutput); std::stringstream out; // Implement transpose operator using consecutive write outputs. @@ -173,13 +171,17 @@ public: return out.str(); } - std::string Generate_GPU_Kernel_ALPAKA(std::string OpName) { + + std::string Generate_GPU_Kernel_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { + if (fIsOutputConstant) return ""; std::string op; OpName = "op_" + OpName; op = "\n//------ TRANSPOSE_KERNEL_ALPAKA\n"; op += SP + "struct TransposeKernel_" + OpName + " {\n"; op += SP + SP + "template\n"; op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const& acc, T const* input, T* output,"; + for (auto &p : dynParamNames) + op += "const std::size_t " + p + ","; op += "const std::size_t totalElements) const {\n"; op += SP + SP + SP + SP + "auto const idx = alpaka::getIdx(acc)[0];\n"; op += SP + SP + SP + SP + "if(idx >= totalElements) return;\n"; @@ -187,18 +189,16 @@ public: op += SP + SP + SP + SP + "std::size_t remaining = idx;\n"; op += SP + SP + SP + SP + "std::size_t coord;\n"; - auto dimShapeData = fDimShapeData.empty() ? ConvertShapeToDim(fShapeData) : fDimShapeData; - auto dimShapeOutput = fDimShapeOutput.empty() ? ConvertShapeToDim(fShapeOutput) : fDimShapeOutput; - auto inputStrides = UTILITY::ComputeStrideFromShape(dimShapeData); - auto outputStrides = UTILITY::ComputeStrideFromShape(dimShapeOutput); - - for (size_t k = 0; k < dimShapeData.size(); k++) { - op += SP + SP + SP + SP + "coord = remaining / " - + outputStrides[k].GetVal() + "u;\n"; - op += SP + SP + SP + SP + "remaining = remaining - coord * " - + outputStrides[k].GetVal() + "u;\n"; - op += SP + SP + SP + SP + "input_idx += coord * " - + inputStrides[fAttrPerm[k]].GetVal() + "u;\n"; + auto inputStrides = UTILITY::ComputeStrideFromShape(fDimShapeData); + auto outputStrides = UTILITY::ComputeStrideFromShape(fDimShapeOutput); + + for (size_t k = 0; k < fDimShapeData.size(); k++) { + op += SP + SP + SP + SP + "coord = remaining / (" + + outputStrides[k].GetVal() + ");\n"; + op += SP + SP + SP + SP + "remaining = remaining - coord * (" + + outputStrides[k].GetVal() + ");\n"; + op += SP + SP + SP + SP + "input_idx += coord * (" + + inputStrides[fAttrPerm[k]].GetVal() + ");\n"; } op += SP + SP + SP + SP + "output[idx] = input[input_idx];\n"; @@ -209,16 +209,17 @@ public: } std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string OpName) override { + if (fIsOutputConstant) return ""; return SP + "TransposeKernel_op_" + OpName + " transposeKernel_" + OpName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string OpName) override { - auto dimShapeOutput = fDimShapeOutput.empty() ? ConvertShapeToDim(fShapeOutput) : fDimShapeOutput; - if (dimShapeOutput.empty()) { + std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector &dynParamNames) override { + if (fIsOutputConstant) return ""; + if (fDimShapeOutput.empty()) { throw std::runtime_error("SOFIE Operator Transpose called to Generate without being initialized first"); } std::stringstream out; - std::string length = ConvertDimShapeToLength(dimShapeOutput); + std::string length = ConvertDimShapeToLength(fDimShapeOutput); out << "\n//------ TRANSPOSE_GPU_ALPAKA\n"; out << SP << "auto const elementsPerThread_"<(1));\n"; @@ -226,7 +227,10 @@ public: out << SP << "auto const workDiv_" << fNOutput << " = sofie_workdiv(elementsPerGrid_" << fNOutput << ");\n"; out << SP << "auto task_" << OpName << " = alpaka::createTaskKernel(workDiv_" << fNOutput << ", transposeKernel_" << OpName << ", alpaka::getPtrNative(deviceBuf_" << fNData - << "), alpaka::getPtrNative(deviceBuf_" << fNOutput << "), static_cast(" << length << "));\n"; + << "), alpaka::getPtrNative(deviceBuf_" << fNOutput << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << length << "));\n"; out << SP <<"alpaka::enqueue(queue, task_" << OpName << ");\n"; return out.str(); } diff --git a/core/inc/SOFIE/ROperator_Trilu.hxx b/core/inc/SOFIE/ROperator_Trilu.hxx index 04e18d52..8d8b73d5 100644 --- a/core/inc/SOFIE/ROperator_Trilu.hxx +++ b/core/inc/SOFIE/ROperator_Trilu.hxx @@ -26,11 +26,10 @@ private: std::string fNK; std::string fNY; - std::vector fShape; - size_t fM = 0; - size_t fN = 0; - size_t fBatch = 1; - size_t fTotal = 0; + std::vector fShape; + Dim fM; + Dim fN; + std::string fTotal; // ConvertDimShapeToLength(fShape) public: ROperator_Trilu() {} @@ -69,7 +68,7 @@ public: throw std::runtime_error("SOFIE Trilu: input tensor '" + fNX + "' not found in model"); - fShape = model.GetTensorShape(fNX); + fShape = model.GetDimTensorShape(fNX); if (fShape.size() < 2) throw std::runtime_error("SOFIE Trilu: input tensor '" + fNX + "' must have at least 2 dimensions, got " + @@ -77,10 +76,7 @@ public: fN = fShape.back(); fM = fShape[fShape.size() - 2]; - fBatch = 1; - for (size_t d = 0; d + 2 < fShape.size(); ++d) - fBatch *= fShape[d]; - fTotal = fBatch * fM * fN; + fTotal = ConvertDimShapeToLength(fShape); if (!fNK.empty()) { if (model.IsInitializedTensor(fNK) || model.IsConstantTensor(fNK)) { @@ -102,7 +98,7 @@ public: if (fKIsStatic) std::cout << fK; else std::cout << "dyn(" << fNK << ")"; std::cout << " -> " << fNY - << " " << ConvertShapeToString(fShape) << std::endl; + << " " << ConvertDimShapeToString(fShape) << std::endl; } } @@ -122,13 +118,13 @@ public: << " = static_cast(tensor_" << fNK << "[0]);\n"; } - out << SP << "for (std::size_t id = 0; id < " << fTotal << "u; ++id) {\n"; - out << SP << SP << "const std::size_t mat_id = id % " - << (fM * fN) << "u;\n"; + out << SP << "for (std::size_t id = 0; id < static_cast(" << fTotal << "); ++id) {\n"; + out << SP << SP << "const std::size_t mat_id = id % static_cast(" + << fM.GetVal() << " * " << fN.GetVal() << ");\n"; out << SP << SP << "const std::ptrdiff_t row = " - << "static_cast(mat_id / " << fN << "u);\n"; + << "static_cast(mat_id / static_cast(" << fN.GetVal() << "));\n"; out << SP << SP << "const std::ptrdiff_t col = " - << "static_cast(mat_id % " << fN << "u);\n"; + << "static_cast(mat_id % static_cast(" << fN.GetVal() << "));\n"; if (fUpper) { out << SP << SP << "const bool keep = (col >= row + k_" << OpName << ");\n"; } else { @@ -156,12 +152,13 @@ public: << "T const* __restrict__ input, " << "T* __restrict__ output, " << "const std::size_t total, " + << "const std::size_t M, " + << "const std::size_t N, " << "const std::ptrdiff_t k) const {\n"; op << SP << SP << "auto const idx = " << "alpaka::getIdx(acc)[0];\n"; op << SP << SP << "if (idx >= total) return;\n"; - op << SP << SP << "constexpr std::size_t N = " << fN << "u;\n"; - op << SP << SP << "constexpr std::size_t MN = " << (fM * fN) << "u;\n"; + op << SP << SP << "const std::size_t MN = M * N;\n"; op << SP << SP << "const std::size_t mat_id = idx % MN;\n"; op << SP << SP << "const std::ptrdiff_t row = " << "static_cast(mat_id / N);\n"; @@ -212,7 +209,7 @@ public: out << SP << "auto const elementsPerThread_" << fNY << " = Vec::all(static_cast(1));\n"; out << SP << "auto const elementsPerGrid_" << fNY - << " = Vec::all(Idx{" << fTotal << "});\n"; + << " = Vec::all(Idx{static_cast(" << fTotal << ")});\n"; out << SP << "auto const workDiv_" << fNY << " = sofie_workdiv(elementsPerGrid_" << fNY << ");\n"; out << SP << "auto task_" << cleanOp @@ -220,7 +217,9 @@ public: << ", triluKernel_" << cleanOp << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" - << ", static_cast(" << fTotal << ")" + << ", static_cast(" << fTotal << ")" + << ", static_cast(" << fM.GetVal() << ")" + << ", static_cast(" << fN.GetVal() << ")" << ", k_" << cleanOp << ");\n"; out << SP << "alpaka::enqueue(queue, task_" << cleanOp << ");\n"; return out.str(); diff --git a/core/inc/SOFIE/ROperator_Where.hxx b/core/inc/SOFIE/ROperator_Where.hxx index b9956e90..2593ec48 100644 --- a/core/inc/SOFIE/ROperator_Where.hxx +++ b/core/inc/SOFIE/ROperator_Where.hxx @@ -38,7 +38,7 @@ private: std::vector fDimShapeY; std::vector fDimShapeZ; - // Broadcast flag: mirrors convention of BasicBinary + // Broadcast flag: // bit 0: broadcast Y->X (Y needs expanding) // bit 1: broadcast X->Y (X needs expanding) // bit 2: broadcast C->Z (C needs expanding) @@ -479,34 +479,43 @@ public: } - std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override { + // fDimShapeX/Y/C are not necessarily rank-padded to fDimShapeZ (only the dynamic + // Initialize() branch pads them; the static branch may leave a shorter, unpadded + // shape when broadcasting is done purely via strides rather than materialized data) + static std::vector PadToRank(const std::vector &shape, std::size_t rank) { + std::vector padded = shape; + if (padded.size() < rank) + padded.insert(padded.begin(), rank - padded.size(), Dim{1}); + return padded; + } + + bool IsPartialBroadcast(const std::vector &shape) const { + return ConvertDimShapeToLength(shape) != "1" && !UTILITY::AreSameShape(shape, fDimShapeZ); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; - if (fShapeZ.empty()) + if (fDimShapeZ.empty()) throw std::runtime_error("SOFIE Where Op called to Generate without being initialized first"); - const std::size_t D = fShapeZ.size(); - std::size_t totalElements = ConvertShapeToLength(fShapeZ); - - std::vector shapeA_padded(D, 1); - std::vector shapeB_padded(D, 1); - std::vector shapeC_padded(D, 1); - { - size_t offA = D - fShapeX.size(); - for (size_t i = 0; i < fShapeX.size(); ++i) shapeA_padded[offA + i] = fShapeX[i]; - size_t offB = D - fShapeY.size(); - for (size_t i = 0; i < fShapeY.size(); ++i) shapeB_padded[offB + i] = fShapeY[i]; - size_t offC = D - fShapeC.size(); - for (size_t i = 0; i < fShapeC.size(); ++i) shapeC_padded[offC + i] = fShapeC[i]; - } + const std::size_t D = fDimShapeZ.size(); + auto shapeX = PadToRank(fDimShapeX, D); + auto shapeY = PadToRank(fDimShapeY, D); + auto shapeC = PadToRank(fDimShapeC, D); + auto stridesX = UTILITY::ComputeStrideFromShape(shapeX); + auto stridesY = UTILITY::ComputeStrideFromShape(shapeY); + auto stridesC = UTILITY::ComputeStrideFromShape(shapeC); + auto stridesZ = UTILITY::ComputeStrideFromShape(fDimShapeZ); - auto stridesA = UTILITY::ComputeStrideFromShape(shapeA_padded); - auto stridesB = UTILITY::ComputeStrideFromShape(shapeB_padded); - auto stridesC = UTILITY::ComputeStrideFromShape(shapeC_padded); - auto stridesZ = UTILITY::ComputeStrideFromShape(fShapeZ); + bool isXScalar = ConvertDimShapeToLength(fDimShapeX) == "1"; + bool isYScalar = ConvertDimShapeToLength(fDimShapeY) == "1"; + bool isCScalar = ConvertDimShapeToLength(fDimShapeC) == "1"; + bool isXPartial = IsPartialBroadcast(fDimShapeX); + bool isYPartial = IsPartialBroadcast(fDimShapeY); + bool isCPartial = IsPartialBroadcast(fDimShapeC); - std::string typeName = TensorType::Name(); - std::string kname = "WhereKernel_" + opName; + std::string kname = "WhereKernel_" + opName; std::string op; op = "\n//------ WHERE_KERNEL_ALPAKA\n"; @@ -517,56 +526,40 @@ public: op += SP + SP + SP + "T const* __restrict__ x,\n"; op += SP + SP + SP + "T const* __restrict__ y,\n"; op += SP + SP + SP + "uint8_t const* __restrict__ cond,\n"; - op += SP + SP + SP + "T* __restrict__ output,\n"; - op += SP + SP + SP + "std::size_t const totalElements) const {\n\n"; + op += SP + SP + SP + "T* __restrict__ output"; + for (auto &p : dynParamNames) + op += ",\n" + SP + SP + SP + "std::size_t const " + p; + op += ",\n" + SP + SP + SP + "std::size_t const totalElements) const {\n\n"; op += SP + SP + SP + "auto const global_thread_idx = alpaka::getIdx(acc)[0];\n"; - op += SP + SP + SP + "if (global_thread_idx >= totalElements) return;\n"; op += SP + SP + SP + "auto const grid_thread_extent = alpaka::getWorkDiv(acc)[0];\n\n"; op += SP + SP + SP + "for (std::size_t elem_idx = global_thread_idx; elem_idx < totalElements; elem_idx += grid_thread_extent) {\n\n"; - for (std::size_t d = 0; d < D; ++d) { - op += SP + SP + SP + SP + "std::size_t const out_" + std::to_string(d) - + " = (elem_idx / " + std::to_string(stridesZ[d]) + "u) % " - + std::to_string(fShapeZ[d]) + "u;\n"; - } - op += "\n"; - - op += SP + SP + SP + SP + "std::size_t const c_idx =\n"; - for (std::size_t d = 0; d < D; ++d) { - if (shapeC_padded[d] == 1) - op += SP + SP + SP + SP + SP + "0u"; - else - op += SP + SP + SP + SP + SP - + "out_" + std::to_string(d) - + " * " + std::to_string(stridesC[d]) + "u"; - op += (d + 1 < D) ? " +\n" : ";\n\n"; - } - - op += SP + SP + SP + SP + "std::size_t const x_idx =\n"; - for (std::size_t d = 0; d < D; ++d) { - if (shapeA_padded[d] == 1) - op += SP + SP + SP + SP + SP + "0u"; - else - op += SP + SP + SP + SP + SP - + "out_" + std::to_string(d) - + " * " + std::to_string(stridesA[d]) + "u"; - op += (d + 1 < D) ? " +\n" : ";\n\n"; - } - - op += SP + SP + SP + SP + "std::size_t const y_idx =\n"; - for (std::size_t d = 0; d < D; ++d) { - if (shapeB_padded[d] == 1) - op += SP + SP + SP + SP + SP + "0u"; - else - op += SP + SP + SP + SP + SP - + "out_" + std::to_string(d) - + " * " + std::to_string(stridesB[d]) + "u"; - op += (d + 1 < D) ? " +\n" : ";\n\n"; + if (isXPartial || isYPartial || isCPartial) { + op += SP + SP + SP + SP + "std::size_t remaining = elem_idx;\n"; + op += SP + SP + SP + SP + "std::size_t coord;\n"; + if (isXPartial) op += SP + SP + SP + SP + "std::size_t idxX = 0;\n"; + if (isYPartial) op += SP + SP + SP + SP + "std::size_t idxY = 0;\n"; + if (isCPartial) op += SP + SP + SP + SP + "std::size_t idxC = 0;\n"; + for (std::size_t d = 0; d < D; d++) { + std::string sZ = "(" + stridesZ[d].GetVal() + ")"; + op += SP + SP + SP + SP + "coord = remaining / " + sZ + ";\n"; + if (d + 1 < D) + op += SP + SP + SP + SP + "remaining -= coord * " + sZ + ";\n"; + if (isXPartial && shapeX[d].GetVal() != "1") + op += SP + SP + SP + SP + "idxX += coord * (" + stridesX[d].GetVal() + ");\n"; + if (isYPartial && shapeY[d].GetVal() != "1") + op += SP + SP + SP + SP + "idxY += coord * (" + stridesY[d].GetVal() + ");\n"; + if (isCPartial && shapeC[d].GetVal() != "1") + op += SP + SP + SP + SP + "idxC += coord * (" + stridesC[d].GetVal() + ");\n"; + } } + std::string indexX = isXScalar ? "0" : (isXPartial ? "idxX" : "elem_idx"); + std::string indexY = isYScalar ? "0" : (isYPartial ? "idxY" : "elem_idx"); + std::string indexC = isCScalar ? "0" : (isCPartial ? "idxC" : "elem_idx"); - op += SP + SP + SP + SP + "output[elem_idx] = cond[c_idx] ? x[x_idx] : y[y_idx];\n"; + op += SP + SP + SP + SP + "output[elem_idx] = cond[" + indexC + "] ? x[" + indexX + "] : y[" + indexY + "];\n"; op += SP + SP + SP + "}\n"; op += SP + SP + "}\n"; op += SP + "};\n"; @@ -581,27 +574,29 @@ public: return SP + kname + " whereKernel_" + opName + ";\n"; } - std::string Generate_GPU_ALPAKA(std::string opName) override { + std::string Generate_GPU_ALPAKA(std::string opName, const std::vector &dynParamNames) override { if (fIsOutputConstant) return ""; opName = "op_" + opName; - if (fShapeZ.empty()) + if (fDimShapeZ.empty()) throw std::runtime_error("SOFIE Where Op called to Generate without being initialized first"); - std::size_t totalElements = ConvertShapeToLength(fShapeZ); + std::string totalElements = ConvertDimShapeToLength(fDimShapeZ); std::string kname = "whereKernel_" + opName; std::stringstream out; out << "\n//------ WHERE_GPU_ALPAKA\n"; out << SP << "auto const elementsPerThread_" << opName << " = Vec::all(static_cast(1));\n"; - out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{" << totalElements << "});\n"; + out << SP << "auto const elementsPerGrid_" << opName << " = Vec::all(Idx{static_cast(" << totalElements << ")});\n"; out << SP << "auto const workDiv_" << opName << " = sofie_workdiv(elementsPerGrid_" << opName << ");\n"; out << SP << "alpaka::exec(queue, workDiv_" << opName << ", " << kname << ", alpaka::getPtrNative(deviceBuf_" << fNX << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")" << ", alpaka::getPtrNative(deviceBuf_" << fNC << ")" - << ", alpaka::getPtrNative(deviceBuf_" << fNZ << ")" - << ", static_cast(" << totalElements << "));\n"; + << ", alpaka::getPtrNative(deviceBuf_" << fNZ << ")"; + for (auto &p : dynParamNames) + out << ", static_cast(" << p << ")"; + out << ", static_cast(" << totalElements << "));\n"; return out.str(); } diff --git a/core/inc/SOFIE/SOFIE_common.hxx b/core/inc/SOFIE/SOFIE_common.hxx index e36df0ab..033a3742 100644 --- a/core/inc/SOFIE/SOFIE_common.hxx +++ b/core/inc/SOFIE/SOFIE_common.hxx @@ -547,6 +547,17 @@ void UnidirectionalBroadcast(const T* data, const std::vector& shape, co std::vector ComputeStrideFromShape(const std::vector & shape); std::vector ComputeStrideFromShape(const std::vector & shape); +/// layout of the independent 1-D slices of a row-major tensor along one axis, as code expressions +struct SliceInfo { + std::string nBefore; // slice groups before the axis ("1" for axis 0) + std::string nAfter; // slices per group after the axis + std::string nSlices; // nBefore * nAfter + std::string nElements; // length of a slice (the axis dimension) + std::string strideAxis; // step along the axis + std::string strideBefore; // step between slice groups ("0" for axis 0) +}; +SliceInfo ComputeSliceInfo(const std::vector & shape, size_t axis); + /// function to check if a >> 0 and a < MAX using a single comparison //// use trick casting to unsigned values so it becomes a single comparison inline bool is_a_ge_zero_and_a_lt_b(int a, int b) { @@ -863,6 +874,11 @@ void ReadTensorFromStream(std::istream &is, T &target, std::string const &expect void EmitNestedLoops(std::stringstream &out, size_t loopRank, const std::vector shape); void CloseNestedLoops(std::stringstream &out, size_t loopRank); +//emit the out_0..out_{D-1} coordinate declarations of the flat kernel index elem_idx +//from the output strides and shape +void EmitOutputCoordsFromThreadIdx(std::string &op, const std::string &indent, + const std::vector &strides, const std::vector &shape); + // code for the memory greeding allocations struct TensorLifeInfo { diff --git a/core/src/RModel.cxx b/core/src/RModel.cxx index 35032a1b..86005ad1 100644 --- a/core/src/RModel.cxx +++ b/core/src/RModel.cxx @@ -1126,7 +1126,11 @@ std::string typeForOutput(ETensorType t) { return ConvertTypeToString(t); } -std::string memberNameForDimShape(std::string name) +} + +// Session member holding the construction-time value of a shape parameter (N -> fN); +// shared by the CPU and GPU session generators +std::string GetMemberNameForDimShape(std::string name) { if (!name.empty()) { name[0] = std::toupper(static_cast(name[0])); @@ -1135,8 +1139,6 @@ std::string memberNameForDimShape(std::string name) return name; } -} - void RModel::GenerateOutput() { size_t outputSize = fOutputTensorNames.size(); @@ -1189,7 +1191,7 @@ void RModel::GenerateOutput() if (dim.isParam && IsIdentifier(dim.param) && !IsInputTensorShapeParam(dim.param)) hasRuntimeParam = true; } - n = hasRuntimeParam ? memberNameForDimShape(dimLen) : dimLen; + n = hasRuntimeParam ? GetMemberNameForDimShape(dimLen) : dimLen; } std::string outputName = "output_tensor_" + name; fGC += SP + "std::vector<" + typeForOutput(GetTensorType(name)) + " > " + outputName + "(" + n + ");\n"; @@ -1215,7 +1217,7 @@ void RModel::GenerateOutput() for (auto &d : shape) { std::string pName = d.param; if (d.isParam && input_params_checked.count(pName) == 0) { - std::string memberName = memberNameForDimShape(d.param); + std::string memberName = GetMemberNameForDimShape(d.param); dynamic_parameters_check += d.param + " > " + memberName + " || "; input_params_checked.insert(pName); fGC += SP + "if (" + d.param + " > " + memberName + ") {\n"; @@ -1338,7 +1340,7 @@ void RModel::GenerateSessionCode() auto dimShapeNames = fDimShapeNames; std::sort(dimShapeNames.begin(), dimShapeNames.end()); for (const auto &p : dimShapeNames) { - fGC += "size_t " + memberNameForDimShape(p) + ";\n"; + fGC += "size_t " + GetMemberNameForDimShape(p) + ";\n"; } } @@ -1388,7 +1390,7 @@ void RModel::GenerateSessionCode() fGC += "\n\n"; std::sort(fDimShapeNames.begin(), fDimShapeNames.end()); for (const auto &p : fDimShapeNames) { - fGC += " " + memberNameForDimShape(p) + " = " + p + ";\n"; + fGC += " " + GetMemberNameForDimShape(p) + " = " + p + ";\n"; } } // add some extra code needed for initialization of dynamic parameters diff --git a/core/src/RModelProfilerGPU.cxx b/core/src/RModelProfilerGPU.cxx index bf946b57..dec760de 100644 --- a/core/src/RModelProfilerGPU.cxx +++ b/core/src/RModelProfilerGPU.cxx @@ -157,12 +157,12 @@ std::string RModelProfilerGPU::GenerateBeginInferCode() return gc; } -std::string RModelProfilerGPU::GenerateOperatorCode(ROperator &op, size_t op_idx) +std::string RModelProfilerGPU::GenerateOperatorCode(ROperator &op, size_t op_idx, const std::vector &dynParamNames) { std::string gc; gc += " // -- GPU Profiling operator: " + op.Name() + " --\n"; gc += " tp_start = std::chrono::steady_clock::now();\n"; - gc += op.Generate_GPU_ALPAKA(std::to_string(op_idx)); + gc += op.Generate_GPU_ALPAKA(std::to_string(op_idx), dynParamNames); // Force synchronisation so chrono measures actual GPU execution time gc += " alpaka::wait(queue);\n"; gc += " fProfilingResults[\"" + op.Name() + "\"].push_back(\n"; diff --git a/core/src/RModel_ALPAKA.cxx b/core/src/RModel_ALPAKA.cxx index 2c14c44a..07cf7acf 100644 --- a/core/src/RModel_ALPAKA.cxx +++ b/core/src/RModel_ALPAKA.cxx @@ -20,6 +20,42 @@ namespace SOFIE { +// device buffer / view alias per tensor type; must match the 'using Buf* / ViewConst*' +// aliases emitted into the generated Session +static std::string GetBufType(ETensorType t) { + switch (t) { + case ETensorType::FLOAT: return "BufF1D"; + case ETensorType::DOUBLE: return "BufD1D"; + case ETensorType::INT32: return "BufI321D"; + case ETensorType::INT64: return "BufI641D"; + case ETensorType::BOOL: + case ETensorType::UINT8: return "BufUI81D"; + default: + throw std::runtime_error("sofie: tensor type " + ConvertTypeToString(t) + + " is not supported on the ALPAKA backend"); + } +} + +static std::string GetViewConstType(ETensorType t) { + switch (t) { + case ETensorType::FLOAT: return "ViewConstF1D"; + case ETensorType::DOUBLE: return "ViewConstD1D"; + case ETensorType::INT32: return "ViewConstI321D"; + case ETensorType::INT64: return "ViewConstI641D"; + case ETensorType::BOOL: + case ETensorType::UINT8: return "ViewConstUI81D"; + default: + throw std::runtime_error("sofie: tensor type " + ConvertTypeToString(t) + + " is not supported on the ALPAKA backend"); + } +} + +// declaration + allocation line for one device buffer +static std::string AllocBufLine(const std::string &name, ETensorType t, const std::string &length) { + return GetBufType(t) + " deviceBuf_" + name + " = alpaka::allocBuf<" + ConvertOutputTypeToString(t) + + ", Idx>(devAcc, Ext1D::all(Idx{" + length + "}));\n"; +} + void RModel::ComputeEltwiseFusionGroups() { fEltwiseFusionGroups.clear(); fOpToFusionGroupIdx.clear(); @@ -53,7 +89,6 @@ void RModel::ComputeEltwiseFusionGroups() { auto firstInputs = fOperators[i]->GetOpInputTensors(); group.inputTensor = firstInputs.empty() ? "" : std::string(firstInputs[0]); - // Extend chain: only if CURRENT op is elementwise and its single output can be fused size_t current = i; while (fOperators[current]->IsElementwise()) { auto curOutputs = fOperators[current]->GetOpOutputTensors(); @@ -78,12 +113,10 @@ void RModel::ComputeEltwiseFusionGroups() { auto lastOutputs = fOperators[current]->GetOpOutputTensors(); group.outputTensor = lastOutputs.empty() ? "" : std::string(lastOutputs[0]); - // Element count from intermediate tensor info (all op outputs are intermediates) - if (!group.outputTensor.empty()) { - auto it = fIntermediateTensorInfos.find(group.outputTensor); - if (it != fIntermediateTensorInfos.end()) - group.numElements = ConvertShapeToLength(it->second.shape); - } + // only fused groups launch a fused kernel and need the element count; + // GetDimTensorShape covers static and dynamic tensors and throws if unknown + if (group.isFused()) + group.iterationLengthExpr = ConvertDimShapeToLength(GetDimTensorShape(group.outputTensor)); size_t gIdx = fEltwiseFusionGroups.size(); for (auto opIdx : group.opIndices) @@ -163,24 +196,7 @@ void RModel::GenerateInitializedTensorInfo_GPU_ALPAKA() { } size_t length = ConvertShapeToLength(i.second.shape()); - if (i.second.type() == ETensorType::FLOAT) { - fGC += "BufF1D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type() == ETensorType::INT32) { - fGC += "BufI321D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type() == ETensorType::INT64) { - fGC += "BufI641D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type() == ETensorType::BOOL || - i.second.type() == ETensorType::UINT8) { - fGC += "BufUI81D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } + fGC += AllocBufLine(i.first, i.second.type(), std::to_string(length)); } } @@ -194,16 +210,8 @@ void RModel::GenerateTemporaryInitializedTensorContainers_GPU_ALPAKA() if (fUseWeightFile && !i.second.IsConstantTensor()) { // case of tensors which are read from a file size_t length = ConvertShapeToLength(i.second.shape()); - if (i.second.type() == ETensorType::FLOAT) { - fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; - } else if (i.second.type() == ETensorType::INT32) { - fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; - } else if (i.second.type() == ETensorType::INT64) { - fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; - } else if (i.second.type() == ETensorType::BOOL || - i.second.type() == ETensorType::UINT8) { - fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; - } + fGC += "std::vector<" + ConvertOutputTypeToString(i.second.type()) + "> tensor_" + i.first + "(" + + std::to_string(length) + ");\n"; } } } @@ -218,27 +226,7 @@ void RModel::GenerateGPU_ALPAKA_Buffers() { size_t length = ConvertShapeToLength(i.second.shape); - if (i.second.type == ETensorType::FLOAT) { - tensor_declaration_block += "BufF1D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type == ETensorType::DOUBLE) { - tensor_declaration_block += "BufD1D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type == ETensorType::INT32) { - tensor_declaration_block += "BufI321D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type == ETensorType::INT64) { - tensor_declaration_block += "BufI641D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } else if (i.second.type == ETensorType::BOOL) { - tensor_declaration_block += "BufUI81D deviceBuf_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + - std::to_string(length) + "}));\n"; - } + tensor_declaration_block += AllocBufLine(i.first, i.second.type, std::to_string(length)); } if (tensor_declaration_block.length()) { @@ -249,18 +237,20 @@ void RModel::GenerateGPU_ALPAKA_Buffers() { // add also the dynamic tensors (only declarations, allocation will be done later) if (!fDynamicTensorInfos.empty()) { fGC += "//--- declare the dynamic tensors\n"; - fGC += "using bufDev_float = alpaka::Buf, size_t>;\n"; - fGC += "using bufDev_double = alpaka::Buf, size_t>;\n"; - fGC += "using bufDev_int64 = alpaka::Buf, size_t>;\n"; - for (auto &i : fDynamicTensorInfos) { - if (i.second.type == ETensorType::FLOAT) { - fGC += "bufDev_float bufDev_" + i.first + ";\n"; - } else if (i.second.type == ETensorType::DOUBLE) { - fGC += "bufDev_double bufDev_" + i.first + ";\n"; - } else if (i.second.type == ETensorType::INT64) { - fGC += "bufDev_int64 bufDev_" + i.first + ";\n"; - } + if (fFusionIntermediateTensors.count(i.first)) continue; + fGC += AllocBufLine(i.first, i.second.type, "1"); + } + } + + if (!fShapeTensors.empty()) { + fGC += "//--- declare the shape tensors\n"; + for (auto &i : fShapeTensors) { + size_t len = i.second.first.size(); + if (len == 0) continue; + fGC += "int64_t tensor_" + i.first + "[" + std::to_string(len) + "];\n"; + fGC += "BufI641D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + + std::to_string(len) + "}));\n"; } } } @@ -270,50 +260,45 @@ void RModel::GenerateDynamicTensorInfo_GPU_ALPAKA() { std::stringstream out; for (auto &i : fDynamicTensorInfos) { + if (fFusionIntermediateTensors.count(i.first)) continue; auto length = ConvertDimShapeToLength(i.second.shape); out << SP << "if (" << length << " > 0) {\n"; - out << "auto bufDev_" + i.first + - " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" << length << "}));\n"; + out << SP << "deviceBuf_" << i.first << " = alpaka::allocBuf<" << ConvertOutputTypeToString(i.second.type) + << ", Idx>(devAcc, Ext1D::all(Idx{" << length << "}));\n"; out << SP << "}\n"; } fGC += out.str(); } -std::string RModel::GenerateInferSignature_GPU_ALPAKA(bool isdecl) { - - auto GetBufType = [this](const std::string& name) -> std::string { - ETensorType type = GetTensorType(name); - if (type == ETensorType::FLOAT) return "BufF1D"; - if (type == ETensorType::DOUBLE) return "BufD1D"; - if (type == ETensorType::INT32) return "BufI321D"; - if (type == ETensorType::INT64) return "BufI641D"; - if (type == ETensorType::BOOL) return "BufUI81D"; - throw std::runtime_error("sofie: input tensor " + name + - " is of a data type which is not yet supported."); - }; - - std::string rGC; - std::unordered_map inputParams; - int i_input = 0; +void RModel::ForEachInferArg_GPU_ALPAKA(const std::function &onParam, + const std::function &onInput) const +{ + std::unordered_map seen; for (auto &name : fInputTensorNames) { - // if is a dynamic tensor pass initial parameters if (IsDimInputTensor(name)) { - auto shape = GetDynamicTensorShape(name); - for (auto &d : shape) { - std::string pName = d.param; - if (d.isParam && inputParams.count(pName) == 0) { - if (isdecl) rGC += "size_t "; - rGC += d.param + ","; - inputParams[pName] = i_input; + for (auto &d : GetDynamicTensorShape(name)) { + if (d.isParam && seen.count(d.param) == 0) { + seen[d.param] = 1; + onParam(d.param); } } } - if (isdecl) { - rGC += GetBufType(name) + " const "; - } - rGC += "deviceBuf_" + name + ","; - i_input++; + onInput(name); } +} + +std::string RModel::GenerateInferSignature_GPU_ALPAKA(bool isdecl) { + + std::string rGC; + ForEachInferArg_GPU_ALPAKA( + [&](const std::string &p) { + if (isdecl) rGC += "size_t "; + rGC += p + ","; + }, + [&](const std::string &name) { + if (isdecl) rGC += GetBufType(GetTensorType(name)) + " const "; + rGC += "deviceBuf_" + name + ","; + }); if (fInputTensorNames.size() > 0) rGC.pop_back(); // remove last "," return rGC; @@ -321,38 +306,16 @@ std::string RModel::GenerateInferSignature_GPU_ALPAKA(bool isdecl) { std::string RModel::GenerateImplSignature_GPU_ALPAKA(bool isdecl) { - auto GetViewConstType = [this](const std::string& name) -> std::string { - ETensorType type = GetTensorType(name); - if (type == ETensorType::FLOAT) return "ViewConstF1D"; - if (type == ETensorType::DOUBLE) return "ViewConstD1D"; - if (type == ETensorType::INT32) return "ViewConstI321D"; - if (type == ETensorType::INT64) return "ViewConstI641D"; - if (type == ETensorType::BOOL) return "ViewConstUI81D"; - throw std::runtime_error("sofie: input tensor " + name + - " is of a data type which is not yet supported."); - }; - std::string rGC; - std::unordered_map inputParams; - int i_input = 0; - for (auto &name : fInputTensorNames) { - if (IsDimInputTensor(name)) { - auto shape = GetDynamicTensorShape(name); - for (auto &d : shape) { - std::string pName = d.param; - if (d.isParam && inputParams.count(pName) == 0) { - if (isdecl) rGC += "size_t "; - rGC += d.param + ","; - inputParams[pName] = i_input; - } - } - } - if (isdecl) { - rGC += GetViewConstType(name) + " const& "; - } - rGC += "deviceBuf_" + name + ","; - i_input++; - } + ForEachInferArg_GPU_ALPAKA( + [&](const std::string &p) { + if (isdecl) rGC += "size_t "; + rGC += p + ","; + }, + [&](const std::string &name) { + if (isdecl) rGC += GetViewConstType(GetTensorType(name)) + " const& "; + rGC += "deviceBuf_" + name + ","; + }); if (fInputTensorNames.size() > 0) rGC.pop_back(); return rGC; @@ -373,38 +336,26 @@ void RModel::GenerateOutput_GPU_ALPAKA() { sameOutputTypes = false; } - auto GetViewConstType = [this](const std::string &name) -> std::string { - ETensorType type = GetTensorType(name); - if (type == ETensorType::FLOAT) return "ViewConstF1D"; - if (type == ETensorType::DOUBLE) return "ViewConstD1D"; - if (type == ETensorType::INT32) return "ViewConstI321D"; - if (type == ETensorType::INT64) return "ViewConstI641D"; - if (type == ETensorType::BOOL) return "ViewConstUI81D"; - throw std::runtime_error("sofie: input tensor " + name + " is of an unsupported data type."); - }; - // Collect deduplicated dynamic dimension parameter names in declaration order std::vector dynParamNames; - { - std::unordered_map seen; - for (auto &name : fInputTensorNames) { - if (IsDimInputTensor(name)) { - auto shape = GetDynamicTensorShape(name); - for (auto &d : shape) { - if (d.isParam && seen.count(d.param) == 0) { - dynParamNames.push_back(d.param); - seen[d.param] = 1; - } - } - } - } - } + ForEachInferArg_GPU_ALPAKA([&](const std::string &p) { dynParamNames.push_back(p); }, + [](const std::string &) {}); fGC += "\n\n"; fGC += "void _infer_impl("; fGC += GenerateImplSignature_GPU_ALPAKA(); fGC += "){\n"; + // device buffers were sized in the ctor from its shape params; refuse larger ones + // (same check and message as the CPU session) + ForEachInferArg_GPU_ALPAKA( + [&](const std::string &p) { + fGC += SP + "if (" + p + " > " + GetMemberNameForDimShape(p) + ") {\n"; + fGC += SP + SP + "throw std::runtime_error(\"sofie: dynamic input tensor shape parameter " + p + + " exceeds the initialized maximum allowed shape.\");\n"; + fGC += SP + "}\n"; + }, + [](const std::string &) {}); // GPU profiling: _infer_impl is a member of Session, so fProfilingResults // is directly accessible without any alias. @@ -433,11 +384,11 @@ void RModel::GenerateOutput_GPU_ALPAKA() { fusedCode += "\n//------ FUSED_ELTWISE_GPU_ALPAKA" + sfx + "\n"; fusedCode += SP + "{\n"; fusedCode += SP + SP + "auto const elementsPerThread_fused" + sfx + " = Vec::all(static_cast(1));\n"; - fusedCode += SP + SP + "auto const elementsPerGrid_fused" + sfx + " = Vec::all(Idx{" + std::to_string(grp.numElements) + "});\n"; + fusedCode += SP + SP + "auto const elementsPerGrid_fused" + sfx + " = Vec::all(Idx{" + grp.iterationLengthExpr + "});\n"; fusedCode += SP + SP + "auto const workDiv_fused" + sfx + " = sofie_workdiv(elementsPerGrid_fused" + sfx + ");\n"; fusedCode += SP + SP + "auto task_fused" + sfx + " = alpaka::createTaskKernel(workDiv_fused" + sfx + ", " + kname + ", alpaka::getPtrNative(deviceBuf_" + grp.inputTensor + "), alpaka::getPtrNative(deviceBuf_" + grp.outputTensor + - "), static_cast(" + std::to_string(grp.numElements) + "));\n"; + "), static_cast(" + grp.iterationLengthExpr + "));\n"; fusedCode += SP + SP + "alpaka::enqueue(queue, task_fused" + sfx + ");\n"; fusedCode += SP + "}\n"; if (fProfile) { @@ -458,9 +409,9 @@ void RModel::GenerateOutput_GPU_ALPAKA() { // Chain followers: skip — their logic is inside the fused kernel } else { if (fProfile) { - fGC += RModelProfilerGPU::GenerateOperatorCode(*fOperators[op_idx], op_idx); + fGC += RModelProfilerGPU::GenerateOperatorCode(*fOperators[op_idx], op_idx, dynParamNames); } else { - fGC += fOperators[op_idx]->Generate_GPU_ALPAKA(std::to_string(op_idx)); + fGC += fOperators[op_idx]->Generate_GPU_ALPAKA(std::to_string(op_idx), dynParamNames); } } } @@ -481,18 +432,18 @@ void RModel::GenerateOutput_GPU_ALPAKA() { fGC += "void infer(std::span inputs, std::span outputs" + spanDynDecl + "){\n"; { + // each symbol goes just before the first input whose shape introduces it fGC += SP + "_infer_impl("; bool first = true; - for (auto &p : dynParamNames) { + size_t i_input = 0; + auto appendArg = [&](const std::string &arg) { if (!first) fGC += ", "; - fGC += p; + fGC += arg; first = false; - } - for (size_t i = 0; i < fInputTensorNames.size(); i++) { - if (!first) fGC += ", "; - fGC += "inputs[" + std::to_string(i) + "]"; - first = false; - } + }; + ForEachInferArg_GPU_ALPAKA(appendArg, [&](const std::string &) { + appendArg("inputs[" + std::to_string(i_input++) + "]"); + }); fGC += ");\n"; } @@ -525,16 +476,15 @@ void RModel::GenerateOutput_GPU_ALPAKA() { fGC += GenerateInferSignature_GPU_ALPAKA(); fGC += "){\n"; - // Wrap each typed input buffer in a ViewConstXX, then call _infer_impl std::vector typedImplArgs; - for (auto &p : dynParamNames) - typedImplArgs.push_back(p); - for (auto &name : fInputTensorNames) { - std::string viewType = GetViewConstType(name); - fGC += SP + viewType + " const view_" + name + - "{alpaka::getPtrNative(deviceBuf_" + name + "), devAcc, alpaka::getExtents(deviceBuf_" + name + ")};\n"; - typedImplArgs.push_back("view_" + name); - } + ForEachInferArg_GPU_ALPAKA( + [&](const std::string &p) { typedImplArgs.push_back(p); }, + [&](const std::string &name) { + std::string viewType = GetViewConstType(GetTensorType(name)); + fGC += SP + viewType + " const view_" + name + + "{alpaka::getPtrNative(deviceBuf_" + name + "), devAcc, alpaka::getExtents(deviceBuf_" + name + ")};\n"; + typedImplArgs.push_back("view_" + name); + }); fGC += SP + "_infer_impl("; for (size_t i = 0; i < typedImplArgs.size(); i++) { @@ -558,6 +508,11 @@ void RModel::GenerateOutput_GPU_ALPAKA() { void RModel::GenerateSessionCode_GPU_ALPAKA() { + // the model's dynamic shape parameters in infer-argument order + std::vector dynParamNames; + ForEachInferArg_GPU_ALPAKA([&](const std::string &p) { dynParamNames.push_back(p); }, + [](const std::string &) {}); + std::set registered_operators; std::set fusedGroupsEmitted; // tracks which fusion groups have had their struct/decl emitted @@ -565,7 +520,6 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { SOFIE::OperatorKind::RELU, SOFIE::OperatorKind::SIGMOID, SOFIE::OperatorKind::TANH, - SOFIE::OperatorKind::SOFTMAX, SOFIE::OperatorKind::LEAKYRELU, SOFIE::OperatorKind::EINSUM, SOFIE::OperatorKind::ELU, @@ -619,13 +573,13 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { if (registered_operators.find(fOperators[id]->GetKind()) == registered_operators.end()) { if (fVerbose) std::cout << "Generating ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); + fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id), dynParamNames); registered_operators.insert(fOperators[id]->GetKind()); } } else { if (fVerbose) std::cout << "Generating ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); + fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id), dynParamNames); } } } @@ -718,15 +672,46 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { if (fOperators[id]->GetKind() == OperatorKind::GEMM || fOperators[id]->GetKind() == OperatorKind::CONV) { for (auto &blasCfg : fOperators[id]->GetBlasConfigs()) { if (!blasCfg.empty()) - fGC += "\nblas.addLayoutConfig(" + blasCfg + ");\n"; + fGC += "\nblas.addOperationConfig(" + blasCfg + ");\n"; } } } fGC += "\nalpaka::wait(queue);\n"; - std::string ctorBody = fGC; + /* + * Constructor shape parameters: first the ones the infer signature introduces, in that + * order, so positional constructor arguments cannot permute on multi-symbol models; then + * the ones an operator registered itself. The same list drives the constructor signature, + * the Session members that keep the construction-time values, and their assignment at + * the top of the constructor body. + */ + std::vector ctorParamNames; + ForEachInferArg_GPU_ALPAKA( + [&](const std::string &p) { ctorParamNames.push_back(p); }, + [](const std::string &) {}); + for (auto &p : fShapeParams) { + if (std::find(ctorParamNames.begin(), ctorParamNames.end(), p.first) == ctorParamNames.end()) + ctorParamNames.push_back(p.first); + } + + std::string ctorBody; + for (auto &p : ctorParamNames) + ctorBody += SP + GetMemberNameForDimShape(p) + " = " + p + ";\n"; + ctorBody += fGC; fGC = savedGC; + std::string ctorParams; + for (auto &p : ctorParamNames) + ctorParams += ",\n size_t " + p + " = " + fShapeParams[p]; + + /* + * One Session member per shape parameter, the same members the CPU session declares. + * The infer arguments are checked against them at the top of _infer_impl; a parameter + * registered by an operator is checked by that operator. + */ + for (auto &p : ctorParamNames) + fGC += "size_t " + GetMemberNameForDimShape(p) + ";\n"; + // ---- public constructors with inlined body ---- fGC += "public:\n"; @@ -735,10 +720,7 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fGC += "\n\n" + sessionName + "(std::string filename = \"" + fileName + "\""; else fGC += "\n\n" + sessionName + "(std::string filename = \"\""; - for (auto &p : fShapeParams) { - fGC += ",\n"; - fGC += " size_t " + p.first + " = " + p.second; - } + fGC += ctorParams; fGC += ") {\n"; fGC += ctorBody; fGC += "}\n\n"; @@ -748,10 +730,7 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fGC += sessionName + "(QueueAcc& extQueue, std::string filename = \"" + fileName + "\""; else fGC += sessionName + "(QueueAcc& extQueue, std::string filename = \"\""; - for (auto &p : fShapeParams) { - fGC += ",\n"; - fGC += " size_t " + p.first + " = " + p.second; - } + fGC += ctorParams; fGC += ")\n : queue(extQueue)"; if (OpNeedsBlas) fGC += ", blas(queue)"; @@ -797,13 +776,19 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { GenerateOutput_GPU_ALPAKA(); - // emit resetState() — zeros all intermediate tensors to reset model state + // zeros all intermediate tensors (static- and dynamic-shape alike) to + // reset model state. A dynamic tensor's device buffer is zeroed over its full allocated + // capacity if (fUseSession) { fGC += "\nvoid resetState(QueueAcc& queue) {\n"; for (auto &i : fIntermediateTensorInfos) { if (fFusionIntermediateTensors.count(i.first)) continue; fGC += SP + "alpaka::memset(queue, deviceBuf_" + i.first + ", 0);\n"; } + for (auto &i : fDynamicTensorInfos) { + if (fFusionIntermediateTensors.count(i.first)) continue; + fGC += SP + "alpaka::memset(queue, deviceBuf_" + i.first + ", 0);\n"; + } fGC += SP + "alpaka::wait(queue);\n"; fGC += "}\n"; } diff --git a/core/src/RModel_Base.cxx b/core/src/RModel_Base.cxx index 793a70dc..ece2e163 100644 --- a/core/src/RModel_Base.cxx +++ b/core/src/RModel_Base.cxx @@ -117,32 +117,6 @@ void RModel_Base::GenerateHeaderInfo_GPU_ALPAKA(std::string& hgname) { fGC += " alpaka::Vec::all(TIdx{1}));\n"; fGC += "}\n"; fGC += "#endif // SOFIE_ALPAKA_WORKDIV_DEFINED\n"; - - // Macros that redirect alpaka two-argument device-math calls to direct C/C++ calls - // when compiled by NVCC's device pass (__CUDA_ARCH__ defined). The normal alpaka - // form is used in host compilation and on non-CUDA backends. - // This prevents the deep template instantiation of alpaka's device-math dispatch - // from corrupting NVCC's namespace tracking (which would misattribute Session to - // when multiple model headers are included in the same translation unit). - fGC += "\n#ifndef SOFIE_ALPAKA_DEVICE_MATH_DEFINED\n"; - fGC += "#define SOFIE_ALPAKA_DEVICE_MATH_DEFINED\n"; - fGC += "#if defined(__CUDA_ARCH__)\n"; - for (const char* fn : {"exp","sqrt","log","sin","cos","abs","pow","atan2","atan", - "floor","ceil","round","cbrt","tanh","sinh","cosh"}) { - std::string sfn = fn; - // pow and atan2 take two numeric args plus acc; handle with variadic macro. - // All others: SOFIE_DEVICE_fn(acc, x) -> ::fn(x) - fGC += "#define SOFIE_DEVICE_" + sfn + "(acc, ...) ::" + sfn + "(__VA_ARGS__)\n"; - } - fGC += "#else\n"; - for (const char* fn : {"exp","sqrt","log","sin","cos","abs","pow","atan2","atan", - "floor","ceil","round","cbrt","tanh","sinh","cosh"}) { - std::string sfn = fn; - fGC += "#define SOFIE_DEVICE_" + sfn + "(acc, ...) " + sfn + "(acc, __VA_ARGS__)\n"; - } - fGC += "#endif // __CUDA_ARCH__\n"; - fGC += "#endif // SOFIE_ALPAKA_DEVICE_MATH_DEFINED\n"; - fGC += "\nnamespace SOFIE_" + fName + "{\n"; } diff --git a/core/src/SOFIE_common.cxx b/core/src/SOFIE_common.cxx index a2bafde2..32fdaea1 100644 --- a/core/src/SOFIE_common.cxx +++ b/core/src/SOFIE_common.cxx @@ -556,6 +556,30 @@ std::vector UTILITY::ComputeStrideFromShape(const std::vector & shape) return strides; } +UTILITY::SliceInfo UTILITY::ComputeSliceInfo(const std::vector & shape, size_t axis) { + auto stride = ComputeStrideFromShape(shape); + std::vector before(shape.begin(), shape.begin() + axis); + std::vector others(shape); + others.erase(others.begin() + axis); + SliceInfo s; + s.nBefore = ConvertDimShapeToLength(before); + s.nAfter = stride[axis].GetVal(); + s.nSlices = ConvertDimShapeToLength(others); + s.nElements = shape[axis].GetVal(); + s.strideAxis = stride[axis].GetVal(); + s.strideBefore = (axis > 0) ? stride[axis - 1].GetVal() : "0"; + return s; +} + +void EmitOutputCoordsFromThreadIdx(std::string &op, const std::string &indent, + const std::vector &strides, const std::vector &shape) { + for (std::size_t d = 0; d < shape.size(); ++d) { + op += indent + "std::size_t const out_" + std::to_string(d) + + " = (elem_idx / (" + strides[d].GetVal() + ")) % (" + + shape[d].GetVal() + ");\n"; + } +} + struct FreeBlock { std::size_t offset; std::size_t size; diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index bcbfbaa4..8dc13d4f 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -181,10 +181,18 @@ if (ENABLE_ALPAKA_TESTS) ${CMAKE_CURRENT_SOURCE_DIR} ) + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "") + set(SOFIE_TEST_CUDA_ARCH "75") # RTX 2070 SUPER default + else() + set(SOFIE_TEST_CUDA_ARCH "${CMAKE_CUDA_ARCHITECTURES}") + endif() + message(STATUS "SOFIE Alpaka CUDA tests: CUDA architectures = ${SOFIE_TEST_CUDA_ARCH}") + set_target_properties( TestCustomModelsFromONNXForAlpakaCuda PROPERTIES CUDA_SEPARABLE_COMPILATION OFF + CUDA_ARCHITECTURES ${SOFIE_TEST_CUDA_ARCH} CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON ) diff --git a/test/TestCustomModelsFromONNXForAlpakaCuda.cxx b/test/TestCustomModelsFromONNXForAlpakaCuda.cxx new file mode 100644 index 00000000..0a0c62c6 --- /dev/null +++ b/test/TestCustomModelsFromONNXForAlpakaCuda.cxx @@ -0,0 +1,4206 @@ +#include +#include + +// ── Trilu ────────────────────────────────────────────────────────────────── +#include "Trilu_upper_FromONNX_GPU_ALPAKA.hxx" +#include "Trilu_lower_FromONNX_GPU_ALPAKA.hxx" +#include "Trilu_k2_FromONNX_GPU_ALPAKA.hxx" +#include "Trilu_kn1_FromONNX_GPU_ALPAKA.hxx" +#include "Trilu_3D_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Trilu_upper.ref.hxx" +#include "input_models/references/Trilu_upper_input.ref.hxx" +#include "input_models/references/Trilu_lower.ref.hxx" +#include "input_models/references/Trilu_lower_input.ref.hxx" +#include "input_models/references/Trilu_k2.ref.hxx" +#include "input_models/references/Trilu_k2_input.ref.hxx" +#include "input_models/references/Trilu_kn1.ref.hxx" +#include "input_models/references/Trilu_kn1_input.ref.hxx" +#include "input_models/references/Trilu_3D.ref.hxx" +#include "input_models/references/Trilu_3D_input.ref.hxx" +// ── Logic ─────────────────────────────────────────────────────────────────── +#include "Logic_And_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_Or_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_Xor_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_BitwiseAnd_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_BitwiseOr_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_BitwiseXor_FromONNX_GPU_ALPAKA.hxx" +#include "Logic_BitwiseNot_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Logic_And.ref.hxx" +#include "input_models/references/Logic_And_input.ref.hxx" +#include "input_models/references/Logic_Or.ref.hxx" +#include "input_models/references/Logic_Or_input.ref.hxx" +#include "input_models/references/Logic_Xor.ref.hxx" +#include "input_models/references/Logic_Xor_input.ref.hxx" +#include "input_models/references/Logic_BitwiseAnd.ref.hxx" +#include "input_models/references/Logic_BitwiseAnd_input.ref.hxx" +#include "input_models/references/Logic_BitwiseOr.ref.hxx" +#include "input_models/references/Logic_BitwiseOr_input.ref.hxx" +#include "input_models/references/Logic_BitwiseXor.ref.hxx" +#include "input_models/references/Logic_BitwiseXor_input.ref.hxx" +#include "input_models/references/Logic_BitwiseNot.ref.hxx" +#include "input_models/references/Logic_BitwiseNot_input.ref.hxx" +// ───────────────────────────────────────────────────────────────────────── + +#include "Linear_64_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Linear_64.ref.hxx" + +#include "AddBroadcast1_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/AddBroadcast1.ref.hxx" + +#include "LinearWithLeakyRelu_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/LinearWithLeakyRelu.ref.hxx" + +#include "LinearWithSigmoid_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/LinearWithSigmoid.ref.hxx" + +#include "Transpose_FromONNX_GPU_ALPAKA.hxx" + +#include "Concat_0D_FromONNX_GPU_ALPAKA.hxx" +#include "ScatterElements_FromONNX_GPU_ALPAKA.hxx" + +#include "Split_0_FromONNX_GPU_ALPAKA.hxx" +#include "Split_1_FromONNX_GPU_ALPAKA.hxx" +#include "Split_2_FromONNX_GPU_ALPAKA.hxx" + +#include "Tile5D_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Tile5D.ref.hxx" +#include "DynamicTile_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicEqual_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicAddBroadcast_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicGather_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicTranspose_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConcat_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicSlice_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicNegRelu_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicLinear_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConv1DNoBias_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConv2DNoBias_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicRange_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicRangeMul_FromONNX_GPU_ALPAKA.hxx" + +#include "GatherAxis0_FromONNX_GPU_ALPAKA.hxx" +#include "GatherAxis1_FromONNX_GPU_ALPAKA.hxx" +#include "GatherAxis2_FromONNX_GPU_ALPAKA.hxx" +#include "GatherAxis3_FromONNX_GPU_ALPAKA.hxx" +#include "Gather2d_FromONNX_GPU_ALPAKA.hxx" +#include "GatherNegativeIndices_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/GatherAxis0.ref.hxx" +#include "input_models/references/GatherAxis1.ref.hxx" +#include "input_models/references/GatherAxis2.ref.hxx" +#include "input_models/references/GatherAxis3.ref.hxx" +#include "input_models/references/Gather2d.ref.hxx" +#include "input_models/references/GatherNegativeIndices.ref.hxx" + +#include "ExpandSameSize_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ExpandSameSize.ref.hxx" + +#include "ExpandDiffSize_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ExpandDiffSize.ref.hxx" + +#include "GatherND_Ex1_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_Ex2_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_Ex3_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_Ex4_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_Ex5_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_NegativeIndices_FromONNX_GPU_ALPAKA.hxx" +#include "GatherND_Batch_FromONNX_GPU_ALPAKA.hxx" + +#include "Equal_FromONNX_GPU_ALPAKA.hxx" +#include "LessOrEqual_FromONNX_GPU_ALPAKA.hxx" +#include "GreaterOrEqual_FromONNX_GPU_ALPAKA.hxx" +#include "Greater_FromONNX_GPU_ALPAKA.hxx" +#include "Less_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Equal.ref.hxx" +#include "input_models/references/LessOrEqual.ref.hxx" +#include "input_models/references/GreaterOrEqual.ref.hxx" +#include "input_models/references/Greater.ref.hxx" +#include "input_models/references/Less.ref.hxx" + +#include "Slice_FromONNX_GPU_ALPAKA.hxx" +#include "Slice_Default_Axis_FromONNX_GPU_ALPAKA.hxx" +#include "Slice_Default_Steps_FromONNX_GPU_ALPAKA.hxx" +#include "Slice_Neg_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Slice.ref.hxx" +#include "input_models/references/Slice_Default_Axis.ref.hxx" +#include "input_models/references/Slice_Default_Steps.ref.hxx" +#include "input_models/references/Slice_Neg.ref.hxx" + +#include "Sin_FromONNX_GPU_ALPAKA.hxx" +#include "Cos_FromONNX_GPU_ALPAKA.hxx" +#include "Abs_FromONNX_GPU_ALPAKA.hxx" +#include "Sqrt_FromONNX_GPU_ALPAKA.hxx" +#include "Reciprocal_FromONNX_GPU_ALPAKA.hxx" +#include "Exp_FromONNX_GPU_ALPAKA.hxx" +#include "Log_FromONNX_GPU_ALPAKA.hxx" +#include "Neg_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Sqrt.ref.hxx" +#include "input_models/references/Reciprocal.ref.hxx" +#include "input_models/references/Exp.ref.hxx" +#include "input_models/references/Log.ref.hxx" +#include "input_models/references/Neg.ref.hxx" + +#include "Where_FromONNX_GPU_ALPAKA.hxx" + +#include "Softplus_FromONNX_GPU_ALPAKA.hxx" + +#include "ReduceMean_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceProd_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceSum_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceSumSquare_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceL2_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceL2Large_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceMax_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceMax_axis0_FromONNX_GPU_ALPAKA.hxx" +#include "ReduceMax_mid_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceSumLast_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceMeanMid_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceMaxFirst_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceSumMulti_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ReduceMean.ref.hxx" +#include "input_models/references/ReduceProd.ref.hxx" +#include "input_models/references/ReduceL2.ref.hxx" +#include "input_models/references/ReduceMax.ref.hxx" +#include "input_models/references/ReduceMax_axis0.ref.hxx" +#include "input_models/references/ReduceMax_mid.ref.hxx" + +#include "ConvWithPadding_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithPadding.ref.hxx" + +#include "ConvWithoutPadding_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithoutPadding.ref.hxx" + +#include "ConvWithAutopadSameLower_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithAutopadSameLower.ref.hxx" + +#include "ConvWithStridesPadding_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithStridesPadding.ref.hxx" + +#include "ConvWithStridesNoPadding_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithStridesNoPadding.ref.hxx" + +#include "ConvWithAsymmetricPadding_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConv1D_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/ConvWithAsymmetricPadding.ref.hxx" + +#include "BatchNorm_FromONNX_GPU_ALPAKA.hxx" +#include "BatchNormRelu_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicBatchNorm4D_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicBatchNormDynSpatialRelu_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicBatchNorm2D_FromONNX_GPU_ALPAKA.hxx" + +#include "LayerNorm_FromONNX_GPU_ALPAKA.hxx" +#include "LayerNormScaleBias_FromONNX_GPU_ALPAKA.hxx" +#include "LayerNorm3D_FromONNX_GPU_ALPAKA.hxx" + +#include "IsInf_FromONNX_GPU_ALPAKA.hxx" +#include "IsNaN_FromONNX_GPU_ALPAKA.hxx" +#include "Clip_FromONNX_GPU_ALPAKA.hxx" +#include "Not_FromONNX_GPU_ALPAKA.hxx" + +#include "TopK_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/TopK.ref.hxx" +#include "Softmax1d_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Softmax1d.ref.hxx" +#include "Softmax2d_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Softmax2d.ref.hxx" +#include "Softmax3d_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Softmax3d.ref.hxx" +#include "Softmax4d_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Softmax4d.ref.hxx" + +#include "GNN_model_FromONNX_GPU_ALPAKA.hxx" + +#include +#include +#include +#include "gtest/gtest.h" + +constexpr float DEFAULT_TOLERANCE = 1e-3f; + +using Idx = std::size_t; +using Dim = alpaka::DimInt<1>; +using Ext1D = alpaka::Vec; + +class SofieAlpakaTest : public ::testing::Test { +protected: + // Shared devices and platforms + alpaka::PlatformCpu hostPlatform; + alpaka::DevCpu host; + alpaka::PlatformCudaRt platform; + alpaka::DevCudaRt device; + alpaka::Queue queue; + + SofieAlpakaTest() + : hostPlatform{} + , host(alpaka::getDevByIdx(hostPlatform, 0u)) + , platform{} + , device(alpaka::getDevByIdx(platform, 0u)) + , queue(device) + { + } + + void SetUp() override { + cudaDeviceSynchronize(); + } + + void TearDown() override { + alpaka::wait(queue); + cudaDeviceSynchronize(); + } + + ~SofieAlpakaTest() override { + cudaDeviceSynchronize(); + } +}; + + +TEST_F(SofieAlpakaTest, Linear64) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + auto A = alpaka::allocBuf(host, Ext1D::all(Idx{6400})); + float *A_ptr = reinterpret_cast(alpaka::getPtrNative(A)); + + for (Idx i = 0; i < 6400; ++i) { + A_ptr[i] = 1.0; + } + + auto A_d = alpaka::allocBuf(device, Ext1D::all(Idx{6400})); + alpaka::memcpy(queue, A_d, A); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{640})); + + { + SOFIE_Linear_64::Session session("Linear_64_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(A_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 = Linear_64_ExpectedOutput::all_ones; + + for (size_t i = 0; i < 640; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, LinearWithLeakyRelu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({ + 0.4369, -0.6882, 1.0309, -1.0263, -0.1519, 1.2237, -0.7054, -0.1762, + -0.6811, -2.2597, 1.0388, -0.7993, 0.1468, 1.3257, -0.4714, -0.0958, + 0.7057, -0.3749, -0.3310, 0.0986, -0.1370, 0.0832, -1.6465, -0.2793 + }); + + auto A = alpaka::allocBuf(host, Ext1D::all(Idx{input.size()})); + float *A_ptr = reinterpret_cast(alpaka::getPtrNative(A)); + + for (Idx i = 0; i < input.size(); ++i) { + A_ptr[i] = input[i]; + } + + auto A_d = alpaka::allocBuf(device, Ext1D::all(Idx{input.size()})); + alpaka::memcpy(queue, A_d, A); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{24})); + + { + SOFIE_LinearWithLeakyRelu::Session session; + auto result = session.infer(A_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 = LinearWithLeakyRelu_ExpectedOutput::outputs; + + for (size_t i = 0; i < 24; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, LinearWithSigmoid) +{ + + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + auto A = alpaka::allocBuf(host, Ext1D::all(Idx{48})); + float *A_ptr = reinterpret_cast(alpaka::getPtrNative(A)); + + for (Idx i = 0; i < 48; ++i) { + A_ptr[i] = 1.0; + } + + auto A_d = alpaka::allocBuf(device, Ext1D::all(Idx{48})); + alpaka::memcpy(queue, A_d, A); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{24})); + + { + SOFIE_LinearWithSigmoid::Session session("LinearWithSigmoid_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(A_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 = LinearWithSigmoid_ExpectedOutput::all_ones; + for (size_t i = 0; i < 24; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, AddBroadcast1) +{ + + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + auto A = alpaka::allocBuf(host, Ext1D::all(Idx{5})); + float *A_ptr = reinterpret_cast(alpaka::getPtrNative(A)); + + auto B = alpaka::allocBuf(host, Ext1D::all(Idx{20})); + float *B_ptr = reinterpret_cast(alpaka::getPtrNative(B)); + + std::vector A_vec({-0.78023305, -1.34029483, -3.01482951, 0.53641361, + -1.22594789}); + std::vector B_vec({1.0626695, 0.43842875, 1.22476468, 0.79763274, 0.98688211, + 0.25267614, 0.44874883, 0.31516773, -0.78771195, 0.64565664, + 0.50450593, -0.41265227, -0.22474539, -0.22362374, 0.00509674, + 0.16927211, 1.06756969, -0.81634773, 0.88467744, 0.78902059}); + + for (Idx i = 0; i < A_vec.size(); ++i) { + A_ptr[i] = A_vec[i]; + } + + for (Idx i = 0; i < B_vec.size(); ++i) { + B_ptr[i] = B_vec[i]; + } + + auto A_d = alpaka::allocBuf(device, Ext1D::all(Idx{5})); + alpaka::memcpy(queue, A_d, A); + alpaka::wait(queue); + + auto B_d = alpaka::allocBuf(device, Ext1D::all(Idx{20})); + alpaka::memcpy(queue, B_d, B); + alpaka::wait(queue); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{20})); + + { + SOFIE_AddBroadcast1::Session session; + auto result = session.infer(A_d, B_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 = AddBroadcast1_ExpectedOutput::output; + for (size_t i = 0; i < 20; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, Transpose) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Input shape: (2, 1, 3, 4) -> 24 elements + constexpr Idx inputSize = 24; + // Output shape: (2, 3, 4, 1) -> 24 elements + constexpr Idx outputSize = 24; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + + std::vector input_vec({ + // shape (2, 1, 3, 4) + 0.f, 1.f, 2.f, 3.f, + 4.f, 5.f, 6.f, 7.f, + 8.f, 9.f, 10.f, 11.f, + + 12.f, 13.f, 14.f, 15.f, + 16.f, 17.f, 18.f, 19.f, + 20.f, 21.f, 22.f, 23.f + }); + + for (Idx i = 0; i < inputSize; ++i) + input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Transpose::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + std::vector expected(outputSize); + std::vector inputShape = {2, 1, 3, 4}; + std::vector perm = {0, 2, 3, 1}; + std::vector outputShape = {2, 3, 4, 1}; + + std::vector inputStrides = {12, 12, 4, 1}; + std::vector outputStrides = {12, 4, 1, 1}; + + for (size_t i = 0; i < outputSize; ++i) + { + size_t remaining = i; + size_t inputIdx = 0; + for (size_t d = 0; d < 4; ++d) + { + size_t const coord = remaining / outputStrides[d]; + remaining = remaining - coord * outputStrides[d]; + inputIdx += coord * inputStrides[perm[d]]; + } + expected[i] = input_vec[inputIdx]; + } + + float* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - expected[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, DynamicTranspose) +{ + // X[N,3,n_pf] perm(0,2,1) -> Y[N,n_pf,3], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicTranspose::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t p = 0; p < P; ++p) + for (std::size_t c = 0; c < C; ++c) { + float expected = in_ptr[n * C * P + c * P + p]; + float got = res[n * P * C + p * C + c]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " p=" << p << " c=" << c; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicRange) +{ + // X[N,K] dynamic -> Shape -> Gather(dim1=K) -> Range(0,K,1) -> Y[K] = arange(K) + const std::size_t Ns[] = {1, 4}; + const std::size_t Ks[] = {3, 7}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], K = Ks[t]; + const std::size_t sz = N * K; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + { + SOFIE_DynamicRange::Session session("", N, K); + auto result = session.infer(N, K, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int64_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < K; ++i) + EXPECT_EQ(res[i], static_cast(i)) << "i=" << i << " N=" << N << " K=" << K; + } +} + +TEST_F(SofieAlpakaTest, DynamicRangeMul) +{ + // X[N,K] dynamic -> Shape -> Gather(N),Gather(K) -> Range(0,K,1) -> Mul(arange, N) -> Y[K] = i*N + // N is a shape tensor read on-device by the Mul (the ParticleNet deviceBuf_168 pattern). + const std::size_t Ns[] = {1, 5}; + const std::size_t Ks[] = {4, 6}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], K = Ks[t]; + const std::size_t sz = N * K; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + { + SOFIE_DynamicRangeMul::Session session("", N, K); + auto result = session.infer(N, K, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int64_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < K; ++i) + EXPECT_EQ(res[i], static_cast(i * N)) << "i=" << i << " N=" << N << " K=" << K; + } +} + +TEST_F(SofieAlpakaTest, Concat0D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({1.40519865e+00, -2.87660856e-01}); + std::vector expected_output({ + 1.40519865e+00, -2.87660856e-01, + 1.40519865e+00, -2.87660856e-01 + }); + + // Host input buffer + 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]; + + // Device input buffer + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{input.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + // Host output buffer + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected_output.size()})); + + { + SOFIE_Concat_0D::Session session("Concat_0D_FromONNX_GPU_ALPAKA.dat"); + + 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)); + + for (size_t i = 0; i < expected_output.size(); ++i) { + EXPECT_LE(std::abs(res_ptr[i] - expected_output[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, DynamicConcat) +{ + // A[N,2,n_pf] + B[N,3,n_pf] axis 1 -> Y[N,5,n_pf], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Ca = 2, Cb = 3, Cy = 5; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t szA = N * Ca * P, szB = N * Cb * P, szY = N * Cy * P; + + auto a_h = alpaka::allocBuf(host, Ext1D::all(Idx{szA})); + auto b_h = alpaka::allocBuf(host, Ext1D::all(Idx{szB})); + float* a_ptr = reinterpret_cast(alpaka::getPtrNative(a_h)); + float* b_ptr = reinterpret_cast(alpaka::getPtrNative(b_h)); + for (Idx i = 0; i < szA; ++i) a_ptr[i] = static_cast(i % 7) - 3.0f; + for (Idx i = 0; i < szB; ++i) b_ptr[i] = static_cast(i % 5) + 10.0f; + + auto a_d = alpaka::allocBuf(device, Ext1D::all(Idx{szA})); + auto b_d = alpaka::allocBuf(device, Ext1D::all(Idx{szB})); + alpaka::memcpy(queue, a_d, a_h); + alpaka::memcpy(queue, b_d, b_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{szY})); + { + SOFIE_DynamicConcat::Session session("", N, P); + auto result = session.infer(N, P, a_d, b_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < Cy; ++j) + for (std::size_t p = 0; p < P; ++p) { + float expected = (j < Ca) ? a_ptr[n * Ca * P + j * P + p] + : b_ptr[n * Cb * P + (j - Ca) * P + p]; + float got = res[n * Cy * P + j * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " j=" << j << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, ScatterElements) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input (9, 0.f); + std::vector indices = { 1, 0, 2, 0, 2, 1 }; + std::vector updates = { 1.f, 1.1f, 1.2f, 2.f, 2.1f, 2.2f }; + std::vector correct = { 2.f, 1.1f, 0.f, 1.f, 0.f, 2.2f, 0.f, 2.1f, 1.2f }; + + // Allocate and fill host buffers + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{input.size()})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{indices.size()})); + auto updates_h = alpaka::allocBuf(host, Ext1D::all(Idx{updates.size()})); + + float* input_ptr = reinterpret_cast (alpaka::getPtrNative(input_h)); + int64_t* indices_ptr = reinterpret_cast(alpaka::getPtrNative(indices_h)); + float* updates_ptr = reinterpret_cast (alpaka::getPtrNative(updates_h)); + + for (Idx i = 0; i < input.size(); ++i) input_ptr[i] = input[i]; + for (Idx i = 0; i < indices.size(); ++i) indices_ptr[i] = indices[i]; + for (Idx i = 0; i < updates.size(); ++i) updates_ptr[i] = updates[i]; + + // Allocate device buffers and copy + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{input.size()})); + auto indices_d = alpaka::allocBuf(device, Ext1D::all(Idx{indices.size()})); + auto updates_d = alpaka::allocBuf(device, Ext1D::all(Idx{updates.size()})); + + alpaka::memcpy(queue, input_d, input_h); + alpaka::memcpy(queue, indices_d, indices_h); + alpaka::memcpy(queue, updates_d, updates_h); + alpaka::wait(queue); + + // Host result buffer + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct.size()})); + + { + SOFIE_ScatterElements::Session session; + auto result = session.infer(input_d, indices_d, updates_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + EXPECT_EQ(correct.size(), 9u); + for (size_t i = 0; i < correct.size(); ++i){ + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, Split_0) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // split in axis 0 in 2 tensors {2,2,3} -> {1,2,3} each + std::vector input {1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.}; + std::vector> correct_output = { {1.,2.,3.,4.,5.,6.}, {7.,8.,9.,10.,11.,12.} }; + + 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 result0_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[0].size()})); + auto result1_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[1].size()})); + + { + SOFIE_Split_0::Session session; + auto [result0, result1] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result0_h, result0); + alpaka::memcpy(queue, result1_h, result1); + alpaka::wait(queue); + } + + float* res0_ptr = reinterpret_cast(alpaka::getPtrNative(result0_h)); + float* res1_ptr = reinterpret_cast(alpaka::getPtrNative(result1_h)); + + for (size_t j = 0; j < correct_output[0].size(); ++j) + EXPECT_LE(std::abs(res0_ptr[j] - correct_output[0][j]), TOLERANCE); + for (size_t j = 0; j < correct_output[1].size(); ++j) + EXPECT_LE(std::abs(res1_ptr[j] - correct_output[1][j]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, Split_1) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // split in axis 1 in 2 tensors {2,2,3} -> {2,1,3} each + std::vector input {1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.}; + std::vector> correct_output = { {1.,2.,3.,7.,8.,9.}, {4.,5.,6.,10.,11.,12.} }; + + 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 result0_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[0].size()})); + auto result1_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[1].size()})); + + { + SOFIE_Split_1::Session session; + auto [result0, result1] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result0_h, result0); + alpaka::memcpy(queue, result1_h, result1); + alpaka::wait(queue); + } + + float* res0_ptr = reinterpret_cast(alpaka::getPtrNative(result0_h)); + float* res1_ptr = reinterpret_cast(alpaka::getPtrNative(result1_h)); + + for (size_t j = 0; j < correct_output[0].size(); ++j) + EXPECT_LE(std::abs(res0_ptr[j] - correct_output[0][j]), TOLERANCE); + for (size_t j = 0; j < correct_output[1].size(); ++j) + EXPECT_LE(std::abs(res1_ptr[j] - correct_output[1][j]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, Split_2) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // split in axis 2 in 2 tensors {2,2,3} -> {2,2,2} and {2,2,1} + std::vector input {1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.}; + std::vector> correct_output = { {1.,2.,4.,5.,7.,8.,10.,11.}, {3.,6.,9.,12.} }; + + 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); + + // outputs have different sizes: {2,2,2}=8 and {2,2,1}=4 + auto result0_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[0].size()})); + auto result1_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct_output[1].size()})); + + { + SOFIE_Split_2::Session session; + auto [result0, result1] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result0_h, result0); + alpaka::memcpy(queue, result1_h, result1); + alpaka::wait(queue); + } + + float* res0_ptr = reinterpret_cast(alpaka::getPtrNative(result0_h)); + float* res1_ptr = reinterpret_cast(alpaka::getPtrNative(result1_h)); + + for (size_t j = 0; j < correct_output[0].size(); ++j) + EXPECT_LE(std::abs(res0_ptr[j] - correct_output[0][j]), TOLERANCE); + for (size_t j = 0; j < correct_output[1].size(); ++j) + EXPECT_LE(std::abs(res1_ptr[j] - correct_output[1][j]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, Tile5D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input_data({ + 0.2386120855808258, 0.5549510717391968, -1.8190287351608276, 0.5724563598632812, -0.6596977710723877, + 0.17560836672782898, 0.7608169317245483, 0.08603227883577347, -0.049375515431165695, 0.2705111503601074, + 1.42119562625885, 0.032626643776893616, -1.212586522102356, -0.5129594802856445, -0.43296414613723755, + -0.1606937050819397, 1.1884371042251587, -0.662174642086029, -2.291109323501587, -0.6852569580078125, + 2.325223922729492, -0.19389064610004425, -0.5784135460853577, -0.39328137040138245, 0.2831517457962036, + 0.4496127665042877, -0.2029038816690445, 0.35477763414382935, 0.4266718924045563, 0.24683749675750732, + 1.90426504611969, -0.4861580729484558, 0.9139055013656616, -0.5031066536903381, 0.9583520293235779, + -0.23210509121418, 1.3183971643447876, 1.7042455673217773, -0.3201166093349457, -0.14444805681705475, + -0.8829464912414551, 1.725736141204834, 0.45657631754875183, 0.4920198321342468, -1.088847041130066, + 0.49437597393989563, -0.006085286382585764, 2.475630760192871, 0.12170185893774033, -0.8953945636749268, + 1.1430096626281738, 1.3278610706329346, 0.3076854348182678, 0.036237504333257675, 0.05180325731635094, + 0.2802475392818451, 0.5289335250854492, 0.9356630444526672, 0.7863689064979553, 0.4239695370197296, + 0.8723016977310181, -0.2248474359512329, 0.3891502320766449, 0.5463842153549194, -0.7782878875732422, + -0.8570080399513245, -2.593783378601074, -0.11392943561077118, 0.5637082457542419, 2.075004816055298, + -1.0598397254943848, 1.0823975801467896 + }); + + const std::size_t inputSize = input_data.size(); + const std::size_t outputSize = sizeof(Tile5D_ExpectedOutput::output) / sizeof(float); + + // Allocate and fill host input buffer + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inputSize; ++i) + input_ptr[i] = input_data[i]; + + // Copy to device + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + // Host result buffer + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Tile5D::Session session; + auto result = session.infer(input_d); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* correct = Tile5D_ExpectedOutput::output; + + EXPECT_EQ(outputSize, sizeof(Tile5D_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, DynamicTile) +{ + // X[N,2] -> Tile([2,3]) -> Y[2N,6], with N dynamic. Run at two batch sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t inCols = 2, outCols = 6; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inRows = N, outRows = 2 * N; + const std::size_t inSize = inRows * inCols, outSize = outRows * outCols; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) + in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicTile::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t r = 0; r < outRows; ++r) + for (std::size_t c = 0; c < outCols; ++c) { + float expected = in_ptr[(r % inRows) * inCols + (c % inCols)]; + EXPECT_LE(std::abs(res[r * outCols + c] - expected), TOLERANCE); + } + } +} + +TEST_F(SofieAlpakaTest, GatherAxis0) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 120; + const std::size_t outputSize = sizeof(GatherAxis0_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GatherAxis0::Session session("GatherAxis0_FromONNX_GPU_ALPAKA.dat"); + 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 = GatherAxis0_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(GatherAxis0_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, GatherAxis1) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 120; + const std::size_t outputSize = sizeof(GatherAxis1_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GatherAxis1::Session session("GatherAxis1_FromONNX_GPU_ALPAKA.dat"); + 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 = GatherAxis1_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(GatherAxis1_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, GatherAxis2) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 120; + const std::size_t outputSize = sizeof(GatherAxis2_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GatherAxis2::Session session("GatherAxis2_FromONNX_GPU_ALPAKA.dat"); + 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 = GatherAxis2_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(GatherAxis2_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, GatherAxis3) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 120; + const std::size_t outputSize = sizeof(GatherAxis3_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GatherAxis3::Session session("GatherAxis3_FromONNX_GPU_ALPAKA.dat"); + 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 = GatherAxis3_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(GatherAxis3_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, Gather2d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 9; + const std::size_t outputSize = sizeof(Gather2d_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Gather2d::Session session("Gather2d_FromONNX_GPU_ALPAKA.dat"); + 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 = Gather2d_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(Gather2d_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, GatherNegativeIndices) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + constexpr Idx inputSize = 10; + const std::size_t outputSize = sizeof(GatherNegativeIndices_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + std::iota(input_ptr, input_ptr + inputSize, 0.f); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GatherNegativeIndices::Session session("GatherNegativeIndices_FromONNX_GPU_ALPAKA.dat"); + 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 = GatherNegativeIndices_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(GatherNegativeIndices_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, ExpandSameSize) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({0.f, 1.f, 2.f}); + const std::size_t outputSize = sizeof(ExpandSameSize_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_ExpandSameSize::Session session("ExpandSameSize_FromONNX_GPU_ALPAKA.dat"); + 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 = ExpandSameSize_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(ExpandSameSize_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, ExpandDiffSize) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({0.f, 1.f, 2.f}); + const std::size_t outputSize = sizeof(ExpandDiffSize_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_ExpandDiffSize::Session session("ExpandDiffSize_FromONNX_GPU_ALPAKA.dat"); + 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 = ExpandDiffSize_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(ExpandDiffSize_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); +} + +TEST_F(SofieAlpakaTest, DynamicGather) +{ + // X[N,3,n_pf], indices [2,0], axis 1 -> Y[N,2,n_pf], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3, K = 2; + const std::size_t idxs[2] = {2, 0}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * C * P, outSize = N * K * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicGather::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t i = 0; i < K; ++i) + for (std::size_t p = 0; p < P; ++p) { + float expected = in_ptr[n * C * P + idxs[i] * P + p]; + float got = res[n * K * P + i * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " i=" << i << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, GatherND_Ex1) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f, 1.f, 2.f, 3.f}; + std::vector expected = {0.f, 3.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Ex1::Session session("GatherND_Ex1_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 2u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_Ex2) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f, 1.f, 2.f, 3.f}; + std::vector expected = {2.f, 3.f, 0.f, 1.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Ex2::Session session("GatherND_Ex2_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 4u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_Ex3) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f,1.f,2.f,3.f,4.f,5.f,6.f,7.f}; + std::vector expected = {2.f, 3.f, 4.f, 5.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Ex3::Session session("GatherND_Ex3_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 4u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_Ex4) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f,1.f,2.f,3.f,4.f,5.f,6.f,7.f}; + std::vector expected = {2.f, 3.f, 4.f, 5.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Ex4::Session session("GatherND_Ex4_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 4u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_Ex5) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f,1.f,2.f,3.f,4.f,5.f,6.f,7.f}; + std::vector expected = {2.f, 3.f, 4.f, 5.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Ex5::Session session("GatherND_Ex5_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 4u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_NegativeIndices) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data = {0.f,1.f,2.f,3.f,4.f,5.f,6.f,7.f,8.f}; + std::vector expected = {6.f, 2.f, 4.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_NegativeIndices::Session session("GatherND_NegativeIndices_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 3u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GatherND_Batch) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector data(24); + std::iota(data.begin(), data.end(), 0.f); + std::vector expected = {4.f,5.f,6.f,7.f, 20.f,21.f,22.f,23.f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{data.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < data.size(); ++i) input_ptr[i] = data[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{data.size()})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{expected.size()})); + + { + SOFIE_GatherND_Batch::Session session("GatherND_Batch_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(expected.size(), 8u); + for (size_t i = 0; i < expected.size(); ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Equal) +{ + std::vector input1 = {1.0f, 2.0f, 3.0f}; + std::vector input2 = {4.0f, 2.0f, 6.0f}; + const std::size_t outputSize = sizeof(Equal_ExpectedOutput::outputs) / sizeof(bool); + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + alpaka::memcpy(queue, input2_d, input2_h); + alpaka::wait(queue); + + // Output is bool — allocate as bool buffer + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Equal::Session session("Equal_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + bool* correct = Equal_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(Equal_ExpectedOutput::outputs) / sizeof(bool)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, DynamicEqual) +{ + // X1[N,3] == X2[N,3] -> Y[N,3] (bool/uint8), N dynamic. Run at two batch sizes. + const std::size_t cols = 3; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * cols; + + auto x1_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + auto x2_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* x1p = reinterpret_cast(alpaka::getPtrNative(x1_h)); + float* x2p = reinterpret_cast(alpaka::getPtrNative(x2_h)); + for (Idx i = 0; i < sz; ++i) { + x1p[i] = static_cast(i % 3); + x2p[i] = static_cast(i % 2); // mix of equal and unequal + } + + auto x1_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + auto x2_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, x1_d, x1_h); + alpaka::memcpy(queue, x2_d, x2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicEqual::Session session("", N); + auto result = session.infer(N, x1_d, x2_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + uint8_t expected = (x1p[i] == x2p[i]) ? 1 : 0; + EXPECT_EQ(res[i], expected); + } + } +} + +TEST_F(SofieAlpakaTest, DynamicAddBroadcast) +{ + // X[N,4,n_pf] + bias[4,1] -> Y[N,4,n_pf], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float bias[4] = {0.5f, -1.0f, 0.25f, 2.0f}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicAddBroadcast::Session session("DynamicAddBroadcast_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / P) % C; + float expected = in_ptr[i] + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i << " N=" << N << " P=" << P; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicNegRelu) +{ + // X[N,3,n_pf] -> Neg -> Relu -> Y, fused eltwise chain, N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicNegRelu::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + float expected = std::max(0.0f, -in_ptr[i]); + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i << " N=" << N << " P=" << P; + } + } +} + +TEST_F(SofieAlpakaTest, LessOrEqual) +{ + std::vector input1 = {1.0f, 2.0f, 3.0f}; + std::vector input2 = {4.0f, 2.0f, 6.0f}; + const std::size_t outputSize = sizeof(LessOrEqual_ExpectedOutput::outputs) / sizeof(bool); + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + alpaka::memcpy(queue, input2_d, input2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_LessOrEqual::Session session("LessOrEqual_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + bool* correct = LessOrEqual_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(LessOrEqual_ExpectedOutput::outputs) / sizeof(bool)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, GreaterOrEqual) +{ + std::vector input1 = {1.0f, 2.0f, 3.0f}; + std::vector input2 = {4.0f, 2.0f, 6.0f}; + const std::size_t outputSize = sizeof(GreaterOrEqual_ExpectedOutput::outputs) / sizeof(bool); + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + alpaka::memcpy(queue, input2_d, input2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_GreaterOrEqual::Session session("GreaterOrEqual_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + bool* correct = GreaterOrEqual_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(GreaterOrEqual_ExpectedOutput::outputs) / sizeof(bool)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Greater) +{ + std::vector input1 = {1.0f, 2.0f, 3.0f}; + std::vector input2 = {4.0f, 2.0f, 6.0f}; + const std::size_t outputSize = sizeof(Greater_ExpectedOutput::outputs) / sizeof(bool); + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + alpaka::memcpy(queue, input2_d, input2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Greater::Session session("Greater_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + bool* correct = Greater_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(Greater_ExpectedOutput::outputs) / sizeof(bool)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Less) +{ + std::vector input1 = {1.0f, 2.0f, 3.0f}; + std::vector input2 = {4.0f, 2.0f, 6.0f}; + const std::size_t outputSize = sizeof(Less_ExpectedOutput::outputs) / sizeof(bool); + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + alpaka::memcpy(queue, input2_d, input2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_Less::Session session("Less_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + bool* correct = Less_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(Less_ExpectedOutput::outputs) / sizeof(bool)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Slice) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = Slice::input; + const std::size_t outputSize = sizeof(Slice::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Slice::Session session("Slice_FromONNX_GPU_ALPAKA.dat"); + 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 = Slice::output; + EXPECT_EQ(outputSize, sizeof(Slice::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Slice_Default_Axis) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = Slice_Default_Axis::input; + const std::size_t outputSize = sizeof(Slice_Default_Axis::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Slice_Default_Axis::Session session("Slice_Default_Axis_FromONNX_GPU_ALPAKA.dat"); + 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 = Slice_Default_Axis::output; + EXPECT_EQ(outputSize, sizeof(Slice_Default_Axis::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Slice_Default_Steps) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = Slice_Default_Steps::input; + const std::size_t outputSize = sizeof(Slice_Default_Steps::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Slice_Default_Steps::Session session("Slice_Default_Steps_FromONNX_GPU_ALPAKA.dat"); + 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 = Slice_Default_Steps::output; + EXPECT_EQ(outputSize, sizeof(Slice_Default_Steps::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Slice_Neg) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = Slice_Neg::input; + const std::size_t outputSize = sizeof(Slice_Neg::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Slice_Neg::Session session("Slice_Neg_FromONNX_GPU_ALPAKA.dat"); + 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 = Slice_Neg::output; + EXPECT_EQ(outputSize, sizeof(Slice_Neg::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, DynamicSlice) +{ + // X[N,4,n_pf] axis 1 [1:3] -> Y[N,2,n_pf], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4, K = 2, start = 1; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * C * P, outSize = N * K * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicSlice::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < K; ++j) + for (std::size_t p = 0; p < P; ++p) { + float expected = in_ptr[n * C * P + (start + j) * P + p]; + float got = res[n * K * P + j * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " j=" << j << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, Sin) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({ + -0.786738f, -0.197796f, -0.187787f, 0.142758f, + 0.876096f, -0.653239f, 0.145444f, -1.107658f, + 2.259171f, -0.947054f, -0.506689f, 1.801250f + }); + + 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{input.size()})); + + { + SOFIE_Sin::Session session; + 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)); + EXPECT_EQ(input.size(), 12u); + for (size_t i = 0; i < input.size(); ++i) + EXPECT_LE(std::abs(res_ptr[i] - std::sin(input[i])), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Cos) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({ + 1.152504f, -1.459324f, 0.691594f, 0.347690f, + -1.307323f, 1.832516f, -1.261772f, 0.014224f, + 1.311477f, 1.147405f, -0.567206f, -0.530606f + }); + + 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{input.size()})); + + { + SOFIE_Cos::Session session; + 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)); + EXPECT_EQ(input.size(), 12u); + for (size_t i = 0; i < input.size(); ++i) + EXPECT_LE(std::abs(res_ptr[i] - std::cos(input[i])), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Abs) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({1.f, -2.f, -3.f, 4.f, -5.f, 6.f}); + + 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{input.size()})); + + { + SOFIE_Abs::Session session; + 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)); + EXPECT_EQ(input.size(), 6u); + for (size_t i = 0; i < input.size(); ++i) + EXPECT_LE(std::abs(res_ptr[i] - std::abs(input[i])), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Sqrt) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({0.8344f, 0.4716f, 0.6226f, 0.8448f, 0.2483f, 0.9467f}); + const std::size_t outputSize = sizeof(Sqrt_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Sqrt::Session session("Sqrt_FromONNX_GPU_ALPAKA.dat"); + 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 = Sqrt_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(Sqrt_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Reciprocal) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({1.2691f, -1.2160f, 0.6393f, -0.4438f, 0.8065f, 0.2011f}); + const std::size_t outputSize = sizeof(Reciprocal_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Reciprocal::Session session("Reciprocal_FromONNX_GPU_ALPAKA.dat"); + 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 = Reciprocal_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(Reciprocal_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Exp) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({ + 1.46566453f, 0.63334515f, 2.4048165f, 0.54468453f, + -1.41271672f, -0.18609187f, 0.2754482f, 1.10615209f, + 0.88474389f, 0.47531232f + }); + const std::size_t outputSize = sizeof(Exp_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Exp::Session session("Exp_FromONNX_GPU_ALPAKA.dat"); + 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 = Exp_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(Exp_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Log) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({1.f, 2.f, 3.f, 4.f}); + const std::size_t outputSize = sizeof(Log_ExpectedOutput::outputs) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Log::Session session("Log_FromONNX_GPU_ALPAKA.dat"); + 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 = Log_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(Log_ExpectedOutput::outputs) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Neg) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({ + -1.9100f, 1.8811f, -1.7269f, -0.1094f, + -0.0145f, 0.2509f, 0.5893f, -2.2733f, + -0.7077f, 1.0645f, -0.8607f, 0.2085f + }); + const std::size_t outputSize = sizeof(Neg_ExpectedOutput::outputs) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_Neg::Session session("Neg_FromONNX_GPU_ALPAKA.dat"); + 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 = Neg_ExpectedOutput::outputs; + EXPECT_EQ(outputSize, sizeof(Neg_ExpectedOutput::outputs) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Softplus) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input({0.1,-0.2,0.3,-0.4,0.5,1.}); + + 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{input.size()})); + + { + SOFIE_Softplus::Session session("Softplus_FromONNX_GPU_ALPAKA.dat"); + 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)); + for (size_t i = 0; i < input.size(); ++i){ + double exp_value = std::log(std::exp(input[i])+1); + EXPECT_LE(std::abs(res_ptr[i] - exp_value), TOLERANCE); + } +} + +TEST_F(SofieAlpakaTest, Where) +{ + std::vector input1 = {1.f, 2.f}; + std::vector input2 = {3.f, 4.f, 5.f, 6.f}; + std::vector cond_vec = {true, false, true}; + std::vector correct = {1.f, 2.f, 5.f, 6.f, 1.f, 2.f}; + + auto input1_h = alpaka::allocBuf(host, Ext1D::all(Idx{input1.size()})); + float* in1_ptr = reinterpret_cast(alpaka::getPtrNative(input1_h)); + for (Idx i = 0; i < input1.size(); ++i) in1_ptr[i] = input1[i]; + + auto input1_d = alpaka::allocBuf(device, Ext1D::all(Idx{input1.size()})); + alpaka::memcpy(queue, input1_d, input1_h); + + auto input2_h = alpaka::allocBuf(host, Ext1D::all(Idx{input2.size()})); + float* in2_ptr = reinterpret_cast(alpaka::getPtrNative(input2_h)); + for (Idx i = 0; i < input2.size(); ++i) in2_ptr[i] = input2[i]; + + auto input2_d = alpaka::allocBuf(device, Ext1D::all(Idx{input2.size()})); + alpaka::memcpy(queue, input2_d, input2_h); + + auto cond_h = alpaka::allocBuf(host, Ext1D::all(Idx{cond_vec.size()})); + uint8_t* cond_ptr = reinterpret_cast(alpaka::getPtrNative(cond_h)); + for (Idx i = 0; i < cond_vec.size(); ++i) cond_ptr[i] = cond_vec[i]; + + auto cond_d = alpaka::allocBuf(device, Ext1D::all(Idx{cond_vec.size()})); + alpaka::memcpy(queue, cond_d, cond_h); + + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{correct.size()})); + + { + SOFIE_Where::Session session("Where_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(input1_d, input2_d, cond_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + EXPECT_EQ(correct.size(), 6u); + for (size_t i = 0; i < correct.size(); ++i) + EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, ReduceMean) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + const std::size_t outputSize = sizeof(ReduceMean_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_ReduceMean::Session session("ReduceMean_FromONNX_GPU_ALPAKA.dat"); + 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 = ReduceMean_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(ReduceMean_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, ReduceProd) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + const std::size_t outputSize = sizeof(ReduceProd_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_ReduceProd::Session session("ReduceProd_FromONNX_GPU_ALPAKA.dat"); + 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 = ReduceProd_ExpectedOutput::output; + EXPECT_EQ(outputSize, sizeof(ReduceProd_ExpectedOutput::output) / sizeof(float)); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, ReduceSum) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + std::vector correct = {24.f}; + + 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{correct.size()})); + + { + SOFIE_ReduceSum::Session session("ReduceSum_FromONNX_GPU_ALPAKA.dat"); + 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)); + EXPECT_EQ(correct.size(), 1u); + for (size_t i = 0; i < correct.size(); ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, ReduceSumSquare) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + std::vector correct = {38.f, 66.f}; + + 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{correct.size()})); + + { + SOFIE_ReduceSumSquare::Session session("ReduceSumSquare_FromONNX_GPU_ALPAKA.dat"); + 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)); + for (size_t i = 0; i < correct.size(); ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +// ReduceL2: input [1,2,3]={5,2,3,5,5,4}, reduce axis=1, keepdims=0 → [1,3] +// Expected: {sqrt(50), sqrt(29), 5.0} +TEST_F(SofieAlpakaTest, ReduceL2) +{ + constexpr float TOLERANCE = 1e-3f; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + const std::size_t outputSize = sizeof(ReduceL2_ExpectedOutput::output) / sizeof(float); + + 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{outputSize})); + + { + SOFIE_ReduceL2::Session session("ReduceL2_FromONNX_GPU_ALPAKA.dat"); + 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 = ReduceL2_ExpectedOutput::output; + EXPECT_EQ(outputSize, 3u); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +// ReduceL2Large: input [4,512], reduce axis=1, keepdims=0 → [4] +// Row i is filled with (i+1), so L2 norm = (i+1)*sqrt(512). +// This test exercises the 256-thread block reduction with reducedLength > BLOCK_SIZE. +TEST_F(SofieAlpakaTest, ReduceL2Large) +{ + constexpr float TOLERANCE = 1e-2f; // slightly looser: large sum, float accumulation + + constexpr std::size_t nrows = 4; + constexpr std::size_t ncols = 512; + const std::size_t inputSize = nrows * ncols; + const std::size_t outputSize = nrows; + + // Fill row i with value (i+1) + std::vector input(inputSize); + for (std::size_t r = 0; r < nrows; ++r) + for (std::size_t c = 0; c < ncols; ++c) + input[r * ncols + c] = static_cast(r + 1); + + // Expected L2 per row: sqrt(ncols) * (row+1) + const float sqrt512 = std::sqrt(static_cast(ncols)); + std::vector correct(nrows); + for (std::size_t r = 0; r < nrows; ++r) + correct[r] = static_cast(r + 1) * sqrt512; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inputSize; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + + { + SOFIE_ReduceL2Large::Session session; + 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)); + EXPECT_EQ(outputSize, nrows); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - correct[i]) / correct[i], TOLERANCE) << "row=" << i; +} + +// ── ReduceMax: [1,2,3] axis=1 keepdims=0 (kLast path) ────────────────────── +TEST_F(SofieAlpakaTest, ReduceMax) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = {5.f, 2.f, 3.f, 5.f, 5.f, 4.f}; + const std::size_t outputSize = sizeof(ReduceMax_ExpectedOutput::output) / sizeof(float); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{input.size()})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t 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{outputSize})); + { + SOFIE_ReduceMax::Session session; + 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 = ReduceMax_ExpectedOutput::output; + for (std::size_t i = 0; i < outputSize; ++i) + EXPECT_NEAR(res_ptr[i], correct[i], TOLERANCE) << " i=" << i; +} + +// ── ReduceMax_axis0: [3,4] axis=0 keepdims=0 (kFirst path) ───────────────── +TEST_F(SofieAlpakaTest, ReduceMax_axis0) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // numpy default_rng(42).standard_normal((3,4)) — same seed/sequence as generator + const std::size_t inputSize = 12; + const std::size_t outputSize = sizeof(ReduceMax_axis0_ExpectedOutput::output) / sizeof(float); + float vals[] = { 0.30471709f, -1.03998411f, 0.75045121f, 0.94056469f, + -1.95103514f, -1.30217946f, 0.12784040f, -0.31624261f, + -0.01680116f, -0.85304391f, 0.87939799f, 0.77779192f}; + std::vector input(vals, vals + inputSize); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < inputSize; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + { + SOFIE_ReduceMax_axis0::Session session; + 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 = ReduceMax_axis0_ExpectedOutput::output; + for (std::size_t i = 0; i < outputSize; ++i) + EXPECT_NEAR(res_ptr[i], correct[i], TOLERANCE) << " i=" << i; +} + +// ── ReduceMax_mid: [2,3,4] axis=1 keepdims=0 (kMiddle path) ──────────────── +TEST_F(SofieAlpakaTest, ReduceMax_mid) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + const std::size_t inputSize = 24; // 2×3×4 + const std::size_t outputSize = sizeof(ReduceMax_mid_ExpectedOutput::output) / sizeof(float); + + // numpy default_rng(42).standard_normal((2,3,4)) — same seed/sequence as generator + float vals[] = { 0.06603070f, 1.12724125f, 0.46750933f, -0.85929245f, + 0.36875078f, -0.95888263f, 0.87845027f, -0.04992591f, + -0.18486236f, -0.68092954f, 1.22254133f, -0.15452948f, + -0.42832783f, -0.35213354f, 0.53230917f, 0.36544406f, + 0.41273260f, 0.43082100f, 2.14164758f, -0.40641502f, + -0.51224273f, -0.81377274f, 0.61597943f, 1.12897229f}; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inputSize})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < inputSize; ++i) input_ptr[i] = vals[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inputSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outputSize})); + { + SOFIE_ReduceMax_mid::Session session; + 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 = ReduceMax_mid_ExpectedOutput::output; + for (std::size_t i = 0; i < outputSize; ++i) + EXPECT_NEAR(res_ptr[i], correct[i], TOLERANCE) << " i=" << i; +} + +// ── Dynamic (symbolic-shape) Reduce tests: each runs at two runtime sizes so a +// baked-in shape would fail one of them. References are computed in-test. + +// DynamicReduceSumLast: X[N,4] -> ReduceSum(axis=-1, keepdims=0) -> Y[N] +// last-axis (kLast), pruned output, negative axis, static reduced-length. +TEST_F(SofieAlpakaTest, DynamicReduceSumLast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t cols = 4; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * cols, outSize = N; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceSumLast::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + float expected = 0.f; + for (std::size_t c = 0; c < cols; ++c) expected += in_ptr[n * cols + c]; + EXPECT_LE(std::abs(res[n] - expected), TOLERANCE); + } + } +} + +// DynamicReduceMeanMid: X[N,4,M] -> ReduceMean(axis=1, keepdims=1) -> Y[N,1,M] +// middle-axis (kMiddle) with TWO symbolic dims. Reconstruct the session per size so +// the returned Y buffer matches outSize; the output length N*M is a symmetric product +// so the unordered_map ctor-arg order is immaterial, and infer args stay in +// declaration order (N, M). +TEST_F(SofieAlpakaTest, DynamicReduceMeanMid) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t K = 4; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ms[] = {1, 3}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], M = Ms[t]; + const std::size_t inSize = N * K * M, outSize = N * M; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceMeanMid::Session session("", N, M); + auto result = session.infer(N, M, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t m = 0; m < M; ++m) { + float sum = 0.f; + for (std::size_t k = 0; k < K; ++k) + sum += in_ptr[n * K * M + k * M + m]; + float expected = sum / static_cast(K); + EXPECT_LE(std::abs(res[n * M + m] - expected), TOLERANCE); + } + } +} + +// DynamicReduceMaxFirst: X[N,4] -> ReduceMax(axis=0, keepdims=0) -> Y[4] +// reduces the DYNAMIC axis, so reducedLength is symbolic (=N); kFirst path, Max op. +TEST_F(SofieAlpakaTest, DynamicReduceMaxFirst) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t cols = 4; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * cols, outSize = cols; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast((i * 7 + 3) % 11); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceMaxFirst::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t c = 0; c < cols; ++c) { + float expected = in_ptr[c]; // n == 0 + for (std::size_t n = 1; n < N; ++n) + if (in_ptr[n * cols + c] > expected) expected = in_ptr[n * cols + c]; + EXPECT_LE(std::abs(res[c] - expected), TOLERANCE); + } + } +} + +// DynamicReduceSumMulti: X[N,3,2] -> ReduceSum(axes=[1,2], keepdims=0) -> Y[N] +// multi-axis reduce -> exercises the redStrides product string in the kernel. +TEST_F(SofieAlpakaTest, DynamicReduceSumMulti) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t slice = 3 * 2; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * slice, outSize = N; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceSumMulti::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + float expected = 0.f; + for (std::size_t j = 0; j < slice; ++j) expected += in_ptr[n * slice + j]; + EXPECT_LE(std::abs(res[n] - expected), TOLERANCE); + } + } +} + +TEST_F(SofieAlpakaTest, ConvWithPadding) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(25); + 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(ConvWithPadding_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithPadding::Session session("ConvWithPadding_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithPadding_ExpectedOutput::all_ones; + + for (size_t i = 0; i < 25; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } +} + + +TEST_F(SofieAlpakaTest, ConvWithoutPadding) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(25); + 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(ConvWithoutPadding_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithoutPadding::Session session("ConvWithoutPadding_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithoutPadding_ExpectedOutput::all_ones; + constexpr size_t nOut_convNoPad = sizeof(ConvWithoutPadding_ExpectedOutput::all_ones) / sizeof(float); + + for (size_t i = 0; i < nOut_convNoPad; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } + +} + + +TEST_F(SofieAlpakaTest, ConvWithAutopadSameLower) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(25); + 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(ConvWithAutopadSameLower_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithAutopadSameLower::Session session("ConvWithAutopadSameLower_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithAutopadSameLower_ExpectedOutput::all_ones; + + for (size_t i = 0; i < 9; ++i) { + std::cout << "res: " << res_ptr[i] << ", correct: " << correct[i] << std::endl; + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } +} + + +TEST_F(SofieAlpakaTest, ConvWithStridesPadding) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(35); + 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(ConvWithStridesPadding_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithStridesPadding::Session session("ConvWithStridesPadding_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithStridesPadding_ExpectedOutput::all_ones; + constexpr size_t nOut_stridesPad = sizeof(ConvWithStridesPadding_ExpectedOutput::all_ones) / sizeof(float); + + for (size_t i = 0; i < nOut_stridesPad; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } +} + + +TEST_F(SofieAlpakaTest, ConvWithStridesNoPadding) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(35); + 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(ConvWithStridesNoPadding_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithStridesNoPadding::Session session("ConvWithStridesNoPadding_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithStridesNoPadding_ExpectedOutput::all_ones; + constexpr size_t nOut_stridesNoPad = sizeof(ConvWithStridesNoPadding_ExpectedOutput::all_ones) / sizeof(float); + + for (size_t i = 0; i < nOut_stridesNoPad; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } +} + + +// Disables test (asymmetric padding not supported) +TEST_F(SofieAlpakaTest, ConvWithAsymmetricPadding) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // Preparing the standard all-ones input + std::vector input(35); + 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(ConvWithAsymmetricPadding_ExpectedOutput::all_ones) / sizeof(float)})); + + { + SOFIE_ConvWithAsymmetricPadding::Session session("ConvWithAsymmetricPadding_FromONNX_GPU_ALPAKA.dat"); + 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 = ConvWithAsymmetricPadding_ExpectedOutput::all_ones; + constexpr size_t nOut_asymPad = sizeof(ConvWithAsymmetricPadding_ExpectedOutput::all_ones) / sizeof(float); + + for (size_t i = 0; i < nOut_asymPad; ++i) { + EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; + } +} + +// dynamic 1D Conv (k=3, pad=1) + bias, run at (N,n_pf) = (1,1) and (8,5) +TEST_F(SofieAlpakaTest, DynamicConv1D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3, K = 3; + const float bias[3] = {0.5f, -0.5f, 1.0f}; + auto Wv = [&](std::size_t oc, std::size_t ic, std::size_t kk) { + return static_cast(oc * (Cin * K) + ic * K + kk); + }; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P, outSize = N * Cout * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv1D::Session session("DynamicConv1D_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t l = 0; l < P; ++l) { + float acc = bias[oc]; + for (std::size_t ic = 0; ic < Cin; ++ic) + for (std::size_t kk = 0; kk < K; ++kk) { + int64_t li = static_cast(l) + static_cast(kk) - 1; // pad=1 + if (li >= 0 && li < static_cast(P)) + acc += in_ptr[n * Cin * P + ic * P + li] * Wv(oc, ic, kk); + } + std::size_t idx = n * Cout * P + oc * P + l; + EXPECT_LE(std::abs(res[idx] - acc), TOLERANCE) << "n=" << n << " oc=" << oc << " l=" << l; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicLinear) +{ + // X[N,4] -> Gemm(W[4,3],B[3]) -> Relu -> Y[N,3], N dynamic. Gemm+Relu fuses to gemmrelu. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t In = 4, Out = 3; + const float W[4][3] = {{0.5f, -1.0f, 0.25f}, + {0.75f, 0.5f, -0.5f}, + {-0.25f, 1.0f, 0.75f}, + {1.5f, -0.75f, 0.5f}}; + const float B[3] = {0.5f, -0.5f, 0.25f}; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * In, outSize = N * Out; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicLinear::Session session("DynamicLinear_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < Out; ++j) { + float acc = B[j]; + for (std::size_t i = 0; i < In; ++i) acc += in_ptr[n * In + i] * W[i][j]; + float expected = acc > 0.0f ? acc : 0.0f; + EXPECT_LE(std::abs(res[n * Out + j] - expected), TOLERANCE) << "n=" << n << " j=" << j << " N=" << N; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicConv1DNoBias) +{ + // X[N,2,n_pf] k=1 no-bias conv -> Y[N,3,n_pf], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3; + const float W[3][2] = {{0.5f, -1.0f}, {0.25f, 0.75f}, {-0.5f, 1.5f}}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P, outSize = N * Cout * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv1DNoBias::Session session("DynamicConv1DNoBias_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t p = 0; p < P; ++p) { + float expected = 0.0f; + for (std::size_t ic = 0; ic < Cin; ++ic) + expected += W[oc][ic] * in_ptr[n * Cin * P + ic * P + p]; + float got = res[n * Cout * P + oc * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " oc=" << oc << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicConv2DNoBias) +{ + // X[N,2,n_pf,4] 1x1 no-bias conv -> Y[N,3,n_pf,4], N and n_pf dynamic. Run at two sizes. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3, Q = 4; + const float W[3][2] = {{0.5f, -1.0f}, {0.25f, 0.75f}, {-0.5f, 1.5f}}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P * Q, outSize = N * Cout * P * Q; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv2DNoBias::Session session("DynamicConv2DNoBias_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t p = 0; p < P; ++p) + for (std::size_t q = 0; q < Q; ++q) { + float expected = 0.0f; + for (std::size_t ic = 0; ic < Cin; ++ic) + expected += W[oc][ic] * in_ptr[n * Cin * P * Q + ic * P * Q + p * Q + q]; + float got = res[n * Cout * P * Q + oc * P * Q + p * Q + q]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " oc=" << oc << " p=" << p << " q=" << q; + } + } +} + +TEST_F(SofieAlpakaTest, BatchNormalization) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = { + 1.f, 2.f, 3.f, 4.f, // channel 0 + 5.f, 6.f, 7.f, 8.f // channel 1 + }; + const std::size_t outputSize = input.size(); + + 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{outputSize})); + + { + SOFIE_BatchNorm::Session session("BatchNorm_FromONNX_GPU_ALPAKA.dat"); + 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 inv_std = 1.f / std::sqrt(1.f + 1e-5f); + ASSERT_EQ(outputSize, 8u); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - input[i] * inv_std), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, BatchNormalizationRelu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + std::vector input = { + -1.f, 2.f, -3.f, 4.f, + 5.f, -6.f, 7.f, -8.f + }; + const std::size_t outputSize = input.size(); + + 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{outputSize})); + + { + SOFIE_BatchNormRelu::Session session("BatchNormRelu_FromONNX_GPU_ALPAKA.dat"); + 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 inv_std = 1.f / std::sqrt(1.f + 1e-5f); + ASSERT_EQ(outputSize, 8u); + for (size_t i = 0; i < outputSize; ++i) { + float expected = std::max(0.f, input[i] * inv_std); + EXPECT_LE(std::abs(res_ptr[i] - expected), TOLERANCE) << "i=" << i; + } +} + +// dynamic BatchNorm, run at two batch sizes. weights are per-channel [C=4] (same values in +// each test); fused scale is scale[c]/sqrt(var[c]+eps). the static BN tests above cover the +// shared code path, these cover the dynamic shapes. + +// X[N,4,2,3]: dynamic batch, static spatial (6) +TEST_F(SofieAlpakaTest, DynamicBatchNorm4D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4, spatial = 2 * 3; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * C * spatial; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNorm4D::Session session("DynamicBatchNorm4D_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / spatial) % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + +// X[N,4,n_pf] + relu: dynamic batch and spatial. rebuilt per size so Y matches its length; +// Y length is a symmetric product so the ctor arg order doesn't matter, infer is (N, n_pf). +TEST_F(SofieAlpakaTest, DynamicBatchNormDynSpatialRelu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNormDynSpatialRelu::Session session("DynamicBatchNormDynSpatialRelu_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / P) % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + expected = expected > 0.0f ? expected : 0.0f; // relu + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + +// X[N,4]: rank 2, spatial is 1 +TEST_F(SofieAlpakaTest, DynamicBatchNorm2D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * C; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNorm2D::Session session("DynamicBatchNorm2D_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = i % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + +TEST_F(SofieAlpakaTest, LayerNorm) +{ + constexpr float TOLERANCE = 1e-4f; + std::vector input = {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}; + const std::size_t outputSize = input.size(); + + 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{outputSize})); + + { + SOFIE_LayerNorm::Session session("LayerNorm_FromONNX_GPU_ALPAKA.dat"); + 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)); + + // Row 0: mean=2.5, std=sqrt(1.25+1e-5) ≈ 1.118034 + // Row 1: mean=6.5, std=sqrt(1.25+1e-5) ≈ 1.118034 + // Y[0] = (1-2.5)/1.118034 ≈ -1.3416 + // Y[1] = (2-2.5)/1.118034 ≈ -0.4472 + // Y[2] = (3-2.5)/1.118034 ≈ 0.4472 + // Y[3] = (4-2.5)/1.118034 ≈ 1.3416 + float inv_std = 1.f / std::sqrt(1.25f + 1e-5f); + std::vector expected = { + (1.f - 2.5f) * inv_std, (2.f - 2.5f) * inv_std, + (3.f - 2.5f) * inv_std, (4.f - 2.5f) * inv_std, + (5.f - 6.5f) * inv_std, (6.f - 6.5f) * inv_std, + (7.f - 6.5f) * inv_std, (8.f - 6.5f) * inv_std + }; + ASSERT_EQ(outputSize, 8u); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, LayerNormScaleBias) +{ + constexpr float TOLERANCE = 1e-4f; + + std::vector input = {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}; + const std::size_t outputSize = input.size(); + + 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{outputSize})); + + { + SOFIE_LayerNormScaleBias::Session session("LayerNormScaleBias_FromONNX_GPU_ALPAKA.dat"); + 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 inv_std = 1.f / std::sqrt(1.25f + 1e-5f); + std::vector expected = { + 2.f * (1.f - 2.5f) * inv_std + 1.f, 2.f * (2.f - 2.5f) * inv_std + 1.f, + 2.f * (3.f - 2.5f) * inv_std + 1.f, 2.f * (4.f - 2.5f) * inv_std + 1.f, + 2.f * (5.f - 6.5f) * inv_std + 1.f, 2.f * (6.f - 6.5f) * inv_std + 1.f, + 2.f * (7.f - 6.5f) * inv_std + 1.f, 2.f * (8.f - 6.5f) * inv_std + 1.f + }; + ASSERT_EQ(outputSize, 8u); + for (size_t i = 0; i < outputSize; ++i) + EXPECT_LE(std::abs(res_ptr[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, LayerNorm3D) +{ + constexpr float TOLERANCE = 1e-4f; + + std::vector input(24); + std::iota(input.begin(), input.end(), 0.f); // 0..23 + const std::size_t outputSize = input.size(); + + 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{outputSize})); + + { + SOFIE_LayerNorm3D::Session session("LayerNorm3D_FromONNX_GPU_ALPAKA.dat"); + 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)); + + auto compute_expected = [](std::vector row) { + float mean = 0.f; + for (float v : row) mean += v; + mean /= row.size(); + float var = 0.f; + for (float v : row) var += (v - mean) * (v - mean); + var /= row.size(); + float inv_std = 1.f / std::sqrt(var + 1e-5f); + std::vector out; + for (float v : row) out.push_back((v - mean) * inv_std); + return out; + }; + + std::vector row0(input.begin(), input.begin() + 12); + std::vector row1(input.begin() + 12, input.end()); + auto exp0 = compute_expected(row0); + auto exp1 = compute_expected(row1); + + ASSERT_EQ(outputSize, 24u); + for (size_t i = 0; i < 12; ++i) + EXPECT_LE(std::abs(res_ptr[i] - exp0[i]), TOLERANCE) << "row0 i=" << i; + for (size_t i = 0; i < 12; ++i) + EXPECT_LE(std::abs(res_ptr[12 + i] - exp1[i]), TOLERANCE) << "row1 i=" << i; +} + +TEST_F(SofieAlpakaTest, IsInf) +{ + // Input contains finite values, +inf, -inf; output is bool (uint8_t). + float pos_inf = std::numeric_limits::infinity(); + float neg_inf = -std::numeric_limits::infinity(); + std::vector input = {1.0f, pos_inf, neg_inf, 0.0f, -1.0f, 2.0f, neg_inf, pos_inf}; + const std::size_t N = input.size(); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + + { + SOFIE_IsInf::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(N, 8u); + for (size_t i = 0; i < N; ++i) + EXPECT_EQ(static_cast(res_ptr[i]), std::isinf(input[i])) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, IsNaN) +{ + // Input contains finite values, +inf, and NaN; output is bool (uint8_t). + float nan_val = std::numeric_limits::quiet_NaN(); + float pos_inf = std::numeric_limits::infinity(); + std::vector input = {1.0f, nan_val, 0.0f, pos_inf, nan_val, 2.0f, -1.0f, nan_val}; + const std::size_t N = input.size(); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + + { + SOFIE_IsNaN::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(N, 8u); + for (size_t i = 0; i < N; ++i) + EXPECT_EQ(static_cast(res_ptr[i]), std::isnan(input[i])) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Clip) +{ + // Model clips to [-1.0, 1.0]. + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr float clip_min = -1.0f; + constexpr float clip_max = 1.0f; + + std::vector input = { + -2.0f, -1.5f, -1.0f, -0.5f, + 0.0f, 0.5f, 1.0f, 1.5f, + 2.0f, -0.3f, 0.7f, 1.2f + }; + const std::size_t N = input.size(); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + + { + SOFIE_Clip::Session session("Clip_FromONNX_GPU_ALPAKA.dat"); + 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)); + ASSERT_EQ(N, 12u); + for (size_t i = 0; i < N; ++i) { + float expected = std::max(clip_min, std::min(clip_max, input[i])); + EXPECT_LE(std::abs(res_ptr[i] - expected), TOLERANCE) << "i=" << i; + } +} + +TEST_F(SofieAlpakaTest, Not) +{ + // Input and output are bool tensors (uint8_t on device). + std::vector input = {1, 0, 1, 1, 0, 0, 1, 0}; + const std::size_t N = input.size(); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + uint8_t* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + + { + SOFIE_Not::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(N, 8u); + for (size_t i = 0; i < N; ++i) + EXPECT_EQ(static_cast(res_ptr[i]), !static_cast(input[i])) << "i=" << i; +} + +// GNN model: 3370 nodes (29 features each), 24126 edges (5 features each), +// edge_index shape [2, 24126]. Output: sigmoid score per edge in [0, 1]. +TEST_F(SofieAlpakaTest, GNN_model) +{ + // ---- sizes ------------------------------------------------------- + constexpr Idx N_x = 97730; // 3370 nodes × 29 features + constexpr Idx N_ef = 120630; // 24126 edges × 5 features + constexpr Idx N_ei = 48252; // 2 rows × 24126 edges (int64) + constexpr Idx N_out = 24126; // one sigmoid score per edge + + // ---- host buffers ------------------------------------------------- + auto x_h = alpaka::allocBuf(host, Ext1D::all(Idx{N_x})); + auto ef_h = alpaka::allocBuf(host, Ext1D::all(Idx{N_ef})); + auto ei_h = alpaka::allocBuf(host, Ext1D::all(Idx{N_ei})); + + float* x_ptr = reinterpret_cast (alpaka::getPtrNative(x_h)); + float* ef_ptr = reinterpret_cast (alpaka::getPtrNative(ef_h)); + int64_t* ei_ptr = reinterpret_cast(alpaka::getPtrNative(ei_h)); + + for (Idx i = 0; i < N_x; ++i) x_ptr[i] = 0.5f; + for (Idx i = 0; i < N_ef; ++i) ef_ptr[i] = 0.5f; + for (Idx i = 0; i < N_ei; ++i) ei_ptr[i] = 0; // all self-loops on node 0 + + // ---- device buffers ----------------------------------------------- + auto x_d = alpaka::allocBuf(device, Ext1D::all(Idx{N_x})); + auto ef_d = alpaka::allocBuf(device, Ext1D::all(Idx{N_ef})); + auto ei_d = alpaka::allocBuf(device, Ext1D::all(Idx{N_ei})); + + alpaka::memcpy(queue, x_d, x_h); + alpaka::memcpy(queue, ef_d, ef_h); + alpaka::memcpy(queue, ei_d, ei_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N_out})); + + { + SOFIE_GNN_model::Session session("GNN_model_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(x_d, ef_d, ei_d); + alpaka::wait(session.queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + ASSERT_EQ(N_out, 24126u); + for (Idx i = 0; i < N_out; ++i) { + EXPECT_GE(res_ptr[i], 0.0f) << "output[" << i << "] < 0"; + EXPECT_LE(res_ptr[i], 1.0f) << "output[" << i << "] > 1"; + } +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Trilu operator tests +// ═══════════════════════════════════════════════════════════════════════════ + +// Helper: copy a host C-array into an Alpaka host buffer then to device. +template +static alpaka::Buf +makeDeviceBuf(alpaka::DevCpu const& host, + alpaka::DevCudaRt const& device, + alpaka::Queue& queue, + const T* src, std::size_t n) +{ + auto hbuf = alpaka::allocBuf(host, Ext1D::all(Idx{n})); + T* hp = reinterpret_cast(alpaka::getPtrNative(hbuf)); + for (std::size_t i = 0; i < n; ++i) hp[i] = src[i]; + auto dbuf = alpaka::allocBuf(device, Ext1D::all(Idx{n})); + alpaka::memcpy(queue, dbuf, hbuf); + alpaka::wait(queue); + return dbuf; +} + +// ── Trilu_upper: 4×4, upper=1, k=0 ───────────────────────────────────────── +TEST_F(SofieAlpakaTest, Trilu_upper) +{ + constexpr std::size_t N = 16; // 4×4 + auto d_input = makeDeviceBuf(host, device, queue, + Trilu_upper_Input::data, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Trilu_upper::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Trilu_upper_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; +} + +// ── Trilu_lower: 4×4, upper=0, k=0 ───────────────────────────────────────── +TEST_F(SofieAlpakaTest, Trilu_lower) +{ + constexpr std::size_t N = 16; + auto d_input = makeDeviceBuf(host, device, queue, + Trilu_lower_Input::data, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Trilu_lower::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Trilu_lower_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; +} + +// ── Trilu_k2: 3×5, upper=1, k=+2 ──────────────────────────────────────────── +TEST_F(SofieAlpakaTest, Trilu_k2) +{ + constexpr std::size_t N = 15; // 3×5 + auto d_input = makeDeviceBuf(host, device, queue, + Trilu_k2_Input::data, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Trilu_k2::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Trilu_k2_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; +} + +// ── Trilu_kn1: 3×5, upper=0, k=-1 ──────────────────────────────────────────── +TEST_F(SofieAlpakaTest, Trilu_kn1) +{ + constexpr std::size_t N = 15; + auto d_input = makeDeviceBuf(host, device, queue, + Trilu_kn1_Input::data, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Trilu_kn1::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Trilu_kn1_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; +} + +// ── Trilu_3D: 2×3×4, upper=1, k=0 (batched) ───────────────────────────────── +TEST_F(SofieAlpakaTest, Trilu_3D) +{ + constexpr std::size_t N = 24; // 2×3×4 + auto d_input = makeDeviceBuf(host, device, queue, + Trilu_3D_Input::data, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Trilu_3D::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Trilu_3D_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Logic / Bitwise operator tests +// ═══════════════════════════════════════════════════════════════════════════ + +// ── Logic_And: 4×4 bool, And ──────────────────────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_And) +{ + constexpr std::size_t N = 16; // 4×4 + auto d_a = makeDeviceBuf(host, device, queue, + Logic_And_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_And_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_And::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + uint8_t* ref = Logic_And_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_Or: 4×4 bool, Or ───────────────────────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_Or) +{ + constexpr std::size_t N = 16; + auto d_a = makeDeviceBuf(host, device, queue, + Logic_Or_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_Or_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_Or::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + uint8_t* ref = Logic_Or_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_Xor: 4×4 bool, Xor ─────────────────────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_Xor) +{ + constexpr std::size_t N = 16; + auto d_a = makeDeviceBuf(host, device, queue, + Logic_Xor_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_Xor_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_Xor::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + uint8_t* ref = Logic_Xor_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_BitwiseAnd: 3×5 int32, BitwiseAnd ──────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_BitwiseAnd) +{ + constexpr std::size_t N = 15; // 3×5 + auto d_a = makeDeviceBuf(host, device, queue, + Logic_BitwiseAnd_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_BitwiseAnd_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_BitwiseAnd::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int32_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + int32_t* ref = Logic_BitwiseAnd_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_BitwiseOr: 3×5 int32, BitwiseOr ────────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_BitwiseOr) +{ + constexpr std::size_t N = 15; + auto d_a = makeDeviceBuf(host, device, queue, + Logic_BitwiseOr_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_BitwiseOr_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_BitwiseOr::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int32_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + int32_t* ref = Logic_BitwiseOr_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_BitwiseXor: 3×5 int32, BitwiseXor ──────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_BitwiseXor) +{ + constexpr std::size_t N = 15; + auto d_a = makeDeviceBuf(host, device, queue, + Logic_BitwiseXor_Input::data_a, N); + auto d_b = makeDeviceBuf(host, device, queue, + Logic_BitwiseXor_Input::data_b, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_BitwiseXor::Session session; + auto result = session.infer(d_a, d_b); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int32_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + int32_t* ref = Logic_BitwiseXor_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +// ── Logic_BitwiseNot: 2×3×4 int32, BitwiseNot ────────────────────────────── +TEST_F(SofieAlpakaTest, Logic_BitwiseNot) +{ + constexpr std::size_t N = 24; // 2×3×4 + auto d_input = makeDeviceBuf(host, device, queue, + Logic_BitwiseNot_Input::data_a, N); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N})); + { + SOFIE_Logic_BitwiseNot::Session session; + auto result = session.infer(d_input); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int32_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + int32_t* ref = Logic_BitwiseNot_ExpectedOutput::outputs; + for (std::size_t i = 0; i < N; ++i) + EXPECT_EQ(res[i], ref[i]) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, TopK) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // axis=-1, largest=1, sorted=1, k=5 (baked); input is a single 9-element row + std::vector input {9.0, 8.0, 4.5, 1.7, 2.9, 3.2, 4.0, 2.6, 7.4}; + constexpr std::size_t K = 5; + + 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 values_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + + { + SOFIE_TopK::Session session; + auto [values, indices] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + + for (std::size_t i = 0; i < K; ++i) { + EXPECT_LE(std::abs(val[i] - TopK_ExpectedOutput::values[i]), TOLERANCE) << " value index=" << i; + EXPECT_EQ(idx[i], static_cast(TopK_ExpectedOutput::indexes[i])) << " index index=" << i; + } +} + +TEST_F(SofieAlpakaTest, Softmax1d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({-1.f, 0.f, 1.f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax1d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax1d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax2d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({-1.f, 0.f, 1.f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax2d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax2d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax3d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({ + -0.8939f, -0.3674f, 0.1763f, 1.5804f, -0.4687f, 1.2253f, -1.3488f, -0.1000f, + -0.1262f, 0.4962f, 1.0870f, 0.6905f, -0.3451f, -1.6981f, -0.4688f, 0.4468f, + -0.5479f, 0.0650f, 1.0446f, -1.6249f, -0.7190f, -1.7520f, 3.7753f, -1.4939f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax3d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax3d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax4d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({ + -0.5869f, -1.4272f, -0.1546f, 0.0096f, 0.1706f, 0.0388f, -0.3484f, -0.7829f, + 1.1138f, -0.5644f, -0.6264f, -1.1890f, 1.6741f, -0.7130f, 0.9592f, 1.7477f, + -0.4775f, 1.3407f, -0.3882f, -0.4560f, 1.0385f, -0.1669f, 0.5540f, -1.0790f, + -0.6153f, -0.6274f, -1.2304f, -0.6757f, 1.0178f, -0.2379f, -0.7912f, -0.0165f, + -0.5423f, 0.1459f, 1.3585f, -0.5005f, -0.2187f, -1.8181f, -0.6642f, 0.0287f, + -1.9103f, 0.7984f, -0.7860f, 1.5134f, 1.3873f, -0.6462f, -0.6354f, -0.1335f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax4d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax4d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} diff --git a/test/alpaka/TestAlpakaConv.cxx b/test/alpaka/TestAlpakaConv.cxx index d0577f5a..c6faf3ee 100644 --- a/test/alpaka/TestAlpakaConv.cxx +++ b/test/alpaka/TestAlpakaConv.cxx @@ -680,3 +680,141 @@ TEST_F(SofieAlpakaTest, ConvWithDilation) EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; } } +<<<<<<< HEAD +======= + +// dynamic 1D Conv (k=3, pad=1) + bias, run at (N,n_pf) = (1,1) and (8,5) +TEST_F(SofieAlpakaTest, DynamicConv1D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3, K = 3; + const float bias[3] = {0.5f, -0.5f, 1.0f}; + auto Wv = [&](std::size_t oc, std::size_t ic, std::size_t kk) { + return static_cast(oc * (Cin * K) + ic * K + kk); + }; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P, outSize = N * Cout * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv1D::Session session("DynamicConv1D_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t l = 0; l < P; ++l) { + float acc = bias[oc]; + for (std::size_t ic = 0; ic < Cin; ++ic) + for (std::size_t kk = 0; kk < K; ++kk) { + int64_t li = static_cast(l) + static_cast(kk) - 1; // pad=1 + if (li >= 0 && li < static_cast(P)) + acc += in_ptr[n * Cin * P + ic * P + li] * Wv(oc, ic, kk); + } + std::size_t idx = n * Cout * P + oc * P + l; + EXPECT_LE(std::abs(res[idx] - acc), TOLERANCE) << "n=" << n << " oc=" << oc << " l=" << l; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicConv1DNoBias) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3; + const float W[3][2] = {{0.5f, -1.0f}, {0.25f, 0.75f}, {-0.5f, 1.5f}}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P, outSize = N * Cout * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv1DNoBias::Session session("DynamicConv1DNoBias_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t p = 0; p < P; ++p) { + float expected = 0.0f; + for (std::size_t ic = 0; ic < Cin; ++ic) + expected += W[oc][ic] * in_ptr[n * Cin * P + ic * P + p]; + float got = res[n * Cout * P + oc * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " oc=" << oc << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicConv2DNoBias) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Cin = 2, Cout = 3, Q = 4; + const float W[3][2] = {{0.5f, -1.0f}, {0.25f, 0.75f}, {-0.5f, 1.5f}}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * Cin * P * Q, outSize = N * Cout * P * Q; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicConv2DNoBias::Session session("DynamicConv2DNoBias_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t oc = 0; oc < Cout; ++oc) + for (std::size_t p = 0; p < P; ++p) + for (std::size_t q = 0; q < Q; ++q) { + float expected = 0.0f; + for (std::size_t ic = 0; ic < Cin; ++ic) + expected += W[oc][ic] * in_ptr[n * Cin * P * Q + ic * P * Q + p * Q + q]; + float got = res[n * Cout * P * Q + oc * P * Q + p * Q + q]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " oc=" << oc << " p=" << p << " q=" << q; + } + } +} +>>>>>>> 9cb225f (chore: in line comments cleanup in the code) diff --git a/test/alpaka/TestAlpakaElementwiseBinaryLogic.cxx b/test/alpaka/TestAlpakaElementwiseBinaryLogic.cxx index 5b053f8a..885469ca 100644 --- a/test/alpaka/TestAlpakaElementwiseBinaryLogic.cxx +++ b/test/alpaka/TestAlpakaElementwiseBinaryLogic.cxx @@ -34,11 +34,23 @@ #include "input_models/references/Greater.ref.hxx" #include "input_models/references/Less.ref.hxx" #include "Where_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicWhere_FromONNX_GPU_ALPAKA.hxx" +#include "MaxMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.hxx" +#include "MinMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.hxx" +#include "MeanMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.hxx" +#include "SumMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/MaxMultidirectionalBroadcast.ref.hxx" +#include "input_models/references/MinMultidirectionalBroadcast.ref.hxx" +#include "input_models/references/MeanMultidirectionalBroadcast.ref.hxx" +#include "input_models/references/SumMultidirectionalBroadcast.ref.hxx" #include "IsInf_FromONNX_GPU_ALPAKA.hxx" #include "IsNaN_FromONNX_GPU_ALPAKA.hxx" #include "Clip_FromONNX_GPU_ALPAKA.hxx" #include "Not_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicEqual_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicAddBroadcast_FromONNX_GPU_ALPAKA.hxx" + TEST_F(SofieAlpakaTest, AddBroadcast1) { @@ -91,6 +103,158 @@ TEST_F(SofieAlpakaTest, AddBroadcast1) } } +TEST_F(SofieAlpakaTest, DynamicAddBroadcast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float bias[4] = {0.5f, -1.0f, 0.25f, 2.0f}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicAddBroadcast::Session session("DynamicAddBroadcast_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / P) % C; + float expected = in_ptr[i] + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i << " N=" << N << " P=" << P; + } + } +} + +// ── BasicNary multidirectional broadcast: A[3,1], B[2,3,1], C[1,4] -> Y[2,3,4] ────── +TEST_F(SofieAlpakaTest, MaxMultidirectionalBroadcast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const float a[] = {0.35974154f, -2.20873388f, 0.95746274f}; + const float b[] = {0.75901985f, -0.46544461f, -0.34920575f, -0.1460754f, 0.08269051f, -0.70045695f}; + const float c[] = {-0.41468981f, -0.46591926f, 0.56172534f, 0.05616931f}; + + auto a_d = makeDeviceBuf(host, device, queue, a, 3); + auto b_d = makeDeviceBuf(host, device, queue, b, 6); + auto c_d = makeDeviceBuf(host, device, queue, c, 4); + + const std::size_t outSize = sizeof(MaxMultidirectionalBroadcast_ExpectedOutput::output) / sizeof(float); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_MaxMultidirectionalBroadcast::Session session("MaxMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(a_d, b_d, c_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* correct = MaxMultidirectionalBroadcast_ExpectedOutput::output; + EXPECT_EQ(outSize, 24u); + for (std::size_t i = 0; i < outSize; ++i) + EXPECT_LE(std::abs(res[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, MinMultidirectionalBroadcast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const float a[] = {0.35974154f, -2.20873388f, 0.95746274f}; + const float b[] = {0.75901985f, -0.46544461f, -0.34920575f, -0.1460754f, 0.08269051f, -0.70045695f}; + const float c[] = {-0.41468981f, -0.46591926f, 0.56172534f, 0.05616931f}; + + auto a_d = makeDeviceBuf(host, device, queue, a, 3); + auto b_d = makeDeviceBuf(host, device, queue, b, 6); + auto c_d = makeDeviceBuf(host, device, queue, c, 4); + + const std::size_t outSize = sizeof(MinMultidirectionalBroadcast_ExpectedOutput::output) / sizeof(float); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_MinMultidirectionalBroadcast::Session session("MinMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(a_d, b_d, c_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* correct = MinMultidirectionalBroadcast_ExpectedOutput::output; + EXPECT_EQ(outSize, 24u); + for (std::size_t i = 0; i < outSize; ++i) + EXPECT_LE(std::abs(res[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, MeanMultidirectionalBroadcast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const float a[] = {0.35974154f, -2.20873388f, 0.95746274f}; + const float b[] = {0.75901985f, -0.46544461f, -0.34920575f, -0.1460754f, 0.08269051f, -0.70045695f}; + const float c[] = {-0.41468981f, -0.46591926f, 0.56172534f, 0.05616931f}; + + auto a_d = makeDeviceBuf(host, device, queue, a, 3); + auto b_d = makeDeviceBuf(host, device, queue, b, 6); + auto c_d = makeDeviceBuf(host, device, queue, c, 4); + + const std::size_t outSize = sizeof(MeanMultidirectionalBroadcast_ExpectedOutput::output) / sizeof(float); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_MeanMultidirectionalBroadcast::Session session("MeanMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(a_d, b_d, c_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* correct = MeanMultidirectionalBroadcast_ExpectedOutput::output; + EXPECT_EQ(outSize, 24u); + for (std::size_t i = 0; i < outSize; ++i) + EXPECT_LE(std::abs(res[i] - correct[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, SumMultidirectionalBroadcast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const float a[] = {0.35974154f, -2.20873388f, 0.95746274f}; + const float b[] = {0.75901985f, -0.46544461f, -0.34920575f, -0.1460754f, 0.08269051f, -0.70045695f}; + const float c[] = {-0.41468981f, -0.46591926f, 0.56172534f, 0.05616931f}; + + auto a_d = makeDeviceBuf(host, device, queue, a, 3); + auto b_d = makeDeviceBuf(host, device, queue, b, 6); + auto c_d = makeDeviceBuf(host, device, queue, c, 4); + + const std::size_t outSize = sizeof(SumMultidirectionalBroadcast_ExpectedOutput::output) / sizeof(float); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_SumMultidirectionalBroadcast::Session session("SumMultidirectionalBroadcast_FromONNX_GPU_ALPAKA.dat"); + auto result = session.infer(a_d, b_d, c_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* correct = SumMultidirectionalBroadcast_ExpectedOutput::output; + EXPECT_EQ(outSize, 24u); + for (std::size_t i = 0; i < outSize; ++i) + EXPECT_LE(std::abs(res[i] - correct[i]), TOLERANCE) << "i=" << i; +} +// ──────────────────────────────────────────────────────────────────────────────────── + TEST_F(SofieAlpakaTest, Equal) { std::vector input1 = {1.0f, 2.0f, 3.0f}; @@ -129,6 +293,44 @@ TEST_F(SofieAlpakaTest, Equal) EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; } +TEST_F(SofieAlpakaTest, DynamicEqual) +{ + const std::size_t cols = 3; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * cols; + + auto x1_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + auto x2_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* x1p = reinterpret_cast(alpaka::getPtrNative(x1_h)); + float* x2p = reinterpret_cast(alpaka::getPtrNative(x2_h)); + for (Idx i = 0; i < sz; ++i) { + x1p[i] = static_cast(i % 3); + x2p[i] = static_cast(i % 2); // mix of equal and unequal + } + + auto x1_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + auto x2_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, x1_d, x1_h); + alpaka::memcpy(queue, x2_d, x2_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicEqual::Session session("", N); + auto result = session.infer(N, x1_d, x2_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + uint8_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + uint8_t expected = (x1p[i] == x2p[i]) ? 1 : 0; + EXPECT_EQ(res[i], expected); + } + } +} + TEST_F(SofieAlpakaTest, LessOrEqual) { std::vector input1 = {1.0f, 2.0f, 3.0f}; @@ -324,6 +526,49 @@ TEST_F(SofieAlpakaTest, Where) EXPECT_EQ(res_ptr[i], correct[i]) << "i=" << i; } +TEST_F(SofieAlpakaTest, DynamicWhere) +{ + constexpr std::size_t Cols = 5; + for (std::size_t N : {std::size_t(1), std::size_t(4)}) { + const std::size_t size = N * Cols; + + auto x_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + float* x_ptr = reinterpret_cast(alpaka::getPtrNative(x_h)); + auto y_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + float* y_ptr = reinterpret_cast(alpaka::getPtrNative(y_h)); + auto c_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + uint8_t* c_ptr = reinterpret_cast(alpaka::getPtrNative(c_h)); + for (Idx i = 0; i < size; ++i) { + x_ptr[i] = static_cast(i) + 1.0f; + y_ptr[i] = -static_cast(i) - 1.0f; + c_ptr[i] = (i % 2 == 0) ? 1 : 0; + } + + auto x_d = alpaka::allocBuf(device, Ext1D::all(Idx{size})); + auto y_d = alpaka::allocBuf(device, Ext1D::all(Idx{size})); + auto c_d = alpaka::allocBuf(device, Ext1D::all(Idx{size})); + alpaka::memcpy(queue, x_d, x_h); + alpaka::memcpy(queue, y_d, y_h); + alpaka::memcpy(queue, c_d, c_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + { + SOFIE_DynamicWhere::Session session("DynamicWhere_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, c_d, x_d, y_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (Idx i = 0; i < size; ++i) { + float expected = c_ptr[i] ? x_ptr[i] : y_ptr[i]; + EXPECT_EQ(res[i], expected) << "i=" << i << " N=" << N; + } + } +} + TEST_F(SofieAlpakaTest, IsInf) { // Input contains finite values, +inf, -inf; output is bool (uint8_t). @@ -635,4 +880,3 @@ TEST_F(SofieAlpakaTest, Logic_BitwiseNot) for (std::size_t i = 0; i < N; ++i) EXPECT_EQ(res[i], ref[i]) << " index=" << i; } - diff --git a/test/alpaka/TestAlpakaElementwiseUnary.cxx b/test/alpaka/TestAlpakaElementwiseUnary.cxx index 61f714bd..4ac6c04c 100644 --- a/test/alpaka/TestAlpakaElementwiseUnary.cxx +++ b/test/alpaka/TestAlpakaElementwiseUnary.cxx @@ -293,6 +293,45 @@ TEST_F(SofieAlpakaTest, Neg) EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; } +<<<<<<< HEAD +======= +TEST_F(SofieAlpakaTest, DynamicNegRelu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicNegRelu::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + float expected = std::max(0.0f, -in_ptr[i]); + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i << " N=" << N << " P=" << P; + } + } +} + +>>>>>>> 9cb225f (chore: in line comments cleanup in the code) TEST_F(SofieAlpakaTest, Softplus) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; diff --git a/test/alpaka/TestAlpakaLinear.cxx b/test/alpaka/TestAlpakaLinear.cxx index d46e0609..39161ea2 100644 --- a/test/alpaka/TestAlpakaLinear.cxx +++ b/test/alpaka/TestAlpakaLinear.cxx @@ -7,6 +7,9 @@ #include "LinearWithSigmoid_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/LinearWithSigmoid.ref.hxx" +#include "DynamicLinear_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicMatMulBatched_FromONNX_GPU_ALPAKA.hxx" + TEST_F(SofieAlpakaTest, Linear64) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -118,3 +121,88 @@ TEST_F(SofieAlpakaTest, LinearWithSigmoid) } } +TEST_F(SofieAlpakaTest, DynamicLinear) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t In = 4, Out = 3; + const float W[4][3] = {{0.5f, -1.0f, 0.25f}, + {0.75f, 0.5f, -0.5f}, + {-0.25f, 1.0f, 0.75f}, + {1.5f, -0.75f, 0.5f}}; + const float B[3] = {0.5f, -0.5f, 0.25f}; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * In, outSize = N * Out; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicLinear::Session session("DynamicLinear_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < Out; ++j) { + float acc = B[j]; + for (std::size_t i = 0; i < In; ++i) acc += in_ptr[n * In + i] * W[i][j]; + float expected = acc > 0.0f ? acc : 0.0f; + EXPECT_LE(std::abs(res[n * Out + j] - expected), TOLERANCE) << "n=" << n << " j=" << j << " N=" << N; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicMatMulBatched) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t M = 4, K = 8, Nc = 6; + + for (std::size_t N : {std::size_t(1), std::size_t(5)}) { + const std::size_t aSize = N * M * K, bSize = N * K * Nc, ySize = N * M * Nc; + + auto a_h = alpaka::allocBuf(host, Ext1D::all(Idx{aSize})); + float* a_ptr = reinterpret_cast(alpaka::getPtrNative(a_h)); + for (Idx i = 0; i < aSize; ++i) a_ptr[i] = (static_cast(i % 11) - 5.0f) * 0.3f; + + auto b_h = alpaka::allocBuf(host, Ext1D::all(Idx{bSize})); + float* b_ptr = reinterpret_cast(alpaka::getPtrNative(b_h)); + for (Idx i = 0; i < bSize; ++i) b_ptr[i] = (static_cast(i % 7) - 3.0f) * 0.5f; + + auto a_d = alpaka::allocBuf(device, Ext1D::all(Idx{aSize})); + auto b_d = alpaka::allocBuf(device, Ext1D::all(Idx{bSize})); + alpaka::memcpy(queue, a_d, a_h); + alpaka::memcpy(queue, b_d, b_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{ySize})); + { + SOFIE_DynamicMatMulBatched::Session session("DynamicMatMulBatched_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, a_d, b_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t i = 0; i < M; ++i) + for (std::size_t j = 0; j < Nc; ++j) { + float expected = 0.0f; + for (std::size_t k = 0; k < K; ++k) + expected += a_ptr[n * M * K + i * K + k] * b_ptr[n * K * Nc + k * Nc + j]; + float got = res[n * M * Nc + i * Nc + j]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) + << "n=" << n << " i=" << i << " j=" << j << " N=" << N; + } + } +} diff --git a/test/alpaka/TestAlpakaNormReduction.cxx b/test/alpaka/TestAlpakaNormReduction.cxx index f9a9d1ca..95746958 100644 --- a/test/alpaka/TestAlpakaNormReduction.cxx +++ b/test/alpaka/TestAlpakaNormReduction.cxx @@ -1,7 +1,9 @@ #include "TestAlpakaCommon.h" #include "RMSNorm_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicRMSNorm_FromONNX_GPU_ALPAKA.hxx" #include "GroupNorm_FromONNX_GPU_ALPAKA.hxx" +#include "GroupNormBig_FromONNX_GPU_ALPAKA.hxx" #include "CumSum_FromONNX_GPU_ALPAKA.hxx" #include "CumSumExclusive_FromONNX_GPU_ALPAKA.hxx" @@ -43,6 +45,41 @@ TEST_F(SofieAlpakaTest, RMSNorm) EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; } +TEST_F(SofieAlpakaTest, DynamicRMSNorm) +{ + constexpr float TOLERANCE = 1e-4f; + constexpr std::size_t D = 4; + const float scale[D] = {1.0f, 1.0f, 1.0f, 1.0f}; + + for (std::size_t N : {std::size_t(1), std::size_t(6)}) { + std::vector input(N * D); + for (std::size_t i = 0; i < input.size(); ++i) input[i] = (static_cast(i % 9) - 4.0f) * 0.5f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), input.size()); + auto scale_d = makeDeviceBuf(host, device, queue, scale, D); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * D})); + + { + SOFIE_DynamicRMSNorm::Session session("DynamicRMSNorm_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d, scale_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + const float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + double sumSq = 0; + for (std::size_t d = 0; d < D; ++d) sumSq += double(input[n * D + d]) * input[n * D + d]; + float invRms = static_cast(1.0 / std::sqrt(sumSq / D + 1e-5)); + for (std::size_t d = 0; d < D; ++d) { + float expected = scale[d] * input[n * D + d] * invRms; + EXPECT_LE(std::abs(res[n * D + d] - expected), TOLERANCE) << "n=" << n << " d=" << d << " N=" << N; + } + } + } +} + TEST_F(SofieAlpakaTest, RMSNormResetState) { SOFIE_RMSNorm::Session session("RMSNorm_FromONNX_GPU_ALPAKA.dat"); @@ -78,6 +115,67 @@ TEST_F(SofieAlpakaTest, GroupNorm) EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; } +TEST_F(SofieAlpakaTest, GroupNormBig) +{ + constexpr float TOLERANCE = 1e-3f; + constexpr std::size_t N = 2, C = 4, H = 20, W = 20, G = 2; + constexpr std::size_t spatial = H * W, gs = C / G; + constexpr std::size_t size = N * C * spatial; + + std::vector x(size); + for (std::size_t i = 0; i < size; ++i) x[i] = (static_cast(i % 23) - 11.0f) * 0.37f; + float scale[C] = {1.1f, 0.9f, 1.3f, 0.8f}; + float bias[C] = {0.1f, -0.2f, 0.3f, -0.1f}; + + auto x_d = makeDeviceBuf(host, device, queue, x.data(), size); + auto scale_d = makeDeviceBuf(host, device, queue, scale, C); + auto bias_d = makeDeviceBuf(host, device, queue, bias, C); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + + { + SOFIE_GroupNormBig::Session session; + auto result = session.infer(x_d, scale_d, bias_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + std::vector expected(size); + for (std::size_t n = 0; n < N; ++n) { + for (std::size_t g = 0; g < G; ++g) { + double sum = 0; + for (std::size_t ci = 0; ci < gs; ++ci) { + std::size_t c = g * gs + ci; + for (std::size_t s = 0; s < spatial; ++s) + sum += x[n * C * spatial + c * spatial + s]; + } + double mean = sum / (gs * spatial); + double varsum = 0; + for (std::size_t ci = 0; ci < gs; ++ci) { + std::size_t c = g * gs + ci; + for (std::size_t s = 0; s < spatial; ++s) { + double d = x[n * C * spatial + c * spatial + s] - mean; + varsum += d * d; + } + } + double var = varsum / (gs * spatial); + double inv_std = 1.0 / std::sqrt(var + 1e-5); + for (std::size_t ci = 0; ci < gs; ++ci) { + std::size_t c = g * gs + ci; + for (std::size_t s = 0; s < spatial; ++s) { + std::size_t idx = n * C * spatial + c * spatial + s; + expected[idx] = static_cast((x[idx] - mean) * inv_std) * scale[c] + bias[c]; + } + } + } + } + + const float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < size; ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + TEST_F(SofieAlpakaTest, CumSumInclusive) { constexpr float TOLERANCE = 1e-4f; diff --git a/test/alpaka/TestAlpakaNormalization.cxx b/test/alpaka/TestAlpakaNormalization.cxx index 7ab0736d..230dd35b 100644 --- a/test/alpaka/TestAlpakaNormalization.cxx +++ b/test/alpaka/TestAlpakaNormalization.cxx @@ -5,6 +5,19 @@ #include "LayerNorm_FromONNX_GPU_ALPAKA.hxx" #include "LayerNormScaleBias_FromONNX_GPU_ALPAKA.hxx" #include "LayerNorm3D_FromONNX_GPU_ALPAKA.hxx" +#include "LayerNormBig_FromONNX_GPU_ALPAKA.hxx" + +#include "DynamicBatchNorm4D_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicBatchNormDynSpatialRelu_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicBatchNorm2D_FromONNX_GPU_ALPAKA.hxx" +#include "Softmax1d_FromONNX_GPU_ALPAKA.hxx" +#include "Softmax2d_FromONNX_GPU_ALPAKA.hxx" +#include "Softmax3d_FromONNX_GPU_ALPAKA.hxx" +#include "Softmax4d_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/Softmax1d.ref.hxx" +#include "input_models/references/Softmax2d.ref.hxx" +#include "input_models/references/Softmax3d.ref.hxx" +#include "input_models/references/Softmax4d.ref.hxx" TEST_F(SofieAlpakaTest, BatchNormalization) { @@ -82,6 +95,127 @@ TEST_F(SofieAlpakaTest, BatchNormalizationRelu) } } +// X[N,4,2,3]: dynamic batch, static spatial (6) +TEST_F(SofieAlpakaTest, DynamicBatchNorm4D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4, spatial = 2 * 3; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * C * spatial; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNorm4D::Session session("DynamicBatchNorm4D_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / spatial) % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + +// X[N,4,n_pf] + relu: dynamic batch and spatial, rebuilt per size so Y matches its length +TEST_F(SofieAlpakaTest, DynamicBatchNormDynSpatialRelu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNormDynSpatialRelu::Session session("DynamicBatchNormDynSpatialRelu_FromONNX_GPU_ALPAKA.dat", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = (i / P) % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + expected = expected > 0.0f ? expected : 0.0f; // relu + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + +// X[N,4]: rank 2, spatial is 1 +TEST_F(SofieAlpakaTest, DynamicBatchNorm2D) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4; + const float scale[4] = {1.0f, 2.0f, 0.5f, 1.5f}, bias[4] = {0.1f, -0.2f, 0.3f, 0.0f}; + const float mean_[4] = {0.5f, 1.0f, -0.5f, 2.0f}, var_[4] = {1.0f, 4.0f, 0.25f, 2.0f}; + const float eps = 1e-5f; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t sz = N * C; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 10) - 5.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicBatchNorm2D::Session session("DynamicBatchNorm2D_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < sz; ++i) { + std::size_t c = i % C; + float fused = scale[c] / std::sqrt(var_[c] + eps); + float expected = (in_ptr[i] - mean_[c]) * fused + bias[c]; + EXPECT_LE(std::abs(res[i] - expected), TOLERANCE) << "i=" << i; + } + } +} + TEST_F(SofieAlpakaTest, LayerNorm) { constexpr float TOLERANCE = 1e-4f; @@ -223,3 +357,186 @@ TEST_F(SofieAlpakaTest, LayerNorm3D) EXPECT_LE(std::abs(res_ptr[12 + i] - exp1[i]), TOLERANCE) << "row1 i=" << i; } +TEST_F(SofieAlpakaTest, LayerNormBig) +{ + constexpr float TOLERANCE = 1e-3f; + const std::size_t N = 3, D = 2050; + const std::size_t total = N * D; + + std::vector input(total); + for (std::size_t r = 0; r < N; ++r) + for (std::size_t j = 0; j < D; ++j) + input[r * D + j] = std::sin((j + 37.f * r) * 0.01f) + 0.15f * std::cos(j * 0.003f + r); + + auto compute_expected = [](std::vector row) { + float mean = 0.f; + for (float v : row) mean += v; + mean /= row.size(); + float var = 0.f; + for (float v : row) var += (v - mean) * (v - mean); + var /= row.size(); + float inv_std = 1.f / std::sqrt(var + 1e-5f); + std::vector out; + for (float v : row) out.push_back((v - mean) * inv_std); + return out; + }; + + std::vector expected(total); + for (std::size_t r = 0; r < N; ++r) { + std::vector row(input.begin() + r * D, input.begin() + (r + 1) * D); + auto exp_row = compute_expected(row); + std::copy(exp_row.begin(), exp_row.end(), expected.begin() + r * D); + } + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{total})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < total; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{total})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{total})); + + { + SOFIE_LayerNormBig::Session session("LayerNormBig_FromONNX_GPU_ALPAKA.dat"); + 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)); + for (std::size_t i = 0; i < total; ++i) + EXPECT_LE(std::abs(res_ptr[i] - expected[i]), TOLERANCE) << "i=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax1d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({-1.f, 0.f, 1.f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax1d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax1d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax2d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({-1.f, 0.f, 1.f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax2d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax2d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax3d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({ + -0.8939f, -0.3674f, 0.1763f, 1.5804f, -0.4687f, 1.2253f, -1.3488f, -0.1000f, + -0.1262f, 0.4962f, 1.0870f, 0.6905f, -0.3451f, -1.6981f, -0.4688f, 0.4468f, + -0.5479f, 0.0650f, 1.0446f, -1.6249f, -0.7190f, -1.7520f, 3.7753f, -1.4939f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax3d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax3d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} + +TEST_F(SofieAlpakaTest, Softmax4d) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + std::vector input_vec({ + -0.5869f, -1.4272f, -0.1546f, 0.0096f, 0.1706f, 0.0388f, -0.3484f, -0.7829f, + 1.1138f, -0.5644f, -0.6264f, -1.1890f, 1.6741f, -0.7130f, 0.9592f, 1.7477f, + -0.4775f, 1.3407f, -0.3882f, -0.4560f, 1.0385f, -0.1669f, 0.5540f, -1.0790f, + -0.6153f, -0.6274f, -1.2304f, -0.6757f, 1.0178f, -0.2379f, -0.7912f, -0.0165f, + -0.5423f, 0.1459f, 1.3585f, -0.5005f, -0.2187f, -1.8181f, -0.6642f, 0.0287f, + -1.9103f, 0.7984f, -0.7860f, 1.5134f, 1.3873f, -0.6462f, -0.6354f, -0.1335f}); + const Idx N = static_cast(input_vec.size()); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(N)); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N; ++i) input_ptr[i] = input_vec[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(N)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(N)); + { + SOFIE_Softmax4d::Session session; + auto result = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + float* ref = Softmax4d_ExpectedOutput::output; + for (Idx i = 0; i < N; ++i) + EXPECT_LE(std::abs(res[i] - ref[i]), TOLERANCE) << " index=" << i; +} diff --git a/test/alpaka/TestAlpakaPool.cxx b/test/alpaka/TestAlpakaPool.cxx index 993920af..b1a9d450 100644 --- a/test/alpaka/TestAlpakaPool.cxx +++ b/test/alpaka/TestAlpakaPool.cxx @@ -14,6 +14,43 @@ #include "input_models/references/AvgPoolCountIncludePad.ref.hxx" #include "GlobalAvgPool2d_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/GlobalAvgPool2d.ref.hxx" +#include "DynamicMaxPool_FromONNX_GPU_ALPAKA.hxx" + +TEST_F(SofieAlpakaTest, DynamicMaxPool) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + for (std::size_t N : {std::size_t(1), std::size_t(3)}) { + std::vector input(N * 16); + for (std::size_t i = 0; i < input.size(); ++i) input[i] = static_cast(i % 17) - 8.0f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), input.size()); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * 4})); + { + SOFIE_DynamicMaxPool::Session session("DynamicMaxPool_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + const float* img = input.data() + n * 16; + float expected[4]; + for (int oh = 0; oh < 2; ++oh) + for (int ow = 0; ow < 2; ++ow) { + float m = -1e30f; + for (int kh = 0; kh < 2; ++kh) + for (int kw = 0; kw < 2; ++kw) + m = std::max(m, img[(oh * 2 + kh) * 4 + (ow * 2 + kw)]); + expected[oh * 2 + ow] = m; + } + for (int i = 0; i < 4; ++i) + EXPECT_LE(std::abs(res[n * 4 + i] - expected[i]), TOLERANCE) << "n=" << n << " i=" << i << " N=" << N; + } + } +} TEST_F(SofieAlpakaTest, MaxPool2d) { @@ -251,8 +288,6 @@ TEST_F(SofieAlpakaTest, AvgPoolCountIncludePad) } // GlobalAveragePool: one output per channel = the mean of the whole channel. -// Input x[1,2,3,3] = iota 0..17, so the channel means are 4 and 13. - TEST_F(SofieAlpakaTest, GlobalAvgPool2d) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; diff --git a/test/alpaka/TestAlpakaReduce.cxx b/test/alpaka/TestAlpakaReduce.cxx index 6005b38b..0dc09a66 100644 --- a/test/alpaka/TestAlpakaReduce.cxx +++ b/test/alpaka/TestAlpakaReduce.cxx @@ -1,5 +1,8 @@ #include "TestAlpakaCommon.h" +#include +#include + #include "ReduceMean_FromONNX_GPU_ALPAKA.hxx" #include "ReduceProd_FromONNX_GPU_ALPAKA.hxx" #include "ReduceSum_FromONNX_GPU_ALPAKA.hxx" @@ -16,6 +19,16 @@ #include "input_models/references/ReduceMax_axis0.ref.hxx" #include "input_models/references/ReduceMax_mid.ref.hxx" +#include "DynamicReduceSumLast_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceMeanMid_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceMaxFirst_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicReduceSumMulti_FromONNX_GPU_ALPAKA.hxx" +#include "TopK_FromONNX_GPU_ALPAKA.hxx" +#include "input_models/references/TopK.ref.hxx" +#include "TopKBig_FromONNX_GPU_ALPAKA.hxx" +#include "TopKTies_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicTopK_FromONNX_GPU_ALPAKA.hxx" + TEST_F(SofieAlpakaTest, ReduceMean) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -338,3 +351,475 @@ TEST_F(SofieAlpakaTest, ReduceMax_mid) EXPECT_NEAR(res_ptr[i], correct[i], TOLERANCE) << " i=" << i; } +// X[N,4] -> ReduceSum(axis=-1, keepdims=0) -> Y[N]: kLast, pruned output, negative axis +TEST_F(SofieAlpakaTest, DynamicReduceSumLast) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t cols = 4; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * cols, outSize = N; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceSumLast::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + float expected = 0.f; + for (std::size_t c = 0; c < cols; ++c) expected += in_ptr[n * cols + c]; + EXPECT_LE(std::abs(res[n] - expected), TOLERANCE); + } + } +} + +// X[N,4,M] -> ReduceMean(axis=1, keepdims=1) -> Y[N,1,M]: kMiddle with two symbolic dims +TEST_F(SofieAlpakaTest, DynamicReduceMeanMid) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t K = 4; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ms[] = {1, 3}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], M = Ms[t]; + const std::size_t inSize = N * K * M, outSize = N * M; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceMeanMid::Session session("", N, M); + auto result = session.infer(N, M, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t m = 0; m < M; ++m) { + float sum = 0.f; + for (std::size_t k = 0; k < K; ++k) + sum += in_ptr[n * K * M + k * M + m]; + float expected = sum / static_cast(K); + EXPECT_LE(std::abs(res[n * M + m] - expected), TOLERANCE); + } + } +} + + +TEST_F(SofieAlpakaTest, DynamicReduceMaxFirst) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t cols = 4; + for (std::size_t N : {std::size_t(1), std::size_t(8), std::size_t(300), std::size_t(50000)}) { + const std::size_t inSize = N * cols, outSize = cols; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast((i * 7 + 3) % 11); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceMaxFirst::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t c = 0; c < cols; ++c) { + float expected = in_ptr[c]; // n == 0 + for (std::size_t n = 1; n < N; ++n) + if (in_ptr[n * cols + c] > expected) expected = in_ptr[n * cols + c]; + EXPECT_LE(std::abs(res[c] - expected), TOLERANCE); + } + } +} + +TEST_F(SofieAlpakaTest, DynamicReduceMaxFirstScratchCacheReuse) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t cols = 4; + const std::size_t N = 50000; + const std::size_t inSize = N * cols, outSize = cols; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + + SOFIE_DynamicReduceMaxFirst::Session session("", N); + for (int rep = 0; rep < 4; ++rep) { + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast((i * 7 + 3 + rep * 13) % 11); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t c = 0; c < cols; ++c) { + float expected = in_ptr[c]; + for (std::size_t n = 1; n < N; ++n) + if (in_ptr[n * cols + c] > expected) expected = in_ptr[n * cols + c]; + EXPECT_LE(std::abs(res[c] - expected), TOLERANCE) << "rep=" << rep << " c=" << c; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicReduceSumMulti) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t slice = 3 * 2; + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inSize = N * slice, outSize = N; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicReduceSumMulti::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + float expected = 0.f; + for (std::size_t j = 0; j < slice; ++j) expected += in_ptr[n * slice + j]; + EXPECT_LE(std::abs(res[n] - expected), TOLERANCE); + } + } +} + +TEST_F(SofieAlpakaTest, TopK) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + + // axis=-1, largest=1, sorted=1, k=5 fixed at codegen; input is a single 9-element row + std::vector input {9.0, 8.0, 4.5, 1.7, 2.9, 3.2, 4.0, 2.6, 7.4}; + constexpr std::size_t K = 5; + + 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 values_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + + { + SOFIE_TopK::Session session; + auto [values, indices] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + + for (std::size_t i = 0; i < K; ++i) { + EXPECT_LE(std::abs(val[i] - TopK_ExpectedOutput::values[i]), TOLERANCE) << " value index=" << i; + EXPECT_EQ(idx[i], static_cast(TopK_ExpectedOutput::indexes[i])) << " index index=" << i; + } +} + +// Regression test for the two-phase parallel GPU TopK path: axis length 10007 is far +// past the point where a single serial thread per slice would leave the GPU mostly +// idle (the scenario this optimization targets, e.g. top-k over vocabulary-sized +// logits), and it's deliberately not a multiple of the chunk size, exercising phase +// 1's tail-chunk handling. Reference is computed by a plain stable sort in the test +// itself: a stable descending sort naturally keeps equal values in ascending original +// index order, matching the tie-break rule the two-phase kernel implements explicitly. +TEST_F(SofieAlpakaTest, TopKBig) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t N = 4, D = 10007, K = 13; + + std::vector input(N * D); + for (std::size_t r = 0; r < N; ++r) + for (std::size_t j = 0; j < D; ++j) + input[r * D + j] = std::sin(j * 0.0173f + r * 1.7f) + 0.3f * std::cos(j * 0.0051f - static_cast(r)); + + std::vector expVal(N * K); + std::vector expInd(N * K); + for (std::size_t r = 0; r < N; ++r) { + std::vector> row(D); + for (std::size_t j = 0; j < D; ++j) row[j] = {input[r * D + j], static_cast(j)}; + std::stable_sort(row.begin(), row.end(), [](auto const &a, auto const &b) { return a.first > b.first; }); + for (std::size_t k = 0; k < K; ++k) { + expVal[r * K + k] = row[k].first; + expInd[r * K + k] = row[k].second; + } + } + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * D})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < N * D; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N * D})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto values_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * K})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * K})); + + { + SOFIE_TopKBig::Session session; + auto [values, indices] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + + for (std::size_t i = 0; i < N * K; ++i) { + EXPECT_LE(std::abs(val[i] - expVal[i]), TOLERANCE) << " i=" << i; + EXPECT_EQ(idx[i], expInd[i]) << " i=" << i; + } +} + +TEST_F(SofieAlpakaTest, TopKTies) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t N = 2, D = 40, K = 6; + + std::vector input(N * D, 5.f); + input[3] = 9.f; input[7] = 9.f; input[15] = 9.f; input[22] = 9.f; // row 0 + input[1] = 1.f; input[30] = 1.f; + for (std::size_t j = 0; j < D; ++j) input[D + j] = static_cast(j % 5); // row 1 + + std::vector expVal(N * K); + std::vector expInd(N * K); + for (std::size_t r = 0; r < N; ++r) { + std::vector> row(D); + for (std::size_t j = 0; j < D; ++j) row[j] = {input[r * D + j], static_cast(j)}; + std::stable_sort(row.begin(), row.end(), [](auto const &a, auto const &b) { return a.first > b.first; }); + for (std::size_t k = 0; k < K; ++k) { + expVal[r * K + k] = row[k].first; + expInd[r * K + k] = row[k].second; + } + } + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * D})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < N * D; ++i) input_ptr[i] = input[i]; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N * D})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto values_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * K})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * K})); + + { + SOFIE_TopKTies::Session session; + auto [values, indices] = session.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + + for (std::size_t i = 0; i < N * K; ++i) { + EXPECT_LE(std::abs(val[i] - expVal[i]), TOLERANCE) << " i=" << i; + EXPECT_EQ(idx[i], expInd[i]) << " i=" << i; + } +} + +namespace DynamicTopKHelpers { + // Stable descending sort per row: ties resolve to ascending original index + inline void referenceTopK(const std::vector &input, std::size_t N, std::size_t D, std::size_t topKCount, + std::vector &expVal, std::vector &expInd) { + expVal.assign(N * topKCount, 0.f); + expInd.assign(N * topKCount, 0); + for (std::size_t r = 0; r < N; ++r) { + std::vector> row(D); + for (std::size_t j = 0; j < D; ++j) row[j] = {input[r * D + j], static_cast(j)}; + std::stable_sort(row.begin(), row.end(), [](auto const &a, auto const &b) { return a.first > b.first; }); + for (std::size_t k = 0; k < topKCount; ++k) { + expVal[r * topKCount + k] = row[k].first; + expInd[r * topKCount + k] = row[k].second; + } + } + } +} + +// Regression test for the dynamic-axis two-phase GPU path: the axis itself (D) is a +// symbolic ONNX dim here, so P/chunk size/the clamped k are all computed at runtime +TEST_F(SofieAlpakaTest, DynamicTopK_SmallerThanRequestedK) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t N = 3, D = 3, K = 5, topKCount = std::min(K, D); + + std::vector input {5.f, 1.f, 9.f, -2.f, 4.f, 0.f, 7.f, 7.f, 3.f}; + std::vector expVal; + std::vector expInd; + DynamicTopKHelpers::referenceTopK(input, N, D, topKCount, expVal, expInd); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * D})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N * D; ++i) input_ptr[i] = input[i]; + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N * D})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto values_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * topKCount})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * topKCount})); + { + SOFIE_DynamicTopK::Session session("", N, D); + auto [values, indices] = session.infer(N, D, input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + for (std::size_t i = 0; i < N * topKCount; ++i) { + EXPECT_LE(std::abs(val[i] - expVal[i]), TOLERANCE) << " i=" << i; + EXPECT_EQ(idx[i], expInd[i]) << " i=" << i; + } +} + +TEST_F(SofieAlpakaTest, DynamicTopK_LargeAxis) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t N = 3, D = 6000, K = 5, topKCount = std::min(K, D); + + std::vector input(N * D); + for (std::size_t r = 0; r < N; ++r) + for (std::size_t j = 0; j < D; ++j) + input[r * D + j] = std::sin(j * 0.021f + r * 2.3f) + 0.2f * std::cos(j * 0.0037f - static_cast(r)); + + std::vector expVal; + std::vector expInd; + DynamicTopKHelpers::referenceTopK(input, N, D, topKCount, expVal, expInd); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * D})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < N * D; ++i) input_ptr[i] = input[i]; + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{N * D})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto values_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * topKCount})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * topKCount})); + { + SOFIE_DynamicTopK::Session session("", N, D); + auto [values, indices] = session.infer(N, D, input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + for (std::size_t i = 0; i < N * topKCount; ++i) { + EXPECT_LE(std::abs(val[i] - expVal[i]), TOLERANCE) << " i=" << i; + EXPECT_EQ(idx[i], expInd[i]) << " i=" << i; + } +} + +TEST_F(SofieAlpakaTest, DynamicTopK_ScratchBufferReuse) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t maxN = 4, maxD = 12000, K = 5; + const std::size_t capacity = maxN * std::min(K, maxD); + + SOFIE_DynamicTopK::Session session("", maxN, maxD); + struct Step { std::size_t N, D; }; + Step steps[] = { {4, 8000}, {2, 100}, {4, 8000}, {3, 12000} }; + + for (auto const &st : steps) { + std::size_t topKCount = std::min(K, st.D); + std::vector input(st.N * st.D); + for (std::size_t r = 0; r < st.N; ++r) + for (std::size_t j = 0; j < st.D; ++j) + input[r * st.D + j] = std::sin(j * 0.013f + r * 0.9f) + 0.1f * std::cos(j * 0.0021f + static_cast(r)); + + std::vector expVal; + std::vector expInd; + DynamicTopKHelpers::referenceTopK(input, st.N, st.D, topKCount, expVal, expInd); + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{st.N * st.D})); + float* input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (std::size_t i = 0; i < st.N * st.D; ++i) input_ptr[i] = input[i]; + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{st.N * st.D})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + // sized to the session's full construction-time capacit + auto values_h = alpaka::allocBuf(host, Ext1D::all(Idx{capacity})); + auto indices_h = alpaka::allocBuf(host, Ext1D::all(Idx{capacity})); + { + auto [values, indices] = session.infer(st.N, st.D, input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, values_h, values); + alpaka::memcpy(queue, indices_h, indices); + alpaka::wait(queue); + } + + float* val = reinterpret_cast(alpaka::getPtrNative(values_h)); + int64_t* idx = reinterpret_cast(alpaka::getPtrNative(indices_h)); + for (std::size_t i = 0; i < st.N * topKCount; ++i) { + EXPECT_LE(std::abs(val[i] - expVal[i]), TOLERANCE) << " N=" << st.N << " D=" << st.D << " i=" << i; + EXPECT_EQ(idx[i], expInd[i]) << " N=" << st.N << " D=" << st.D << " i=" << i; + } + } +} diff --git a/test/alpaka/TestAlpakaSequenceOps.cxx b/test/alpaka/TestAlpakaSequenceOps.cxx index 93859ad9..258ddd06 100644 --- a/test/alpaka/TestAlpakaSequenceOps.cxx +++ b/test/alpaka/TestAlpakaSequenceOps.cxx @@ -1,6 +1,7 @@ #include "TestAlpakaCommon.h" #include "SDPA_FromONNX_GPU_ALPAKA.hxx" +#include "SDPABig_FromONNX_GPU_ALPAKA.hxx" #include "MambaScan_FromONNX_GPU_ALPAKA.hxx" #include "RWKV_WKV6_FromONNX_GPU_ALPAKA.hxx" #include "GriffinRGLRU_FromONNX_GPU_ALPAKA.hxx" @@ -45,6 +46,75 @@ TEST_F(SofieAlpakaTest, SDPA) EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; } +TEST_F(SofieAlpakaTest, SDPABig) +{ + constexpr float TOLERANCE = 1e-3f; + constexpr std::size_t B = 1, H = 2, S = 16, D = 8; + constexpr std::size_t qkvSize = B * H * S * D; + constexpr std::size_t maskSize = B * H * S * S; + + unsigned rngState = 12345u; + auto rnd = [&]() { + rngState = 1103515245u * rngState + 12345u; + return ((rngState & 0x7fffffffu) / static_cast(0x7fffffffu)) * 2.0f - 1.0f; + }; + std::vector q(qkvSize), k(qkvSize), v(qkvSize), mask(maskSize); + for (std::size_t i = 0; i < qkvSize; ++i) { + q[i] = rnd(); + k[i] = rnd(); + v[i] = rnd(); + } + for (std::size_t i = 0; i < maskSize; ++i) + mask[i] = ((i % 5) == 0) ? -1e9f : 0.0f; // mask out every 5th key + + auto q_d = makeDeviceBuf(host, device, queue, q.data(), qkvSize); + auto k_d = makeDeviceBuf(host, device, queue, k.data(), qkvSize); + auto v_d = makeDeviceBuf(host, device, queue, v.data(), qkvSize); + auto mask_d = makeDeviceBuf(host, device, queue, mask.data(), maskSize); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{qkvSize})); + + { + SOFIE_SDPABig::Session session; + auto result = session.infer(q_d, k_d, v_d, mask_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + const float scale = 1.0f / std::sqrt(static_cast(D)); + std::vector expected(qkvSize); + for (std::size_t b = 0; b < B; ++b) + for (std::size_t h = 0; h < H; ++h) + for (std::size_t s = 0; s < S; ++s) { + std::vector scores(S); + double maxScore = -1e300; + for (std::size_t j = 0; j < S; ++j) { + double dot = 0; + for (std::size_t d = 0; d < D; ++d) + dot += q[b*H*S*D + h*S*D + s*D + d] * k[b*H*S*D + h*S*D + j*D + d]; + double sc = dot * scale + mask[b*H*S*S + h*S*S + s*S + j]; + scores[j] = sc; + maxScore = std::max(maxScore, sc); + } + double sumExp = 0; + for (std::size_t j = 0; j < S; ++j) { + scores[j] = std::exp(scores[j] - maxScore); + sumExp += scores[j]; + } + for (std::size_t d = 0; d < D; ++d) { + double acc = 0; + for (std::size_t j = 0; j < S; ++j) + acc += scores[j] * v[b*H*S*D + h*S*D + j*D + d]; + expected[b*H*S*D + h*S*D + s*D + d] = static_cast(acc / sumExp); + } + } + + const float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < qkvSize; ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; +} + TEST_F(SofieAlpakaTest, MambaScan) { constexpr float TOLERANCE = 1e-4f; diff --git a/test/alpaka/TestAlpakaShapeIndexing.cxx b/test/alpaka/TestAlpakaShapeIndexing.cxx index f445b3a7..c7b835b0 100644 --- a/test/alpaka/TestAlpakaShapeIndexing.cxx +++ b/test/alpaka/TestAlpakaShapeIndexing.cxx @@ -5,6 +5,7 @@ #include "Trilu_k2_FromONNX_GPU_ALPAKA.hxx" #include "Trilu_kn1_FromONNX_GPU_ALPAKA.hxx" #include "Trilu_3D_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicTrilu_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/Trilu_upper.ref.hxx" #include "input_models/references/Trilu_upper_input.ref.hxx" #include "input_models/references/Trilu_lower.ref.hxx" @@ -17,14 +18,20 @@ #include "input_models/references/Trilu_3D_input.ref.hxx" #include "Transpose_FromONNX_GPU_ALPAKA.hxx" #include "Concat_0D_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicExpand_FromONNX_GPU_ALPAKA.hxx" #include "ScatterElements_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicScatterElements_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicCumSum_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConstantOfShape_FromONNX_GPU_ALPAKA.hxx" #include "ScatterND_Ex1_FromONNX_GPU_ALPAKA.hxx" #include "ScatterND_Ex2_FromONNX_GPU_ALPAKA.hxx" #include "ScatterND_NegativeIndices_FromONNX_GPU_ALPAKA.hxx" #include "ScatterND_2D_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicScatterND_FromONNX_GPU_ALPAKA.hxx" #include "Split_0_FromONNX_GPU_ALPAKA.hxx" #include "Split_1_FromONNX_GPU_ALPAKA.hxx" #include "Split_2_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicSplit_FromONNX_GPU_ALPAKA.hxx" #include "Tile5D_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/Tile5D.ref.hxx" #include "GatherAxis0_FromONNX_GPU_ALPAKA.hxx" @@ -50,6 +57,7 @@ #include "GatherND_Ex5_FromONNX_GPU_ALPAKA.hxx" #include "GatherND_NegativeIndices_FromONNX_GPU_ALPAKA.hxx" #include "GatherND_Batch_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicGatherND_FromONNX_GPU_ALPAKA.hxx" #include "Slice_FromONNX_GPU_ALPAKA.hxx" #include "Slice_Default_Axis_FromONNX_GPU_ALPAKA.hxx" #include "Slice_Default_Steps_FromONNX_GPU_ALPAKA.hxx" @@ -59,6 +67,14 @@ #include "input_models/references/Slice_Default_Steps.ref.hxx" #include "input_models/references/Slice_Neg.ref.hxx" +#include "DynamicTranspose_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicConcat_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicTile_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicGather_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicSlice_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicRange_FromONNX_GPU_ALPAKA.hxx" +#include "DynamicRangeMul_FromONNX_GPU_ALPAKA.hxx" + TEST_F(SofieAlpakaTest, Transpose) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -127,6 +143,45 @@ TEST_F(SofieAlpakaTest, Transpose) EXPECT_LE(std::abs(res_ptr[i] - expected[i]), TOLERANCE); } +TEST_F(SofieAlpakaTest, DynamicTranspose) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t sz = N * C * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + { + SOFIE_DynamicTranspose::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t p = 0; p < P; ++p) + for (std::size_t c = 0; c < C; ++c) { + float expected = in_ptr[n * C * P + c * P + p]; + float got = res[n * P * C + p * C + c]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " p=" << p << " c=" << c; + } + } +} + TEST_F(SofieAlpakaTest, Concat0D) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -170,6 +225,51 @@ TEST_F(SofieAlpakaTest, Concat0D) } } +TEST_F(SofieAlpakaTest, DynamicConcat) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t Ca = 2, Cb = 3, Cy = 5; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t szA = N * Ca * P, szB = N * Cb * P, szY = N * Cy * P; + + auto a_h = alpaka::allocBuf(host, Ext1D::all(Idx{szA})); + auto b_h = alpaka::allocBuf(host, Ext1D::all(Idx{szB})); + float* a_ptr = reinterpret_cast(alpaka::getPtrNative(a_h)); + float* b_ptr = reinterpret_cast(alpaka::getPtrNative(b_h)); + for (Idx i = 0; i < szA; ++i) a_ptr[i] = static_cast(i % 7) - 3.0f; + for (Idx i = 0; i < szB; ++i) b_ptr[i] = static_cast(i % 5) + 10.0f; + + auto a_d = alpaka::allocBuf(device, Ext1D::all(Idx{szA})); + auto b_d = alpaka::allocBuf(device, Ext1D::all(Idx{szB})); + alpaka::memcpy(queue, a_d, a_h); + alpaka::memcpy(queue, b_d, b_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{szY})); + { + SOFIE_DynamicConcat::Session session("", N, P); + auto result = session.infer(N, P, a_d, b_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < Cy; ++j) + for (std::size_t p = 0; p < P; ++p) { + float expected = (j < Ca) ? a_ptr[n * Ca * P + j * P + p] + : b_ptr[n * Cb * P + (j - Ca) * P + p]; + float got = res[n * Cy * P + j * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " j=" << j << " p=" << p; + } + } +} + TEST_F(SofieAlpakaTest, ScatterElements) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -222,6 +322,103 @@ TEST_F(SofieAlpakaTest, ScatterElements) } } +TEST_F(SofieAlpakaTest, DynamicScatterElements) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 4; + + for (std::size_t N : {std::size_t(1), std::size_t(3)}) { + const std::size_t size = N * Cols; + std::vector input(size, 0.f); + std::vector indices(size); + std::vector updates(size); + for (std::size_t i = 0; i < size; ++i) { + indices[i] = static_cast((i + 1) % Cols); + updates[i] = static_cast(i) + 0.5f; + } + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), size); + auto indices_d = makeDeviceBuf(host, device, queue, indices.data(), size); + auto updates_d = makeDeviceBuf(host, device, queue, updates.data(), size); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + { + SOFIE_DynamicScatterElements::Session session("DynamicScatterElements_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d, indices_d, updates_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + std::vector expected = input; + for (std::size_t n = 0; n < N; ++n) + for (std::size_t c = 0; c < Cols; ++c) { + std::size_t srcIdx = n * Cols + c; + std::size_t dstIdx = n * Cols + static_cast(indices[srcIdx]); + expected[dstIdx] = updates[srcIdx]; + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < size; ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i << " N=" << N; + } +} + +TEST_F(SofieAlpakaTest, DynamicCumSum) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 5; + + for (std::size_t N : {std::size_t(1), std::size_t(4)}) { + const std::size_t size = N * Cols; + std::vector input(size); + for (std::size_t i = 0; i < size; ++i) input[i] = static_cast(i % 6) * 0.5f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), size); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + { + SOFIE_DynamicCumSum::Session session("DynamicCumSum_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) { + float acc = 0.f; + for (std::size_t c = 0; c < Cols; ++c) { + acc += input[n * Cols + c]; + EXPECT_LE(std::abs(res[n * Cols + c] - acc), TOLERANCE) << "n=" << n << " c=" << c << " N=" << N; + } + } + } +} + +TEST_F(SofieAlpakaTest, DynamicConstantOfShape) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 3; + + for (std::size_t N : {std::size_t(1), std::size_t(5)}) { + std::vector input(N * Cols, 0.0f); + auto input_d = makeDeviceBuf(host, device, queue, input.data(), input.size()); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * Cols})); + + { + SOFIE_DynamicConstantOfShape::Session session("DynamicConstantOfShape_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < N * Cols; ++i) + EXPECT_LE(std::abs(res[i] - 7.0f), TOLERANCE) << "i=" << i << " N=" << N; + } +} + // ── ScatterND GPU tests ──────────────────────────────────────────────────── TEST_F(SofieAlpakaTest, ScatterND_Ex1) @@ -349,6 +546,38 @@ TEST_F(SofieAlpakaTest, ScatterND_2D) EXPECT_LE(std::abs(res[i] - correct[i]), TOLERANCE) << " index=" << i; } +TEST_F(SofieAlpakaTest, DynamicScatterND) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 3; + std::vector updates = {5.f, 6.f, 7.f, 8.f, 9.f, 10.f}; + + for (std::size_t N : {std::size_t(2), std::size_t(5)}) { + std::vector data(N * Cols, 0.f); + auto data_d = makeDeviceBuf(host, device, queue, data.data(), data.size()); + auto updates_d = makeDeviceBuf(host, device, queue, updates.data(), updates.size()); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * Cols})); + { + SOFIE_DynamicScatterND::Session session("DynamicScatterND_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, data_d, updates_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + std::vector expected = data; + for (std::size_t c = 0; c < Cols; ++c) { + expected[(N - 1) * Cols + c] = updates[c]; // index -1 -> row N-1 + expected[c] = updates[Cols + c]; // index 0 + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < N * Cols; ++i) + EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i << " N=" << N; + } +} + TEST_F(SofieAlpakaTest, Split_0) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -467,6 +696,40 @@ TEST_F(SofieAlpakaTest, Split_2) EXPECT_LE(std::abs(res1_ptr[j] - correct_output[1][j]), TOLERANCE); } +TEST_F(SofieAlpakaTest, DynamicSplit) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 6, Split0 = 2, Split1 = 4; + + for (std::size_t N : {std::size_t(1), std::size_t(3)}) { + const std::size_t size = N * Cols; + std::vector input(size); + for (std::size_t i = 0; i < size; ++i) input[i] = static_cast(i) + 1.0f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), size); + auto result0_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * Split0})); + auto result1_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * Split1})); + + { + SOFIE_DynamicSplit::Session session("DynamicSplit_FromONNX_GPU_ALPAKA.dat", N); + auto [result0, result1] = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result0_h, result0); + alpaka::memcpy(queue, result1_h, result1); + alpaka::wait(queue); + } + + float* res0 = reinterpret_cast(alpaka::getPtrNative(result0_h)); + float* res1 = reinterpret_cast(alpaka::getPtrNative(result1_h)); + for (std::size_t n = 0; n < N; ++n) { + for (std::size_t c = 0; c < Split0; ++c) + EXPECT_LE(std::abs(res0[n * Split0 + c] - input[n * Cols + c]), TOLERANCE) << "n=" << n << " c=" << c << " N=" << N; + for (std::size_t c = 0; c < Split1; ++c) + EXPECT_LE(std::abs(res1[n * Split1 + c] - input[n * Cols + Split0 + c]), TOLERANCE) << "n=" << n << " c=" << c << " N=" << N; + } + } +} + TEST_F(SofieAlpakaTest, Tile5D) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -523,6 +786,43 @@ TEST_F(SofieAlpakaTest, Tile5D) EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); } +// X[N,2] -> Tile([2,3]) -> Y[2N,6], N dynamic +TEST_F(SofieAlpakaTest, DynamicTile) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t inCols = 2, outCols = 6; + + for (std::size_t N : {std::size_t(1), std::size_t(8)}) { + const std::size_t inRows = N, outRows = 2 * N; + const std::size_t inSize = inRows * inCols, outSize = outRows * outCols; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) + in_ptr[i] = static_cast(i + 1); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicTile::Session session("", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t r = 0; r < outRows; ++r) + for (std::size_t c = 0; c < outCols; ++c) { + float expected = in_ptr[(r % inRows) * inCols + (c % inCols)]; + EXPECT_LE(std::abs(res[r * outCols + c] - expected), TOLERANCE); + } + } +} + TEST_F(SofieAlpakaTest, GatherAxis0) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -727,6 +1027,72 @@ TEST_F(SofieAlpakaTest, GatherNegativeIndices) EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE); } +TEST_F(SofieAlpakaTest, DynamicGather) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 3, K = 2; + const std::size_t idxs[2] = {2, 0}; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * C * P, outSize = N * K * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicGather::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t i = 0; i < K; ++i) + for (std::size_t p = 0; p < P; ++p) { + float expected = in_ptr[n * C * P + idxs[i] * P + p]; + float got = res[n * K * P + i * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " i=" << i << " p=" << p; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicExpand) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 4; + + for (std::size_t N : {std::size_t(1), std::size_t(4)}) { + std::vector input(N); + for (std::size_t i = 0; i < N; ++i) input[i] = static_cast(i) + 1.0f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), N); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{N * Cols})); + { + SOFIE_DynamicExpand::Session session("DynamicExpand_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t c = 0; c < Cols; ++c) + EXPECT_LE(std::abs(res[n * Cols + c] - input[n]), TOLERANCE) << "n=" << n << " c=" << c << " N=" << N; + } +} + TEST_F(SofieAlpakaTest, ExpandSameSize) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -1022,6 +1388,33 @@ TEST_F(SofieAlpakaTest, GatherND_Batch) EXPECT_LE(std::abs(res[i] - expected[i]), TOLERANCE) << "i=" << i; } +TEST_F(SofieAlpakaTest, DynamicGatherND) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Cols = 3; + + for (std::size_t N : {std::size_t(2), std::size_t(5)}) { + std::vector data(N * Cols); + for (std::size_t i = 0; i < data.size(); ++i) data[i] = static_cast(i) + 1.0f; + + auto data_d = makeDeviceBuf(host, device, queue, data.data(), data.size()); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{2 * Cols})); + { + SOFIE_DynamicGatherND::Session session("DynamicGatherND_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, data_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t c = 0; c < Cols; ++c) { + EXPECT_LE(std::abs(res[c] - data[(N - 1) * Cols + c]), TOLERANCE) << "row -1, c=" << c << " N=" << N; + EXPECT_LE(std::abs(res[Cols + c] - data[c]), TOLERANCE) << "row 0, c=" << c << " N=" << N; + } + } +} + TEST_F(SofieAlpakaTest, Slice) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; @@ -1154,6 +1547,45 @@ TEST_F(SofieAlpakaTest, Slice_Neg) EXPECT_LE(std::abs(res_ptr[i] - correct[i]), TOLERANCE) << "i=" << i; } +TEST_F(SofieAlpakaTest, DynamicSlice) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + const std::size_t C = 4, K = 2, start = 1; + + const std::size_t Ns[] = {1, 8}; + const std::size_t Ps[] = {1, 5}; // n_pf + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], P = Ps[t]; + const std::size_t inSize = N * C * P, outSize = N * K * P; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{inSize})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inSize; ++i) in_ptr[i] = static_cast(i % 7) - 3.0f; + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{inSize})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{outSize})); + { + SOFIE_DynamicSlice::Session session("", N, P); + auto result = session.infer(N, P, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t j = 0; j < K; ++j) + for (std::size_t p = 0; p < P; ++p) { + float expected = in_ptr[n * C * P + (start + j) * P + p]; + float got = res[n * K * P + j * P + p]; + EXPECT_LE(std::abs(got - expected), TOLERANCE) << "n=" << n << " j=" << j << " p=" << p; + } + } +} + TEST_F(SofieAlpakaTest, Trilu_upper) { constexpr std::size_t N = 16; // 4×4 @@ -1267,3 +1699,96 @@ TEST_F(SofieAlpakaTest, Trilu_3D) for (std::size_t i = 0; i < N; ++i) EXPECT_NEAR(res[i], ref[i], DEFAULT_TOLERANCE) << " index=" << i; } + +TEST_F(SofieAlpakaTest, DynamicTrilu) +{ + constexpr float TOLERANCE = DEFAULT_TOLERANCE; + constexpr std::size_t Rows = 4, Cols = 4; + + for (std::size_t N : {std::size_t(1), std::size_t(3)}) { + const std::size_t size = N * Rows * Cols; + std::vector input(size); + for (std::size_t i = 0; i < size; ++i) input[i] = static_cast(i % 13) + 1.0f; + + auto input_d = makeDeviceBuf(host, device, queue, input.data(), size); + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{size})); + { + SOFIE_DynamicTrilu::Session session("DynamicTrilu_FromONNX_GPU_ALPAKA.dat", N); + auto result = session.infer(N, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + float* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t n = 0; n < N; ++n) + for (std::size_t r = 0; r < Rows; ++r) + for (std::size_t c = 0; c < Cols; ++c) { + std::size_t idx = n * Rows * Cols + r * Cols + c; + float expected = (c >= r) ? input[idx] : 0.0f; + EXPECT_LE(std::abs(res[idx] - expected), TOLERANCE) << "n=" << n << " r=" << r << " c=" << c << " N=" << N; + } + } +} + +TEST_F(SofieAlpakaTest, DynamicRange) +{ + const std::size_t Ns[] = {1, 4}; + const std::size_t Ks[] = {3, 7}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], K = Ks[t]; + const std::size_t sz = N * K; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + { + SOFIE_DynamicRange::Session session("", N, K); + auto result = session.infer(N, K, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int64_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < K; ++i) + EXPECT_EQ(res[i], static_cast(i)) << "i=" << i << " N=" << N << " K=" << K; + } +} + +TEST_F(SofieAlpakaTest, DynamicRangeMul) +{ + const std::size_t Ns[] = {1, 5}; + const std::size_t Ks[] = {4, 6}; + for (int t = 0; t < 2; ++t) { + const std::size_t N = Ns[t], K = Ks[t]; + const std::size_t sz = N * K; + + auto input_h = alpaka::allocBuf(host, Ext1D::all(Idx{sz})); + float* in_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < sz; ++i) in_ptr[i] = static_cast(i); + + auto input_d = alpaka::allocBuf(device, Ext1D::all(Idx{sz})); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(Idx{K})); + { + SOFIE_DynamicRangeMul::Session session("", N, K); + auto result = session.infer(N, K, input_d); + cudaDeviceSynchronize(); + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + int64_t* res = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (std::size_t i = 0; i < K; ++i) + EXPECT_EQ(res[i], static_cast(i * N)) << "i=" << i << " N=" << N << " K=" << K; + } +} diff --git a/test/input_models/DynamicAddBroadcast.onnx b/test/input_models/DynamicAddBroadcast.onnx new file mode 100644 index 00000000..12b7cbcb Binary files /dev/null and b/test/input_models/DynamicAddBroadcast.onnx differ diff --git a/test/input_models/DynamicBatchNorm2D.onnx b/test/input_models/DynamicBatchNorm2D.onnx new file mode 100644 index 00000000..a8219ddd Binary files /dev/null and b/test/input_models/DynamicBatchNorm2D.onnx differ diff --git a/test/input_models/DynamicBatchNorm4D.onnx b/test/input_models/DynamicBatchNorm4D.onnx new file mode 100644 index 00000000..8aecf1a2 Binary files /dev/null and b/test/input_models/DynamicBatchNorm4D.onnx differ diff --git a/test/input_models/DynamicBatchNormDynSpatialRelu.onnx b/test/input_models/DynamicBatchNormDynSpatialRelu.onnx new file mode 100644 index 00000000..626e3fad Binary files /dev/null and b/test/input_models/DynamicBatchNormDynSpatialRelu.onnx differ diff --git a/test/input_models/DynamicConcat.onnx b/test/input_models/DynamicConcat.onnx new file mode 100644 index 00000000..bef89046 Binary files /dev/null and b/test/input_models/DynamicConcat.onnx differ diff --git a/test/input_models/DynamicConstantOfShape.onnx b/test/input_models/DynamicConstantOfShape.onnx new file mode 100644 index 00000000..7915f73c Binary files /dev/null and b/test/input_models/DynamicConstantOfShape.onnx differ diff --git a/test/input_models/DynamicConv1D.onnx b/test/input_models/DynamicConv1D.onnx new file mode 100644 index 00000000..b97f8f86 Binary files /dev/null and b/test/input_models/DynamicConv1D.onnx differ diff --git a/test/input_models/DynamicConv1DNoBias.onnx b/test/input_models/DynamicConv1DNoBias.onnx new file mode 100644 index 00000000..387ebd41 Binary files /dev/null and b/test/input_models/DynamicConv1DNoBias.onnx differ diff --git a/test/input_models/DynamicConv2DNoBias.onnx b/test/input_models/DynamicConv2DNoBias.onnx new file mode 100644 index 00000000..86480de8 Binary files /dev/null and b/test/input_models/DynamicConv2DNoBias.onnx differ diff --git a/test/input_models/DynamicCumSum.onnx b/test/input_models/DynamicCumSum.onnx new file mode 100644 index 00000000..763a2b05 Binary files /dev/null and b/test/input_models/DynamicCumSum.onnx differ diff --git a/test/input_models/DynamicEqual.onnx b/test/input_models/DynamicEqual.onnx new file mode 100644 index 00000000..8627ba11 Binary files /dev/null and b/test/input_models/DynamicEqual.onnx differ diff --git a/test/input_models/DynamicExpand.onnx b/test/input_models/DynamicExpand.onnx new file mode 100644 index 00000000..861fd55f Binary files /dev/null and b/test/input_models/DynamicExpand.onnx differ diff --git a/test/input_models/DynamicGather.onnx b/test/input_models/DynamicGather.onnx new file mode 100644 index 00000000..b9ee64e9 Binary files /dev/null and b/test/input_models/DynamicGather.onnx differ diff --git a/test/input_models/DynamicGatherND.onnx b/test/input_models/DynamicGatherND.onnx new file mode 100644 index 00000000..1baf5b72 Binary files /dev/null and b/test/input_models/DynamicGatherND.onnx differ diff --git a/test/input_models/DynamicLinear.onnx b/test/input_models/DynamicLinear.onnx new file mode 100644 index 00000000..b4526797 Binary files /dev/null and b/test/input_models/DynamicLinear.onnx differ diff --git a/test/input_models/DynamicMatMulBatched.onnx b/test/input_models/DynamicMatMulBatched.onnx new file mode 100644 index 00000000..e8c32a1c Binary files /dev/null and b/test/input_models/DynamicMatMulBatched.onnx differ diff --git a/test/input_models/DynamicMaxPool.onnx b/test/input_models/DynamicMaxPool.onnx new file mode 100644 index 00000000..d5c4e17f Binary files /dev/null and b/test/input_models/DynamicMaxPool.onnx differ diff --git a/test/input_models/DynamicNegRelu.onnx b/test/input_models/DynamicNegRelu.onnx new file mode 100644 index 00000000..6c07396d Binary files /dev/null and b/test/input_models/DynamicNegRelu.onnx differ diff --git a/test/input_models/DynamicRMSNorm.onnx b/test/input_models/DynamicRMSNorm.onnx new file mode 100644 index 00000000..419f2f84 Binary files /dev/null and b/test/input_models/DynamicRMSNorm.onnx differ diff --git a/test/input_models/DynamicRange.onnx b/test/input_models/DynamicRange.onnx new file mode 100644 index 00000000..10cf9578 Binary files /dev/null and b/test/input_models/DynamicRange.onnx differ diff --git a/test/input_models/DynamicRangeMul.onnx b/test/input_models/DynamicRangeMul.onnx new file mode 100644 index 00000000..cc6a9622 Binary files /dev/null and b/test/input_models/DynamicRangeMul.onnx differ diff --git a/test/input_models/DynamicReduceMaxFirst.onnx b/test/input_models/DynamicReduceMaxFirst.onnx new file mode 100644 index 00000000..13fa588d Binary files /dev/null and b/test/input_models/DynamicReduceMaxFirst.onnx differ diff --git a/test/input_models/DynamicReduceMeanMid.onnx b/test/input_models/DynamicReduceMeanMid.onnx new file mode 100644 index 00000000..cb0c0cb3 Binary files /dev/null and b/test/input_models/DynamicReduceMeanMid.onnx differ diff --git a/test/input_models/DynamicReduceSumLast.onnx b/test/input_models/DynamicReduceSumLast.onnx new file mode 100644 index 00000000..03f7bd26 Binary files /dev/null and b/test/input_models/DynamicReduceSumLast.onnx differ diff --git a/test/input_models/DynamicReduceSumMulti.onnx b/test/input_models/DynamicReduceSumMulti.onnx new file mode 100644 index 00000000..0a4dd9f3 Binary files /dev/null and b/test/input_models/DynamicReduceSumMulti.onnx differ diff --git a/test/input_models/DynamicScatterElements.onnx b/test/input_models/DynamicScatterElements.onnx new file mode 100644 index 00000000..c6d7e7ba Binary files /dev/null and b/test/input_models/DynamicScatterElements.onnx differ diff --git a/test/input_models/DynamicScatterND.onnx b/test/input_models/DynamicScatterND.onnx new file mode 100644 index 00000000..05bda628 Binary files /dev/null and b/test/input_models/DynamicScatterND.onnx differ diff --git a/test/input_models/DynamicSlice.onnx b/test/input_models/DynamicSlice.onnx new file mode 100644 index 00000000..a5451af3 Binary files /dev/null and b/test/input_models/DynamicSlice.onnx differ diff --git a/test/input_models/DynamicSplit.onnx b/test/input_models/DynamicSplit.onnx new file mode 100644 index 00000000..34eb4ad4 Binary files /dev/null and b/test/input_models/DynamicSplit.onnx differ diff --git a/test/input_models/DynamicTile.onnx b/test/input_models/DynamicTile.onnx new file mode 100644 index 00000000..b1f557a6 Binary files /dev/null and b/test/input_models/DynamicTile.onnx differ diff --git a/test/input_models/DynamicTopK.onnx b/test/input_models/DynamicTopK.onnx new file mode 100644 index 00000000..68a5c432 Binary files /dev/null and b/test/input_models/DynamicTopK.onnx differ diff --git a/test/input_models/DynamicTranspose.onnx b/test/input_models/DynamicTranspose.onnx new file mode 100644 index 00000000..3f534524 Binary files /dev/null and b/test/input_models/DynamicTranspose.onnx differ diff --git a/test/input_models/DynamicTrilu.onnx b/test/input_models/DynamicTrilu.onnx new file mode 100644 index 00000000..de574b9f Binary files /dev/null and b/test/input_models/DynamicTrilu.onnx differ diff --git a/test/input_models/DynamicWhere.onnx b/test/input_models/DynamicWhere.onnx new file mode 100644 index 00000000..be90b900 Binary files /dev/null and b/test/input_models/DynamicWhere.onnx differ diff --git a/test/input_models/GroupNormBig.onnx b/test/input_models/GroupNormBig.onnx new file mode 100644 index 00000000..c2f43545 Binary files /dev/null and b/test/input_models/GroupNormBig.onnx differ diff --git a/test/input_models/LayerNormBig.onnx b/test/input_models/LayerNormBig.onnx new file mode 100644 index 00000000..60f0a1c4 Binary files /dev/null and b/test/input_models/LayerNormBig.onnx differ diff --git a/test/input_models/SDPABig.onnx b/test/input_models/SDPABig.onnx new file mode 100644 index 00000000..8259200a Binary files /dev/null and b/test/input_models/SDPABig.onnx differ diff --git a/test/input_models/TopKBig.onnx b/test/input_models/TopKBig.onnx new file mode 100644 index 00000000..cb42bded Binary files /dev/null and b/test/input_models/TopKBig.onnx differ diff --git a/test/input_models/TopKTies.onnx b/test/input_models/TopKTies.onnx new file mode 100644 index 00000000..9861245f Binary files /dev/null and b/test/input_models/TopKTies.onnx differ