Back to home page

EIC code displayed by LXR

 
 

    


Warning, file /include/root/TMVA/ROperator_Slice.hxx was not indexed or was modified since last indexation (in which case cross-reference links may be missing, inaccurate or erroneous).

0001 #ifndef TMVA_SOFIE_ROPERATOR_SLICE
0002 #define TMVA_SOFIE_ROPERATOR_SLICE
0003 
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007 
0008 #include <cassert>
0009 #include <sstream>
0010 #include <numeric>
0011 
0012 namespace TMVA{
0013 namespace Experimental{
0014 namespace SOFIE{
0015 
0016 // slice operator
0017 
0018 template <typename IType>
0019 class ROperator_Slice final : public ROperator
0020 {
0021 
0022 private:
0023 
0024    // flags to indicate if start/end and steps are not defined at compiled time
0025    bool fIsStartUndef = false;
0026    bool fIsEndUndef = false;
0027    bool fIsStepUndef = false;
0028    std::string fNData;        // input data tensor name
0029    std::string fNOutput;      // output data name
0030    std::vector<std::string> fNames;       // tensor names for meta(axis) information
0031    std::vector<Dim> fShapeInput;     // input shape data
0032    std::vector<Dim> fShapeOutput;   // output shape data
0033    // saved Start/End.Steps are corrected from initial ONNX for negative/default values
0034    // and are available for each axis
0035    std::vector<Dim> fStart;         // starting values of slices for all axes
0036    std::vector<Dim> fEnd;           // End values of slices for all axes
0037    std::vector<Dim> fSteps;         // step values of slices for all axes
0038    std::vector<Dim> fStartDims;         // input starting values of slices
0039    std::vector<Dim> fEndDims;           // input End values of slices
0040    std::vector<Dim> fStepDims;         // input step values of slices
0041    std::vector<IType> fAxes;           // axes for input start/emd/step values
0042 
0043    std::vector<std::vector<IType>> fAttributes; // attributes for the version <=10 case
0044 
0045 
0046 public:
0047 
0048    ROperator_Slice(){}
0049 
0050    // ctor for versions >= 10
0051    ROperator_Slice(std::string nameData, std::vector<std::string> names, std::string nameOutput)
0052       : fNData(UTILITY::Clean_name(nameData)),
0053       fNOutput(UTILITY::Clean_name(nameOutput))
0054    {
0055     fNames.resize(4);
0056     // axes and steps can be optional
0057     for (size_t i = 0; i < names.size(); ++i) {
0058         fNames[i] = UTILITY::Clean_name(names[i]);
0059     }
0060 
0061     fInputTensorNames = { fNData };
0062     fOutputTensorNames = { fNOutput };
0063    }
0064    // ctor for versions < 10
0065    ROperator_Slice(std::string nameData, std::vector<IType> starts, std::vector<IType> ends, std::vector<IType> axes, std::string nameOutput)
0066       : fNData(UTILITY::Clean_name(nameData)),
0067       fNOutput(UTILITY::Clean_name(nameOutput))
0068    {
0069      fAttributes.push_back(starts);
0070      fAttributes.push_back(ends);
0071      fAttributes.push_back(axes);
0072     }
0073 
0074 
0075 
0076    void Initialize(RModel& model) override {
0077       if (model.CheckIfTensorAlreadyExist(fNData) == false){   //input must be a graph input, or already initialized intermediate tensor
0078          throw std::runtime_error("TMVA Slice Op Input Tensor is not found in model");
0079       }
0080 
0081       std::vector<std::vector<Dim>> shapes;
0082       fShapeInput = model.GetDimTensorShape(fNData);
0083       shapes.push_back(fShapeInput);
0084 
0085       std::vector<std::vector<IType>> itensors(4);
0086 
0087       if (fNames.size() > 0) {  // size has to be equal to 4
0088          // loop on the extra 2 or 3 or 4 inputs
0089          for (size_t i = 0; i < 4; ++i) {
0090             if (!fNames[i].empty()) {
0091                if (model.IsInitializedTensor(fNames[i])) {
0092                   auto dptr = model.GetInitializedTensorData(fNames[i]);
0093                   auto tensor = static_cast<IType *>(dptr.get());
0094                   auto vec = model.GetTensorShape(fNames[i]);
0095                   assert(vec.size() == 1);
0096                   itensors[i] = std::vector<IType>(tensor, tensor + vec[0]);
0097 
0098                } else if (model.IsShapeTensor(fNames[i])) {
0099                   // case is a shape tensor
0100                   if (i == 0) {
0101                      fStartDims = model.GetShapeTensorValues(fNames[i]);
0102                   } else if (i == 1) {
0103                      fEndDims = model.GetShapeTensorValues(fNames[i]);
0104                   } else if (i == 3) {
0105                      fStepDims = model.GetShapeTensorValues(fNames[i]);
0106                   }
0107                } else {
0108                   // case is an intermediate tensor
0109                   auto shape = model.GetTensorShape(fNames[i]);
0110                   size_t s = shape[0];
0111                   for (size_t k = 0; k < s; k++) {
0112                      if (i == 0) {
0113                         fStartDims.push_back( Dim{std::string("start_") + fNOutput + "_" + std::to_string(k)});
0114                         fIsStartUndef = true;
0115                      } else if (i == 1) {
0116                         fEndDims.push_back(Dim{std::string("end_") + fNOutput + "_" + std::to_string(k)});
0117                         fIsEndUndef = true;
0118                      } else if (i == 3) {
0119                         fStepDims.push_back(Dim{std::string("step_") + fNOutput + "_" + std::to_string(k)});
0120                         fIsStepUndef = true;
0121                      }
0122                   }
0123                }
0124             }
0125          }
0126       } else {
0127          // old slice versions
0128          assert(fAttributes.size() > 1);
0129          for (size_t i = 0; i < fAttributes.size(); i++) {
0130             itensors[i] = fAttributes[i];
0131          }
0132       }
0133       size_t dim = fShapeInput.size();
0134 
0135       // default values
0136       fSteps = std::vector<Dim>(dim, Dim{1});
0137       fStart = std::vector<Dim>(dim, Dim{0});
0138       fEnd = fShapeInput;
0139 
0140       // default axes
0141       if (itensors[2].empty()) {
0142          fAxes.resize(dim);
0143          std::iota(fAxes.begin(), fAxes.end(), 0);
0144       } else {
0145          fAxes = itensors[2];
0146          for (size_t i = 0; i < fAxes.size(); i++) {
0147             // negative axes - they count from the back
0148             if (fAxes[i] < 0) fAxes[i] = dim + fAxes[i];
0149             if (fAxes[i] < 0 || fAxes[i] >= static_cast<IType>(dim))
0150                throw std::runtime_error("TMVA Slice Op : invalid axis value " + std::to_string(fAxes[i]) +
0151                   " for  " + std::to_string(i));
0152          }
0153       }
0154       // Loop on axis to get start/end/step values
0155       for (size_t i = 0; i < fAxes.size(); i++) {
0156          if (!itensors[0].empty() )
0157             fStartDims.push_back(Dim{ static_cast<size_t>(itensors[0][i])});
0158          if (fStartDims.empty())
0159             throw std::runtime_error("TMVA Slice Op : Missing start input tensor");
0160 
0161          if (!itensors[1].empty())
0162             fEndDims.push_back(Dim{ static_cast<size_t>(itensors[1][i])});
0163          else if (fEndDims.empty())
0164             throw std::runtime_error("TMVA Slice Op : Missing end input tensor");
0165 
0166          if (!itensors[3].empty()) {
0167             fStepDims.push_back(Dim{ static_cast<size_t>(itensors[3][i])});
0168          }
0169          else if (fStepDims.size() < fAxes.size())  // this can happen since it is optional
0170             fStepDims.push_back(Dim{size_t(1)});
0171 
0172          if (!fShapeInput[fAxes[i]].isParam) {
0173             size_t iAxisDim = fShapeInput[fAxes[i]].dim;
0174             //correct values if too large or too small
0175             IType istart = 0;
0176             if (!fStartDims[i].isParam) {
0177                istart = static_cast<IType>(fStartDims[i].dim);
0178                if (istart < 0) istart = iAxisDim + istart;
0179             }
0180             IType iend = static_cast<IType>(iAxisDim);
0181             if (!fEndDims[i].isParam) {
0182                iend = static_cast<IType>(fEndDims[i].dim);
0183                if (iend < 0) iend = iAxisDim + iend;
0184             }
0185             //steps
0186             IType istep = 1;
0187             if (!fStepDims[i].isParam) {
0188                istep = static_cast<IType>(fStepDims[i].dim);
0189             } else {
0190                throw std::runtime_error("TMVA Slice Op : parametric step inputs are not supported");
0191             }
0192             // clamp start end values depending on steps
0193             // start must be [0,N] for positive steps or [0,N-1] for negative
0194             // end   must be [0,N] for positive steps or [-1, N-1] for negative
0195             if (istart < 0) istart = 0;
0196             if (istep > 0) {
0197                if (istart > static_cast<IType>(iAxisDim)) istart = static_cast<IType>(iAxisDim);
0198                if (iend < 0) iend = 0;
0199                if (iend > static_cast<IType>(iAxisDim)) iend = static_cast<IType>(iAxisDim);
0200             } else if (istep < 0) {
0201                if (istart > static_cast<IType>(iAxisDim)-1) istart = static_cast<IType>(iAxisDim) -1;
0202                if (iend < -1) iend = -1;
0203                if (iend > static_cast<IType>(iAxisDim)-1) iend = static_cast<IType>(iAxisDim) -1;
0204             } else {
0205                throw std::runtime_error("TMVA Slice Op : invalid step value " + std::to_string(istep) +
0206                   " for  " + std::to_string(i));
0207             }
0208             // for parametric values clamping we will done at run time
0209             if (fStartDims[i].isParam)
0210                fStart[fAxes[i]] = fStartDims[i];
0211             else
0212                fStart[fAxes[i]] = Dim{size_t(istart)};
0213             if (fStartDims[i].isParam)
0214                fEnd[fAxes[i]] = fEndDims[i];
0215             else
0216                fEnd[fAxes[i]] = Dim{size_t(iend)};
0217 
0218             fSteps[fAxes[i]] = Dim{size_t(istep)};
0219          } else {
0220             //std::cout << i << " Param dim for " << fAxes[i] << "  " <<  fShapeInput[fAxes[i]] << std::endl;
0221             // correct only negative values
0222             if (!fStartDims[i].isParam) {
0223                IType istart = static_cast<IType>(fStartDims[i].dim);
0224                if (istart < 0) {
0225                   std::string sstart = std::string("(") + fShapeInput[fAxes[i]].param + "-" + std::to_string(-istart) +")";
0226                   fStart[fAxes[i]] = Dim{sstart,size_t(-1)};
0227                } else {
0228                  fStart[fAxes[i]] = Dim{size_t(istart)};
0229                }
0230             } else {
0231                fStart[fAxes[i]] = fStartDims[i];
0232             }
0233             if (!fEndDims[i].isParam) {
0234                IType iend = static_cast<IType>(fEndDims[i].dim);
0235                if (iend < 0) {
0236                   std::string send = std::string("(") + fShapeInput[fAxes[i]].param + "-" + std::to_string(-iend) +")";
0237                   fEnd[fAxes[i]] = Dim{send,size_t(-1)};
0238                } else {
0239                  fEnd[fAxes[i]] = Dim{size_t(iend)};
0240                }
0241             } else {
0242                fEnd[fAxes[i]] = fEndDims[i];
0243             }
0244 
0245             fSteps[fAxes[i]] = fStepDims[i];
0246          }
0247 
0248       }
0249       //  find output shape
0250       fShapeOutput.resize(dim);
0251       for (size_t i = 0; i < dim; i++) {
0252          if (!fEnd[i].isParam && !fStart[i].isParam && !fSteps[i].isParam) {
0253             int64_t istart = static_cast<int64_t>(fStart[i].dim);
0254             int64_t iend = static_cast<int64_t>(fEnd[i].dim);
0255             int64_t istep= static_cast<int64_t>(fSteps[i].dim);
0256             int64_t s = (iend-istart)/istep;
0257             fShapeOutput[i] = Dim{static_cast<size_t>(s)};
0258          } else {
0259             std::string s;
0260             if (fStart[i].GetVal() != "0")
0261                s = "(" + fEnd[i].GetVal() + "-" + fStart[i].GetVal() + ")";
0262             else
0263                s = fEnd[i].GetVal();
0264             if (fSteps[i].GetVal() != "1") {
0265                s.insert(0,"(");
0266                s += ")/" + fSteps[i].GetVal() + ")";
0267             }
0268             fShapeOutput[i] = Dim{s,size_t(-1)};
0269             // add also the shape parameters to RModel to declare them when
0270             // allocating output tensor
0271             if (fEnd[i].isParam && fEnd[i].dim != size_t(-1))
0272                model.AddShapeParam(fEnd[i].param,fEnd[i].dim );
0273             if (fStart[i].isParam && fStart[i].dim != size_t(-1))
0274                model.AddShapeParam(fStart[i].param,fStart[i].dim );
0275             if (fSteps[i].isParam && fSteps[i].dim != size_t(-1))
0276                model.AddShapeParam(fSteps[i].param,fSteps[i].dim );
0277 
0278          }
0279       }
0280       // case input is a constant tensor and of int64 type
0281       if (model.IsInitializedTensor(fNData) && model.GetTensorType(fNData) == ETensorType::INT64) {
0282          fIsOutputConstant = true;
0283          auto inputData = static_cast<int64_t*>(model.GetInitializedTensorData(fNData).get());
0284          size_t outputSize = ConvertShapeToLength(ConvertShapeToInt(fShapeOutput));
0285          std::vector<int64_t> outputData(outputSize);
0286          std::vector<size_t> inputStride = UTILITY::ComputeStrideFromShape(ConvertShapeToInt(fShapeInput));
0287          std::cout << "slice " << ConvertDimShapeToString(fShapeInput) << " output size " << outputSize << "  " << ConvertDimShapeToString(fShapeOutput) << std::endl;
0288          std::cout << " start - end -steps \n";
0289          for (size_t ii = 0; ii< fStart.size(); ii++)
0290             std::cout << fStart[ii] << "  " << fEnd[ii] << "  " << fSteps[ii] << std::endl;
0291           // perform slice using a recursive function- need to use two lambda functions for this
0292          auto sliceRecursive = [&](size_t iaxis, size_t & outIdx, size_t & inOffset) {
0293             auto slice_impl = [&](size_t iax, size_t & outputIdx, size_t & inputOffset, auto & sliceRecImpl) {
0294                std::cout << "SLice_impl " << fStart.size() << "  " << fEnd.size() << " " << fSteps.size() << "  " << iax << std::endl;
0295                if (fStart[iax].isParam || fEnd[iax].isParam || fSteps[iax].isParam)
0296                   throw std::runtime_error("TMVA Slice Op : cannot have parametric values when input is constant");
0297                // compute indices
0298                std::vector<IType> indices;
0299                for (IType i = (IType) fStart[iax].dim; (IType(fSteps[iax].dim) > 0) ? i < IType(fEnd[iax].dim) : i > IType(fEnd[iax].dim); i += IType(fSteps[iax].dim) )
0300                   indices.push_back(i);
0301                if (iax == dim-1) { // last axis
0302                   std::cout << "SLice_impl last axis: " << indices.size() << " : ";
0303                   for (size_t i = 0; i < indices.size(); i++) {
0304                      std::cout << outputIdx << " , " << indices[i] << " " << inputOffset << " ; ";
0305                      outputData[outputIdx] = inputData[inputOffset + indices[i]];
0306                      outputIdx++;
0307                   }
0308                   std::cout << std::endl;
0309                   return;
0310                } else {
0311                   std::cout << "SLice_impl else : " << indices.size() << " : ";
0312                   for (size_t i = 0; i < indices.size(); i++) {
0313                      std::cout << inputStride[iax] << " , " << indices[i] << " " << inputOffset << "  ";
0314                      size_t offset = inputOffset + inputStride[iax]*indices[i];
0315                      sliceRecImpl(iax+1, outputIdx, offset,sliceRecImpl);
0316                   }
0317                   std::cout << std::endl;
0318                }
0319             };
0320             slice_impl(iaxis, outIdx, inOffset,slice_impl);
0321          };
0322          size_t idx = 0;
0323          size_t offset = 0;
0324          sliceRecursive(0, idx, offset);
0325 
0326          model.AddConstantTensor<int64_t>(fNOutput, ConvertShapeToInt(fShapeOutput), outputData.data());
0327          if (model.Verbose()) {
0328             std::cout << "Slice: output is a constant tensor " << ConvertShapeToString(fShapeOutput) << " : "
0329                      << ConvertValuesToString(outputData) << std::endl;
0330          }
0331       }
0332       else {
0333          model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNData), fShapeOutput);
0334          if (model.Verbose()) {
0335             std::cout << "Slice ---> " << fNOutput << " " <<  ConvertShapeToString(fShapeOutput) << std::endl;
0336          }
0337       }
0338    }
0339 
0340    std::string Generate(std::string OpName) override {
0341       if (fIsOutputConstant) return "";  //no op for constant tensors
0342 
0343       OpName = "op_" + OpName;
0344       if (fShapeInput.empty() || fShapeOutput.empty()){
0345          throw std::runtime_error("TMVA SOFIE Slice Op called to Generate without being initialized first");
0346       }
0347 
0348       std::stringstream out;
0349       //std::string opName = "Slice";
0350 
0351       out << SP << "///------- Slice operator\n" << std::endl;
0352       // loop on the dimensions depending no the orders
0353       size_t ndim = fShapeInput.size();
0354       auto strides = UTILITY::ComputeStrideFromShape(fShapeInput);
0355 
0356 
0357       out << SP << "{\n"; // define operator scope
0358       for (size_t i = 0; i < fStepDims.size(); i++) {
0359          if (fStepDims[i].isParam) {
0360             if (fIsStepUndef)
0361                out << SP << "size_t " << fStepDims[i] << " = tensor_" << fNames[3] << "[" << i << "];\n";
0362          }
0363       }
0364       // special case for parametric  values for start/end. Need to do clipping
0365       for (size_t i = 0; i < fStartDims.size(); i++) {
0366          if (fStartDims[i].isParam && fStartDims[i].param != fShapeInput[fAxes[i]].param) {
0367             std::string s_start = "start_" + std::to_string(i);
0368             if (fIsStartUndef) {
0369                s_start = fStartDims[i].param;
0370                out << SP << "size_t " << s_start << " = tensor_" << fNames[0] << "[" << i << "];\n";
0371             } else {
0372                out << SP << "size_t " << s_start << " = " <<  fStartDims[i] << ";\n";
0373                fStart[fAxes[i]] = s_start; // need to use this value later when slicing
0374             }
0375             out << SP << "if (" << s_start << " < 0) " << s_start << " += " << fShapeInput[fAxes[i]] <<";\n";
0376             out << SP << "if (" << s_start << " < 0) " << s_start << " = 0;\n";
0377             if (!fStepDims[i].isParam) {
0378                if (static_cast<IType>(fStepDims[i].dim) > 0 )
0379                   out << SP << "if (" << s_start << " > " << fShapeInput[fAxes[i]] << " ) " << s_start << " = " << fShapeInput[fAxes[i]] <<";\n";
0380                else
0381                   out << SP << "if (" << s_start << " > " << fShapeInput[fAxes[i]] << " - 1" << " ) " << s_start << " = " << fShapeInput[fAxes[i]] << " - 1;\n";
0382             }
0383          }
0384          // special case if step is negative and shape are equal and step is negative
0385          else if (fStartDims[i].isParam && fStartDims[i].param == fShapeInput[fAxes[i]].param && !fStepDims[i].isParam && static_cast<IType>(fStepDims[i].dim) < 0 ) {
0386             fStart[fAxes[i]] = Dim{ fStartDims[i].param + "-1" };
0387          }
0388       }
0389       // now to for end
0390       for (size_t i = 0; i < fEndDims.size(); i++) {
0391          if (fEndDims[i].isParam && fEndDims[i].param != fShapeInput[fAxes[i]].param) {
0392             std::string s_end = "end_" + std::to_string(i);
0393             if (fIsEndUndef) {
0394                s_end = fEndDims[i].param;
0395                out << SP << "size_t " << s_end << " = tensor_" << fNames[1] << "[" << i << "];\n";
0396             } else {
0397                out << SP << "size_t " << s_end << " = " <<  fEndDims[i] << ";\n";
0398                fEnd[fAxes[i]] = s_end; // need to use this value later when slicing
0399             }
0400             out << SP << "if (" << s_end << " < 0) " << s_end << " += " << fShapeInput[fAxes[i]] <<";\n";
0401             if (!fStepDims[i].isParam) {
0402                if (static_cast<IType>(fStepDims[i].dim) > 0 ) {
0403                   out << SP << "if (" << s_end << " < 0) " << s_end << " = 0;\n";
0404                   out << SP << "if (" << s_end << " > " << fShapeInput[fAxes[i]] << " ) " << s_end << " = " << fShapeInput[fAxes[i]] <<";\n";
0405                } else {
0406                   out << SP << "if (" << s_end << " < -1) " << s_end << " = -1;\n";
0407                   out << SP << "if (" << s_end << " > " << fShapeInput[fAxes[i]] << " - 1" << " ) " << s_end << " = " << fShapeInput[fAxes[i]] << " - 1;\n";
0408                }
0409             }
0410          }
0411          // special case if step is negative and shape are equal and step is negative
0412          else if (fEndDims[i].isParam && fEndDims[i].param == fShapeInput[fAxes[i]].param && !fStepDims[i].isParam && static_cast<IType>(fStepDims[i].dim) < 0 ) {
0413             fEnd[fAxes[i]] = Dim{ fEndDims[i].param + "-1" };
0414          }
0415       }
0416 
0417       out << SP << "size_t iOut = 0;\n";
0418       std::string MSP = SP;
0419       for (size_t idim = 0; idim < ndim; idim++) {
0420         out << MSP << "for (size_t i" << idim << " = " << fStart[idim] <<  "; i" << idim << " < " << fEnd[idim]
0421             << "; i" << idim << "+= " << fSteps[idim] << ") {\n";
0422         MSP += SP;
0423         if (idim < ndim-1) out << MSP << "size_t stride" << idim << " = " << strides[idim] << "*i" << idim << ";\n";
0424       }
0425       out << MSP << "size_t iInput = ";
0426       for (size_t idim = 0; idim < ndim-1; idim++) out << " stride" << idim << " + ";
0427       // here should be step size ?
0428       out << "i" << ndim-1 << ";\n";
0429       out << MSP << "tensor_" << fNOutput << "[iOut++] = tensor_" <<fNData << "[iInput];\n";
0430       for (size_t idim = 0; idim < ndim; idim++) {
0431           MSP = MSP.replace(0,SP.length(),"");
0432           out << MSP << "}\n";
0433       }
0434       out << SP << "}\n"; // end operator scope
0435 
0436       return out.str();
0437    }
0438 
0439 };
0440 
0441 }//SOFIE
0442 }//Experimental
0443 }//TMVA
0444 
0445 
0446 #endif //TMVA_SOFIE_ROPERATOR_SLICE