root/TMVA/ROperator_Transpose.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_TRANSPOSE
0002 #define TMVA_SOFIE_ROPERATOR_TRANSPOSE
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <sstream>
0009 #include <cassert>
0010
0011 namespace TMVA{
0012 namespace Experimental{
0013 namespace SOFIE{
0014
0015
0016
0017
0018 template <typename T>
0019 class ROperator_Transpose final : public ROperator
0020 {
0021
0022 private:
0023    std::vector<int_t> fAttrPerm;
0024
0025    std::string fNData;
0026    std::string fNOutput;
0027    std::vector<Dim> fShapeData;
0028    std::vector<Dim> fShapeOutput;
0029
0030 public:
0031
0032    ROperator_Transpose(){}
0033    ROperator_Transpose(std::vector<int_t> attr_perm, std::string nameData, std::string nameOutput):
0034       fAttrPerm(attr_perm), fNData(UTILITY::Clean_name(nameData)), fNOutput(UTILITY::Clean_name(nameOutput)) {
0035             fInputTensorNames = { fNData };
0036             fOutputTensorNames = { fNOutput };
0037    }
0038
0039    ROperator_Transpose(std::string nameData, std::string nameOutput):
0040       fNData(UTILITY::Clean_name(nameData)), fNOutput(UTILITY::Clean_name(nameOutput)) {
0041          fInputTensorNames = { fNData };
0042          fOutputTensorNames = { fNOutput };
0043    }
0044
0045    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
0046       return input;
0047    }
0048
0049    std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
0050       if (input.size() > 1) throw std::runtime_error("TMVA SOFIE Tranpose Op Shape Inference only need 1 input tensor");
0051       auto& data = input[0];
0052       if (fAttrPerm.size() != data.size() )
0053          throw std::runtime_error("TMVA SOFIE Tranpose Op - Invalid axes attributes");
0054
0055       std::vector<size_t> output_shape(fAttrPerm.size());
0056       for (size_t i = 0; i < fAttrPerm.size(); i++){
0057          output_shape[i] = data[fAttrPerm[i]];
0058       }
0059       std::vector<std::vector<size_t>> ret;
0060       ret.push_back(output_shape);
0061       return ret;
0062    }
0063
0064
0065    void Initialize(RModel& model) override {
0066       if (model.CheckIfTensorAlreadyExist(fNData) == false){   //input must be a graph input, or already initialized intermediate tensor
0067          std::cout<<"Input tensor for transpose: "<<fNData<<'\n';
0068          throw std::runtime_error("TMVA SOFIE Tranpose Op Input Tensor is not found in model");
0069       }
0070       fShapeData = model.GetDimTensorShape(fNData);
0071       if (fAttrPerm.empty()){
0072          fAttrPerm.reserve(fShapeData.size());
0073          for (int i = fShapeData.size() - 1; i >= 0; i--){
0074             fAttrPerm.push_back(i);
0075          }
0076       }
0077
0078       // inference of output shape
0079       if (fAttrPerm.size() != fShapeData.size() )
0080          throw std::runtime_error("TMVA SOFIE Tranpose Op - Invalid axes attributes");
0081
0082       fShapeOutput.resize(fAttrPerm.size());
0083       for (size_t i = 0; i < fAttrPerm.size(); i++){
0084          fShapeOutput[i] = fShapeData[fAttrPerm[i]];
0085       }
0086
0087       if (model.IsInitializedTensor(fNData) ) {
0088          // here we know the shape
0089          auto shapeX = ConvertShapeToInt(fShapeData);
0090          auto shapeY = ConvertShapeToInt(fShapeOutput);
0091          fIsOutputConstant = true;
0092          // case input is a constant or initialized tensor we perform here the transpose
0093          auto inStrides = UTILITY::ComputeStrideFromShape(shapeX);
0094          auto outStrides = UTILITY::ComputeStrideFromShape(shapeY);
0095          size_t length = ConvertShapeToLength(shapeY);
0096          auto inputData = static_cast<T*>(model.GetInitializedTensorData(fNData).get());
0097          size_t dim = fShapeData.size();
0098          std::vector<size_t> outputIdx(dim);
0099          std::vector<T> outputData(length);
0100          for (size_t i = 0; i < length; i++) {
0101             outputIdx[0] = i / outStrides[0];
0102             for (size_t j = 1; j < dim; j++) {
0103                outputIdx[j] = (i % outStrides[j-1]) / outStrides[j];
0104             }
0105             // compute input index
0106             size_t inputIndex = 0;
0107             for (size_t j = 0; j < dim; j++) {
0108                // find value in fAtrrPerm corresponding to j
0109                int k = std::find(fAttrPerm.begin(), fAttrPerm.end(), j) - fAttrPerm.begin();
0110                inputIndex += outputIdx[k] * inStrides[j];
0111             }
0112             outputData[i] = inputData[inputIndex];
0113          }
0114          model.AddConstantTensor<T>(fNOutput, shapeY, outputData.data());
0115          if (model.Verbose()) {
0116             std::cout << "Transpose: output is a constant tensor " << ConvertShapeToString(shapeY) << " : "
0117                << ConvertValuesToString(outputData) << std::endl;
0118          }
0119       } else {
0120          model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNData), fShapeOutput);
0121          if (model.Verbose()) {
0122             std::cout << "Transpose ---> " << fNOutput << " " <<  ConvertShapeToString(fShapeOutput) << std::endl;
0123          }
0124       }
0125    }
0126
0127    std::string Generate(std::string OpName) override {
0128       if (fIsOutputConstant) return "";  //no op for constant tensors
0129       OpName = "op_" + OpName;
0130       if (fShapeData.empty() || fShapeOutput.empty()){
0131          throw std::runtime_error("TMVA SOFIE Transpose Op called to Generate without being initialized first");
0132       }
0133       int dim = fShapeData.size();
0134       auto inStrides = UTILITY::ComputeStrideFromShape(fShapeData);
0135       auto outStrides = UTILITY::ComputeStrideFromShape(fShapeOutput);
0136       auto length = ConvertDimShapeToLength(fShapeData);
0137       //assert (length == outStrides[0]*fShapeOutput[0]);
0138
0139       std::stringstream out;
0140       // Implement transpose operator using consecutive read inputs.
0141       // But
0142       // tensorOut[id] = tensorInput[ inStrides[0]*i0 + inStrides[1]*i1 + inStrides[2]*i2 + ...]
0143       // now if (j0,j1,j2) are the output indices
0144       // j0 =  id / outStrides[0]
0145       // j1 =  (id % outStrides[0])/outStrides[1]
0146       // j2 =  (id % outStrides[1])/outStrides[2]
0147       //......
0148       // and we have j_k = i_fAttrPerm[k]
0149       // since we are using consecutive writes we should find the inverse of fAttrPerm
0150       out << SP << "///------- Transpose operator " << OpName << ConvertDimShapeToString(fShapeData)
0151                   << " --> " << ConvertDimShapeToString(fShapeOutput) << std::endl;
0152       out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
0153       out << SP << SP << "tensor_" << fNOutput << "[id] = tensor_" << fNData << "[ ";
0154       // compute output j indices
0155       std::vector<std::string> i_out(dim);
0156       for (int k =0; k < dim; k++){
0157          if (k == 0)
0158             i_out[k] = "id";
0159          else
0160             i_out[k] = "(id % (" + outStrides[k-1].GetVal() + "))";
0161          if (k < dim-1)
0162             i_out[k] += " / (" + outStrides[k].GetVal() + ")";
0163       }
0164       // use now them for input tensors
0165       // need to invert the fAttrPerm[k]
0166       for (int k =0; k < dim; k++){
0167          // find value in fAtrrPerm corresponding to k
0168          int l = std::find(fAttrPerm.begin(), fAttrPerm.end(), k) - fAttrPerm.begin();
0169          assert(l >= 0 && l < dim);
0170          out << "( " << i_out[l] << " )";
0171          if (k < dim-1) {
0172             out << " * " << inStrides[k];
0173             out << " + ";
0174          }
0175       }
0176       out << "];\n";
0177       out << SP << "}\n";
0178       return out.str();
0179    }
0180
0181
0182 };
0183
0184 }//SOFIE
0185 }//Experimental
0186 }//TMVA
0187
0188
0189 #endif //TMVA_SOFIE_ROPERATOR_TRANSPOSE