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<size_t> fShapeData;
0028    std::vector<size_t> 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.GetTensorShape(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       std::vector<std::vector<size_t>> inputs = { fShapeData };
0078       fShapeOutput = ShapeInference(inputs).front();
0079       if (model.IsInitializedTensor(fNData)) {
0080          fIsOutputConstant = true;
0081          // case input is a constant or initialized tensor we perform here the transpose
0082          auto inStrides = UTILITY::ComputeStrideFromShape(fShapeData);
0083          auto outStrides = UTILITY::ComputeStrideFromShape(fShapeOutput);
0084          size_t length = ConvertShapeToLength(fShapeOutput);
0085          auto inputData = static_cast<T*>(model.GetInitializedTensorData(fNData).get());
0086          size_t dim = fShapeData.size();
0087          std::vector<size_t> outputIdx(dim);
0088          std::vector<T> outputData(length);
0089          for (size_t i = 0; i < length; i++) {
0090             outputIdx[0] = i / outStrides[0];
0091             for (size_t j = 1; j < dim; j++) {
0092                outputIdx[j] = (i % outStrides[j-1]) / outStrides[j];
0093             }
0094             // compute input index
0095             size_t inputIndex = 0;
0096             for (size_t j = 0; j < dim; j++) {
0097                // find value in fAtrrPerm corresponding to j
0098                int k = std::find(fAttrPerm.begin(), fAttrPerm.end(), j) - fAttrPerm.begin();
0099                inputIndex += outputIdx[k] * inStrides[j];
0100             }
0101             outputData[i] = inputData[inputIndex];
0102          }
0103          model.AddConstantTensor<T>(fNOutput, fShapeOutput, outputData.data());
0104          if (model.Verbose()) {
0105             std::cout << "Transpose: output is a constant tensor " << ConvertShapeToString(fShapeOutput) << " : "
0106                << ConvertValuesToString(outputData) << std::endl;
0107          }
0108       } else {
0109          model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNData), fShapeOutput);
0110          if (model.Verbose()) {
0111             std::cout << "Transpose ---> " << fNOutput << " " <<  ConvertShapeToString(fShapeOutput) << std::endl;
0112          }
0113       }
0114    }
0115
0116    std::string Generate(std::string OpName) override {
0117       if (fIsOutputConstant) return "";  //no op for constant tensors
0118       OpName = "op_" + OpName;
0119       if (fShapeData.empty() || fShapeOutput.empty()){
0120          throw std::runtime_error("TMVA SOFIE Transpose Op called to Generate without being initialized first");
0121       }
0122       int dim = fShapeData.size();
0123       auto inStrides = UTILITY::ComputeStrideFromShape(fShapeData);
0124       auto outStrides = UTILITY::ComputeStrideFromShape(fShapeOutput);
0125       size_t length = inStrides[0]*fShapeData[0];  // total tensor size
0126       assert (length == outStrides[0]*fShapeOutput[0]);
0127
0128       std::stringstream out;
0129       // Implement transpose operator using consecutive read inputs.
0130       // But
0131       // tensorOut[id] = tensorInput[ inStrides[0]*i0 + inStrides[1]*i1 + inStrides[2]*i2 + ...]
0132       // now if (j0,j1,j2) are the output indices
0133       // j0 =  id / outStrides[0]
0134       // j1 =  (id % outStrides[0])/outStrides[1]
0135       // j2 =  (id % outStrides[1])/outStrides[2]
0136       //......
0137       // and we have j_k = i_fAttrPerm[k]
0138       // since we are using consecutive writes we should find the inverse of fAttrPerm
0139       out << SP << "///------- Transpose operator\n" << std::endl;
0140       out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
0141       out << SP << SP << "tensor_" << fNOutput << "[id] = tensor_" << fNData << "[ ";
0142       // compute output j indices
0143       std::vector<std::string> i_out(dim);
0144       for (int k =0; k < dim; k++){
0145          if (k == 0)
0146             i_out[k] = "id";
0147          else
0148             i_out[k] = "(id % " + std::to_string(outStrides[k-1]) + ")";
0149          if (k < dim-1)
0150             i_out[k] += " / " + std::to_string(outStrides[k]);
0151       }
0152       // use now them for input tensors
0153       // need to invert the fAttrPerm[k]
0154       for (int k =0; k < dim; k++){
0155          // find value in fAtrrPerm corresponding to k
0156          int l = std::find(fAttrPerm.begin(), fAttrPerm.end(), k) - fAttrPerm.begin();
0157          assert(l >= 0 && l < dim);
0158          out << "( " << i_out[l] << " )";
0159          if (k < dim-1) {
0160             out << " * " << inStrides[k];
0161             out << " + ";
0162          }
0163       }
0164       out << "];\n";
0165       out << SP << "}\n";
0166       return out.str();
0167    }
0168
0169
0170 };
0171
0172 }//SOFIE
0173 }//Experimental
0174 }//TMVA
0175
0176
0177 #endif //TMVA_SOFIE_ROPERATOR_TRANSPOSE