Add Alpaka backend support for HardSigmoid, HardSwish, and Softplus

core/CMakeLists.txt

@@ -21,6 +21,8 @@ set(sources_headers
    SOFIE/RModelProfilerGPU.hxx
    SOFIE/ROperator.hxx
    SOFIE/ROperator_BasicUnary.hxx
+    ROperator_HardSigmoid.hxx
+    ROperator_HardSwish.hxx
    SOFIE/ROperator_BasicBinary.hxx
    SOFIE/ROperator_BasicNary.hxx
    SOFIE/ROperator_BatchNormalization.hxx

core/inc/SOFIE/ROperator.hxx

@@ -39,7 +39,10 @@ enum class OperatorKind {
    UNARY_ABS=23,
    CLIP=24,
    NOT=25,
-   POOL=26
+   POOL=26,
+   HARDSIGMOID=27,
+   HARDSWISH=28,
+   SOFTPLUS=29
 };
 
 inline const char* toString(OperatorKind kind) {
@@ -52,6 +55,9 @@ inline const char* toString(OperatorKind kind) {
        case OperatorKind::BATCHNORM:       return "BATCHNORM";  
        case OperatorKind::CONV:       return "CONV";
        case OperatorKind::UNDEFINED:  return "UNDEFINED";
+       case OperatorKind::HARDSIGMOID:return "HARDSIGMOID";
+       case OperatorKind::HARDSWISH:  return "HARDSWISH";
+       case OperatorKind::SOFTPLUS:   return "SOFTPLUS";
        default:                       return "UNKNOWN";
    }
 }

core/inc/SOFIE/ROperator_BasicUnary.hxx

@@ -65,7 +65,7 @@ struct UnaryOpTraits<T, EBasicUnaryOperator::kAbs> {
 template <typename T>
 struct UnaryOpTraits<T, EBasicUnaryOperator::kSoftplus> {
    static std::string Name() { return "Softplus"; }
-   static std::string Op(const std::string &X) { return "std::log(std::exp(" + X + ") + 1)"; }
+   static std::string Op(const std::string &X) { return "((" + X + " >= 0x1.4000000000000p+4f) ? " + X + " : std::log1p(std::exp(" + X + ")))"; }
 };
 
 template <typename T>
@@ -121,6 +121,9 @@ public:
             case EBasicUnaryOperator::kAbs:
                fKind = OperatorKind::UNARY_ABS;
                break;
+            case EBasicUnaryOperator::kSoftplus:
+               fKind = OperatorKind::SOFTPLUS;
+               break;
          }
          fInputTensorNames =  { fNX };
          fOutputTensorNames = { fNY };

core/inc/SOFIE/ROperator_HardSigmoid.hxx

@@ -0,0 +1,101 @@
+#ifndef SOFIE_ROPERATOR_HARDSIGMOID
+#define SOFIE_ROPERATOR_HARDSIGMOID
+
+#include <SOFIE/SOFIE_common.hxx>
+#include <SOFIE/ROperator.hxx>
+#include <SOFIE/RModel.hxx>
+
+#include <sstream>
+
+namespace SOFIE {
+
+template <typename T>
+class ROperator_HardSigmoid final : public ROperator
+{
+
+private:
+
+   std::string fNX;
+   std::string fNY;
+   std::vector<size_t> fShape;
+   float fAlpha;
+   float fBeta;
+
+public:
+   ROperator_HardSigmoid(){}
+   ROperator_HardSigmoid(std::string nameX, std::string nameY, float alpha, float beta):
+      fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)), fAlpha(alpha), fBeta(beta){
+         fInputTensorNames = { fNX };
+         fOutputTensorNames = { fNY };
+         fKind = OperatorKind::HARDSIGMOID;
+      }
+
+   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
+      return input;
+   }
+
+   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
+      return input;
+   }
+
+   void Initialize(RModel& model) override {
+      if (!model.CheckIfTensorAlreadyExist(fNX)){
+         throw std::runtime_error("SOFIE HardSigmoid Op Input Tensor " + fNX + " is not found in model");
+      }
+      fShape = model.GetTensorShape(fNX);
+      model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShape);
+   }
+
+   std::string Generate(std::string OpName) override {
+      OpName = "op_" + OpName;
+      if (fShape.empty()){
+         throw std::runtime_error("SOFIE HardSigmoid operator called to Generate without being initialized first");
+      }
+      std::stringstream out;
+      size_t length = ConvertShapeToLength(fShape);
+
+      // HardSigmoid: y = max(0, min(1, alpha * x + beta))
+      out << "\n//------ HardSigmoid\n";
+      out << SP << "for (int id = 0; id < " << length << " ; id++){\n";
+      out << SP << SP << "tensor_" << fNY << "[id] = std::fmax(0x0p+0f, std::fmin(0x1p+0f, "
+          << fAlpha << "f * tensor_" << fNX << "[id] + " << fBeta << "f));\n";
+      out << SP << "}\n";
+      return out.str();
+   }
+   std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) override {
+      std::string op = "\n//------ HARDSIGMOID_KERNEL_ALPAKA\n";
+      op += SP + "struct HardSigmoidKernel{\n";
+      op += SP + SP + "template<typename TAcc, typename T>\n";
+      op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * data, T * out, std::size_t numElements, T const alpha, T const beta) const {\n";
+      op += SP + SP + SP + "const auto idx = alpaka::getIdx<alpaka::Grid, alpaka::Threads>(acc)[0];\n";
+      op += SP + SP + SP + "if (idx < numElements) {\n";
+      op += SP + SP + SP + SP + "T x = data[idx];\n";
+      op += SP + SP + SP + SP + "T h = alpha * x + beta;\n";
+      op += SP + SP + SP + SP + "out[idx] = (h < T(0)) ? T(0) : ((h > T(1)) ? T(1) : h);\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 + "HardSigmoidKernel hardSigmoidKernel;\n";
+   }
+
+   std::string Generate_GPU_ALPAKA(std::string OpName) override {
+      std::stringstream out;
+      auto length = ConvertShapeToLength(fShape);
+      out << "\n//------ op_" << OpName << "_ALPAKA\n";
+      out << SP << "auto const elementsPerThread_" << fNX << " = alpaka::Vec<Dim, Idx>::all(static_cast<Idx>(1));\n";
+      out << SP << "auto const elementsPerGrid_" << fNX << " = alpaka::Vec<Dim, Idx>::all(static_cast<Idx>(" << length << "));\n";
+      out << SP << "auto const workDiv_" << fNX << " = sofie_workdiv(elementsPerGrid_" << fNX << ");\n";
+      out << SP << "auto task_op_" << OpName << " = alpaka::createTaskKernel<Acc>(workDiv_" << fNX << ", hardSigmoidKernel, alpaka::getPtrNative(deviceBuf_" << fNX << "), alpaka::getPtrNative(deviceBuf_" << fNY << "), static_cast<std::size_t>(" << length << "), static_cast<float>(" << fAlpha << "), static_cast<float>(" << fBeta << "));\n";
+      out << SP << "alpaka::enqueue(queue, task_op_" << OpName << ");\n";
+      return out.str();
+   }
+   std::vector<std::string> GetStdLibs() override { return { std::string("cmath") };}
+};
+
+} // namespace SOFIE
+
+#endif

