ROperator_BasicNary.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_ROPERATOR_BASICNARY
#define TMVA_SOFIE_ROPERATOR_BASICNARY
 
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TMVA/RModel.hxx"
 
#include <vector>
#include <sstream>
#include <algorithm>
 
namespace TMVA{
namespace Experimental{
namespace SOFIE{
 
enum class EBasicNaryOperator {Max, Min, Mean, Sum};
 
template<typename T, EBasicNaryOperator Op>
struct NaryOperatorTraits {};
 
template<typename T>
struct NaryOperatorTraits<T, EBasicNaryOperator::Max> {
   static const std::string Name() {return "Max";}
   static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
      std::stringstream out;
      out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
      for (size_t i = 1; i < inputs.size(); i++) {
         out << "\t" << "\t" << res << " = std::max(" << res << ", " << inputs[i] << ");\n";
      }
      return out.str();
   }
};
 
template<typename T>
struct NaryOperatorTraits<T, EBasicNaryOperator::Min> {
   static const std::string Name() {return "Min";}
   static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
      std::stringstream out;
      out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
      for (size_t i = 1; i < inputs.size(); i++) {
         out << "\t" << "\t" << res << " = std::min(" << res << ", " << inputs[i] << ");\n";
      }
      return out.str();
   }
};
 
template<typename T>
struct NaryOperatorTraits<T, EBasicNaryOperator::Mean> {};
 
template<>
struct NaryOperatorTraits<float, EBasicNaryOperator::Mean> {
   static const std::string Name() {return "Mean";}
   static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
      std::stringstream out;
      out << "\t" << "\t" << res << " = (" << inputs[0];
      for (size_t i = 1; i < inputs.size(); i++) {
         out << " + " << inputs[i];
      }
      out << ") / float(" << inputs.size() << ");\n";
      return out.str();
   }
};
 
template<typename T>
struct NaryOperatorTraits<T, EBasicNaryOperator::Sum> {
   static const std::string Name() {return "Sum";}
   static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
      std::stringstream out;
      out << "\t" << "\t" << res << " = " << inputs[0];
      for (size_t i = 1; i < inputs.size(); i++) {
         out << " + " << inputs[i];
      }
      out << ";\n";
      return out.str();
   }
};
 
template <typename T, EBasicNaryOperator Op>
class ROperator_BasicNary final : public ROperator
{
 
private:
 
   std::vector<std::string> fNInputs;
   std::string fNY;
   std::vector<std::vector<size_t>> fShapeInputs;
 
   std::vector<std::string> fNBroadcastedInputs;
   std::vector<size_t> fShapeY;
 
   bool fBroadcast = false;
 
   std::string fType;
 
public:
   ROperator_BasicNary(){}
 
   ROperator_BasicNary( const std::vector<std::string> & inputNames, const std::string& nameY):
   fNY(UTILITY::Clean_name(nameY)){
      fNInputs.reserve(inputNames.size());
      for (auto & name : inputNames)
         fNInputs.push_back(UTILITY::Clean_name(name));
 
      fInputTensorNames.resize(fNInputs.size());
      std::transform(fNInputs.begin(), fNInputs.end(), fInputTensorNames.begin(),
                  [](const std::string& s) -> std::string_view { return s; });
      fOutputTensorNames = { fNY };
   }
 
   // type of output given input
   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
      return input;
   }
 
   // shape of output tensors given input tensors
   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
      auto ret = std::vector<std::vector<size_t>>(1, input[0]);
      return ret;
   }
 
   void Initialize(RModel& model) override {
      for (auto &it : fNInputs) {
         if (!model.CheckIfTensorAlreadyExist(it)) {
            throw std::runtime_error("TMVA SOFIE BasicNary Op Input Tensor " + it + " is not found in model");
         }
         fShapeInputs.push_back(model.GetTensorShape(it));
      }
      // Find the common shape of the input tensors
      fShapeY = UTILITY::MultidirectionalBroadcastShape(fShapeInputs);
      model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fShapeY);
      // Broadcasting
      size_t N = fNInputs.size();
      fNBroadcastedInputs.reserve(N);
      for (size_t i = 0; i < N; i++) {
         if (!UTILITY::AreSameShape(model.GetTensorShape(fNInputs[i]), fShapeY)) {
            fBroadcast = true;
            std::string name = "Broadcasted"  + fNInputs[i];
            model.AddIntermediateTensor(name, model.GetTensorType(fNInputs[0]), fShapeY);
            fNBroadcastedInputs.emplace_back("tensor_" + name);
         } else {
            fNBroadcastedInputs.emplace_back("tensor_" + fNInputs[i]);
         }
      }
      fType = ConvertTypeToString(model.GetTensorType(fNInputs[0]));
   }
 
   std::string Generate(std::string OpName) override {
      OpName = "op_" + OpName;
      if (fShapeY.empty()) {
         throw std::runtime_error("TMVA SOFIE BasicNary called to Generate without being initialized first");
      }
      std::stringstream out;
      size_t length = ConvertShapeToLength(fShapeY);
      out << SP << "\n//------ BasicNary operator\n";
      if (fBroadcast) {
         for (size_t i = 0; i < fNInputs.size(); i++) {
            if (fNBroadcastedInputs[i] != fNInputs[i]) {
               out << SP << SP << "// Broadcasting " << fNInputs[i] << " to " << ConvertShapeToString(fShapeY) << "\n";
               out << SP << SP << "{\n";
               out << SP << SP << SP << fType << "* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << fType << ">(tensor_" + fNInputs[i] << ", " << ConvertShapeToString(fShapeInputs[i]);
               out << ", " << ConvertShapeToString(fShapeY) << ");\n";
               out << SP << SP << SP << "std::copy(data, data + " << length << ", " << fNBroadcastedInputs[i] << ");\n";
               out << SP << SP << SP << "delete[] data;\n";
               out << SP << SP << "}\n";
            }
         }
      }
 
      if (fNInputs.size() == 1) {
         out << SP << "std::copy(tensor_" << fNInputs[0] << ", tensor_" << fNInputs[0] << " + ";
         out << length << ", tensor_" << fNY << ");\n";
      } else {
         std::vector<std::string> inputs(fNBroadcastedInputs.size());
         for (size_t i = 0; i < fNBroadcastedInputs.size(); i++) {
            inputs[i] = fNBroadcastedInputs[i] + "[id]";
         }
         out << SP << "for (size_t id = 0; id < " << length << "; id++) {\n";
         out << NaryOperatorTraits<T,Op>::Op("tensor_" + fNY + "[id]", inputs);
         out << SP << "}\n";
      }
      return out.str();
   }
 
   std::vector<std::string> GetStdLibs() override {return { std::string("cmath") }; }
};
 
}//SOFIE
}//Experimental
}//TMVA
 
 
#endif //TMVA_SOFIE_ROPERATOR_BasicNary