SWDEV-532479 - Add tracking of hostcall memory allocations (#416)

* SWDEV-532479 - Add tracking of hostcall memory allocations

* SWDEV-532479 - Remove hostcall allocations if request is received

* SWDEV-532479 - Cleanup

* SWDEV-532479 - Naming fix

* SWDEV-532479 - Add new separator after each new function

rocclr/device/devhostcall.cpp

@@ -94,6 +94,7 @@ static void handlePayload(MessageHandler& messages, uint32_t service, uint64_t*
       if (payload[0]) {
         amd::Memory* mem = amd::MemObjMap::FindMemObj(reinterpret_cast<void*>(payload[0]));
         if (mem) {
+          const_cast<amd::Device*>(&dev)->RemoveHostcallMemory(mem);
           amd::MemObjMap::RemoveMemObj(reinterpret_cast<void*>(payload[0]));
           mem->release();
         } else {
@@ -110,6 +111,7 @@ static void handlePayload(MessageHandler& messages, uint32_t service, uint64_t*
             device::Memory* dm = buf->getDeviceMemory(dev);
             va = dm->virtualAddress();
             amd::MemObjMap::AddMemObj(reinterpret_cast<void*>(va), buf);
+            const_cast<amd::Device*>(&dev)->TrackHostcallMemory(buf);
           } else {
             buf->release();
           }

rocclr/device/device.cpp

@@ -807,6 +807,16 @@ Device::~Device() {
     delete vaCacheMap_;
   }
 
+  for (auto memory : hostcall_allocated_memories_) {
+    if (memory != nullptr) {
+      amd::MemObjMap::RemoveMemObj(
+          reinterpret_cast<void*>(memory->getDeviceMemory(*this, false)->virtualAddress()));
+      memory->release();
+    }
+  }
+
+  hostcall_allocated_memories_.clear();
+
   delete vaCacheAccess_;
   delete settings_;
   delete[] info_.extensions_;
@@ -1205,6 +1215,20 @@ bool Device::GetHandleForAddressRange(void* dev_ptr, size_t size, void* handle)
   return dev_mem->GetFDHandleForMem(dev_ptr, size, VmmPtr, handle);
 }
 
+// ================================================================================================
+void Device::TrackHostcallMemory(amd::Memory* memory) {
+  hostcall_allocated_memories_.push_back(memory);
+}
+
+// ================================================================================================
+void Device::RemoveHostcallMemory(amd::Memory* memory) {
+  auto it =
+      std::find(hostcall_allocated_memories_.begin(), hostcall_allocated_memories_.end(), memory);
+  if (it != hostcall_allocated_memories_.end()) {
+    hostcall_allocated_memories_.erase(it);
+  }
+}
+
 }  // namespace amd
 
 namespace amd::device {

rocclr/device/device.hpp

@@ -2226,7 +2226,12 @@ class Device : public RuntimeObject {
 
   bool GetHandleForAddressRange(void* dev_ptr, size_t size, void* handle);
 
- protected:
+  // Registers a memory object allocated via hostcall for later cleanup.
+  void TrackHostcallMemory(amd::Memory* memory);
+
+  // Removes a memory object from hostcall tracking.
+  void RemoveHostcallMemory(amd::Memory* memory);
+
   //! Enable the specified extension
   char* getExtensionString();
 
@@ -2274,6 +2279,9 @@ class Device : public RuntimeObject {
   Monitor* vaCacheAccess_;                            //!< Lock to serialize VA caching access
   std::map<uintptr_t, device::Memory*>* vaCacheMap_;  //!< VA cache map
   uint32_t index_;  //!< Unique device index
+
+  // Tracks all amd::Memory objects allocated via hostcall for this device.
+  std::vector<amd::Memory*> hostcall_allocated_memories_;
 };
 
 /*! @}