ROperator_Cast.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_ROPERATOR_Cast
#define TMVA_SOFIE_ROPERATOR_Cast
 
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TMVA/RModel.hxx"
 
#include <sstream>
 
namespace TMVA{
namespace Experimental{
namespace SOFIE{
 
 
class ROperator_Cast final : public ROperator
{
 
private:
 
   std::string fNX;
   std::string fNY;
   std::vector<Dim> fShape;
   ETensorType fType;
 
public:
   ROperator_Cast(){}
   ROperator_Cast(ETensorType type,std::string nameX, std::string nameY):
      fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)),
      fType(type)
   {
      fInputTensorNames = { fNX };
      fOutputTensorNames = { fNY };
   }
 
   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
      return input;
   }
 
   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
      auto ret = input; //suggest copy to compiler
      return ret;
   }
 
   void Initialize(RModel& model) override {
       //input must be a graph input, or already initialized intermediate tensor
      if (model.CheckIfTensorAlreadyExist(fNX) == false){
        throw std::runtime_error("TMVA SOFIE Cast Op Input Tensor is not found in model");
      }
      fShape = model.GetDimTensorShape(fNX);
      // should we add a check if the same type
      auto inputType = model.GetTensorType(fNX);
      if (model.IsInitializedTensor(fNX)) {
         fIsOutputConstant = true;
         auto inputData = model.GetInitializedTensorData(fNX);
         if (fType == ETensorType::INT64) {
            model.AddConstantTensor<int64_t>(fNY, ConvertShapeToInt(fShape), static_cast<int64_t*>(inputData.get()));
            model.SetNotWritableInitializedTensor(fNX);
         }
         else
            fIsOutputConstant = false;
      } else if (model.IsShapeTensor(fNX) && fType == ETensorType::INT64) {
         auto shapeData = model.GetShapeTensorValues(fNX);
         model.AddShapeTensor(fNY, shapeData, fShape.size() == 0);
         fIsOutputConstant = true;
      }
      if (!fIsOutputConstant)
         model.AddIntermediateTensor(fNY, fType, fShape);
      if (model.Verbose()) {
         std::cout << "Cast : " << ConvertTypeToString(inputType) << " " << fNX << " -> " << ConvertTypeToString(fType);
         if (fType == ETensorType::BOOL) std::cout << " (converted from BOOL) ";
         std::cout << " for " << fNY << " shape " << ConvertDimShapeToString(fShape);
         if (fIsOutputConstant) std::cout << " (constant) ";
         std::cout << std::endl;
      }
   }
 
 
   std::string Generate(std::string opName) override {
 
      // output shape can be empty if is a scalar
 
      std::stringstream out;
      auto length = ConvertDimShapeToLength(fShape);
 
      out << "\n//------ CAST " << opName << " ---> " << fNY << "  " << ConvertDimShapeToString(fShape) << "\n";
       // no generated code for constant outputs
      if (fIsOutputConstant) return out.str();
 
      out << SP << "for (int id = 0; id < " << length << " ; id++){\n";
 
      // need to handle bool case separatly since casting to uint8 will not give right result
      if (fType == ETensorType::BOOL)
         out << SP << SP << "tensor_" << fNY << "[id] = (tensor_" << fNX << "[id] != 0) ? 1 : 0;\n";
      else
         out << SP << SP << "tensor_" << fNY << "[id] = static_cast<"<< ConvertTypeToString(fType) << ">(tensor_" << fNX << "[id]);\n";
 
      out << SP << "}\n";
      return out.str();
   }
 
};
 
}//SOFIE
}//Experimental
}//TMVA
 
 
#endif //TMVA_SOFIE_ROPERATOR_Cast