SWDEV-311271 - Enable MGPU support for memory pool

Change-Id: I36850de282d62139b5cfe342df97bb1204cb7869

hipamd/src/hip_mempool.cpp

@@ -115,9 +115,6 @@ hipError_t hipMemPoolGetAttribute(hipMemPool_t mem_pool, hipMemPoolAttr attr, vo
   if (mem_pool == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  if (mem_pool == nullptr) {
-    HIP_RETURN(hipErrorInvalidValue);
-  }
   auto hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
   HIP_RETURN(hip_mem_pool->GetAttribute(attr, value));
 }
@@ -128,6 +125,21 @@ hipError_t hipMemPoolSetAccess(
     const hipMemAccessDesc* desc_list,
     size_t count) {
   HIP_INIT_API(hipMemPoolSetAccess, mem_pool, desc_list, count);
+  if ((mem_pool == nullptr) || (desc_list == nullptr)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  for (int i = 0; i < count; ++i) {
+    if (desc_list[i].location.type == hipMemLocationTypeDevice) {
+      if (desc_list[i].location.id >= g_devices.size()) {
+        HIP_RETURN(hipErrorInvalidValue);
+      }
+      auto device = g_devices[desc_list[i].location.id];
+      hip_mem_pool->SetAccess(device, desc_list[i].flags);
+    } else {
+      HIP_RETURN(hipErrorInvalidValue);
+    }
+  }
   HIP_RETURN(hipSuccess);
 }
 
@@ -137,6 +149,19 @@ hipError_t hipMemPoolGetAccess(
     hipMemPool_t mem_pool,
     hipMemLocation* location) {
   HIP_INIT_API(hipMemPoolGetAccess, flags, mem_pool, location);
+  if ((mem_pool == nullptr) || (location == nullptr) || (flags == nullptr)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  if (location->type == hipMemLocationTypeDevice) {
+    if (location->id >= g_devices.size()) {
+      HIP_RETURN(hipErrorInvalidValue);
+    }
+    auto device = g_devices[location->id];
+    hip_mem_pool->GetAccess(device, flags);
+  } else {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
   HIP_RETURN(hipSuccess);
 }
 

hipamd/src/hip_mempool_impl.cpp

@@ -127,6 +127,25 @@ void Heap::RemoveStream(hip::Stream* stream) {
   }
 }
 
+// ================================================================================================
+void Heap::SetAccess(hip::Device* device, bool enable) {
+  for (const auto& it : allocations_) {
+    auto peer_device = device->asContext()->devices()[0];
+    device::Memory* mem = it.first->getDeviceMemory(*peer_device);
+    if (mem != nullptr) {
+      if (!mem->getAllowedPeerAccess() && enable) {
+        // Enable p2p access for the specified device
+        peer_device->deviceAllowAccess(reinterpret_cast<void*>(mem->virtualAddress()));
+        mem->setAllowedPeerAccess(true);
+      } else if (mem->getAllowedPeerAccess() && !enable) {
+        mem->setAllowedPeerAccess(false);
+      }
+    } else {
+      LogError("Couldn't find device memory for P2P access");
+    }
+  }
+}
+
 // ================================================================================================
 void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream) {
   amd::ScopedLock lock(lock_pool_ops_);
@@ -155,6 +174,15 @@ void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream) {
     memory = getMemoryObject(dev_ptr, offset);
     // Saves the current device id so that it can be accessed later
     memory->getUserData().deviceId = device_->deviceId();
+
+    // Update access for the new allocation from other devices
+    for (const auto& it : access_map_) {
+      auto vdi_device = it.first->asContext()->devices()[0];
+      device::Memory* mem = memory->getDeviceMemory(*vdi_device);
+      if ((mem != nullptr) && (it.second != hipMemAccessFlagsProtNone)) {
+        mem->setAllowedPeerAccess(true);
+      }
+    }
   } else {
     free_heap_.RemoveMemory(memory);
     const device::Memory* dev_mem = memory->getDeviceMemory(*device_->devices()[0]);
@@ -310,4 +338,47 @@ hipError_t MemoryPool::GetAttribute(hipMemPoolAttr attr, void* value) {
   return hipSuccess;
 }
 
+// ================================================================================================
+void MemoryPool::SetAccess(hip::Device* device, hipMemAccessFlags flags) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  // Check if the requested device is the pool device where memory was allocated
+  if (device == device_) {
+    return;
+  }
+
+  hipMemAccessFlags current_flags = hipMemAccessFlagsProtNone;
+
+  // Check if access was enabled before
+  if (access_map_.find(device) != access_map_.end()) {
+    current_flags = access_map_[device];
+  }
+
+  if (current_flags != flags) {
+    bool enable_access = false;
+    // Save the access state  in the device map
+    access_map_[device] = flags;
+    // Check if access is enabled
+    if ((flags == hipMemAccessFlagsProtRead) || (flags == hipMemAccessFlagsProtReadWrite)) {
+      enable_access = true;
+    }
+    // Update device access on the both pools
+    busy_heap_.SetAccess(device, enable_access);
+    free_heap_.SetAccess(device, enable_access);
+  }
+}
+
+// ================================================================================================
+void MemoryPool::GetAccess(hip::Device* device, hipMemAccessFlags* flags) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  // Current pool device has full access to memory allocation
+  *flags = (device == device_) ? hipMemAccessFlagsProtReadWrite : hipMemAccessFlagsProtNone;
+
+  // Check if access was enabled before
+  if (access_map_.find(device) != access_map_.end()) {
+    *flags = access_map_[device];
+  }
+}
+
 }

hipamd/src/hip_mempool_impl.hpp

@@ -107,6 +107,9 @@ public:
   /// Remove the provided stream from the safe list
   void RemoveStream(hip::Stream* stream);
 
+  /// Enables P2P access to the provided device
+  void SetAccess(hip::Device* device, bool enable);
+
   /// Heap doesn't have any allocations
   bool IsEmpty() const { return (allocations_.size() == 0) ? true : false; }
 
@@ -193,6 +196,12 @@ public:
   /// Get memory pool control attributes
   hipError_t GetAttribute(hipMemPoolAttr attr, void* value);
 
+  /// Set memory pool access by different devices
+  void SetAccess(hip::Device* device, hipMemAccessFlags flags);
+
+  /// Set memory pool access by different devices
+  void GetAccess(hip::Device* device, hipMemAccessFlags* flags);
+
   /// Accessors for the pool state
   bool EventDependencies() const { return (state_.event_dependencies_) ? true : false; }
   bool Opportunistic() const { return (state_.opportunistic_) ? true : false; }
@@ -213,6 +222,7 @@ private:
   } state_;
 
   amd::Monitor  lock_pool_ops_;  //!< Access to the pool must be lock protected
+  std::map<hip::Device*, hipMemAccessFlags> access_map_;  //!< Map of access to the pool from devices
   hip::Device*  device_;    //!< Hip device the heap will reside
 };