SWDEV-485904 - Fix virtual,physical mem obj leaks

Change-Id: Ie0456b5dcfec206ae54a6aabfc2a15a620cac693

hipamd/src/hip_memory.cpp

@@ -426,16 +426,16 @@ hipError_t ihipMemcpy_validate(void* dst, const void* src, size_t sizeBytes,
       return hipErrorUnknown;
     }
 
-    // Size validation
-    if (sizeBytes > (srcMemory->getSize() - sOffset)) {
-      return hipErrorInvalidValue;
-    }
-
     // If the mem object is a VMM sub buffer (subbuffer has parent set),
     // then use parent's size for validation.
     if (srcMemory->parent() && (srcMemory->getMemFlags() & CL_MEM_VA_RANGE_AMD)) {
       srcMemory = srcMemory->parent();
     }
+
+    // Size validation
+    if (sizeBytes > (srcMemory->getSize() - sOffset)) {
+      return hipErrorInvalidValue;
+    }
   }
 
   if (dstMemory != nullptr) {
@@ -444,16 +444,16 @@ hipError_t ihipMemcpy_validate(void* dst, const void* src, size_t sizeBytes,
       return hipErrorUnknown;
     }
 
-    // Size validation
-    if (sizeBytes > (dstMemory->getSize() - dOffset)) {
-      return hipErrorInvalidValue;
-    }
-
     // If the mem object is a VMM sub buffer (subbuffer has parent set),
     // then use parent's size for validation.
     if (dstMemory->parent() && (dstMemory->getMemFlags() & CL_MEM_VA_RANGE_AMD)) {
       dstMemory = dstMemory->parent();
     }
+
+    // Size validation
+    if (sizeBytes > (dstMemory->getSize() - dOffset)) {
+      return hipErrorInvalidValue;
+    }
   }
 
   //If src and dst ptr are null then kind must be either h2h or def.

hipamd/src/hip_vm.cpp

@@ -39,10 +39,13 @@ hipError_t hipMemAddressFree(void* devPtr, size_t size) {
   if (devPtr == nullptr || size == 0) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-
+  amd::Memory* memObj = amd::MemObjMap::FindVirtualMemObj(devPtr);
+  if (memObj == nullptr) {
+    LogPrintfError("Cannot find the Virtual MemObj entry for this addr 0x%x", devPtr);
+  }
   // Single call frees address range for all devices.
   g_devices[0]->devices()[0]->virtualFree(devPtr);
-
+  memObj->release();
   HIP_RETURN(hipSuccess);
 }
 
@@ -126,10 +129,6 @@ hipError_t hipMemCreate(hipMemGenericAllocationHandle_t* handle, size_t size,
   phys_mem_obj->getUserData().data = new hip::GenericAllocation(*phys_mem_obj, size, *prop);
   *handle = reinterpret_cast<hipMemGenericAllocationHandle_t>(phys_mem_obj->getUserData().data);
 
-  // Remove because the entry of 0x1 is not needed in MemObjMap.
-  // We save the copy of Phy mem obj in virtual mem obj during mapping.
-  amd::MemObjMap::RemoveMemObj(ptr);
-
   HIP_RETURN(hipSuccess);
 }
 
@@ -261,7 +260,6 @@ hipError_t hipMemMap(void* ptr, size_t size, size_t offset, hipMemGenericAllocat
   ga->retain();
 
   auto& queue = *g_devices[ga->GetProperties().location.id]->NullStream();
-
   // Map the physical address to virtual address
   amd::Command* cmd = new amd::VirtualMapCommand(queue, amd::Command::EventWaitList{}, ptr, size,
                                                  &ga->asAmdMemory());
@@ -366,6 +364,7 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
   cmd->enqueue();
   cmd->awaitCompletion();
   cmd->release();
+  vaddr_sub_obj->release();
 
   // restore the original pa of the generic allocation
   hip::GenericAllocation* ga

hipamd/src/hip_vm.hpp

@@ -38,7 +38,10 @@ class GenericAllocation : public amd::RuntimeObject {
 public:
   GenericAllocation(amd::Memory& phys_mem_ref, size_t size, const hipMemAllocationProp& prop)
                     : phys_mem_ref_(phys_mem_ref), size_(size), properties_(prop) {}
-  ~GenericAllocation() {}
+  ~GenericAllocation() {
+    amd::Context* amdContext = g_devices[properties_.location.id]->asContext();
+    amd::SvmBuffer::free(*amdContext, phys_mem_ref_.getSvmPtr());
+  }
 
   const hipMemAllocationProp& GetProperties() const { return properties_; }
   hipMemGenericAllocationHandle_t asMemGenericAllocationHandle() {

rocclr/device/device.cpp

@@ -479,11 +479,7 @@ bool Device::DestroyVirtualBuffer(amd::Memory* vaddr_mem_obj) {
     return false;
   }
 
-  if (vaddr_mem_obj->parent() == nullptr) {
-    // If parent is nullptr, then vaddr_mem_obj is the parent.
-    amd::MemObjMap::RemoveVirtualMemObj(vaddr_mem_obj->getSvmPtr());
-    return true;
-  } else {
+  if (vaddr_mem_obj->parent() != nullptr) {
     // If parent is not nullptr, this is the sub-buffer object.
     amd::Memory* vaddr_base_obj = amd::MemObjMap::FindVirtualMemObj(vaddr_mem_obj->getSvmPtr());
     if (vaddr_base_obj == nullptr) {

rocclr/platform/memory.cpp

@@ -456,7 +456,9 @@ Memory::~Memory() {
   }
   hostMemRef_.deallocateMemory(context_());
   if (getMemFlags() & CL_MEM_VA_RANGE_AMD) {
-    amd::MemObjMap::RemoveVirtualMemObj(getSvmPtr());
+    if (parent_ == nullptr) {
+      amd::MemObjMap::RemoveVirtualMemObj(getSvmPtr());
+    }
     // If runtime executes graph mempool with VM, then VA can be mapped in space
     // for graph validation logic during execution. And the reason it's not unmaped
     // in graph itself because the app can have a graph without a free node