From 5e6ad190ff1de3a3388fc3528db0ebf690ec2653 Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Thu, 8 Feb 2024 20:09:51 +0000
Subject: [PATCH] SWDEV-445221 - Fix the memory reported for
 hipMemPoolAttrReservedMemHigh

Change-Id: I894a836643612c73feb322f29a9236ada4d8cbda
---
 hipamd/src/hip_mempool_impl.cpp | 8 ++++----
 hipamd/src/hip_mempool_impl.hpp | 4 +++-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/hipamd/src/hip_mempool_impl.cpp b/hipamd/src/hip_mempool_impl.cpp
index 3afe065bb0..9e2ec4ca28 100644
--- a/hipamd/src/hip_mempool_impl.cpp
+++ b/hipamd/src/hip_mempool_impl.cpp
@@ -210,6 +210,8 @@ void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream, void* dptr) {
   // Place the allocated memory into the busy heap
   busy_heap_.AddMemory(memory, stream);
 
+  max_total_size_ = std::max(max_total_size_, busy_heap_.GetTotalSize() +
+                                                  free_heap_.GetTotalSize());
   // Increment the reference counter on the pool
   retain();
 
@@ -345,8 +347,7 @@ hipError_t MemoryPool::SetAttribute(hipMemPoolAttr attr, void* value) {
       if (reset != 0) {
         return hipErrorInvalidValue;
       }
-      free_heap_.SetMaxTotalSize(reset);
-      busy_heap_.SetMaxTotalSize(reset);
+      max_total_size_ = reset;
       break;
     case hipMemPoolAttrUsedMemCurrent:
       // Should be GetAttribute only
@@ -392,8 +393,7 @@ hipError_t MemoryPool::GetAttribute(hipMemPoolAttr attr, void* value) {
       break;
     case hipMemPoolAttrReservedMemHigh:
       // High watermark of all allocated memory in OS, since the last reset
-      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetMaxTotalSize() +
-                                            free_heap_.GetMaxTotalSize();
+      *reinterpret_cast<uint64_t*>(value) = max_total_size_;
       break;
     case hipMemPoolAttrUsedMemCurrent:
       // Total currently used memory by the pool
diff --git a/hipamd/src/hip_mempool_impl.hpp b/hipamd/src/hip_mempool_impl.hpp
index f625a79119..c330a8a0f2 100644
--- a/hipamd/src/hip_mempool_impl.hpp
+++ b/hipamd/src/hip_mempool_impl.hpp
@@ -187,7 +187,8 @@ class MemoryPool : public amd::ReferenceCountedObject {
         free_heap_(device),
         lock_pool_ops_("Pool operations", true),
         device_(device),
-        shared_(nullptr) {
+        shared_(nullptr),
+        max_total_size_(0) {
     device_->AddMemoryPool(this);
     state_.value_ = 0;
     state_.event_dependencies_ = 1;
@@ -286,6 +287,7 @@ private:
   std::map<hip::Device*, hipMemAccessFlags> access_map_;  //!< Map of access to the pool from devices
   hip::Device*  device_;    //!< Hip device the heap will reside
   SharedMemPool* shared_;   //!< Pointer to shared memory for IPC
+  uint64_t max_total_size_; //!< Max of total reserved memory in the pool since last reset
 };