From 8b132beb9243dee2ff06e00f11f0181549b29a19 Mon Sep 17 00:00:00 2001 From: German Andryeyev Date: Tue, 3 May 2022 12:25:46 -0400 Subject: [PATCH] SWDEV-311271 - Enable MGPU support for memory pool Change-Id: I36850de282d62139b5cfe342df97bb1204cb7869 --- hipamd/src/hip_mempool.cpp | 31 ++++++++++++-- hipamd/src/hip_mempool_impl.cpp | 71 +++++++++++++++++++++++++++++++++ hipamd/src/hip_mempool_impl.hpp | 10 +++++ 3 files changed, 109 insertions(+), 3 deletions(-) diff --git a/hipamd/src/hip_mempool.cpp b/hipamd/src/hip_mempool.cpp index 5992b018e4..960845f1e0 100644 --- a/hipamd/src/hip_mempool.cpp +++ b/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(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(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(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); } diff --git a/hipamd/src/hip_mempool_impl.cpp b/hipamd/src/hip_mempool_impl.cpp index 288359d4ad..84e8e4a526 100644 --- a/hipamd/src/hip_mempool_impl.cpp +++ b/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(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]; + } +} + } diff --git a/hipamd/src/hip_mempool_impl.hpp b/hipamd/src/hip_mempool_impl.hpp index a7547252c6..69ac92a50d 100644 --- a/hipamd/src/hip_mempool_impl.hpp +++ b/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 access_map_; //!< Map of access to the pool from devices hip::Device* device_; //!< Hip device the heap will reside };