root/TMVA/ROperator_Slice.hxx

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