1 files changed, 22 insertions, 10 deletions

diff --git a/runtime/contrib/heap_trace/src/trace.cc b/runtime/contrib/heap_trace/src/trace.cc
index 82f2915cb..020aeb90e 100644
--- a/runtime/contrib/heap_trace/src/trace.cc
+++ b/runtime/contrib/heap_trace/src/trace.cc
@@ -48,7 +48,7 @@ void Trace::logAllocationEvent(void *memory_ptr, size_t size_of_allocated_space_
     _peak_heap_usage_on_cpu = _total_allocated_bytes_on_cpu - _total_deallocated_bytes_on_cpu;
   }
   _memory_in_use_on_cpu[memory_ptr] = size_of_allocated_space_in_bytes;
-  Guard{}.signalizeThatDangerOfRecursionHAsPassed();
+  Guard{}.signalizeThatDangerOfRecursionHasPassed();
 }
 
 void Trace::logDeallocationEvent(void *memory_ptr)
@@ -61,20 +61,29 @@ void Trace::logDeallocationEvent(void *memory_ptr)
     _total_deallocated_bytes_on_cpu += found_memory_space_description->second;
     _memory_in_use_on_cpu.erase(found_memory_space_description);
   }
-  Guard{}.signalizeThatDangerOfRecursionHAsPassed();
+  Guard{}.signalizeThatDangerOfRecursionHasPassed();
 }
 
 void Trace::logAllocationEvent(cl_mem memory_ptr, size_t size_of_allocated_space_in_bytes)
 {
   Guard{}.signalizeAboutPossibleRecursion();
   std::lock_guard<std::mutex> guard(_lock);
-  _total_allocated_bytes_on_gpu += size_of_allocated_space_in_bytes;
-  if (_peak_heap_usage_on_gpu < _total_allocated_bytes_on_gpu - _total_deallocated_bytes_on_gpu)
+  auto found_memory_space_description = _memory_in_use_on_gpu.find(memory_ptr);
+  if (found_memory_space_description == _memory_in_use_on_gpu.end())
+  {
+    _memory_in_use_on_gpu.insert(
+        std::make_pair(memory_ptr, MemoryTraits(1, size_of_allocated_space_in_bytes)));
+    _total_allocated_bytes_on_gpu += size_of_allocated_space_in_bytes;
+    if (_peak_heap_usage_on_gpu < _total_allocated_bytes_on_gpu - _total_deallocated_bytes_on_gpu)
+    {
+      _peak_heap_usage_on_gpu = _total_allocated_bytes_on_gpu - _total_deallocated_bytes_on_gpu;
+    }
+  }
+  else
   {
-    _peak_heap_usage_on_gpu = _total_allocated_bytes_on_gpu - _total_deallocated_bytes_on_gpu;
+    ++found_memory_space_description->second.ref_counter;
   }
-  _memory_in_use_on_gpu[memory_ptr] = size_of_allocated_space_in_bytes;
-  Guard{}.signalizeThatDangerOfRecursionHAsPassed();
+  Guard{}.signalizeThatDangerOfRecursionHasPassed();
 }
 
 void Trace::logDeallocationEvent(cl_mem memory_ptr)
@@ -84,10 +93,13 @@ void Trace::logDeallocationEvent(cl_mem memory_ptr)
   auto found_memory_space_description = _memory_in_use_on_gpu.find(memory_ptr);
   if (found_memory_space_description != _memory_in_use_on_gpu.end())
   {
-    _total_deallocated_bytes_on_gpu += found_memory_space_description->second;
-    _memory_in_use_on_gpu.erase(found_memory_space_description);
+    if (--found_memory_space_description->second.ref_counter == 0)
+    {
+      _total_deallocated_bytes_on_gpu += found_memory_space_description->second.size;
+      _memory_in_use_on_gpu.erase(found_memory_space_description);
+    }
   }
-  Guard{}.signalizeThatDangerOfRecursionHAsPassed();
+  Guard{}.signalizeThatDangerOfRecursionHasPassed();
 }
 
 Trace::~Trace()