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

tmva/sofie/inc/TMVA/RModel.hxx

@@ -34,6 +34,7 @@ private:
 
 
    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;
@@ -68,12 +69,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,
@@ -159,11 +161,15 @@ public:
    std::string GenerateInferSignature(bool isdecl = true);
 
    // calculate total intermediate memory and position intermediate tensor addresses
-   std::string AllocateIntermediateMemory(std::span<const std::string_view> op_output_tensors);
-   void CheckAndFlushIntermediateMemory(std::span<const std::string_view> op_output_tensors, const size_t& op_idx);
+   std::string AllocateIntermediateMemory(std::span<const std::string> op_output_tensors, std::set<std::string>& allocated_tensors);
+   void CheckAndFlushIntermediateMemory(std::span<const std::string> op_output_tensors, const size_t& op_idx);
 
    void SetOptimizationLevel(const OptimizationLevel &optim_level) { fOptimizationLevel = optim_level; }
 
+   void RemoveIntermediateTensor(const std::string& tensor_name){
+      fIntermediateTensorInfos.erase(tensor_name);
+   }
+
 protected:
    // internal functions
    // generate code for the initialized tensors
@@ -180,6 +186,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,32 @@
 
 class RModel;
 
+enum class OperatorKind {
+   GEMM = 0,
+   LAYERNORM = 1,
+   RELU = 2,
+   CONSTANT = 3,
+   CONSTANTOFSHAPE = 4,
+   UNDEFINED = 5,
+   CONV=6,
+   BATCHNORM=7
+};
+
+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::BATCHNORM:       return "batchnorm";  
+                   case OperatorKind::CONV:       return "conv";
+       case OperatorKind::UNDEFINED:  return "UNDEFINED";
+       default:                       return "UNKNOWN";
+   }
+}
+inline std::set<OperatorKind> FusableKinds = { OperatorKind::RELU, OperatorKind::LAYERNORM, OperatorKind::BATCHNORM};
+
 class ROperator{
 
 
@@ -32,30 +59,44 @@
    // 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, const std::function<void(const std::string&)>& removal_func){ 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)
    bool fIsOutputParamShape = false;     ///< flag to identify of the output represents a parametric shape (can be knwon at compile time)
 
-   mutable std::vector<std::string_view> fInputTensorNames;
-   mutable std::vector<std::string_view> fOutputTensorNames;
+   mutable std::vector<std::string> fInputTensorNames;
+   mutable std::vector<std::string> fOutputTensorNames;
 
 public:
-   std::span<const std::string_view> GetOpInputTensors() const {
+   std::span<const std::string> GetOpInputTensors() const {
       return fInputTensorNames;
    }
 
-   std::span<const std::string_view> GetOpOutputTensors() const {
+   std::span<const std::string> 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_BasicNary.hxx

@@ -103,7 +103,7 @@ public:
 
       fInputTensorNames.resize(fNInputs.size());
       std::transform(fNInputs.begin(), fNInputs.end(), fInputTensorNames.begin(),
-                  [](const std::string& s) -> std::string_view { return s; });
+                  [](const std::string& s) -> std::string { return s; });
       fOutputTensorNames = { fNY };
    }
 

tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx

@@ -53,6 +53,7 @@ public:
    fNB(UTILITY::Clean_name(nameB)), fNMean(UTILITY::Clean_name(nameMean)),
    fNVar(UTILITY::Clean_name(nameVar)), fNY(UTILITY::Clean_name(nameY)), fActivation(activation)
    {
+               fKind = OperatorKind::BATCHNORM;
       fInputTensorNames = { fNX };
       fOutputTensorNames = { fNY };
 
@@ -233,6 +234,19 @@ public:
    }
 
    std::vector<std::string> GetBlasRoutines() override { return { std::string("Copy"), std::string("Axpy") }; }
+   std::string GetFusableOutputTensorName() override {
+       return fNY;
+   }
+
+   void UpdateFusableTensorName(std::string fusable_tensor_name, const std::function<void(const std::string&)>& removal_func){
+      removal_func(fNX);
+      removal_func(fNY);      
+      fNX = fusable_tensor_name;
+      fNY = fusable_tensor_name;
+      fInputTensorNames[0] = fNX;
+      fOutputTensorNames[0] = fNY;
+   }
+
 };
 
 }//SOFIE

tmva/sofie/inc/TMVA/ROperator_Concat.hxx

@@ -37,7 +37,7 @@
 
          fInputTensorNames.resize(fInputs.size());
          std::transform(fInputs.begin(), fInputs.end(), fInputTensorNames.begin(),
-                   [](const std::string& s) -> std::string_view { return s; });
+                   [](const std::string& s) -> std::string { return s; });
          fOutputTensorNames = { fOutput };
          }
 

tmva/sofie/inc/TMVA/ROperator_Constant.hxx

@@ -36,7 +36,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 = { };
       }
