[tmva][sofie] Implement multi-operator fusion for efficient memory usage

tmva/sofie/inc/TMVA/RModel.hxx

@@ -31,6 +31,7 @@ private:
    std::vector<std::string> fInputTensorNames; // input tensor names using ONNX order
 
    std::vector<std::unique_ptr<ROperator>> fOperators;
+   std::vector<std::unique_ptr<ROperator>> fConstantOperators;
 
    std::vector<std::shared_ptr<RModel>> fSubGraphs;    ///<!  sub-graph models (transient)
    RModel * fParentGraph = nullptr;
@@ -59,12 +60,13 @@ public:
    bool CheckIfTensorAlreadyExist(std::string tensor_name);
    void AddInputTensorInfo(std::string input_name, ETensorType type, std::vector<Dim> shape);
    void AddInputTensorInfo(std::string input_name, ETensorType type, std::vector<size_t> shape);
-   void AddOperator(std::unique_ptr<ROperator> op, int order_execution = -1);
-   void AddOperatorReference(ROperator *op, int order_execution = -1)
+   void AddOperator(std::unique_ptr<ROperator> op, size_t order_execution = -1);
+   void AddOperatorReference(ROperator *op, size_t order_execution = -1)
    {
       std::unique_ptr<ROperator> tmp(op);
       AddOperator(std::move(tmp), order_execution);
    }
+   void AddConstantOperator(std::unique_ptr<ROperator> op);
    void AddInitializedTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
                              std::shared_ptr<void> data);
    void AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
@@ -161,6 +163,7 @@ protected:
    void GenerateIntermediateMemoryPool();
    // Generate all session code
    void GenerateSessionCode();
+   void CheckAndFuseOperators();
 
 public:
    const std::vector<std::string> &GetInputTensorNames() const { return fInputTensorNames; }

tmva/sofie/inc/TMVA/ROperator.hxx

@@ -2,6 +2,7 @@
 #define TMVA_SOFIE_ROPERATOR
 
 #include <vector>
+#include <set>
 #include <memory>
 
 #include "TMVA/SOFIE_common.hxx"
