Back to home page

EIC code displayed by LXR
 
 

    

File indexing completed on 2026-08-23 09:25:43 

0001 #ifndef TMVA_SOFIE_ROperator_BasicBinary
0002 #define TMVA_SOFIE_ROperator_BasicBinary
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 enum EBasicBinaryOperator { Add, Sub, Mul, Div, Pow, Mod, FMod };
0015 
0016 template <typename T, EBasicBinaryOperator Op1>
0017 struct BinaryOperatorTrait {};
0018 
0019 template <typename T>
0020 struct BinaryOperatorTrait<T, Add> {
0021    static const std::string Name() { return "Add"; }
0022    static std::string Op(const std::string &t1, const std::string t2) { return t1 + " + " + t2; }
0023    static T Func(T t1, T t2) { return t1 + t2; }
0024 };
0025 
0026 template <typename T>
0027 struct BinaryOperatorTrait<T, Sub> {
0028    static const std::string Name() { return "Sub"; }
0029    static std::string Op(const std::string &t1, const std::string t2) { return t1 + " - " + t2; }
0030    static T Func(T t1, T t2) { return t1 - t2; }
0031 };
0032 
0033 template <typename T>
0034 struct BinaryOperatorTrait<T, Mul> {
0035    static const std::string Name() { return "Mul"; }
0036    static std::string Op(const std::string &t1, const std::string t2) { return t1 + " * " + t2; }
0037    static T Func(T t1, T t2) { return t1 * t2; }
0038 };
0039 
0040 template <typename T>
0041 struct BinaryOperatorTrait<T, Div> {
0042    static const std::string Name() { return "Div"; }
0043    static std::string Op(const std::string &t1, const std::string t2) { return t1 + " / " + t2; }
0044    static T Func(T t1, T t2) { return t1 / t2; }
0045 };
0046 
0047 template <typename T>
0048 struct BinaryOperatorTrait<T, Pow> {
0049    static const std::string Name() { return "Pow"; }
0050    static std::string Op(const std::string &t1, const std::string t2) { return "std::pow(" + t1 + "," + t2 + ")"; }
0051    static T Func(T t1, T t2) { return std::pow(t1, t2); }
0052 };
0053 template <typename T>
0054 struct BinaryOperatorTrait<T, Mod> {
0055    static const std::string Name() { return "Mod"; }
0056    static std::string Op(const std::string & t1, const std::string t2) { return "(" + t1 + " % " + t2 + ")"; }
0057    static T Func(T t1, T t2) { return t1 % t2; }
0058 };
0059 template <typename T>
0060 struct BinaryOperatorTrait<T, FMod> {
0061    static const std::string Name() { return "FMod"; }
0062    static std::string Op(const std::string & t1, const std::string t2) { return "std::fmod(" + t1 + "," + t2 + ")"; }
0063    static T Func(T t1, T t2) { return std::fmod(t1, t2); }
0064 };
0065 
0066 template <typename T, EBasicBinaryOperator Op>
0067 class ROperator_BasicBinary final : public ROperator {
0068 private:
0069    int fBroadcastFlag = 0;
0070    std::string fNA;
0071    std::string fNB;
0072    std::string fNBroadcastedA;
0073    std::string fNBroadcastedB;
0074    std::string fNY;
0075 
0076    std::vector<size_t> fShapeA;
0077    std::vector<size_t> fShapeB;
0078    std::vector<size_t> fShapeY;
0079 
0080    std::vector<Dim> fDimShapeA;
0081    std::vector<Dim> fDimShapeB;
0082    std::vector<Dim> fDimShapeY;
0083 
0084 public:
0085    ROperator_BasicBinary() {}
0086    ROperator_BasicBinary(std::string nameA, std::string nameB, std::string nameY)
0087       : fNA(UTILITY::Clean_name(nameA)), fNB(UTILITY::Clean_name(nameB)), fNY(UTILITY::Clean_name(nameY))
0088    {
0089       fInputTensorNames = {fNA, fNB};
0090       fOutputTensorNames = {fNY};
0091    }
0092 
0093    // type of output given input
0094    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override { return input; }
0095 
0096    // shape of output tensors given input tensors
0097    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override
0098    {
0099       // assume now inputs have same shape (no broadcasting)
0100       auto ret = std::vector<std::vector<size_t>>(1, input[0]); // return vector size 1 with first input
0101       return ret;
0102    }
0103 
0104    void Initialize(RModel &model) override
0105    {
0106       // input must be a graph input, or already initialized intermediate tensor
0107       if (!model.CheckIfTensorAlreadyExist(fNA)) {
0108          throw std::runtime_error(std::string("TMVA SOFIE Binary Op Input Tensor ") + fNA + "is not found in model");
0109       }
0110       if (!model.CheckIfTensorAlreadyExist(fNB)) {
0111          throw std::runtime_error(std::string("TMVA SOFIE Binary Op Input Tensor ") + fNB + "is not found in model");
0112       }
0113       int dynamicInputs = 0;
0114       if (model.IsDynamicTensor(fNA)) {
0115          fDimShapeA = model.GetDynamicTensorShape(fNA);
0116          dynamicInputs |= 1;
0117       } else {
0118          fShapeA = model.GetTensorShape(fNA);
0119          fDimShapeA = ConvertShapeToDim(fShapeA);
0120       }
0121       if (model.IsDynamicTensor(fNB)) {
0122          dynamicInputs |= 2;
0123          fDimShapeB = model.GetDynamicTensorShape(fNB);
0124       } else {
0125          fShapeB = model.GetTensorShape(fNB);
0126          fDimShapeB = ConvertShapeToDim(fShapeB);
0127       }
0128       if (dynamicInputs & 1 && model.Verbose())
0129          std::cout << BinaryOperatorTrait<T, Op>::Name() << " : input " << fNA << " is dynamic "
0130                    << ConvertDimShapeToString(fDimShapeA) << std::endl;
0131       if (dynamicInputs & 2 && model.Verbose())
0132          std::cout << BinaryOperatorTrait<T, Op>::Name() << " : input " << fNB << " is dynamic "
0133                    << ConvertDimShapeToString(fDimShapeB) << std::endl;
0134 
0135       // check if need to broadcast at initialization time if shapes are known and different
0136       // (we could broadcast the tensor tensor to maximum values of dynamic shapes - to be done)
0137       // case of known shapes
0138       // if shapes are known find the output shape from broadcasting
0139       if (dynamicInputs == 0) {
0140          auto ret = UTILITY::MultidirectionalBroadcastShape(fShapeA, fShapeB);
0141          fBroadcastFlag = ret.first;
0142          fShapeY = ret.second;
0143          auto  lengthY = ConvertShapeToLength(fShapeY);
0144          if (model.IsConstantTensor(fNA) && model.IsConstantTensor(fNB)) {
0145             bool broadcast = fBroadcastFlag > 0;
0146             if (broadcast) {
0147                // Y is the common shape of A and B
0148                bool broadcastA = fBroadcastFlag & 2;
0149                bool broadcastB = fBroadcastFlag & 1;
0150                // Broadcast A to Y
0151                if (broadcastA) {
0152                   fNBroadcastedA = "Broadcasted" + fNA + "to" + fNY;
0153                   auto data = model.GetInitializedTensorData(fNA);
0154                   std::shared_ptr<void> broadcastedData(
0155                      UTILITY::UnidirectionalBroadcast(static_cast<T *>(data.get()), fShapeA, fShapeY),
0156                      std::default_delete<T[]>());
0157                   if (model.Verbose())
0158                      std::cout << "broadcasted data A " << ConvertShapeToString(fShapeY) << " : "
0159                                << ConvertValuesToString(ConvertShapeToLength(fShapeY),
0160                                                         static_cast<T *>(broadcastedData.get()))
0161                                << std::endl;
0162                   // Update the data and the shape of A
0163                   model.AddConstantTensor(fNBroadcastedA, model.GetTensorType(fNA), fShapeY, broadcastedData);
0164                   fShapeA = fShapeY;
0165                   fDimShapeA = ConvertShapeToDim(fShapeA);
0166                }
0167                // Broadcast B to Y
0168                if (broadcastB) {
0169                   fNBroadcastedB = "Broadcasted" + fNB + "to" + fNY;
0170                   auto data = model.GetInitializedTensorData(fNB);
0171                   if (model.Verbose())
0172                      std::cout << "data B " << ConvertShapeToString(fShapeB) << " : "
0173                                << ConvertValuesToString(ConvertShapeToLength(fShapeB), static_cast<T *>(data.get()))
0174                                << std::endl;
0175                   std::shared_ptr<void> broadcastedData(
0176                      UTILITY::UnidirectionalBroadcast(static_cast<T *>(data.get()), fShapeB, fShapeY),
0177                      std::default_delete<T[]>());
0178                   // do not update tensor B but add broadcasted one (since it can be input to some other operators)
0179                   if (model.Verbose())
0180                      std::cout << "broadcasted data B " << ConvertShapeToString(fShapeY) << " : "
0181                                << ConvertValuesToString(ConvertShapeToLength(fShapeY),
0182                                                         static_cast<T *>(broadcastedData.get()))
0183                                << std::endl;
0184                   model.AddConstantTensor(fNBroadcastedB, model.GetTensorType(fNB), fShapeY, broadcastedData);
0185                   fShapeB = fShapeY;
0186                   fDimShapeB = ConvertShapeToDim(fShapeB);
0187                }
0188             } else {
0189                fShapeY = fShapeA;
0190             }
0191             // tensors are constant: perform here the binary operation
0192 
0193             const std::string &nameA = fNBroadcastedA.empty() ? fNA : fNBroadcastedA;
0194             const std::string &nameB = fNBroadcastedB.empty() ? fNB : fNBroadcastedB;
0195             auto dataA = static_cast<T *>(model.GetInitializedTensorData(nameA).get());
0196             auto dataB = static_cast<T *>(model.GetInitializedTensorData(nameB).get());
0197             std::vector<T> dataY(lengthY);
0198             for (size_t i = 0; i < dataY.size(); i++) {
0199                dataY[i] = BinaryOperatorTrait<T, Op>::Func(dataA[i], dataB[i]);
0200             }
0201             model.AddConstantTensor<T>(fNY, fShapeY, dataY.data());
0202             // flag tensors to not be written in the generated code or weight file
0203             model.SetNotWritableInitializedTensor(nameA);
0204             model.SetNotWritableInitializedTensor(nameB);
0205             fIsOutputConstant = true;
0206             if (model.Verbose()) {
0207                std::cout << BinaryOperatorTrait<T, Op>::Name() << " : " << fNA << "  " << ConvertShapeToString(fShapeA)
0208                          << " , " << fNB << "  " << ConvertShapeToString(fShapeB) << " ---> " << fNY << "  "
0209                          << ConvertShapeToString(fShapeY) << " : " << ConvertValuesToString(dataY) << std::endl;
0210             }
0211          } else if (((model.IsShapeTensor(fNA) && model.IsShapeTensor(fNB)) ||
0212                     (model.IsShapeTensor(fNA) && model.IsConstantTensor(fNB)) ||
0213                     (model.IsShapeTensor(fNB) && model.IsConstantTensor(fNA)))
0214                      && (fShapeA.size() <=1 && fShapeB.size() <=1 &&  model.GetTensorType(fNA) == ETensorType::INT64)) {
0215             // case of shape tensors ( tensors are of rank 0 or 1  )
0216             std::vector<Dim> dimValA;
0217             std::vector<Dim> dimValB;
0218             if (model.IsShapeTensor(fNA))
0219                dimValA = model.GetShapeTensorValues(fNA);
0220             if (model.IsShapeTensor(fNB))
0221                dimValB = model.GetShapeTensorValues(fNB);
0222             // adjust for broadcasting - repet values until it reaches shapes of Y
0223             if (!fShapeY.empty() && fShapeY[0] > 1) {
0224                if (dimValA.size() == 1) dimValA = std::vector<Dim>( fShapeY[0], dimValA[0]);
0225                if (dimValB.size() == 1) dimValB = std::vector<Dim>( fShapeY[0], dimValB[0]);
0226             }
0227 
0228             auto convertDataToDim = [&](const std::string & name, const std::vector<size_t> & shape, std::vector<Dim> & dimValues) {
0229                auto data = static_cast<int64_t *>(model.GetInitializedTensorData(name).get());
0230                dimValues.resize(lengthY);
0231                for (size_t i = 0; i < lengthY; i++) {
0232                   if (!shape.empty() && lengthY == shape[0])
0233                      dimValues[i] = Dim{ static_cast<size_t>(data[i])};
0234                   else // case dataA is a scalar
0235                      dimValues[i] = Dim{ static_cast<size_t>(data[0])};
0236                }
0237             };
0238             if (model.IsConstantTensor(fNA)) {
0239                convertDataToDim(fNA,fShapeA,dimValA);
0240             } else if (model.IsConstantTensor(fNB)) {
0241                convertDataToDim(fNB,fShapeB,dimValB);
0242             }
0243 
0244             //perform binary operations on shape tensors
0245             std::vector<Dim> dimValY(lengthY);
0246             for (size_t i = 0; i < lengthY; i++) {
0247                if (!dimValA[i].isParam && !dimValB[i].isParam) {
0248                   size_t d = BinaryOperatorTrait<size_t, Op>::Func(dimValA[i].dim, dimValB[i].dim);
0249                   dimValY[i] = Dim{d};
0250                } else {
0251                   auto res =  BinaryOperatorTrait<T, Op>::Op(dimValA[i].GetVal(), dimValB[i].GetVal());
0252                   dimValY[i] = Dim{res, static_cast<size_t>(-1)};
0253                }
0254             }
0255             model.AddShapeTensor(fNY,dimValY, fShapeY.empty()); // cannot be a  scalar
0256             if (model.Verbose()) {
0257                std::cout << BinaryOperatorTrait<T, Op>::Name() << " : " << fNA << "  " << ConvertShapeToString(fShapeA)
0258                          << " , " << fNB << "  " << ConvertShapeToString(fShapeB) << " ---> " << fNY << "  "
0259                          << ConvertShapeToString(fShapeY) << " : " << ConvertDimShapeToString(dimValY) << " (shape)" <<  std::endl;
0260             }
0261             // no code needs to be generated (flag this as a constant output tensor)
0262             fIsOutputConstant = true;
0263 
0264          } else {
0265             // case of defined and non-constant tensors
0266             model.AddIntermediateTensor(fNY, model.GetTensorType(fNA), fShapeY);
0267             if (model.Verbose()) {
0268                std::cout << BinaryOperatorTrait<T, Op>::Name() << " : " << fNA << "  " << ConvertShapeToString(fShapeA)
0269                          << " , " << fNB << "  " << ConvertShapeToString(fShapeB) << " ---> " << fNY << "  "
0270                          << ConvertShapeToString(fShapeY) << std::endl;
0271             }
0272             // we convert non-dim shapes to Dim shapes
0273             fDimShapeY = ConvertShapeToDim(fShapeY);
0274          }
0275       } else {
0276          // case A or B have dynamic shapes. We need to broadcast if shape are not same
0277          auto ret = UTILITY::MultidirectionalBroadcastShape(fDimShapeA, fDimShapeB);
0278          fBroadcastFlag = ret.first;
0279          fDimShapeY = ret.second;
0280          // case of all parametric shapes and MultiDirectionalBroadcastShape  return the max of the 2
0281          // need to do before we declare the output tensor shape and the broadcasted ones
0282          if (ret.first & 4) {
0283             // check if one of the parameter is an input dimension
0284             // define function to find this
0285             auto IsInputDimParam = [&](const std::string &p) {
0286                auto inputNames = model.GetInputTensorNames();
0287                for (auto &input : inputNames) {
0288                   for (auto &i_s : model.GetDimTensorShape(input)) {
0289                      if (i_s.isParam && i_s.param == p)
0290                         return true;
0291                   }
0292                }
0293                return false;
0294             };
0295             for (size_t i = 0; i < fDimShapeY.size(); i++) {
0296                auto &s = fDimShapeY[i];
0297                if (s.isParam && s.param.find("std::max") != std::string::npos) {
0298                   if (IsInputDimParam(fDimShapeA[i].param)) {
0299                      // case dim is 1 we indicate that the input parameter is equal to 1
0300                      if (fDimShapeA[i].dim != 1)
0301                         s = fDimShapeA[i];
0302                      else
0303                         s = fDimShapeB[i];
0304                   } else if (IsInputDimParam(fDimShapeB[i].param)) {
0305                      if (fDimShapeB[i].dim != 1)
0306                         s = fDimShapeB[i];
0307                      else
0308                         s = fDimShapeA[i];
0309                   }
0310                }
0311             }
0312          }
0313 
0314          model.AddIntermediateTensor(fNY, model.GetTensorType(fNA), fDimShapeY);
0315          if (model.Verbose()) {
0316             std::cout << BinaryOperatorTrait<T, Op>::Name() << " : " << ConvertDimShapeToString(fDimShapeA) << " , "
0317                       << ConvertDimShapeToString(fDimShapeB) << " --> " << ConvertDimShapeToString(fDimShapeY) << std::endl;
0318          }
0319       }
0320    }
0321 
0322    std::string GenerateInitCode() override
0323    {
0324       std::stringstream out;
0325       return out.str();
0326    }
0327 
0328    std::string Generate(std::string opName) override
0329    {
0330 
0331       if (fIsOutputConstant)
0332          return "";
0333 
0334       opName = "op_" + opName;
0335 
0336       std::stringstream out;
0337       out << SP << "\n//------ " << opName << "  " << BinaryOperatorTrait<T, Op>::Name() << " --> "
0338           << ConvertDimShapeToString(fDimShapeY) << "\n";
0339       auto length = ConvertDimShapeToLength(fDimShapeY);
0340       std::string typeName = TensorType<T>::Name();
0341 
0342       // we need to check if we can broadcast (case flag has bit 4 set)
0343 
0344       if (fBroadcastFlag & 4) {
0345          // need to check if shapes are the same
0346          auto lengthA = ConvertDimShapeToLength(fDimShapeA);
0347          auto lengthB = ConvertDimShapeToLength(fDimShapeB);
0348          out << SP << "if (" << lengthA << "!=" << lengthB << ") {\n";
0349          // check if A->B or B->A
0350          // bool broadcastable = true;
0351          for (size_t i = 0; i < fDimShapeY.size(); i++) {
0352             if (fBroadcastFlag & 5 && fDimShapeY[i] == fDimShapeA[i] && fDimShapeA[i].dim > 1 &&
0353                 fDimShapeB[i].isParam) {
0354                // B->A B[i] needs to be 1
0355                out << SP << SP << "if (" << fDimShapeB[i] << "!= 1)\n";
0356                out << SP << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast B->A in operator "
0357                    << opName << "\");\n";
0358             }
0359             if (fBroadcastFlag & 6 && fDimShapeY[i] == fDimShapeB[i] && fDimShapeB[i].dim > 1 &&
0360                 fDimShapeA[i].isParam) {
0361                // A-> B A[i] needs to be 1
0362                out << SP << SP << "if (" << fDimShapeA[i] << "!= 1)\n";
0363                out << SP << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast A->B in operator "
0364                    << opName << "\");\n";
0365             } else if (fDimShapeA[i].isParam && fDimShapeB[i].isParam) {
0366                // both shapes are parametric and we broadcast to maximum
0367                // we allocate here output vector
0368                out << SP << SP << "if (" << fDimShapeA[i] << " != " << fDimShapeB[i] << " && (" << fDimShapeA[i]
0369                    << " != 1 || " << fDimShapeB[i] << " != 1))\n";
0370                out << SP << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast shapes in operator " << opName
0371                    << "\");\n";
0372             }
0373          }
0374          out << SP << "}\n";
0375       }
0376 
0377       auto stridesA = UTILITY::ComputeStrideFromShape(fDimShapeA);
0378       auto stridesB = UTILITY::ComputeStrideFromShape(fDimShapeB);
0379       auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY);
0380 
0381       std::string compute_idx_A, compute_idx_B, compute_idx_Y;
0382       if (fDimShapeA.empty() ||
0383           std::all_of(fDimShapeA.begin(), fDimShapeA.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
0384          compute_idx_A = "0";
0385       } else {
0386          for (size_t i = 0; i < fDimShapeA.size(); ++i) {
0387             if (fDimShapeA[i].dim == 1 || fDimShapeA[i].GetVal() == "1")
0388                continue;
0389             compute_idx_A += "idx_" + std::to_string(i + (fDimShapeY.size() - fDimShapeA.size()));
0390             if (stridesA[i].GetVal() != "1")
0391                compute_idx_A += " * " + stridesA[i].GetVal();
0392             compute_idx_A += " + ";
0393          }
0394          // remove last 3 character " + "
0395          for (int j = 0; j < 3; j++)
0396             compute_idx_A.pop_back();
0397       }
0398       if (fDimShapeB.empty() ||
0399           std::all_of(fDimShapeB.begin(), fDimShapeB.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
0400          compute_idx_B = "0";
0401       } else {
0402          for (size_t i = 0; i < fDimShapeB.size(); ++i) {
0403             if (fDimShapeB[i].dim == 1 || fDimShapeB[i].GetVal() == "1")
0404                continue;
0405             compute_idx_B += "idx_" + std::to_string(i + (fDimShapeY.size() - fDimShapeB.size()));
0406             if (stridesB[i].GetVal() != "1")
0407                compute_idx_B += " * " + stridesB[i].GetVal();
0408             compute_idx_B += " + ";
0409          }
0410           // remove last 3 character " + "
0411          for (int j = 0; j < 3; j++)
0412             compute_idx_B.pop_back();
0413       }
0414       int nloop = 0;
0415       if (fDimShapeY.empty() ||
0416           std::all_of(fDimShapeY.begin(), fDimShapeY.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
0417          compute_idx_Y = "0";
0418       } else {
0419          for (size_t i = 0; i < fDimShapeY.size(); ++i) {
0420             if (fDimShapeY[i].dim != 1 && fDimShapeY[i].GetVal() != "1") {
0421                nloop++;
0422                for (int j = 0; j < nloop; j++) out << SP;
0423                out << "for (size_t idx_" << i << " = 0; idx_" << i << " < " << fDimShapeY[i]
0424                    << "; ++idx_" << i << "){\n";
0425                compute_idx_Y += "idx_" + std::to_string(i);
0426                if (stridesY[i].GetVal() != "1")
0427                   compute_idx_Y += " * " + stridesY[i].GetVal();
0428                compute_idx_Y += " + ";
0429             }
0430          }
0431          // remove last 3 characters " + "
0432          for (int j = 0; j < 3; j++)
0433             compute_idx_Y.pop_back();
0434       }
0435       for (int j = 0; j < nloop + 1; j++) out << SP;
0436       out << "tensor_" << fNY << "[" << compute_idx_Y << "] = "
0437           << BinaryOperatorTrait<T, Op>::Op("tensor_" + fNA + "[" + compute_idx_A + "]",
0438                                             "tensor_" + fNB + "[" + compute_idx_B + "]")
0439           << " ;\n";
0440 
0441       for (int i = nloop; i > 0; i--) {
0442          for (int j = 0; j < i; j++) out << SP;
0443          out << "}\n";
0444       }
0445       return out.str();
0446    }
0447 
0448    std::vector<std::string> GetStdLibs() override
0449    {
0450       if (Op == EBasicBinaryOperator::Pow) {
0451          return {std::string("cmath")};
0452       } else {
0453          return {};
0454       }
0455    }
0456 };
0457 
0458 } // namespace SOFIE
0459 } // namespace Experimental
0460 } // namespace TMVA
0461 
0462 #endif // TMVA_SOFIE_ROperator_BasicBinary
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!