core/inc/SOFIE/ROperator_HardSwish.hxx

@@ -0,0 +1,103 @@
+#ifndef SOFIE_ROPERATOR_HARDSWISH
+#define SOFIE_ROPERATOR_HARDSWISH
+
+#include <SOFIE/SOFIE_common.hxx>
+#include <SOFIE/ROperator.hxx>
+#include <SOFIE/RModel.hxx>
+
+#include <sstream>
+
+namespace SOFIE {
+
+template <typename T>
+class ROperator_HardSwish final : public ROperator
+{
+
+private:
+
+   std::string fNX;
+   std::string fNY;
+   std::vector<size_t> fShape;
+
+public:
+   ROperator_HardSwish(){}
+   ROperator_HardSwish(std::string nameX, std::string nameY):
+      fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)){
+         fInputTensorNames = { fNX };
+         fOutputTensorNames = { fNY };
+         fKind = OperatorKind::HARDSWISH;
+      }
+
+   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
+      return input;
+   }
+
+   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
+      return input;
+   }
+
+   void Initialize(RModel& model) override {
+      if (!model.CheckIfTensorAlreadyExist(fNX)){
+         throw std::runtime_error("SOFIE HardSwish Op Input Tensor " + fNX + " is not found in model");
+      }
+      fShape = model.GetTensorShape(fNX);
+      model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShape);
+   }
+
+   std::string Generate(std::string OpName) override {
+      OpName = "op_" + OpName;
+      if (fShape.empty()){
+         throw std::runtime_error("SOFIE HardSwish operator called to Generate without being initialized first");
+      }
+      std::stringstream out;
+      size_t length = ConvertShapeToLength(fShape);
+
+      // HardSwish: y = x * max(0, min(1, x/6 + 0.5))
+      // Split topology for debuggability
+      out << "\n//------ HardSwish\n";
+      out << SP << "for (int id = 0; id < " << length << " ; id++){\n";
+      out << SP << SP << "float h = 0x1.5555555555555p-3f * tensor_" << fNX << "[id] + 0x1p-1f;\n";
+      out << SP << SP << "tensor_" << fNY << "[id] = tensor_" << fNX
+          << "[id] * std::fmax(0x0p+0f, std::fmin(0x1p+0f, h));\n";
+      out << SP << "}\n";
+      return out.str();
+   }
+
+   std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) override {
+      std::string op = "\n//------ HARDSWISH_KERNEL_ALPAKA\n";
+      op += SP + "struct HardSwishKernel{\n";
+      op += SP + SP + "template<typename TAcc, typename T>\n";
+      op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * data, T * out, std::size_t numElements) const {\n";
+      op += SP + SP + SP + "const auto idx = alpaka::getIdx<alpaka::Grid, alpaka::Threads>(acc)[0];\n";
+      op += SP + SP + SP + "if (idx < numElements) {\n";
+      op += SP + SP + SP + SP + "T x = data[idx];\n";
+      op += SP + SP + SP + SP + "T h = T(0x1.5555555555555p-3) * x + T(0.5);\n";
+      op += SP + SP + SP + SP + "out[idx] = x * ((h < T(0)) ? T(0) : ((h > T(1)) ? T(1) : h));\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 + "HardSwishKernel hardSwishKernel;\n";
+   }
+
+   std::string Generate_GPU_ALPAKA(std::string OpName) override {
+      std::stringstream out;
+      auto length = ConvertShapeToLength(fShape);
+      out << "\n//------ op_" << OpName << "_ALPAKA\n";
+      out << SP << "auto const elementsPerThread_" << fNX << " = alpaka::Vec<Dim, Idx>::all(static_cast<Idx>(1));\n";
+      out << SP << "auto const elementsPerGrid_" << fNX << " = alpaka::Vec<Dim, Idx>::all(static_cast<Idx>(" << length << "));\n";
+      out << SP << "auto const workDiv_" << fNX << " = sofie_workdiv(elementsPerGrid_" << fNX << ");\n";
+      out << SP << "auto task_op_" << OpName << " = alpaka::createTaskKernel<Acc>(workDiv_" << fNX << ", hardSwishKernel, alpaka::getPtrNative(deviceBuf_" << fNX << "), alpaka::getPtrNative(deviceBuf_" << fNY << "), static_cast<std::size_t>(" << length << "));\n";
+      out << SP << "alpaka::enqueue(queue, task_op_" << OpName << ");\n";
+      return out.str();
+   }
+
+   std::vector<std::string> GetStdLibs() override { return { std::string("cmath") };}
+};
+
+} // namespace SOFIE
+
+#endif