@@ -53,10 +59,9 @@ public:
    void Initialize(RModel& model) override {
        //input must be a graph input, or already initialized intermediate tensor
       size_t length = 1;
-      /// ConstantOfShape-------------
+
+      /// ------- ConstantOfShape ---------
       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_Conv.hxx

@@ -58,7 +58,8 @@ public:
       fAttrPads(pads), fAttrStrides(strides),
       fNX(UTILITY::Clean_name(nameX)), fNW(UTILITY::Clean_name(nameW)),
       fNB(UTILITY::Clean_name(nameB)), fNY(UTILITY::Clean_name(nameY))
-   {
+   {  
+      fKind = OperatorKind::CONV;
       if(std::is_same<T, float>::value) {
          fType = "float";
       } else {
@@ -77,6 +78,7 @@ public:
       fAttrPads(pads), fAttrStrides(strides),
       fNX(UTILITY::Clean_name(nameX)), fNW(UTILITY::Clean_name(nameW)), fNY(UTILITY::Clean_name(nameY))
    {
+      fKind = OperatorKind::CONV;
       if(std::is_same<T, float>::value) {
          fType = "float";
       } else {
@@ -569,6 +571,14 @@ public:
    /*! \brief Returns the blas routines needed to compile the generated code
     */
    std::vector<std::string> GetBlasRoutines() override { return { std::string("Gemm"), std::string("Axpy") }; }
+         std::string GetFusableOutputTensorName() override {
+         return fNY;
+      }
+   void UpdateFusableTensorName(std::string fusable_tensor_name, const std::function<void(const std::string&)>& removal_func) override {
+      removal_func(fNY);
+      fNY = fusable_tensor_name;
+      fOutputTensorNames[0] = fNY;
+   }
 };
 
 } // namespace SOFIE

tmva/sofie/inc/TMVA/ROperator_Einsum.hxx

@@ -51,7 +51,7 @@ public:
 
       fInputTensorNames.resize(fNInputs.size());
       std::transform(fNInputs.begin(), fNInputs.end(), fInputTensorNames.begin(),
-                  [](const std::string& s) -> std::string_view { return s; });
+                  [](const std::string& s) -> std::string { return s; });
       fOutputTensorNames = { fNY };
    }
 

tmva/sofie/inc/TMVA/ROperator_Gemm.hxx

@@ -50,6 +50,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>,
@@ -62,10 +63,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};
+         fInputTensorNames = { fNA, fNB, fNC };
          fOutputTensorNames = { fNY };
       }
 
@@ -402,7 +404,16 @@
       }
 
       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, const std::function<void(const std::string&)>& removal_func){
+         removal_func(fNY);
+         fNY = fusable_tensor_name;
+         fOutputTensorNames[0] = fNY;
+      }
+
    };
 
 

tmva/sofie/inc/TMVA/ROperator_LayerNormalization.hxx

@@ -58,7 +58,8 @@ public:
       : 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,20 @@ public:
    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, const std::function<void(const std::string&)>& removal_func){
+      removal_func(fNX);
+      removal_func(fNY);
+      fNX = fusable_tensor_name;
+      fNY = fusable_tensor_name;
+      fInputTensorNames[0] = fNX;
+      fOutputTensorNames[0] = fNY;
+   }
+
 };
 
 } // namespace SOFIE

tmva/sofie/inc/TMVA/ROperator_Range.hxx

@@ -28,7 +28,7 @@ public:
    ROperator_Range(){}
 
    ROperator_Range(std::string start, std::string limit, std::string delta, std::string nameOutput):
-      fNStart(start), fNLimit(limit), fNDelta(delta),
+      fNStart(UTILITY::Clean_name(start)), fNLimit(UTILITY::Clean_name(limit)), fNDelta(UTILITY::Clean_name(delta)),
       fNOutput(UTILITY::Clean_name(nameOutput)) {
       if (std::is_same<T, float>::value) {
           fType = "float";
@@ -37,6 +37,8 @@ public:
       }
       static_assert( (std::is_same_v<T, float> || std::is_same_v<T, int64_t>),
                   "TMVA::SOFIE - Unsupported type by Range operator");
+      fInputTensorNames = {fNStart, fNLimit, fNDelta};
+      fOutputTensorNames = {fNOutput};
    }
 
    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {

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,20 @@
       return out.str();
    }
 
+
+   std::string GetFusableOutputTensorName() override {
+         return fNY;
+   }
+
+   void UpdateFusableTensorName(std::string fusable_tensor_name, const std::function<void(const std::string&)>& removal_func){
+      removal_func(fNX);
+      removal_func(fNY);
+      fNX = fusable_tensor_name;
+      fNY = fusable_tensor_name;
+      fInputTensorNames[0] =  fNX;
+      fOutputTensorNames[0] = fNY;
+   }
+
 };
 
 }//SOFIE

tmva/sofie/inc/TMVA/ROperator_Split.hxx

@@ -38,7 +38,7 @@ public:
          fInputTensorNames = { fNX };
          fOutputTensorNames.resize(fNYs.size());
          std::transform(fNYs.begin(), fNYs.end(), fOutputTensorNames.begin(),
-                   [](const std::string& s) -> std::string_view { return s; });
+                   [](const std::string& s) -> std::string { return s; });
       }
 
    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {

tmva/sofie/inc/TMVA/ROperator_SubGraph.hxx

@@ -36,7 +36,7 @@ public:
 
          fInputTensorNames = { fNX };
          std::transform(fNYs.begin(), fNYs.end(), fOutputTensorNames.begin(),
-                   [](const std::string& s) -> std::string_view { return s; });
+                   [](const std::string& s) -> std::string { return s; });
       }
 
    std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {