root/TMVA/ROperator_Expand.hxx

0001 #ifndef TMVA_SOFIE_ROperator_Expand
0002 #define TMVA_SOFIE_ROperator_Expand
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <sstream>
0009
0010 namespace TMVA{
0011 namespace Experimental{
0012 namespace SOFIE{
0013
0014 template<typename T>
0015 class ROperator_Expand final : public ROperator{
0016 private:
0017
0018    std::vector<Dim> fShapeX;
0019    std::vector<size_t> fShape;
0020    std::vector<Dim> fShapeY;
0021    std::vector<Dim> fShapeDim;
0022
0023    std::string fNX;
0024    std::string fNShape;
0025    std::string fNY;
0026    std::string fType;
0027
0028    bool fInitialized = false;
0029    bool fInitializedShape = false;
0030    bool fDimShapeValues = false;
0031    bool fInitBroadcast = false;
0032
0033 public:
0034    ROperator_Expand(){}
0035    ROperator_Expand(std::string nameX, std::string nameShape, std::string nameY):
0036       fNX(UTILITY::Clean_name(nameX)), fNShape(UTILITY::Clean_name(nameShape)), fNY(UTILITY::Clean_name(nameY)){
0037          fInputTensorNames = { fNX };
0038          fOutputTensorNames = { fNY };
0039       }
0040
0041
0042    void Initialize(RModel& model) override {
0043       // input must be a graph input, or already initialized intermediate tensor
0044       if (!model.CheckIfTensorAlreadyExist(fNX)) {
0045         throw std::runtime_error("TMVA SOFIE Expand Op Input Tensor " + fNX + " is not found in model");
0046       }
0047       fShapeX = model.GetDimTensorShape(fNX);
0048       if (model.IsInitializedTensor(fNShape)) {
0049          fInitializedShape = true;
0050          int64_t *shapeData =
0051            static_cast<int64_t *>(model.GetInitializedTensorData(fNShape).get());
0052          fShape = model.GetTensorShape(fNShape);
0053          if (fShape.size() != 1) {
0054             throw std::runtime_error("TMVA::SOFIE - Expand operator shape must be a 1d tensor.");
0055          }
0056          size_t N = fShape[0];
0057          // what do we do if shapeData contains negative values?
0058          for (size_t i = 0; i < N; i++) {
0059             if ( shapeData[i] < 0)
0060                throw std::runtime_error("TMVA::SOFIE - Expand: invalid shape value " + std::to_string(shapeData[i]));
0061          }
0062          std::vector<size_t> shape(shapeData, shapeData + N);
0063          fShapeDim = ConvertShapeToDim(shape);
0064       } else if (model.IsShapeTensor(fNShape)) {
0065          // case input shape is a shape tensor
0066          fShapeDim = model.GetShapeTensorValues(fNShape);
0067          fDimShapeValues = true;
0068       } else {
0069          // assume shape of input shape is known (size is 1)
0070          auto shapeOfInputShape = model.GetTensorShape(fNShape);
0071          fShapeDim.resize(shapeOfInputShape[0]);
0072          for (size_t i = 0; i < fShapeDim.size(); i++) {
0073             fShapeDim[i] = Dim{std::string("v_") + fNShape + "_" + std::to_string(i)};
0074             model.AddShapeParam(fShapeDim[i].param);
0075          }
0076       }
0077       // Y is the common shape of fShapeX and shape
0078       auto ret  = TMVA::Experimental::SOFIE::UTILITY::MultidirectionalBroadcastShape(fShapeX, fShapeDim);
0079       fShapeY = ret.second;
0080       fInitialized = model.IsInitializedTensor(fNX) && fInitializedShape;
0081       std::vector<size_t> shapeX;
0082       std::vector<size_t> shapeY;
0083       // case shape tensor and input shape are known
0084       if (!model.IsDynamicTensor(fNX) && !model.IsDimInputTensor(fNX) && fInitializedShape) {
0085          shapeX = ConvertShapeToInt(fShapeX);
0086          shapeY = ConvertShapeToInt(fShapeY);
0087          if (!UTILITY::AreSameShape(shapeX, shapeY))
0088             fInitBroadcast = true;
0089       }
0090       if (fInitialized) {
0091          // cannot have Dim initialized tensors
0092          assert(!shapeX.empty() && !shapeY.empty());
0093          // Broadcast X to the common shape shapeY
0094          // If X is an initialized tensor (constant)
0095          auto data = model.GetInitializedTensorData(fNX);
0096          if (fInitBroadcast) {
0097             std::shared_ptr<void> broadcastedData(
0098                UTILITY::UnidirectionalBroadcast(static_cast<T *>(data.get()), shapeX, shapeY),
0099                std::default_delete<T[]>());
0100             // Update the data and the shape of X
0101             model.UpdateInitializedTensor(fNX, model.GetTensorType(fNX), shapeY, broadcastedData);
0102             fShapeX = fShapeY;
0103             // need to set as a not writable tensor
0104             model.SetNotWritableInitializedTensor(fNX);
0105             data = broadcastedData;
0106          }
0107          if (fInitBroadcast || model.IsConstantTensor(fNX)) {
0108             fIsOutputConstant = true; // constant output in this case
0109             model.AddConstantTensor(fNY, model.GetTensorType(fNX), shapeY, data);
0110             fOutputTensorNames.pop_back();
0111          } else {
0112             model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), shapeY);
0113          }
0114       } else {
0115          // // case input is not initialized
0116          model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
0117       }
0118       fType = ConvertTypeToString(model.GetTensorType(fNX));
0119       if (model.Verbose()) {
0120          std::cout << "Expand - input " << fNX << " shape " << ConvertDimShapeToString(fShapeX) << " --> " << fNY << " shape "
0121                   << ConvertDimShapeToString(fShapeY) << (fIsOutputConstant ? ConvertValuesToString(model.GetTensorData<T>(fNY)) + " (constant)" : "") << std::endl;
0122       }
0123    }
0124
0125    std::string GenerateInitCode() override {
0126       std::stringstream out;
0127       if (!fIsOutputConstant && fInitialized && !fInitBroadcast) {
0128          // shapeX and shapeY are the same in this case
0129          auto length = ConvertDimShapeToLength(fShapeY);
0130          out << "// Copying initialized tensor " << fNX << " to " << fNY << "\n";
0131          out << SP << "std::copy(tensor_" << fNX << ", " << "tensor_" << fNX << " + " << length << ", tensor_" << fNY << ");\n";
0132       }
0133       return out.str();
0134    }
0135
0136    std::string Generate(std::string opName) override {
0137       if (fIsOutputConstant) return "";
0138       opName = "op_" + opName;
0139       if (fShapeY.empty()) {
0140          throw std::runtime_error("TMVA SOFIE Expand Op called to Generate without being initialized first");
0141       }
0142       std::stringstream out;
0143       out << SP << "\n//------ Expand " << opName << " --> " << ConvertDimShapeToString(fShapeY) << "\n";
0144       // need to declare shape parameters for non initialized shapes
0145       if (!fInitializedShape && !fDimShapeValues) {
0146          for (size_t i = 0; i < fShapeDim.size(); i++) {
0147             out << SP << "size_t " << fShapeDim[i] << " = " << "tensor_" << fNShape << "[" << i << "];\n";
0148          }
0149       }
0150       // No need to broadcast A if it's an initialized tensor or shapes are the same
0151       auto lengthX = ConvertDimShapeToLength(fShapeX);
0152       auto lengthY = ConvertDimShapeToLength(fShapeY);
0153       if (lengthX != lengthY) {
0154          out << SP << "if ( (" << lengthX << ") < (" << lengthY << ") ) {\n";
0155          out << SP << SP << "// Broadcasting uninitialized tensor " << fNX << "\n";
0156          out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast(tensor_" << fNX << ", " << ConvertDimShapeToString(fShapeX) << ", " << ConvertDimShapeToString(fShapeY)
0157                    << ", tensor_"<<fNY<<");\n";
0158          out << SP << "} else {\n";
0159          out << SP << SP << "std::copy(tensor_" << fNX << ", " << "tensor_" << fNX << " + (" << lengthX << "), tensor_" << fNY << ");\n";
0160          out << SP << "}\n";
0161       } else {
0162          // case of equal length even if shapes are dims
0163          out << SP << "std::copy(tensor_" << fNX << ", " << "tensor_" << fNX << " + (" << lengthX << "), tensor_" << fNY << ");\n";
0164       }
0165
0166       return out.str();
0167    }
0168
0169 };
0170
0171 }//SOFIE
0172 }//Experimental
0173 }//TMVA
0174
0175
0176 #endif //TMVA_SOFIE_ROperator_Expand