parsers/CMakeLists.txt

@@ -33,6 +33,8 @@ target_include_directories(SOFIE_parsers
 set(sources_cxx
     src/RModelParser_ONNX.cxx
     src/ParseBasicUnary.cxx
+    ParseHardSigmoid.cxx
+    ParseHardSwish.cxx
     src/ParseBasicBinary.cxx
     src/ParseBasicIs.cxx
     src/ParseBatchNormalization.cxx

parsers/src/ParseHardSigmoid.cxx

@@ -0,0 +1,47 @@
+#include "SOFIE/RModelParser_ONNX.hxx"
+#include "SOFIE/ROperator_HardSigmoid.hxx"
+#include "onnx_proto3.pb.h"
+
+namespace SOFIE {
+
+ParserFuncSignature ParseHardSigmoid = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) {
+   ETensorType input_type;
+
+   // ONNX spec defaults: alpha=0.2, beta=0.5
+   float alpha = 0.2f;
+   float beta = 0.5f;
+
+   for (int_t i = 0; i < nodeproto.attribute_size(); i++) {
+      std::string attribute_name = nodeproto.attribute(i).name();
+      if (attribute_name == "alpha")
+         alpha = nodeproto.attribute(i).f();
+      else if (attribute_name == "beta")
+         beta = nodeproto.attribute(i).f();
+   }
+
+   auto input_name = nodeproto.input(0);
+   if (parser.IsRegisteredTensorType(input_name)) {
+      input_type = parser.GetTensorType(input_name);
+   } else {
+      throw std::runtime_error("TMVA::SOFIE ONNX Parser HardSigmoid op has input tensor " + input_name +
+                               " but its type is not yet registered");
+   }
+
+   std::unique_ptr<ROperator> op;
+   std::string output_name = nodeproto.output(0);
+
+   switch (input_type) {
+   case ETensorType::FLOAT: op.reset(new ROperator_HardSigmoid<float>(input_name, output_name, alpha, beta)); break;
+   default:
+      throw std::runtime_error("TMVA::SOFIE - Unsupported - Operator HardSigmoid does not yet support input type " +
+                               std::to_string(static_cast<int>(input_type)));
+   }
+
+   if (!parser.IsRegisteredTensorType(output_name)) {
+      parser.RegisterTensorType(output_name, input_type);
+   }
+
+   return op;
+};
+
+} // namespace SOFIE

parsers/src/ParseHardSwish.cxx

@@ -0,0 +1,35 @@
+#include "SOFIE/RModelParser_ONNX.hxx"
+#include "SOFIE/ROperator_HardSwish.hxx"
+#include "onnx_proto3.pb.h"
+
+namespace SOFIE {
+
+ParserFuncSignature ParseHardSwish = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) {
+   ETensorType input_type;
+
+   auto input_name = nodeproto.input(0);
+   if (parser.IsRegisteredTensorType(input_name)) {
+      input_type = parser.GetTensorType(input_name);
+   } else {
+      throw std::runtime_error("TMVA::SOFIE ONNX Parser HardSwish op has input tensor " + input_name +
+                               " but its type is not yet registered");
+   }
+
+   std::unique_ptr<ROperator> op;
+   std::string output_name = nodeproto.output(0);
+
+   switch (input_type) {
+   case ETensorType::FLOAT: op.reset(new ROperator_HardSwish<float>(input_name, output_name)); break;
+   default:
+      throw std::runtime_error("TMVA::SOFIE - Unsupported - Operator HardSwish does not yet support input type " +
+                               std::to_string(static_cast<int>(input_type)));
+   }
+
+   if (!parser.IsRegisteredTensorType(output_name)) {
+      parser.RegisterTensorType(output_name, input_type);
+   }
+
+   return op;
+};
+
+} // namespace SOFIE

parsers/src/RModelParser_ONNX.cxx

@@ -200,6 +200,8 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un
    RegisterOperator("Cos", ParseCos);
    RegisterOperator("Abs", ParseAbs);
    RegisterOperator("Softplus", ParseSoftplus);
+   RegisterOperator("HardSigmoid", ParseHardSigmoid);
+   RegisterOperator("HardSwish", ParseHardSwish);
    RegisterOperator("Atan", ParseAtan);
    RegisterOperator("Floor", ParseFloor);