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