SWDEV-505795 - Return the same ptr from hipIpcOpenMemHandle if it is called multiple times (#93)

* SWDEV-505795 - Return the same ptr from hipIpcOpenMemHandle if it is called multiple times

* Move initialization outside of if statement

hipamd/src/hip_memory.cpp

@@ -3274,14 +3274,24 @@ hipError_t hipIpcOpenMemHandle(void** dev_ptr, hipIpcMemHandle_t handle, unsigne
     HIP_RETURN(hipErrorInvalidContext);
   }
 
-  if(!device->IpcAttach(&(ihandle->ipc_handle), ihandle->psize,
-                        ihandle->poffset, flags, dev_ptr)) {
-    LogPrintfError("Cannot attach ipc_handle: with ipc_size: %u"
-                      "ipc_offset: %u flags: %u", ihandle->psize, flags);
-    HIP_RETURN(hipErrorInvalidDevicePointer);
+  amd_mem_obj = amd::MemObjMap::FindIpcHandleMemObj(ihandle);
+  if (amd_mem_obj == nullptr) {
+    if (!device->IpcAttach(&(ihandle->ipc_handle), ihandle->psize, ihandle->poffset, flags,
+                           dev_ptr)) {
+      LogPrintfError(
+          "Cannot attach ipc_handle: with ipc_size: %u"
+          "ipc_offset: %u flags: %u",
+          ihandle->psize, flags);
+      HIP_RETURN(hipErrorInvalidDevicePointer);
+    }
+    amd_mem_obj = getMemoryObject(*dev_ptr, offset);
+    amd::MemObjMap::AddIpcHandleMemObj(ihandle, amd_mem_obj);
+  } else {
+    // Handle was already opened by the same process
+    *dev_ptr = amd_mem_obj->getSvmPtr();
+    amd_mem_obj->retain();
   }
 
-  amd_mem_obj = getMemoryObject(*dev_ptr, offset);
   amd_mem_obj->getUserData().deviceId = hip::getCurrentDevice()->deviceId();
 
   HIP_RETURN(hipSuccess, ReturnPtrValue(dev_ptr));
@@ -3299,23 +3309,36 @@ hipError_t hipIpcCloseMemHandle(void* dev_ptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
 
-  amd_mem_obj = amd::MemObjMap::FindMemObj(dev_ptr);
-  if (amd_mem_obj != nullptr) {
-    auto device_id = amd_mem_obj->getUserData().deviceId;
-    g_devices[device_id]->SyncAllStreams();
-  }
-
   /* Call IPC Detach from Device class */
   device = hip::getCurrentDevice()->devices()[0];
   if (device == nullptr) {
     HIP_RETURN(hipErrorNoDevice);
   }
 
-  /* detach the memory */
-  if (!device->IpcDetach(dev_ptr)){
+  amd_mem_obj = amd::MemObjMap::FindMemObj(dev_ptr);
+  if (amd_mem_obj == nullptr) {
+    DevLogPrintfError("Memory object for the ptr: 0x%x cannot be null \n", dev_ptr);
     HIP_RETURN(hipErrorInvalidValue);
   }
 
+  if (!amd_mem_obj->ipcShared()) {
+    DevLogPrintfError("Memory object for the ptr: 0x%x is not ipcShared \n", dev_ptr);
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  // If handle is opened more than once, do detach on last release call
+  if (amd_mem_obj->referenceCount() == 1) {
+    auto device_id = amd_mem_obj->getUserData().deviceId;
+    g_devices[device_id]->SyncAllStreams();
+
+    // Remove entry from the map
+    amd::MemObjMap::RemoveIpcHandleMemObj(amd_mem_obj);
+    /* detach the memory */
+    device->IpcDetach(dev_ptr);
+  } else {
+    amd_mem_obj->release();
+  }
+
   HIP_RETURN(hipSuccess);
 }
 

rocclr/device/device.cpp

@@ -339,6 +339,7 @@ Memory* Device::p2p_stage_ = nullptr;
 std::shared_mutex MemObjMap::AllocatedLock_ ROCCLR_INIT_PRIORITY(101);
 std::map<uintptr_t, amd::Memory*> MemObjMap::MemObjMap_ ROCCLR_INIT_PRIORITY(101);
 std::map<uintptr_t, amd::Memory*> MemObjMap::VirtualMemObjMap_ ROCCLR_INIT_PRIORITY(101);
+std::unordered_map<uintptr_t, amd::Memory*> MemObjMap::IpcHandleMemObjMap_ ROCCLR_INIT_PRIORITY(101);
 
 void MemObjMap::AddMemObj(const void* k, amd::Memory* v) {
   std::unique_lock lock(AllocatedLock_);
@@ -446,6 +447,38 @@ amd::Memory* MemObjMap::FindVirtualMemObj(const void* k) {
   }
 }
 
+void MemObjMap::AddIpcHandleMemObj(const void* k, amd::Memory* v) {
+  std::unique_lock lock(AllocatedLock_);
+  auto rval = IpcHandleMemObjMap_.insert({reinterpret_cast<uintptr_t>(k), v});
+  if (!rval.second) {
+    DevLogPrintfError("IpcHandle Memobj map already has an entry for ptr: 0x%x",
+                      reinterpret_cast<uintptr_t>(k));
+  }
+}
+
+void MemObjMap::RemoveIpcHandleMemObj(amd::Memory* v) {
+  std::unique_lock lock(AllocatedLock_);
+
+  for (const auto it : IpcHandleMemObjMap_) {
+    if (it.second == v) {
+      IpcHandleMemObjMap_.erase(it.first);
+      break;
+    }
+  }
+}
+
+amd::Memory* MemObjMap::FindIpcHandleMemObj(const void* k) {
+  std::shared_lock lock(AllocatedLock_);
+  uintptr_t key = reinterpret_cast<uintptr_t>(k);
+
+  auto it = IpcHandleMemObjMap_.find(key);
+  if (it == IpcHandleMemObjMap_.cend()) {
+    return nullptr;
+  }
+
+  return it->second;
+}
+
 //==================================================================================================
 bool Device::ValidateVirtualAddressRange(amd::Memory* vaddr_base_obj, amd::Memory* vaddr_sub_obj) {
 
@@ -1056,17 +1089,8 @@ bool Device::IpcAttach(const void* handle, size_t mem_size, size_t mem_offset, u
 }
 
 // ================================================================================================
-bool Device::IpcDetach(void* dev_ptr) const {
+void Device::IpcDetach(void* dev_ptr) const {
   amd::Memory* amd_mem_obj = amd::MemObjMap::FindMemObj(dev_ptr);
-  if (amd_mem_obj == nullptr) {
-    DevLogPrintfError("Memory object for the ptr: 0x%x cannot be null \n", dev_ptr);
-    return false;
-  }
-
-  if (!amd_mem_obj->ipcShared()) {
-    DevLogPrintfError("Memory object for the ptr: 0x%x is not ipcShared \n", dev_ptr);
-    return false;
-  }
 
   // Get the original pointer from the amd::Memory object
   void* orig_dev_ptr = nullptr;
@@ -1081,8 +1105,6 @@ bool Device::IpcDetach(void* dev_ptr) const {
   if (amd_mem_obj->release() == 0) {
     amd::MemObjMap::RemoveMemObj(orig_dev_ptr);
   }
-
-  return true;
 }
 
 }  // namespace amd

rocclr/device/device.hpp

@@ -1393,6 +1393,13 @@ class MemObjMap : public AllStatic {
   //!< Same as FindMemObj but for virtual addressing
   static amd::Memory* FindVirtualMemObj(const void* k);
 
+  //!< Same as AddMemObj but for virtual ipc handle to MemObj mapping
+  static void AddIpcHandleMemObj(const void* k, amd::Memory* v);
+  //!< Remove entry from the map by searching values
+  static void RemoveIpcHandleMemObj(amd::Memory* v);
+  //!< Same as FindMemObj but for ipc handle to MemObj mapping
+  static amd::Memory* FindIpcHandleMemObj(const void* k);
+
  private:
   //!< the mem object<->hostptr information container
   static std::map<uintptr_t, amd::Memory*> MemObjMap_;
@@ -1400,6 +1407,8 @@ class MemObjMap : public AllStatic {
   static std::map<uintptr_t, amd::Memory*> VirtualMemObjMap_;
   //!< Shared read/write lock
   static std::shared_mutex AllocatedLock_;
+  //!< the ipc handle<->mem object information container
+  static std::unordered_map<uintptr_t, amd::Memory*> IpcHandleMemObjMap_;
 };
 
 /// @brief Instruction Set Architecture properties.
@@ -2065,7 +2074,7 @@ class Device : public RuntimeObject {
   bool IpcAttach(const void* handle, size_t mem_size, size_t mem_offset, unsigned int flags,
                  void** dev_ptr) const;
 
-  bool IpcDetach(void* dev_ptr) const;
+  void IpcDetach(void* dev_ptr) const;
 
   //! Return context
   amd::Context& context() const { return *context_; }