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
0114 std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0115 return input;
0116 }
0117
0118
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
0139 if (fShapeInputs.size() > 2 ) {
0140
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
0146 fBroadcast = ret.first;
0147 fDimShapeY = ret.second;
0148
0149
0150 if (ret.first & 4) {
0151
0152
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
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
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
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
0246 for (int j = 0; j < 3; j++)
0247 compute_idx_Y.pop_back();
0248 }
0249
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
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
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 }
0291 }
0292 }
0293
0294
0295 #endif