SWDEV-482130 - Fix release of virtual mem obj

Change-Id: I893a8353aa1a25d00e36c8e601caf31cc0fc1f22

hipamd/src/hip_graph_internal.hpp

@@ -2391,6 +2391,13 @@ class GraphMemAllocNode final : public GraphNode {
 
   virtual ~GraphMemAllocNode() final {
     if (va_ != nullptr) {
+      if (va_->referenceCount() == 1) {
+        auto graph = GetParentGraph();
+        if (graph != nullptr) {
+          graph->FreeAddress(va_->getSvmPtr());
+        }
+      }
+
       va_->release();
     }
   }
@@ -2498,6 +2505,7 @@ class GraphMemFreeNode : public GraphNode {
         hip::setCurrentDevice(device_id_);
       }
       // Free virtual address
+      vaddr_sub_obj->release();
       vaddr_mem_obj->release();
       // Release the allocation back to graph's pool
       graph_->FreeMemory(phys_mem_obj->getSvmPtr(), static_cast<hip::Stream*>(queue()));

hipamd/src/hip_mempool_impl.cpp

@@ -95,8 +95,14 @@ bool Heap::RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts) {
 Heap::SortedMap::iterator Heap::EraseAllocaton(Heap::SortedMap::iterator& it) {
   auto memory = it->first.second;
   const device::Memory* dev_mem = memory->getDeviceMemory(*device_->devices()[0]);
+  void* dev_mem_vaddr = reinterpret_cast<void*>(dev_mem->virtualAddress());
   total_size_ -= it->first.first;
-  amd::SvmBuffer::free(memory->getContext(), reinterpret_cast<void*>(dev_mem->virtualAddress()));
+  
+  if (dev_mem_vaddr != nullptr) {
+    amd::SvmBuffer::free(memory->getContext(), dev_mem_vaddr);
+  } else {
+    amd::SvmBuffer::free(memory->getContext(), memory->getSvmPtr());
+  }
   // Clear HIP event
   it->second.SetEvent(nullptr);
   // Remove the allocation from the map

rocclr/platform/memory.cpp

@@ -416,21 +416,21 @@ device::Memory* Memory::getDeviceMemory(const Device& dev, bool alloc) {
 Memory::~Memory() {
   // For_each destructor callback:
   DestructorCallBackEntry* entry;
-  for (entry = destructorCallbacks_; entry != NULL; entry = entry->next_) {
+  for (entry = destructorCallbacks_; entry != nullptr; entry = entry->next_) {
     // invoke the callback function.
     entry->callback_(const_cast<cl_mem>(as_cl(this)), entry->data_);
   }
 
   // Release the parent.
-  if (NULL != parent_) {
+  if (parent_ != nullptr) {
     // Update cache if runtime destroys a subbuffer
-    if (NULL != parent_->getHostMem() && (vDev_ == NULL)) {
+    if (parent_->getHostMem() != nullptr && (vDev_ == nullptr)) {
       cacheWriteBack(nullptr);
     }
     parent_->removeSubBuffer(this);
   }
 
-  if (NULL != deviceMemories_) {
+  if (deviceMemories_ != nullptr) {
     // Destroy all device memory objects
     for (uint i = 0; i < numDevices_; ++i) {
       delete deviceMemories_[i].value_;
@@ -444,19 +444,20 @@ Memory::~Memory() {
 
   // Destroy the destructor callback entries
   DestructorCallBackEntry* callback = destructorCallbacks_;
-  while (callback != NULL) {
+  while (callback != nullptr) {
     DestructorCallBackEntry* next = callback->next_;
     delete callback;
     callback = next;
   }
 
   // Make sure runtime destroys the parent only after subbuffer destruction
-  if (NULL != parent_) {
+  if (parent_ != nullptr) {
     parent_->release();
   }
   hostMemRef_.deallocateMemory(context_());
-  if (getMemFlags() & CL_MEM_VA_RANGE_AMD) {
-    if (parent_ == nullptr) {
+  if (parent_ == nullptr && (getMemFlags() & CL_MEM_VA_RANGE_AMD)) {
+    // The mapping may manually be removed prior to objects destruction
+    if (amd::MemObjMap::FindVirtualMemObj(getSvmPtr())) {
       amd::MemObjMap::RemoveVirtualMemObj(getSvmPtr());
     }
     // If runtime executes graph mempool with VM, then VA can be mapped in space