root/TMVA/ROperator_Gelu.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_GELU
0002 #define TMVA_SOFIE_ROPERATOR_GELU
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <sstream>
0009
0010 namespace TMVA{
0011 namespace Experimental{
0012 namespace SOFIE{
0013
0014 class ROperator_Gelu final : public ROperator
0015 {
0016
0017 private:
0018
0019    std::string fNX;
0020    std::string fNY;
0021    std::string fApproximate; // "none" (exact) or "tanh" (approximate)
0022    std::vector<size_t> fShape;
0023
0024 public:
0025    ROperator_Gelu(){}
0026    ROperator_Gelu(std::string nameX, std::string nameY, std::string approximate = "none"):
0027       fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)), fApproximate(approximate){
0028          fInputTensorNames = { fNX };
0029          fOutputTensorNames = { fNY };
0030       }
0031
0032
0033    void Initialize(RModel& model) override {
0034       if (model.CheckIfTensorAlreadyExist(fNX) == false){
0035          throw std::runtime_error("TMVA SOFIE Gelu Op Input Tensor is not found in model");
0036       }
0037       fShape = model.GetTensorShape(fNX);
0038       model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShape);
0039    }
0040
0041    std::string Generate(std::string OpName) override {
0042       OpName = "op_" + OpName;
0043       if (fShape.empty()){
0044          throw std::runtime_error("TMVA SOFIE Operator Gelu called to Generate without being initialized first");
0045       }
0046       std::stringstream out;
0047       int length = 1;
0048       for(auto& i: fShape){
0049          length *= i;
0050       }
0051       out << SP << "for (int id = 0; id < " << length << " ; id++){\n";
0052       if (fApproximate == "tanh") {
0053          // Tanh approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
0054          out << SP << SP << "float x = tensor_" << fNX << "[id];\n";
0055          out << SP << SP << "tensor_" << fNY << "[id] = 0.5f * x * (1.0f + std::tanh(0.7978845608028654f * (x + 0.044715f * x * x * x)));\n";
0056       } else {
0057          // Exact: 0.5 * x * (1 + erf(x / sqrt(2)))
0058          out << SP << SP << "tensor_" << fNY << "[id] = 0.5f * tensor_" << fNX << "[id] * (1.0f + std::erf(tensor_" << fNX << "[id] * 0.7071067811865475f));\n";
0059       }
0060       out << SP << "}\n";
0061       return out.str();
0062    }
0063
0064    std::vector<std::string> GetStdLibs() override { return { std::string("cmath") };}
0065 };
0066
0067 }//SOFIE
0068 }//Experimental
0069 }//TMVA
0070
0071
0072 #endif //TMVA_SOFIE_ROPERATOR_GELU