From 87e68b3c6c904f3e03023eca9bb5b215cf706ea4 Mon Sep 17 00:00:00 2001
From: Sanjiban Sengupta <sanjiban.sg@gmail.com>
Date: Thu, 30 Jan 2025 13:39:25 +0100
Subject: [PATCH] fix: using available memory for initializing output tensors

---
 tmva/sofie/src/RModel.cxx                    | 28 +++++++++-----------
 tmva/sofie_parsers/src/RModelParser_ONNX.cxx |  2 +-
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/tmva/sofie/src/RModel.cxx b/tmva/sofie/src/RModel.cxx
index 3e00b6a5aba723..3bba9b0e683765 100644
--- a/tmva/sofie/src/RModel.cxx
+++ b/tmva/sofie/src/RModel.cxx
@@ -174,10 +174,10 @@ void RModel::AddOperator(std::unique_ptr<ROperator> op, int order_execution) {
       // check if the tensor is already in the lookup table
       if (fIntermediateTensorFrequencyLookup.find(op_input_tensors[index]) == fIntermediateTensorFrequencyLookup.end()) {
          // first time seeing this tensor: initialize the first and last index with the current index
-         fIntermediateTensorFrequencyLookup[op_input_tensors[index]] = std::make_pair(index, index);
+         fIntermediateTensorFrequencyLookup[op_input_tensors[index]] = std::make_pair(order_execution, order_execution);
       } else {
          // tensor already seen: update last index
-         fIntermediateTensorFrequencyLookup[op_input_tensors[index]].second = index;
+         fIntermediateTensorFrequencyLookup[op_input_tensors[index]].second = order_execution;
       }
    }
 }
@@ -356,15 +356,14 @@ std::string RModel::AllocateIntermediateMemory(std::span<const std::string_view>
                // check if available memory chunks are capable of accomodating the tensor
                if (chunk->second >= tensor_size) {
                   chunk->second -= tensor_size;
-
                   allocated = true;
-
+                  memory_allocation_string += "\nfloat* tensor_"+ std::string(it) + "= reinterpret_cast<float*>(fIntermediateMemoryPool+" + chunk->first + ");\n";
+                  chunk->first += tensor_size;
+                  
                   if (chunk->second == 0) {
                         chunk = fIntermediateMemoryInfo.available_memory.erase(chunk);
                         continue;
                   }
-
-                  memory_allocation_string += "\nfloat* tensor_"+ std::string(it) + "= reinterpret_cast<float*>(fIntermediateMemoryPool+" + chunk->first + ");\n";
                }
                ++chunk;
             }
@@ -398,21 +397,21 @@ std::string RModel::AllocateIntermediateMemory(std::span<const std::string_view>
 
 void RModel::CheckAndFlushIntermediateMemory(std::span<const std::string_view> op_input_tensors, const size_t& op_idx){
    for (auto &it : op_input_tensors){
-      
       // last occurence of the tensor is reached => flush it from memory
       if (fIntermediateTensorFrequencyLookup[it].second == op_idx) {
          for (auto chunk = fIntermediateMemoryInfo.total_memory.begin(); 
                chunk != fIntermediateMemoryInfo.total_memory.end(); ) {
-            if (chunk->tensor_name == it) {
-                  fIntermediateMemoryInfo.available_memory.push_back({chunk->chunk_idx, chunk->tensor_size});
-                  
-                  chunk = fIntermediateMemoryInfo.total_memory.erase(chunk);
-            } else {
-                  ++chunk;
-            }
+               if (chunk->tensor_name == std::string(it)) {
+                     fIntermediateMemoryInfo.available_memory.push_back({chunk->chunk_idx, chunk->tensor_size});
+                     
+                     chunk = fIntermediateMemoryInfo.total_memory.erase(chunk);
+               } else {
+                     ++chunk;
+               }
          }
       }
    }
+   std::cout<<"\n\n\n";
 }
 
 
@@ -557,7 +556,6 @@ void RModel::GenerateInitializedTensorInfo() {
          size_t length = ConvertShapeToLength(i.second.shape());
             // in case we are not using weight files or for tensor created from Constant operator
          if (!fUseWeightFile || i.second.IsConstantTensor() ) {
-            //std::cout << "write tensor " << i.first << std::endl;
             std::stringstream strs;
             if (i.second.type() == ETensorType::FLOAT) {
                strs << "float tensor_" << i.first << "[" << length << "] = {";
diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx
index 9184316bf43051..e8288152c2d871 100644
--- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx
+++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx
@@ -610,7 +610,7 @@ void RModelParser_ONNX::ParseONNXGraph(RModel & rmodel, const onnx::GraphProto &
          // for skipping the fused nodes like Add after MatMul
          continue;
       }
-      rmodel.AddOperator(std::move(op));
+      rmodel.AddOperator(std::move(op), i);
    }
 
    std::vector<std::string> outputnames;