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    std::string fNData;        // input data tensor name
0025    std::string fNOutput;      // output data name
0026    std::vector<std::string> fNames;       // tensor names for meta(axis) information
0027    std::vector<size_t> fShapeInput;     // input shape data
0028    std::vector<size_t> fShapeOutput;   // output shape data
0029    // saved Start/End.Steps are corrected from initial ONNX for negative/default values
0030    // and are available for each axis
0031    std::vector<IType> fStart;         // starting values of slices
0032    std::vector<IType> fEnd;           // End values of slices
0033    std::vector<IType> fSteps;         // step values of slices
0034
0035    std::vector<std::vector<IType>> fAttributes; // attributes for the version <=10 case
0036
0037
0038 public:
0039
0040    ROperator_Slice(){}
0041
0042    // ctor for versions >= 10
0043    ROperator_Slice(std::string nameData, std::vector<std::string> names, std::string nameOutput)
0044       : fNData(UTILITY::Clean_name(nameData)),
0045       fNOutput(UTILITY::Clean_name(nameOutput))
0046    {
0047     fNames.resize(4);
0048     // axes and steps can be optional
0049     for (size_t i = 0; i < names.size(); ++i) {
0050         fNames[i] = UTILITY::Clean_name(names[i]);
0051     }
0052
0053     fInputTensorNames = { fNData };
0054     fOutputTensorNames = { fNOutput };
0055    }
0056    // ctor for versions < 10
0057    ROperator_Slice(std::string nameData, std::vector<IType> starts, std::vector<IType> ends, std::vector<IType> axes, std::string nameOutput)
0058       : fNData(UTILITY::Clean_name(nameData)),
0059       fNOutput(UTILITY::Clean_name(nameOutput))
0060    {
0061      fAttributes.push_back(starts);
0062      fAttributes.push_back(ends);
0063      fAttributes.push_back(axes);
0064     }
0065
0066    // output type is same as input
0067    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0068       auto ret = std::vector<ETensorType>(1, input[0]);
0069       return ret;
0070    }
0071
0072    // output shape
0073    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0074       auto & input_shape = input[0];
0075        // assume dimension of output shape is SAME AS INPUT !
0076       std::vector<std::vector<size_t>> ret(1, input_shape);
0077       auto & output_shape = ret[0];
0078       for (size_t i = 0; i < input_shape.size(); i++) {
0079           output_shape[i] = (fEnd[i]-fStart[i])/ fSteps[i];
0080       }
0081       return ret;
0082    }
0083
0084
0085    void Initialize(RModel& model) override {
0086       if (model.CheckIfTensorAlreadyExist(fNData) == false){   //input must be a graph input, or already initialized intermediate tensor
0087          throw std::runtime_error("TMVA Slice Op Input Tensor is not found in model");
0088       }
0089
0090       std::vector<std::vector<size_t>> shapes;
0091       fShapeInput = model.GetTensorShape(fNData);
0092       shapes.push_back(fShapeInput);
0093
0094       std::vector<std::vector<IType>> itensors(4);
0095       if (fNames.size() > 0) {
0096          // loop on the extra 2 or 3 or 4 inputs
0097          for (size_t i = 0; i < fNames.size(); ++i) {
0098             if (!fNames[i].empty()) {
0099                // std::cout << " i " << i << " getting data for tensor " << fNames[i] << std::endl;
0100                auto dptr = model.GetInitializedTensorData(fNames[i]);
0101                auto tensor = static_cast<IType *>(dptr.get());
0102                auto vec = model.GetTensorShape(fNames[i]);
0103                assert(vec.size() == 1);
0104                itensors[i] = std::vector<IType>(tensor, tensor + vec[0]);
0105             } else {
0106                switch (i) {
0107                case 2: // missing axes
0108                   itensors[2] = std::vector<IType>(fShapeInput.size());
0109                   std::iota(itensors[2].begin(), itensors[2].end(), 0);
0110                   break;
0111                case 3: // missing steps
0112                   itensors[3] = std::vector<IType>(itensors[0].size(), 1);
0113                default: break;
0114                }
0115             }
0116          }
0117       } else {
0118          assert(fAttributes.size() > 1);
0119          for (size_t i = 0; i < fAttributes.size(); i++) {
0120             itensors[i] = fAttributes[i];
0121          }
0122       }
0123       size_t dim = fShapeInput.size();
0124
0125       fSteps = std::vector<IType>(dim, 1);
0126       fStart = std::vector<IType>(dim, 0);
0127       fEnd = std::vector<IType>(dim, 0);
0128       std::copy(fShapeInput.begin(), fShapeInput.end(), fEnd.begin());
0129
0130       auto istart = itensors[0];
0131       auto iend = itensors[1];
0132       auto iaxes = itensors[2];
0133       auto isteps  = itensors[3];
0134
0135       // make tensor axis
0136       // if iaxes.size is =0 tensor axis is missing and use defaults
0137       if (iaxes.size() > 0) {
0138          for (size_t i = 0; i < iaxes.size(); i++) {
0139             // negative axes - they count from the back
0140             if (iaxes[i] < 0) iaxes[i] = dim + iaxes[i];
0141             if (iaxes[i] < 0 || iaxes[i] >= static_cast<IType>(dim))
0142                throw std::runtime_error("TMVA Slice Op : invalid axis value " + std::to_string(iaxes[i]) +
0143                   " for  " + std::to_string(i));
0144
0145             size_t iAxisDim = fShapeInput[iaxes[i]];
0146             // find start/end/step for given axis
0147             // check step size for clamping starting/end value
0148             if (istart[i] < 0) istart[i] = iAxisDim + istart[i];
0149             if (iend[i] < 0) iend[i] = iAxisDim + iend[i];
0150             if (istart[i] < 0) istart[i] = 0;
0151             if (isteps[i] > 0) {
0152                if (istart[i] > static_cast<IType>(iAxisDim)) istart[i] = static_cast<IType>(iAxisDim);
0153                if (iend[i] < 0) iend[i] = 0;
0154                if (iend[i] > static_cast<IType>(iAxisDim)) iend[i] = static_cast<IType>(iAxisDim);
0155             } else if (isteps[i] < 0) {
0156                if (istart[i] > static_cast<IType>(iAxisDim)-1) istart[i] = static_cast<IType>(iAxisDim) -1;
0157                if (iend[i] < -1) iend[i] = -1;
0158                if (iend[i] > static_cast<IType>(iAxisDim)-1) iend[i] = static_cast<IType>(iAxisDim) -1;
0159             } else {
0160                throw std::runtime_error("TMVA Slice Op : invalid step value " + std::to_string(isteps[i]) +
0161                   " for  " + std::to_string(i));
0162             }
0163             fStart[iaxes[i]] = istart[i];
0164             fEnd[iaxes[i]] = iend[i];
0165             fSteps[iaxes[i]] = isteps[i];
0166          }
0167       }
0168
0169       fShapeOutput = ShapeInference({fShapeInput})[0];
0170       // case input is a constant tensor and of int64 type
0171       if (model.IsInitializedTensor(fNData) && model.GetTensorType(fNData) == ETensorType::INT64) {
0172          fIsOutputConstant = true;
0173          auto inputData = static_cast<int64_t*>(model.GetInitializedTensorData(fNData).get());
0174          size_t outputSize = ConvertShapeToLength(fShapeOutput);
0175          std::vector<int64_t> outputData(outputSize);
0176          std::vector<size_t> inputStride = UTILITY::ComputeStrideFromShape(fShapeInput);
0177          // perform slice using a recursive function- need to use two lambda functions for this
0178          auto sliceRecursive = [&](size_t iaxis, size_t & outIdx, size_t & inOffset) {
0179             auto slice_impl = [&](size_t iax, size_t & outputIdx, size_t & inputOffset, auto & sliceRecImpl) {
0180                // compute indices
0181                std::vector<IType> indices;
0182                for (IType i = fStart[iax]; (fSteps[iax] > 0) ? i < fEnd[iax] : i > fEnd[iax]; i += fSteps[iax] )
0183                   indices.push_back(i);
0184                if (iax == dim-1) { // last axis
0185                   for (size_t i = 0; i < indices.size(); i++) {
0186                      outputData[outputIdx] = inputData[inputOffset + indices[i]];
0187                      outputIdx++;
0188                   }
0189                   return;
0190                } else {
0191                   for (size_t i = 0; i < indices.size(); i++) {
0192                      size_t offset = inputOffset + inputStride[iax]*indices[i];
0193                      sliceRecImpl(iax+1, outputIdx, offset,sliceRecImpl);
0194                   }
0195                }
0196             };
0197             slice_impl(iaxis, outIdx, inOffset,slice_impl);
0198          };
0199          size_t idx = 0;
0200          size_t offset = 0;
0201          sliceRecursive(0, idx, offset);
0202
0203          model.AddConstantTensor<int64_t>(fNOutput, fShapeOutput, outputData.data());
0204          if (model.Verbose()) {
0205             std::cout << "Slice: output is a constant tensor " << ConvertShapeToString(fShapeOutput) << " : "
0206                      << ConvertValuesToString(outputData) << std::endl;
0207          }
0208       }
0209       else {
0210          model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNData), fShapeOutput);
0211          if (model.Verbose()) {
0212             std::cout << "Slice ---> " << fNOutput << " " <<  ConvertShapeToString(fShapeOutput) << std::endl;
0213          }
0214       }
0215    }
0216
0217    std::string Generate(std::string OpName) override {
0218       if (fIsOutputConstant) return "";  //no op for constant tensors
0219
0220       OpName = "op_" + OpName;
0221       if (fShapeInput.empty() || fShapeOutput.empty()){
0222          throw std::runtime_error("TMVA SOFIE Slice Op called to Generate without being initialized first");
0223       }
0224
0225       std::stringstream out;
0226       //std::string opName = "Slice";
0227
0228       out << SP << "///------- Slice operator\n" << std::endl;
0229       // loop on the dimensions depending no the orders
0230       size_t ndim = fShapeInput.size();
0231       std::vector<size_t> strides(ndim,1);
0232       for (int i = int(ndim-2); i >=0 ; i--) {
0233           strides[i] = strides[i+1]*fShapeInput[i+1];
0234       }
0235
0236       out << SP << "{\n"; // define operator scope
0237       out << SP << "size_t iOut = 0;\n";
0238       std::string MSP = SP;
0239       for (size_t idim = 0; idim < ndim; idim++) {
0240         out << MSP << "for (size_t i" << idim << " = " << fStart[idim] <<  "; i" << idim << " < " << fEnd[idim]
0241             << "; i" << idim << "+= " << fSteps[idim] << ") {\n";
0242         MSP += SP;
0243         if (idim < ndim-1) out << MSP << "size_t stride" << idim << " = " << strides[idim] << "*i" << idim << ";\n";
0244       }
0245       out << MSP << "size_t iInput = ";
0246       for (size_t idim = 0; idim < ndim-1; idim++) out << " stride" << idim << " + ";
0247       // here should be step size ?
0248       out << "i" << ndim-1 << ";\n";
0249       out << MSP << "tensor_" << fNOutput << "[iOut++] = tensor_" <<fNData << "[iInput];\n";
0250       for (size_t idim = 0; idim < ndim; idim++) {
0251           MSP = MSP.replace(0,SP.length(),"");
0252           out << MSP << "}\n";
0253       }
0254       out << SP << "}\n"; // end operator scope
0255
0256       return out.str();
0257    }
0258
0259 };
0260
0261 }//SOFIE
0262 }//Experimental
0263 }//TMVA
0264
0265
0266 #endif //TMVA_SOFIE_ROPERATOR_SLICE