Back to home page

EIC code displayed by LXR
 
 

    

File indexing completed on 2026-08-16 09:21:40 

0001 #ifndef TMVA_SOFIE_ROPERATOR_GatherND
0002 #define TMVA_SOFIE_ROPERATOR_GatherND
0003 
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007 
0008 #include <sstream>
0009 #include <stdexcept>
0010 #include <string>
0011 
0012 namespace TMVA{
0013 namespace Experimental{
0014 namespace SOFIE{
0015 
0016 class ROperator_GatherND final : public ROperator
0017 {
0018 private:
0019 
0020    size_t fBatchDims = 0;
0021    std::string fNX;
0022    std::string fNIndices;
0023    std::string fNY;
0024 
0025    std::vector<Dim> fShapeX;
0026    std::vector<Dim> fShapeIndices;
0027    std::vector<Dim> fShapeY;
0028 
0029    std::vector<int64_t> fIndices;  // indices vector in case they are known at initialization
0030 
0031    std::string fType;
0032 
0033 public:
0034    ROperator_GatherND(){}
0035    ROperator_GatherND(int batch_dims, std::string nameX, std::string nameIndices, std::string nameY):
0036       fBatchDims(batch_dims), fNX(UTILITY::Clean_name(nameX)), fNIndices(UTILITY::Clean_name(nameIndices)), fNY(UTILITY::Clean_name(nameY)) {
0037          fInputTensorNames = { fNX, fNIndices };
0038          fOutputTensorNames = { fNY };
0039    }
0040 
0041    void Initialize(RModel& model) override {
0042       if (!model.CheckIfTensorAlreadyExist(fNX)) {
0043          throw std::runtime_error("TMVA SOFIE GatherND Op Input Tensor " + fNX + " is not found in model");
0044       }
0045       fShapeX = model.GetDimTensorShape(fNX);
0046       if (model.Verbose())
0047          std::cout << "GatherND - initial shape " << ConvertDimShapeToString(fShapeX) << " shape of indices "
0048                << ConvertDimShapeToString(model.GetDimTensorShape(fNIndices)) << std::endl;
0049       //  fShapeIndices can be  dynamic
0050       fShapeIndices = model.GetDimTensorShape(fNIndices);
0051       size_t q = fShapeIndices.size();
0052       // Axis in range [0, r) where r=rank(X)
0053       size_t r = fShapeX.size();
0054 
0055       if (q < 1) {
0056          throw std::runtime_error("TMVA SOFIE GatherND : rank of Indices is < 1");
0057       }
0058       if (r < 1) {
0059          throw std::runtime_error("TMVA SOFIE GatherND : rank of input tensor is < 1");
0060       }
0061       if (fBatchDims >= std::min(q,r)) {
0062          throw std::runtime_error("TMVA SOFIE GatherND : invalid batch dim value");
0063       }
0064       if (fBatchDims > 0) {
0065          for (size_t i = 0; i < fBatchDims; i++) {
0066             if (fShapeX[i] != fShapeIndices[i]) {
0067                std::cout << " input shape " << ConvertDimShapeToString(fShapeX) << " "
0068                          << " index shape " << ConvertDimShapeToString(fShapeIndices) << std::endl;
0069                throw std::runtime_error("TMVA SOFIE GatherND : invalid input or index shape for " + std::to_string(i));
0070             }
0071          }
0072       }
0073 
0074       //general case. Assumption is that last dimension of index shape is known (is not dynamic)
0075       if (fShapeIndices.back().isParam)
0076          throw std::runtime_error("TMVA SOFIE GatherND : Index_shape(-1) is not known");
0077 
0078       // output shape size (output rank)
0079       // is (q-1)+r -index_shape[-1]
0080       size_t last_index_shape = fShapeIndices.back().dim;
0081       if (last_index_shape < 1 || last_index_shape > r - fBatchDims) {
0082          throw std::runtime_error("TMVA SOFIE GatherND : Index_shape(-1) has wrong value " +
0083             std::to_string(last_index_shape));
0084       }
0085 
0086       size_t output_rank = r + q -1 - last_index_shape - fBatchDims;
0087       //fShapeY.resize(output_rank);
0088       // first index shape dimensions are same in output
0089       fShapeY = std::vector<Dim>(fShapeIndices.begin(), fShapeIndices.end() - 1);
0090       fShapeY.insert(fShapeY.end(), fShapeX.begin() + fBatchDims + last_index_shape, fShapeX.end());
0091       if (fShapeY.size() != output_rank) {
0092          std::cout << " input shape " << ConvertDimShapeToString(fShapeX) << " "
0093                          << " index shape " << ConvertDimShapeToString(fShapeIndices)
0094                          << " output shape " << ConvertDimShapeToString(fShapeY)
0095                          << " and output rank should be " << output_rank << std::endl;
0096          throw std::runtime_error("TMVA SOFIE GatherND : Something is wrong in initialization ");
0097       }
0098 
0099       if (!fIsOutputConstant) {
0100          // Add output tensor
0101          model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
0102          fType = ConvertTypeToString(model.GetTensorType(fNX));
0103          if (model.Verbose())
0104                std::cout <<  "GatherND: input " << fNX << " " << ConvertDimShapeToString(fShapeX) << " indices " << fNIndices << ConvertDimShapeToString(fShapeIndices)
0105                          << " -> " << fNY << " with shape " << ConvertDimShapeToString(fShapeY) << std::endl;
0106       }
0107 
0108 
0109 
0110       // // case indices tensor is initialized
0111       // if (model.IsInitializedTensor(fNIndices)) {
0112       //     // empty shape Indices is a scalar value for the indices
0113       //    size_t indicesLength = ConvertShapeToLength(model.GetTensorShape(fNIndices));
0114       //    int64_t* indicesData = static_cast<int64_t*>(model.GetInitializedTensorData(fNIndices).get());
0115       //    //flag index tensor as not writable (not sure this is needed since index tensor might be used in generated code)
0116       //    model.SetNotWritableInitializedTensor(fNIndices);
0117       //    // update indices data in case of negative dim values
0118       //    for (size_t i = 0; i < indicesLength; i++) {
0119       //       // move this at generation time?
0120       //       if (!fShapeX[fAttrAxis].isParam) {
0121       //          if (indicesData[i] < 0) {
0122       //             indicesData[i] += fShapeX[fAttrAxis].dim;
0123       //          }
0124       //       }
0125       //    }
0126       //    // Save in a vector GatherND Indices of size q
0127       //    fIndices = std::vector<int64_t>(indicesData, indicesData + indicesLength);
0128       // }
0129 
0130       // case input is known (type is an integer) and input indices is a scalar (or vector of size 1)
0131       // if (model.IsInitializedTensor(fNX) && q <= 1 && r == 1 && fIndices.size() > 0) {
0132       //    auto shapeX = ConvertShapeToInt(fShapeX);  // we assume model is not dynamic
0133       //    auto shapeY = ConvertShapeToInt(fShapeY);
0134       //    if (model.GetTensorType(fNX) == ETensorType::INT64) {
0135       //       auto inputData = static_cast<int64_t*>(model.GetInitializedTensorData(fNX).get());
0136       //       // if q <=1 and r = 1 output length = 1 (it is a scalar)
0137       //       std::vector<int64_t> outputData(1); //ConvertShapeToLength(shapeY));
0138       //       outputData[0] = inputData[fIndices[0]];
0139       //       model.AddConstantTensor(fNY, shapeY, outputData.data());
0140       //       if (model.Verbose())
0141       //          std::cout << "GatherND: " << fNX << " " << ConvertShapeToString(shapeX) << " -> " << fNY << " with shape " << ConvertShapeToString(shapeY)
0142       //              << " and values " << ConvertValuesToString(outputData) << " (constant) " << std::endl;
0143       //       fIsOutputConstant = true;
0144       //    }
0145       // }
0146       // // case input is a shape tensor  (r is == 1 by definition) and indices are known
0147       // else if (model.IsShapeTensor(fNX) && q <=1  && fIndices.size() > 0) {
0148       //    auto inputData = model.GetShapeTensorValues(fNX);
0149       //    // if r == 1 and q<=1 then output length is 1 (is a scalar or tensor of size1)
0150       //    std::vector<Dim> outputData(1);
0151       //    outputData[0] = inputData[fIndices[0]];
0152       //    if (outputData[0].isParam) {
0153       //       fIsOutputConstant = true;
0154       //       // shapeY can be scalar or vector of size1
0155       //       model.AddShapeTensor(fNY, outputData, fShapeY.size() == 0);
0156       //       if (model.Verbose())
0157       //          std::cout << "GatherND: " << fNX << " " << ConvertShapeToString(fShapeX) << " -> " << fNY << " with shape " << ConvertShapeToString(fShapeY)
0158       //              << " and values " << ConvertShapeToString(outputData) << " (shape) " << std::endl;
0159       //    } else {
0160       //       int64_t value = static_cast<int64_t>(outputData[0].dim);
0161       //       auto shapeY = ConvertShapeToInt(fShapeY);
0162       //       model.AddConstantTensor(fNY, shapeY, &value);
0163       //       fIsOutputConstant = true;
0164       //       if (model.Verbose())
0165       //          std::cout << "GatherND: " << fNX << " " << ConvertShapeToString(fShapeX) << " -> " << fNY << " with shape " << ConvertShapeToString(fShapeY)
0166       //              << " and values {" << value <<  "} (constant) " << std::endl;
0167       //    }
0168       // }
0169 
0170    }
0171 
0172    std::string Generate(std::string opName) override {
0173       if (fIsOutputConstant) {
0174          // no code to generate here for constant output. Tensor output is defined in Session constructor
0175          return "//---------------------------------------\n";
0176       }
0177       opName = "op_" + opName;
0178       std::stringstream out;
0179       out << "//--------- GatherND " << opName << " --> " << ConvertDimShapeToString(fShapeY) << "\n";
0180       // The shape of the output is q + r - 1
0181       size_t r = fShapeX.size();
0182       // Indices of shape q
0183       size_t q = fShapeIndices.size();
0184       // Strides
0185       auto stridesX = UTILITY::ComputeStrideFromShape(fShapeX);
0186       auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
0187       auto stridesIndices = UTILITY::ComputeStrideFromShape(fShapeIndices);
0188 
0189       // case input_index_shape == rank of input
0190       size_t ss = fShapeIndices.back().dim;
0191 
0192       // check for negative indices
0193       auto indicesLength = ConvertDimShapeToLength(fShapeIndices);
0194       out << SP << "for (size_t i = 0; i < " << indicesLength << "; i++) {\n";
0195       out << SP << SP << "if (tensor_" << fNIndices << "[i] < 0 ) {\n";
0196       // corresponding input shape is  i % strides[N-1]
0197       out << SP << SP << SP << "size_t s_i = " << fShapeX[fBatchDims] << ";\n";
0198       for (size_t j = 1; j < ss; j++) {
0199          out << SP << SP << SP << "if (i % " << ss << " == " << j << ") s_i = " <<  fShapeX[fBatchDims+j] << ";\n";
0200       }
0201       out << SP << SP << SP << "const_cast<int64_t &>(tensor_" << fNIndices << "[i]) +=  s_i;\n";
0202       out << SP << SP << "}\n";
0203       out << SP << "}\n";
0204       // loop on batch dims
0205       std::string outIndex;
0206       std::string inIndex;
0207       std::string idIndex;
0208       for (size_t j = 0; j < fBatchDims; j++) {
0209          std::string index = "i_" + std::to_string(j);
0210          for (size_t k = 0; k <= j; k++) out << SP;
0211          out << "for (size_t " << index << " = 0; " << index << " < " << fShapeY[j] << "; " << index << "++) {\n";
0212          if (j > 0) {
0213             outIndex += " + ";
0214             inIndex += " + ";
0215             idIndex += " + ";
0216          }
0217          outIndex += index;
0218          if (stridesY[j].GetVal() != "1")
0219             outIndex += " * " + stridesY[j].GetVal();
0220          inIndex += index;
0221          if (stridesX[j].GetVal() != "1")
0222             inIndex += " * " + stridesX[j].GetVal();
0223          idIndex += index;
0224          if (stridesIndices[j].GetVal() != "1")
0225             idIndex += " * " + stridesIndices[j].GetVal();
0226       }
0227       // loop between b and q-1
0228       for (size_t j = fBatchDims; j < q - 1; j++) {
0229          std::string index = "i_" + std::to_string(j);
0230          for (size_t k = 0; k <= j; k++) out << SP;
0231          out << "for (size_t " << index << " = 0; " << index << " < " << fShapeY[j] << "; " << index << "++) {\n";
0232          if (j > 0) {
0233             outIndex += " + ";
0234             idIndex += " + ";
0235          }
0236          outIndex += index;
0237          if (stridesY[j].GetVal() != "1")
0238             outIndex += " * " + stridesY[j].GetVal();
0239          idIndex += index;
0240          if (stridesIndices[j].GetVal() != "1")
0241             idIndex += " * " + stridesIndices[j].GetVal();
0242       }
0243       for (size_t k = 0; k <= q - 1; k++) out << SP;
0244       out << "size_t inputIndex = " << inIndex;
0245       std::string indexIndex = idIndex;
0246       for (size_t l = 0; l < ss; l++) {
0247          if (l > 0)
0248             indexIndex = idIndex + " + " + std::to_string(l);
0249          // compute input index using index tensors
0250          if (!indexIndex.empty() || l>0)
0251             out << " + ";
0252          out << "tensor_" << fNIndices << "[" << indexIndex << "]";
0253          if (stridesX[fBatchDims + l].GetVal() != "1") out
0254              << " * " << stridesX[fBatchDims + l];
0255       }
0256       out << ";\n";
0257       for (size_t k = 0; k <= q - 1; k++) out << SP;
0258       // case slice is a scalar
0259       if (ss == r - fBatchDims) {
0260          out << "tensor_" << fNY << "[" << outIndex << "] = "
0261              << "tensor_" << fNX << "[inputIndex];\n";
0262       } else {
0263          // we make a copy of slice
0264          out << "std::copy(tensor_" << fNX << " + inputIndex, tensor_" << fNX << " + inputIndex + "
0265              << stridesX[fBatchDims + ss - 1] << ","
0266              << "tensor_" << fNY << "+" << outIndex << ");\n";
0267       }
0268       // close the loops
0269 
0270       // end loops j_k, j_{k + 1}, ..., j_{r - 2}
0271       for (size_t j = q-1; j > 0; j--) {
0272          for (size_t k = 0; k <j; k++) out << SP;
0273          out << "}\n";
0274       }
0275 
0276       return out.str();
0277    }
0278 
0279 };
0280 
0281 }//SOFIE
0282 }//Experimental
0283 }//TMVA
0284 
0285 
0286 #endif //TMVA_SOFIE_ROPERATOR_RELU
[ 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!