root/TMVA/ROperator_BasicNary.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_BASICNARY
0002 #define TMVA_SOFIE_ROPERATOR_BASICNARY
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <vector>
0009 #include <sstream>
0010 #include <algorithm>
0011
0012 namespace TMVA{
0013 namespace Experimental{
0014 namespace SOFIE{
0015
0016 enum class EBasicNaryOperator {Max, Min, Mean, Sum};
0017
0018 template<typename T, EBasicNaryOperator Op>
0019 struct NaryOperatorTraits {};
0020
0021 template<typename T>
0022 struct NaryOperatorTraits<T, EBasicNaryOperator::Max> {
0023    static const std::string Name() {return "Max";}
0024    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0025       std::stringstream out;
0026       out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
0027       for (size_t i = 1; i < inputs.size(); i++) {
0028          out << "\t" << "\t" << res << " = std::max(" << res << ", " << inputs[i] << ");\n";
0029       }
0030       return out.str();
0031    }
0032 };
0033
0034 template<typename T>
0035 struct NaryOperatorTraits<T, EBasicNaryOperator::Min> {
0036    static const std::string Name() {return "Min";}
0037    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0038       std::stringstream out;
0039       out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
0040       for (size_t i = 1; i < inputs.size(); i++) {
0041          out << "\t" << "\t" << res << " = std::min(" << res << ", " << inputs[i] << ");\n";
0042       }
0043       return out.str();
0044    }
0045 };
0046
0047 template<typename T>
0048 struct NaryOperatorTraits<T, EBasicNaryOperator::Mean> {};
0049
0050 template<>
0051 struct NaryOperatorTraits<float, EBasicNaryOperator::Mean> {
0052    static const std::string Name() {return "Mean";}
0053    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0054       std::stringstream out;
0055       out << "\t" << "\t" << res << " = (" << inputs[0];
0056       for (size_t i = 1; i < inputs.size(); i++) {
0057          out << " + " << inputs[i];
0058       }
0059       out << ") / float(" << inputs.size() << ");\n";
0060       return out.str();
0061    }
0062 };
0063
0064 template<typename T>
0065 struct NaryOperatorTraits<T, EBasicNaryOperator::Sum> {
0066    static const std::string Name() {return "Sum";}
0067    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0068       std::stringstream out;
0069       out << "\t" << "\t" << res << " = " << inputs[0];
0070       for (size_t i = 1; i < inputs.size(); i++) {
0071          out << " + " << inputs[i];
0072       }
0073       out << ";\n";
0074       return out.str();
0075    }
0076 };
0077
0078 template <typename T, EBasicNaryOperator Op>
0079 class ROperator_BasicNary final : public ROperator
0080 {
0081
0082 private:
0083
0084    std::vector<std::string> fNInputs;
0085    std::string fNY;
0086    std::vector<std::vector<size_t>> fShapeInputs;
0087
0088    std::vector<std::string> fNBroadcastedInputs;
0089    std::vector<size_t> fShapeY;
0090
0091    bool fBroadcast = false;
0092
0093    std::string fType;
0094
0095 public:
0096    ROperator_BasicNary(){}
0097
0098    ROperator_BasicNary( const std::vector<std::string> & inputNames, const std::string& nameY):
0099    fNY(UTILITY::Clean_name(nameY)){
0100       fNInputs.reserve(inputNames.size());
0101       for (auto & name : inputNames)
0102          fNInputs.push_back(UTILITY::Clean_name(name));
0103
0104       fInputTensorNames.resize(fNInputs.size());
0105       std::transform(fNInputs.begin(), fNInputs.end(), fInputTensorNames.begin(),
0106                   [](const std::string& s) -> std::string_view { return s; });
0107       fOutputTensorNames = { fNY };
0108    }
0109
0110    // type of output given input
0111    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0112       return input;
0113    }
0114
0115    // shape of output tensors given input tensors
0116    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0117       auto ret = std::vector<std::vector<size_t>>(1, input[0]);
0118       return ret;
0119    }
0120
0121    void Initialize(RModel& model) override {
0122       for (auto &it : fNInputs) {
0123          if (!model.CheckIfTensorAlreadyExist(it)) {
0124             throw std::runtime_error("TMVA SOFIE BasicNary Op Input Tensor " + it + " is not found in model");
0125          }
0126          fShapeInputs.push_back(model.GetTensorShape(it));
0127       }
0128       // Find the common shape of the input tensors
0129       fShapeY = UTILITY::MultidirectionalBroadcastShape(fShapeInputs);
0130       model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fShapeY);
0131       // Broadcasting
0132       size_t N = fNInputs.size();
0133       fNBroadcastedInputs.reserve(N);
0134       for (size_t i = 0; i < N; i++) {
0135          if (!UTILITY::AreSameShape(model.GetTensorShape(fNInputs[i]), fShapeY)) {
0136             fBroadcast = true;
0137             std::string name = "Broadcasted"  + fNInputs[i];
0138             model.AddIntermediateTensor(name, model.GetTensorType(fNInputs[0]), fShapeY);
0139             fNBroadcastedInputs.emplace_back("tensor_" + name);
0140          } else {
0141             fNBroadcastedInputs.emplace_back("tensor_" + fNInputs[i]);
0142          }
0143       }
0144       fType = ConvertTypeToString(model.GetTensorType(fNInputs[0]));
0145    }
0146
0147    std::string Generate(std::string OpName) override {
0148       OpName = "op_" + OpName;
0149       if (fShapeY.empty()) {
0150          throw std::runtime_error("TMVA SOFIE BasicNary called to Generate without being initialized first");
0151       }
0152       std::stringstream out;
0153       size_t length = ConvertShapeToLength(fShapeY);
0154       out << SP << "\n//------ BasicNary operator\n";
0155       if (fBroadcast) {
0156          for (size_t i = 0; i < fNInputs.size(); i++) {
0157             if (fNBroadcastedInputs[i] != fNInputs[i]) {
0158                out << SP << SP << "// Broadcasting " << fNInputs[i] << " to " << ConvertShapeToString(fShapeY) << "\n";
0159                out << SP << SP << "{\n";
0160                out << SP << SP << SP << fType << "* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << fType << ">(tensor_" + fNInputs[i] << ", " << ConvertShapeToString(fShapeInputs[i]);
0161                out << ", " << ConvertShapeToString(fShapeY) << ");\n";
0162                out << SP << SP << SP << "std::copy(data, data + " << length << ", " << fNBroadcastedInputs[i] << ");\n";
0163                out << SP << SP << SP << "delete[] data;\n";
0164                out << SP << SP << "}\n";
0165             }
0166          }
0167       }
0168
0169       if (fNInputs.size() == 1) {
0170          out << SP << "std::copy(tensor_" << fNInputs[0] << ", tensor_" << fNInputs[0] << " + ";
0171          out << length << ", tensor_" << fNY << ");\n";
0172       } else {
0173          std::vector<std::string> inputs(fNBroadcastedInputs.size());
0174          for (size_t i = 0; i < fNBroadcastedInputs.size(); i++) {
0175             inputs[i] = fNBroadcastedInputs[i] + "[id]";
0176          }
0177          out << SP << "for (size_t id = 0; id < " << length << "; id++) {\n";
0178          out << NaryOperatorTraits<T,Op>::Op("tensor_" + fNY + "[id]", inputs);
0179          out << SP << "}\n";
0180       }
0181       return out.str();
0182    }
0183
0184    std::vector<std::string> GetStdLibs() override {return { std::string("cmath") }; }
0185 };
0186
0187 }//SOFIE
0188 }//Experimental
0189 }//TMVA
0190
0191
0192 #endif //TMVA_SOFIE_ROPERATOR_BasicNary