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