RModel_Base.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_RMODEL_BASE
#define TMVA_SOFIE_RMODEL_BASE
 
#include <type_traits>
#include <unordered_set>
#include <vector>
#include <unordered_map>
#include <memory>
#include <ctime>
#include <set>
#include <iomanip>
#include <fstream>
#include <sstream>
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TBuffer.h"
 
namespace TMVA {
namespace Experimental {
namespace SOFIE {
 
enum class Options {
   kDefault = 0x0,
   kNoSession = 0x1,
   kNoWeightFile = 0x2,
   kRootBinaryWeightFile = 0x4,
   kGNN = 0x8,
   kGNNComponent = 0x10,
};
 
// Optimization levels inspired by ONNXRuntime.
// We only get Operator Fusion with the Basic, and
// memory reuse with Extended. kExtended is enabled
// by default
enum class OptimizationLevel {
   kBasic = 0x0,
   kExtended = 0x1,
};
 
enum class WeightFileType { None, RootBinary, Text };
 
std::underlying_type_t<Options> operator|(Options opA, Options opB);
std::underlying_type_t<Options> operator|(std::underlying_type_t<Options> opA, Options opB);
 
class RModel_Base {
 
protected:
   std::string fFileName;  // file name of original model file for identification
   std::string fParseTime; // UTC date and time string at parsing
 
   WeightFileType fWeightFile = WeightFileType::Text;
 
   std::unordered_set<std::string> fNeededBlasRoutines;
   // Set to true once GenerateHeaderInfo has emitted the extern "C" declaration
   // of the BLAS sgemm_ routine (from fNeededBlasRoutines). It lets the
   // standalone Gemm_Call helper skip emitting a second, duplicate declaration.
   bool fBlasSgemmDeclared = false; //!
 
   std::unordered_set<std::string> fNeededStdLib = {"vector"};
   std::unordered_set<std::string> fCustomOpHeaders;
 
   // Inference helper functions (from SOFIE_common) that the generated code
   // needs. Their standalone definitions are emitted into the generated header
   // so that it does not depend on including TMVA/SOFIE_common.hxx.
   std::set<std::string> fNeededHelperFunctions;
 
   std::string fName = "UnnamedModel";
   std::string fGC; // generated code
   bool fUseWeightFile = true;
   bool fUseSession = true;
   bool fIsGNN = false;
   bool fIsGNNComponent = false;
 
public:
   /**
       Default constructor. Needed to allow serialization of ROOT objects. See
       https://root.cern/manual/io_custom_classes/#restrictions-on-types-root-io-can-handle
   */
   RModel_Base() = default;
 
   RModel_Base(std::string name, std::string parsedtime);
 
   // For GNN Functions usage
   RModel_Base(std::string function_name) : fName(function_name) {}
 
   void AddBlasRoutines(std::vector<std::string> routines)
   {
      for (auto &routine : routines) {
         fNeededBlasRoutines.insert(routine);
      }
   }
   void AddNeededStdLib(std::string libname)
   {
      // if the library is already in the set, insert does nothing, so we don't need to check before inserting
      fNeededStdLib.insert(libname);
   }
   void AddNeededCustomHeader(std::string filename)
   {
       fCustomOpHeaders.insert(filename);
   }
   // Register an inference helper function that the generated code needs. See
   // GenerateHelperFunctionsCode for the list of recognised keys.
   void AddNeededHelperFunction(std::string name)
   {
      fNeededHelperFunctions.insert(std::move(name));
   }
   const std::set<std::string> &GetNeededHelperFunctions() const { return fNeededHelperFunctions; }
 
   // Placeholder tokens emitted by GenerateHeaderInfo and later replaced by
   // EmitHelperFunctionsCode with the actual helper includes / definitions.
   // This two-step approach is needed because the full set of required helpers
   // is only known once all operators (and sub-graphs) have been generated.
   static constexpr const char *kHelperIncludesMarker = "//@SOFIE_HELPER_INCLUDES@\n";
   static constexpr const char *kHelperFunctionsMarker = "//@SOFIE_HELPER_FUNCTIONS@\n";
 
   void GenerateHeaderInfo(std::string &hgname);
   // Replace the helper markers in the generated code with the standalone
   // definitions of the helper functions collected in fNeededHelperFunctions.
   void EmitHelperFunctionsCode();
   void PrintGenerated(std::ostream &os=std::cout) { os << fGC; }
 
   std::string ReturnGenerated() { return fGC; }
   void OutputGenerated(std::string filename = "", bool append = false);
   void SetFilename(std::string filename) { fName = filename; }
   std::string GetFilename() { return fName; }
   const std::string & GetName() const { return fName;}
};
 
enum class GraphType { INVALID = 0, GNN = 1, GraphIndependent = 2 };
 
enum class FunctionType { UPDATE = 0, AGGREGATE = 1 };
enum class FunctionTarget { INVALID = 0, NODES = 1, EDGES = 2, GLOBALS = 3 };
enum class FunctionReducer { INVALID = 0, SUM = 1, MEAN = 2 };
enum class FunctionRelation { INVALID = 0, NODES_EDGES = 1, NODES_GLOBALS = 2, EDGES_GLOBALS = 3 };
 
class RModel_GNNBase : public RModel_Base {
public:
   virtual void Generate() = 0;
   virtual ~RModel_GNNBase() = default;
};
 
} // namespace SOFIE
} // namespace Experimental
} // namespace TMVA
 
#endif // TMVA_SOFIE_RMODEL_BASE