Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-08-14 09:29:34

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 << res << " = std::max({ " << inputs[0];
0027       for (size_t i = 1; i < inputs.size(); i++) {
0028          out << ", " << inputs[i];
0029       }
0030       out << "});\n";
0031       return out.str();
0032    }
0033 };
0034 
0035 template<typename T>
0036 struct NaryOperatorTraits<T, EBasicNaryOperator::Min> {
0037    static const std::string Name() {return "Min";}
0038    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0039       std::stringstream out;
0040        out << res << " = std::min({ " << inputs[0];
0041       for (size_t i = 1; i < inputs.size(); i++) {
0042          out << ", " << inputs[i];
0043       }
0044       out << "});\n";
0045       return out.str();
0046    }
0047 };
0048 
0049 template<typename T>
0050 struct NaryOperatorTraits<T, EBasicNaryOperator::Mean> {};
0051 
0052 template<>
0053 struct NaryOperatorTraits<float, EBasicNaryOperator::Mean> {
0054    static const std::string Name() {return "Mean";}
0055    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0056       std::stringstream out;
0057       out << res << " = (" << inputs[0];
0058       for (size_t i = 1; i < inputs.size(); i++) {
0059          out << " + " << inputs[i];
0060       }
0061       out << ") / float(" << inputs.size() << ");\n";
0062       return out.str();
0063    }
0064 };
0065 
0066 template<typename T>
0067 struct NaryOperatorTraits<T, EBasicNaryOperator::Sum> {
0068    static const std::string Name() {return "Sum";}
0069    static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
0070       std::stringstream out;
0071       out << res << " = " << inputs[0];
0072       for (size_t i = 1; i < inputs.size(); i++) {
0073          out << " + " << inputs[i];
0074       }
0075       out << ";\n";
0076       return out.str();
0077    }
0078 };
0079 
0080 template <typename T, EBasicNaryOperator Op>
0081 class ROperator_BasicNary final : public ROperator
0082 {
0083 
0084 private:
0085 
0086    std::vector<std::string> fNInputs;
0087    std::string fNY;
0088    std::vector<std::vector<Dim>> fShapeInputs;
0089 
0090    std::vector<std::string> fNBroadcastedInputs;
0091    std::vector<size_t> fShapeY;
0092    std::vector<Dim> fDimShapeY;
0093 
0094    bool fBroadcast = false;
0095 
0096    std::string fType;
0097 
0098 public:
0099    ROperator_BasicNary(){}
0100 
0101    ROperator_BasicNary( const std::vector<std::string> & inputNames, const std::string& nameY):
0102    fNY(UTILITY::Clean_name(nameY)){
0103       fNInputs.reserve(inputNames.size());
0104       for (auto & name : inputNames)
0105          fNInputs.push_back(UTILITY::Clean_name(name));
0106 
0107       fInputTensorNames.resize(fNInputs.size());
0108       std::transform(fNInputs.begin(), fNInputs.end(), fInputTensorNames.begin(),
0109                   [](const std::string& s) -> std::string_view { return s; });
0110       fOutputTensorNames = { fNY };
0111    }
0112 
0113    // type of output given input
0114    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0115       return input;
0116    }
0117 
0118    // shape of output tensors given input tensors
0119    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0120       auto ret = std::vector<std::vector<size_t>>(1, input[0]);
0121       return ret;
0122    }
0123 
0124    void Initialize(RModel& model) override {
0125       std::vector<std::vector<size_t>> inputShapes;
0126       for (auto &it : fNInputs) {
0127          if (!model.CheckIfTensorAlreadyExist(it)) {
0128             throw std::runtime_error("TMVA SOFIE BasicNary Op Input Tensor " + it + " is not found in model");
0129          }
0130          fShapeInputs.push_back(model.GetDimTensorShape(it));
0131          if (fNInputs.size()> 2) {
0132             if (model.IsDimInputTensor(it))
0133                throw std::runtime_error("TMVA SOFIE BasicNary : supports only 2 inputs for dynamic tensors");
0134             else
0135                inputShapes.push_back(model.GetTensorShape(it));
0136          }
0137       }
0138       // Find the common shape of the input tensors
0139       if (fShapeInputs.size() > 2 ) {
0140          // support dynamic tensors now for input list of size=2
0141          auto shapeY = UTILITY::MultidirectionalBroadcastShape(inputShapes);
0142          fDimShapeY = ConvertShapeToDim(shapeY);
0143       } else if (fShapeInputs.size() == 2 ) {
0144          auto ret  = UTILITY::MultidirectionalBroadcastShape(fShapeInputs[0], fShapeInputs[1]);
0145          // use same code as in BinaryOperator (need to extend for input sizes > 2)
0146          fBroadcast = ret.first;
0147          fDimShapeY = ret.second;
0148          // case of all parametric shapes and MultiDirectionalBroadcastShape  return the max of the 2
0149          // need to do before we declare the output tensor shape and the broadcasted ones
0150          if (ret.first & 4) {
0151             // check if one of the parameter is an input dimension
0152             // define function to find this
0153             auto IsInputDimParam = [&](const std::string &p) {
0154                auto inputNames = model.GetInputTensorNames();
0155                for (auto &input : inputNames) {
0156                   for (auto &i_s : model.GetDimTensorShape(input)) {
0157                      if (i_s.isParam && i_s.param == p)
0158                         return true;
0159                   }
0160                }
0161                return false;
0162             };
0163             auto & shapeA = fShapeInputs[0];
0164             auto & shapeB = fShapeInputs[1];
0165             for (size_t i = 0; i < fDimShapeY.size(); i++) {
0166                auto &s = fDimShapeY[i];
0167                if (s.isParam && s.param.find("std::max") != std::string::npos) {
0168                   if (IsInputDimParam(shapeA[i].param)) {
0169                      // case dim is 1 we indicate that the input parameter is equal to 1
0170                      if (shapeA[i].dim != 1)
0171                         s = shapeA[i];
0172                      else
0173                         s = shapeB[i];
0174                   } else if (IsInputDimParam(shapeB[i].param)) {
0175                      if (shapeB[i].dim != 1)
0176                         s = shapeB[i];
0177                      else
0178                         s = shapeA[i];
0179                   }
0180                }
0181             }
0182          }
0183       } else if  (fShapeInputs.size() == 1 ) {
0184          fDimShapeY = fShapeInputs[0];
0185       }
0186       if (!fShapeY.empty())
0187          model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fShapeY);
0188       else
0189          model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fDimShapeY);
0190 
0191 
0192       fType = ConvertTypeToString(model.GetTensorType(fNInputs[0]));
0193 
0194       if (model.Verbose()) {
0195          std::cout << NaryOperatorTraits<T, Op>::Name() << " : ";
0196          if (fNInputs.size() == 2)
0197             std::cout << ConvertDimShapeToString(fShapeInputs[0]) << " , "
0198                       << ConvertDimShapeToString(fShapeInputs[1]);
0199          std::cout << " --> " << ConvertDimShapeToString(fDimShapeY) << std::endl;
0200       }
0201    }
0202 
0203    std::string Generate(std::string OpName) override {
0204       OpName = "op_" + OpName;
0205       if (fDimShapeY.empty()) {
0206          throw std::runtime_error("TMVA SOFIE BasicNary called to Generate without being initialized first");
0207       }
0208       std::stringstream out;
0209       auto length = ConvertDimShapeToLength(fDimShapeY);
0210       out << SP << "\n//------ BasicNary operator\n";
0211 
0212       int nInputs = fNInputs.size();
0213 
0214       if (nInputs == 1) {
0215          out << SP << "std::copy(tensor_" << fNInputs[0] << ", tensor_" << fNInputs[0] << " + ";
0216          out << length << ", tensor_" << fNY << ");\n";
0217       } else {
0218 
0219          // implement operator without broadcasting, but using loos on all indices
0220          std::vector<std::vector<Dim>> inputStrides(nInputs);
0221          for (int i = 0; i < nInputs; i++)
0222             inputStrides[i] = UTILITY::ComputeStrideFromShape(fShapeInputs[i]);
0223 
0224          auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY);
0225 
0226          // make loop on output indices
0227          std::string compute_idx_Y;
0228          int nloop = 0;
0229          if (fDimShapeY.empty() ||
0230                std::all_of(fDimShapeY.begin(), fDimShapeY.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
0231             compute_idx_Y = "0";
0232          } else {
0233             for (size_t i = 0; i < fDimShapeY.size(); ++i) {
0234                if (fDimShapeY[i].dim != 1 && fDimShapeY[i].GetVal() != "1") {
0235                   nloop++;
0236                   for (int j = 0; j < nloop; j++) out << SP;
0237                   out << "for (size_t idx_" << i << " = 0; idx_" << i << " < " << fDimShapeY[i]
0238                       << "; ++idx_" << i << "){\n";
0239                   compute_idx_Y += "idx_" + std::to_string(i);
0240                   if (stridesY[i].GetVal() != "1")
0241                      compute_idx_Y += " * " + stridesY[i].GetVal();
0242                   compute_idx_Y += " + ";
0243                }
0244             }
0245             // remove last 3 characters " + "
0246             for (int j = 0; j < 3; j++)
0247                compute_idx_Y.pop_back();
0248          }
0249          // find indices for input tensors
0250          std::vector<std::string> inputs(nInputs);
0251          for (int ipt = 0; ipt < nInputs; ipt++ ) {
0252             std::string compute_idx_X;
0253             auto & shape = fShapeInputs[ipt];
0254             auto & stride = inputStrides[ipt];
0255             if (shape.empty() ||
0256                 std::all_of(shape.begin(), shape.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
0257                compute_idx_X = "0";
0258             } else {
0259                for (size_t i = 0; i < shape.size(); ++i) {
0260                   if (shape[i].dim == 1 || shape[i].GetVal() == "1")
0261                      continue;
0262                   compute_idx_X += "idx_" + std::to_string(i + (fDimShapeY.size() - shape.size()));
0263                   if (stride[i].GetVal() != "1")
0264                      compute_idx_X += " * " + stride[i].GetVal();
0265                   compute_idx_X += " + ";
0266                }
0267                // remove last 3 character " + "
0268                for (int j = 0; j < 3; j++)
0269                   compute_idx_X.pop_back();
0270             }
0271             inputs[ipt] = "tensor_" + fNInputs[ipt] + "[" + compute_idx_X + "]";
0272          }
0273 
0274          // perform the operation
0275          for (int j = 0; j < nloop + 1; j++) out << SP;
0276          std::string output = "tensor_" + fNY + "[" + compute_idx_Y + "]";
0277          out << NaryOperatorTraits<T,Op>::Op(output, inputs);
0278 
0279          for (int i = nloop; i > 0; i--) {
0280             for (int j = 0; j < i; j++) out << SP;
0281             out << "}\n";
0282          }
0283       }
0284       return out.str();
0285    }
0286 
0287    std::vector<std::string> GetStdLibs() override {return { std::string("cmath") }; }
0288 };
0289 
0290 }//SOFIE
0291 }//Experimental
0292 }//TMVA
0293 
0294 
0295 #endif //TMVA_SOFIE_ROPERATOR_BasicNary