root/TMVA/RModel.hxx

0001 #ifndef TMVA_SOFIE_RMODEL
0002 #define TMVA_SOFIE_RMODEL
0003
0004 #include "TMVA/RModel_Base.hxx"
0005 #include "TMVA/SOFIE_common.hxx"
0006 #include "TMVA/ROperator.hxx"
0007
0008 namespace TMVA {
0009 namespace Experimental {
0010 namespace SOFIE {
0011
0012 class RModel final : public RModel_Base {
0013
0014 private:
0015    bool fIsInitialized = false;
0016    bool fIsSubGraph = false;
0017    bool fUseVDT = false;
0018    int fVerbose = 0;
0019    int fBatchSize = -1;
0020    long fReadPos = 0;  // reading file position
0021    size_t fConstantTensorSize = 0; // size  (in Bytes) of the allocated constant tensors
0022    size_t fWeightsTensorSize = 0;  // size  (in Bytes) of the allocated weight tensors
0023    size_t fOtherTensorSize = 0;    // size  (in Bytes) of intermediate tensors which are not managed by the memory pool
0024
0025    OptimizationLevel fOptimizationLevel = OptimizationLevel::kExtended;
0026
0027    std::unordered_map<std::string, InputTensorInfo> fInputTensorInfos; // input tensors where shape may not fully defined or other graph inputs?
0028    std::unordered_map<std::string, TensorInfo> fReadyInputTensorInfos; // input tensors where shape is full defined
0029    std::unordered_map<std::string, InitializedTensor> fInitializedTensors;
0030    std::unordered_map<std::string, TensorInfo> fIntermediateTensorInfos;
0031    std::unordered_map<std::string, DynamicTensorInfo> fDynamicTensorInfos;
0032    std::unordered_map<std::string, std::pair<std::vector<Dim>, bool>> fShapeTensors; // constant tensors describing a shape
0033    std::unordered_map<std::string, std::string> fShapeParams; // parameters defining the dynamic shape (e.g. batch size), store also its default value
0034    std::unordered_map<std::string, std::string> fAliasTensors;   // list of alias tensors
0035    std::vector<std::string> fDimShapeNames; // parameter names used to define the shapes
0036    std::vector<std::string> fOutputTensorNames;
0037    std::vector<std::string> fInputTensorNames; // input tensor names using ONNX order
0038
0039    std::vector<std::unique_ptr<ROperator>> fOperators;
0040
0041    std::vector<std::shared_ptr<RModel>> fSubGraphs;    ///<!  sub-graph models (transient)
0042    RModel * fParentGraph = nullptr;
0043
0044    // memory pool information for intermediate tensors
0045    MemoryPoolInfo fIntermediateMemoryInfo;    ///<!  intermediate memory info (transient)
0046    std::unordered_map<std::string_view, size_t> fIntermediateTensorFrequencyLookup;    ///<!  lookup table for intermediate tensor frequency (transient)
0047
0048 public:
0049    /**
0050        Default constructor. Needed to allow serialization of ROOT objects. See
0051        https://root.cern/manual/io_custom_classes/#restrictions-on-types-root-io-can-handle
0052    */
0053    RModel() = default;
0054    RModel(std::string name, std::string parsedtime) : RModel_Base(name, parsedtime) {}
0055
0056    // For GNN Functions usage
0057    RModel(std::string function_name) : RModel_Base(function_name) {}
0058
0059    int Verbose() const { return fVerbose;}
0060
0061    std::vector<size_t> GetTensorShape(const std::string & name) const;
0062    std::vector<Dim> GetDimTensorShape(const std::string & name) const;
0063    std::vector<Dim> GetDynamicTensorShape(const std::string & name) const ;
0064
0065    // get the values for the tensor representing a shape
0066    const std::vector<Dim> & GetShapeTensorValues(const std::string & tensor_name) const;
0067
0068    ETensorType GetTensorType(std::string name) const;
0069
0070
0071    bool CheckIfTensorAlreadyExist(std::string tensor_name);
0072    void AddInputTensorInfo(std::string input_name, ETensorType type, std::vector<Dim> shape);
0073    void AddInputTensorInfo(std::string input_name, ETensorType type, std::vector<size_t> shape);
0074    void AddOperator(std::unique_ptr<ROperator> op, int order_execution = -1);
0075    void AddInitializedTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
0076                              std::shared_ptr<void> data);
0077    void AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
0078                              std::shared_ptr<void> data);
0079
0080    void AddAliasTensor(const std::string & tensor_name, const std::string & orig_tensor_name);
0081
0082
0083    template<class T>
0084    void AddConstantTensor(const std::string & name, const std::vector<size_t> & shape, const T * data) {
0085       size_t length = ConvertShapeToLength(shape);
0086       std::shared_ptr<void> data_ptr(malloc(length * sizeof(T)), free);
0087       std::memcpy(data_ptr.get(), (void*) data, length * sizeof(T));
0088       AddConstantTensor(name, GetTemplatedType<T>(T()), shape, data_ptr);
0089    }
0090    // for boolean can be more convenient passing an std::vector
0091    template<class T>
0092    void AddConstantTensor(const std::string & name, const std::vector<size_t> & shape, const std::vector<T> & data) {
0093       size_t length = data.size();
0094       std::shared_ptr<void> data_ptr(malloc(length * sizeof(T)), free);
0095       std::copy(data.begin(), data.end(), (T*) data_ptr.get());
0096       //std::memcpy(data_ptr.get(), (void*) data, length * sizeof(T));
0097       AddConstantTensor(name, GetTemplatedType<T>(T()), shape, data_ptr);
0098    }
0099
0100    template <typename T>
0101    void AddInitializedTensor(const std::string & tensor_name, const std::vector<std::size_t> & shape, T *raw_data)
0102    {
0103       size_t size = ConvertShapeToLength(shape);
0104       std::shared_ptr<void> data(malloc(size * sizeof(T)), free);
0105       std::memcpy(data.get(), raw_data, size * sizeof(T));
0106       AddInitializedTensor(tensor_name,  GetTemplatedType(T()), shape, data);
0107    }
0108
0109    void AddShapeTensor(const std::string & name, const std::vector<Dim> & shapeValues, bool scalar = false);
0110
0111
0112    // add and initialize subgraph to the model
0113    void InitializeSubGraph(std::shared_ptr<RModel>  graph);
0114
0115    // set a flag to indicate tensor does not need to be written in a weight file
0116    // (e.g. shape tensors used as input to define a shape (in Reshape))
0117    void SetNotWritableInitializedTensor(const std::string & tensor_name);
0118
0119    // Check if a tensor is initialized
0120    bool IsInitializedTensor(const std::string &name) const;
0121    // Check if a tensor is Constant (note a Constant tensor is also initialized)
0122    bool IsConstantTensor(const std::string &name) const;
0123    bool IsDynamicTensor(const std::string &name) const;
0124    // Check if tensor is a input dynamic tensor (without a specified shape, based on Sim structure
0125    bool IsDimInputTensor(const std::string &name) const;
0126    // check if tensor is a fully specified input tensor
0127    bool IsReadyInputTensor(const std::string &name) const;
0128    /// check if a tensor is a shape tensor
0129    bool IsShapeTensor(const std::string & name) const;
0130    /// check if a tensor is a alias tensor
0131    bool IsAliasTensor(const std::string & name) const;
0132
0133    // Add intermediate tensor
0134    void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<Dim> dim_shape);
0135    void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape);
0136    // Add an intermediate dynamic tensor
0137    void AddDynamicTensor(std::string tensor_name, ETensorType type, std::vector<Dim> shape);
0138    // void Add a shape parameter
0139    void AddShapeParam(const std::string & name, size_t def_value = 0);
0140    void AddInputTensorName(std::string name);
0141    void AddOutputTensorNameList(std::vector<std::string> output_tensor_names);
0142    void
0143    UpdateOutputTensorList(std::vector<std::string> curr_output_tensor, std::vector<std::string> modify_output_tensor);
0144    void UpdateInitializedTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
0145                                 std::shared_ptr<void> data);
0146    std::shared_ptr<void> GetInitializedTensorData(std::string tensor_name);
0147
0148    template<class T>
0149    std::vector<T> GetTensorData(const std::string & name);
0150
0151    void Initialize(int batchSize = -1, bool verbose = false);
0152    void Initialize(const std::map<std::string,size_t> & inputParams, bool verbose = false);
0153
0154    void Generate(std::underlying_type_t<Options> options, int batchSize = -1, long pos = 0, bool verbose = false);
0155    void Generate(Options options = Options::kDefault, int batchSize = -1, int pos = 0, bool verbose = false)
0156    {
0157       Generate(static_cast<std::underlying_type_t<Options>>(options), batchSize, pos, verbose);
0158    }
0159    // generate the infer function signature. If isdecl= false generate the calling infer function
0160    // used to infer the sub-graphs
0161    std::string GenerateInferSignature(bool isdecl = true);
0162
0163    // calculate total intermediate memory and position intermediate tensor addresses
0164    std::string AllocateIntermediateMemory(std::span<const std::string_view> op_output_tensors);
0165    void CheckAndFlushIntermediateMemory(std::span<const std::string_view> op_output_tensors, const size_t& op_idx);
0166
0167    void SetOptimizationLevel(OptimizationLevel optim_level) { fOptimizationLevel = optim_level; }
0168
0169    // get the size in bytes of the constant tensors
0170    size_t GetConstantTensorSize() const { return fConstantTensorSize; }
0171    // get the size in bytes of the weight tensors
0172    size_t GetWeightsTensorSize() const { return fWeightsTensorSize; }
0173    // get the size in bytes of the intermediate tensors which are not part of the memory pool
0174    size_t GetOtherTensorSize() const { return fOtherTensorSize; }
0175    // get the size in bytes of the intermediate tensors managed by the memory pool
0176    size_t GetIntermediateTensorSize() const {
0177       return (!fIntermediateMemoryInfo.total_stack.empty())
0178                 ? fIntermediateMemoryInfo.total_stack.rbegin()->first + fIntermediateMemoryInfo.total_stack.rbegin()->second.tensor_size
0179                 : 0;
0180    }
0181
0182 protected:
0183    // internal functions
0184    // generate code for the initialized tensors
0185    void GenerateInitializedTensorInfo();
0186    // generate code for the intermediate tensors
0187    void GenerateIntermediateTensorInfo();
0188    // generate code for the dynamic tensors
0189    void GenerateDynamicTensorInfo();
0190    // generate code for declarations needed by operators
0191    void GenerateOperatorDeclarations();
0192    // generate code for inference
0193    void GenerateOutput();
0194    // generate code for initializing memory pool for intermediate tensors
0195    void GenerateIntermediateMemoryPool();
0196    // Generate all session code
0197    void GenerateSessionCode();
0198    bool IsInputTensorShapeParam(std::string const &name) const;
0199    std::vector<std::string> CollectTensorMemberNames(const std::string &input);
0200    void GenerateRequiredInputTensorInfo();
0201
0202 public:
0203    const std::vector<std::string> & GetInputTensorNames() const { return fInputTensorNames; }
0204    const std::vector<std::string> & GetOutputTensorNames() const { return fOutputTensorNames; }
0205    const std::vector<std::string> & GetDimShapeNames() const { return fDimShapeNames; }
0206
0207    void ReadInitializedTensorsFromFile(long);
0208    long WriteInitializedTensorsToFile(std::string filename = "");
0209
0210    void PrintSummary() const;
0211    void PrintIntermediateTensors() const;
0212    void PrintOutputTensors() const;
0213    void OutputGenerated(std::string filename = "", bool append = false);
0214    void SetFilename(std::string filename) { fName = filename; }
0215
0216    /*
0217       template <typename T>
0218       void AddInitializedTensor(std::string tensor_name, RTensor<T> new_tensor){
0219          //a view only
0220          T obj;
0221          if (fInitializedTensors.find(tensor_name) != fInitializedTensors.end()){
0222             throw std::runtime_error("TMVA-SOFIE: initialized tensor with name " + tensor_name + " already exists \n");
0223          }
0224          InitializedTensor new_tensor_ {GetTemplatedType(obj), new_tensor.GetShape() ,
0225       static_cast<void>(new_tensor.GetData())}; fInitializedTensors[tensor_name] = new_tensor_;
0226       }
0227    */
0228
0229    void PrintRequiredInputTensors() const;
0230    void PrintInitializedTensors() const;
0231    void PrintDynamicTensors() const;
0232    void HeadInitializedTensors(std::string name, int n_print = 50);
0233
0234    bool UseSession() const { return fUseSession; }
0235    // flag to use vdt for fast math functions (e.g. exp in softmax)
0236    void SetUseVDT(bool on) {
0237       fUseVDT = on;
0238    }
0239    bool UseVDT() const { return fUseVDT;}
0240
0241    // Use the ClassDef macro to allow definition of custom streaming
0242    ClassDefNV(RModel, 3);
0243 };
0244
0245 // need to implement here templated member functions and its specialization
0246
0247
0248 template<class T>
0249 inline std::vector<T> RModel::GetTensorData(const std::string & name) {
0250    if (!IsInitializedTensor(name)) return std::vector<T>{};
0251    T * data = static_cast<T*>(GetInitializedTensorData(name).get());
0252    size_t size = ConvertShapeToLength(GetTensorShape(name));
0253    return std::vector<T>(data, data+size);
0254 }
0255
0256 template<>
0257 inline std::vector<Dim> RModel::GetTensorData<Dim>(const std::string & name) {
0258    if (!IsShapeTensor(name)) return std::vector<Dim>{};
0259    return GetShapeTensorValues(name);
0260 }
0261
0262 } // namespace SOFIE
0263 } // namespace Experimental
0264 } // namespace TMVA
0265
0266 #endif // TMVA_SOFIE_RMODEL