@@ -15,6 +16,28 @@ namespace SOFIE{
 
 class RModel;
 
+enum class OperatorKind {
+   GEMM = 0,
+   LAYERNORM = 1,
+   RELU = 2,
+   CONSTANT = 3,
+   CONSTANTOFSHAPE = 4,
+   UNDEFINED = 5
+};
+
+inline const char* toString(OperatorKind kind) {
+   switch (kind) {
+       case OperatorKind::GEMM:       return "GEMM";
+       case OperatorKind::LAYERNORM:  return "LAYERNORM";
+       case OperatorKind::RELU:       return "RELU";
+       case OperatorKind::CONSTANT:       return "CONSTANT";
+       case OperatorKind::CONSTANTOFSHAPE:       return "CONSTANTOFSHAPE";
+       case OperatorKind::UNDEFINED:  return "UNDEFINED";
+       default:                       return "UNKNOWN";
+   }
+}
+inline std::set<OperatorKind> FusableKinds = { OperatorKind::RELU, OperatorKind::LAYERNORM };
+
 class ROperator{
 
 
@@ -32,13 +55,17 @@ public:
    // generate session data members specific to operator
    virtual std::string GenerateSessionMembersCode(std::string /*opName*/) { return ""; }
    virtual std::string Header() { return "";}
+   virtual std::string GetFusableOutputTensorName() { return "";}
+   virtual void UpdateFusableTensorName(std::string){ return;};
+
 
    //virtual void Forward_reference() = 0;
    //virtual void Forward_blas() = 0;
    virtual ~ROperator(){}
 
 protected:
-
+   OperatorKind fKind = OperatorKind::UNDEFINED;
+   size_t fOpOrder = 0;
    const std::string SP = "   ";    ///< space used to correctly indent the generated C++ code
    bool fUseSession = false;        ///< flag to identify if using the session class
    bool fIsOutputConstant = false;  ///< flag to identify if operator has a constant output (no need to generate code)
@@ -54,7 +81,17 @@ public:
    std::span<const std::string_view> GetOpOutputTensors() const {
       return fOutputTensorNames;
    }
-
+
+   OperatorKind GetOpKind(){
+            return fKind;
+   }
+   void RegisterOperatorOrder(const size_t ord){
+      fOpOrder = ord;
+   }
+   size_t GetOpOrder(){
+      return fOpOrder;
+   }
+
 };
 
 

tmva/sofie/inc/TMVA/ROperator_Constant.hxx

@@ -33,7 +33,13 @@ public:
       fShape(shape),
       fValues(values),
       fAttrType(type)
-      {
+      {  
+         fKind = OperatorKind::CONSTANT;
+         if (!fNX.empty()) {
+            // case of ConstantOfShape (since no inputs in case of Constant operator)
+            fIsConstantOfShape  = true;
+            fKind = OperatorKind::CONSTANTOFSHAPE;
+         }
          fInputTensorNames = { };
          fOutputTensorNames = { };
       }
@@ -50,9 +56,9 @@ public:
    void Initialize(RModel& model) override {
        //input must be a graph input, or already initialized intermediate tensor
       size_t length = 1;
+
+      // constant of shape case
       if (!fNX.empty()) {
-         // case of ConstantOfShape (since no inputs in case of Constant operator)
-         fIsConstantOfShape  = true;
          if (model.CheckIfTensorAlreadyExist(fNX) == false){
            throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor is not found in model");
          }

tmva/sofie/inc/TMVA/ROperator_Gemm.hxx

@@ -49,6 +49,7 @@
          fAttrAlpha(alpha), fAttrBeta(beta), fAttrTransA(transA), fAttrTransB(transB), fNA(UTILITY::Clean_name(nameA)),
          fNB(UTILITY::Clean_name(nameB)), fNY(UTILITY::Clean_name(nameY))
       {
+         fKind = OperatorKind::GEMM;
          fActivation = activation;
          fType = "float";
          static_assert(std::is_same_v<T, float>,
@@ -61,9 +62,11 @@
          fAttrAlpha(alpha), fAttrBeta(beta), fAttrTransA(transA), fAttrTransB(transB), fNA(UTILITY::Clean_name(nameA)),
          fNB(UTILITY::Clean_name(nameB)), fNC(UTILITY::Clean_name(nameC)), fNY(UTILITY::Clean_name(nameY)), fActivation(activation)
       {
+         fKind = OperatorKind::GEMM;
          fActivation = activation;
          fType = "float";
 
+         fInputTensorNames = { fNA, fNB, fNC };
          fOutputTensorNames = { fNY };
       }
 
@@ -383,7 +386,13 @@
       }
 
       std::vector<std::string> GetBlasRoutines() override { return { std::string("Gemm"), std::string("Gemv") }; }
-
+      std::string GetFusableOutputTensorName() override {
+         return fNY;
+      }
+
+      void UpdateFusableTensorName(std::string fusable_tensor_name){
+         fNY = UTILITY::Clean_name(fusable_tensor_name);
+      }
    };
 
 

tmva/sofie/inc/TMVA/ROperator_LayerNormalization.hxx

@@ -58,7 +58,8 @@
       : fAttrAxis(axis), fAttrEpsilon(epsilon), fAttrStashType(stashType), fNX(UTILITY::Clean_name(nameX)),
         fNScale(UTILITY::Clean_name(nameScale)), fNB(UTILITY::Clean_name(nameB)),
         fNY(UTILITY::Clean_name(nameY)), fNMean(UTILITY::Clean_name(nameMean)), fNInvStdDev(UTILITY::Clean_name(nameInvStdDev))
-   {
+   {     
+         fKind = OperatorKind::LAYERNORM;
          fInputTensorNames = { fNX, fNScale };
          if (!fNB.empty()){
             fInputTensorNames.emplace_back(fNB);
@@ -336,6 +337,15 @@
    std::vector<std::string> GetBlasRoutines() override { return { std::string("Axpy") }; }
 
    std::vector<std::string> GetStdLibs() override { return { std::string("cmath") }; }
+
+   std::string GetFusableOutputTensorName() override {
+       return fNY;
+   }
+
+   void UpdateFusableTensorName(std::string fusable_tensor_name){
+            fNX = UTILITY::Clean_name(fusable_tensor_name);
+            fNY = UTILITY::Clean_name(fusable_tensor_name);
+   }
 };
 
 } // namespace SOFIE

tmva/sofie/inc/TMVA/ROperator_Relu.hxx

@@ -25,6 +25,7 @@
    ROperator_Relu(){}
    ROperator_Relu(std::string nameX, std::string nameY):
       fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)){
+         fKind = OperatorKind::RELU;
          fInputTensorNames = { fNX };
          fOutputTensorNames = { fNY };
       }
@@ -66,6 +67,16 @@
       return out.str();
    }
 
+
+   std::string GetFusableOutputTensorName() override {
+         return fNY;
+   }
+
+   void UpdateFusableTensorName(std::string fusable_tensor_name){
+         fNX = UTILITY::Clean_name(fusable_tensor_name);
+         fNY = UTILITY::Clean_name(fusable_tensor_name);
+   }
+
 };
 
 }//SOFIE