[OpenMP][libomptarget] Centralize host pinned buffers map to NextGen's PluginInterface

[kevinsala]

Summary

This patch moves the management/tracking of host pinned buffers to the common PluginInterface in NextGen plugins. For the moment, the management consists of tracking the host pinned allocations into a map in each device.

[jdoerfert]

LG

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

 
 #include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <deque>
 #include <hsa.h>
 #include <hsa_ext_amd.h>
 #include <mutex>
-#include <shared_mutex>
 #include <string>
 #include <unistd.h>
 #include <unordered_map>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
 #include "PluginInterface.h"
if (!MemoryPool) {
       return OFFLOAD_FAIL;
     }
 
     if (Error Err = MemoryPool->deallocate(TgtPtr)) {
       REPORT("%s\n", toString(std::move(Err)).data());
       return OFFLOAD_FAIL;
     }
 
-    if (Kind == TARGET_ALLOC_HOST) {
-      std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
-      size_t Erased = HostAllocations.erase(TgtPtr);
-      if (!Erased) {
-        REPORT("Cannot find a host allocation in the map\n");
-        return OFFLOAD_FAIL;
-      }
-    }
-
     return OFFLOAD_SUCCESS;
   }
 
   /// Synchronize current thread with the pending operations on the async info.
   Error synchronizeImpl(__tgt_async_info &AsyncInfo) override {
     AMDGPUStreamTy *Stream =
         reinterpret_cast<AMDGPUStreamTy *>(AsyncInfo.Queue);
     assert(Stream && "Invalid stream");
Error queryAsyncImpl(__tgt_async_info &AsyncInfo) override {
     return Plugin::success();
   }
 
   /// Submit data to the device (host to device transfer).
   Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                        AsyncInfoWrapperTy &AsyncInfoWrapper) override {
 
     // Use one-step asynchronous operation when host memory is already pinned.
-    if (isHostPinnedMemory(HstPtr)) {
+    if (isHostPinnedMemoryBuffer(HstPtr)) {
       AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
       return Stream.pushPinnedMemoryCopyAsync(TgtPtr, HstPtr, Size);
     }
 
     void *PinnedHstPtr = nullptr;
 
     // For large transfers use synchronous behavior.
     if (Size >= OMPX_MaxAsyncCopyBytes) {
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
     AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
     return Stream.pushMemoryCopyH2DAsync(TgtPtr, HstPtr, PinnedHstPtr, Size,
                                          PinnedMemoryManager);
   }
 
   /// Retrieve data from the device (device to host transfer).
   Error dataRetrieveImpl(void *HstPtr, const void *TgtPtr, int64_t Size,
                          AsyncInfoWrapperTy &AsyncInfoWrapper) override {
-    if (isHostPinnedMemory(HstPtr)) {
+
+    // Use one-step asynchronous operation when host memory is already pinned.
+    if (isHostPinnedMemoryBuffer(HstPtr)) {
       // Use one-step asynchronous operation when host memory is already pinned.
       AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
       return Stream.pushPinnedMemoryCopyAsync(HstPtr, TgtPtr, Size);
     }
 
     void *PinnedHstPtr = nullptr;
 
     // For large transfers use synchronous behavior.
struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// Get the next queue in a round-robin fashion.
   AMDGPUQueueTy &getNextQueue() {
     static std::atomic<uint32_t> NextQueue(0);
 
     uint32_t Current = NextQueue.fetch_add(1, std::memory_order_relaxed);
     return Queues[Current % Queues.size()];
   }
 
-  /// Check whether a buffer is a host pinned buffer.
-  bool isHostPinnedMemory(const void *Ptr) const {
-    bool Found = false;
-    HostAllocationsMutex.lock_shared();
-    if (!HostAllocations.empty()) {
-      auto It = HostAllocations.lower_bound((const void *)Ptr);
-      if (It != HostAllocations.end() && It->first == Ptr) {
-        Found = true;
-      } else if (It != HostAllocations.begin()) {
-        --It;
-        Found = ((const char *)It->first + It->second > (const char *)Ptr);
-      }
-    }
-    HostAllocationsMutex.unlock_shared();
-    return Found;
-  }
-
 private:
   using AMDGPUStreamRef = AMDGPUResourceRef<AMDGPUStreamTy>;
   using AMDGPUEventRef = AMDGPUResourceRef<AMDGPUEventTy>;
 
   using AMDGPUStreamManagerTy = GenericDeviceResourceManagerTy<AMDGPUStreamRef>;
   using AMDGPUEventManagerTy = GenericDeviceResourceManagerTy<AMDGPUEventRef>;
 
   /// Envar for controlling the number of HSA queues per device. High number of
private:
   /// The agent handler corresponding to the device.
   hsa_agent_t Agent;
 
   /// Reference to the host device.
   AMDHostDeviceTy &HostDevice;
 
   /// List of device packet queues.
   std::vector<AMDGPUQueueTy> Queues;
-
-  /// Map of host pinned allocations. We track these pinned allocations so that
-  /// memory transfers involving these allocations do not need a two-step copy
-  /// with an intermediate pinned buffer.
-  std::map<const void *, size_t> HostAllocations;
-  mutable std::shared_mutex HostAllocationsMutex;
 };
 
 Error AMDGPUDeviceImageTy::loadExecutable(const AMDGPUDeviceTy &Device) {
   hsa_status_t Status;
   Status = hsa_code_object_deserialize(getStart(), getSize(), "", &CodeObject);
   if (auto Err =
           Plugin::check(Status, "Error in hsa_code_object_deserialize: %s"))
     return Err;
void *AMDGPUDeviceTy::allocate(size_t Size, void *, TargetAllocTy Kind) {
   if (Kind == TARGET_ALLOC_HOST && Alloc) {
     auto &KernelAgents = Plugin::get<AMDGPUPluginTy>().getKernelAgents();
 
     // Enable all kernel agents to access the host pinned buffer.
     if (auto Err = MemoryPool->enableAccess(Alloc, Size, KernelAgents)) {
       REPORT("%s\n", toString(std::move(Err)).data());
       return nullptr;
     }
-
-    // Keep track of the host pinned allocations for optimizations in transfers.
-    std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
-    HostAllocations.insert({Alloc, Size});
   }
 
   return Alloc;
 }
 
 } // namespace plugin
 } // namespace target
 } // namespace omp
 } // namespace llvm

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

 
 #ifndef OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_PLUGININTERFACE_H
 #define OPENMP_LIBOMPTARGET_PLUGINS_NEXTGEN_COMMON_PLUGININTERFACE_H
 
 #include <cstddef>
 #include <cstdint>
 #include <list>
 #include <map>
+#include <shared_mutex>
 #include <vector>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
 #include "GlobalHandler.h"
 #include "MemoryManager.h"
 #include "Utilities.h"
 #include "omptarget.h"
private:
   virtual Error getDeviceHeapSize(uint64_t &V) = 0;
   virtual Error setDeviceHeapSize(uint64_t V) = 0;
 
   /// Indicate whether the device should setup the device environment. Notice
   /// that returning false in this function will change the behavior of the
   /// setupDeviceEnvironment() function.
   virtual bool shouldSetupDeviceEnvironment() const { return true; }
 
+  /// Register a host buffer as host pinned allocation.
+  Error registerHostPinnedMemoryBuffer(const void *Buffer, size_t Size);
+
+  /// Unregister a host pinned allocations.
+  Error unregisterHostPinnedMemoryBuffer(const void *Buffer);
+
   /// Pointer to the memory manager or nullptr if not available.
   MemoryManagerTy *MemoryManager;
 
   /// Environment variables defined by the OpenMP standard.
   Int32Envar OMP_TeamLimit;
   Int32Envar OMP_NumTeams;
   Int32Envar OMP_TeamsThreadLimit;
 
   /// Environment variables defined by the LLVM OpenMP implementation.
   Int32Envar OMPX_DebugKind;
   UInt32Envar OMPX_SharedMemorySize;
   UInt64Envar OMPX_TargetStackSize;
   UInt64Envar OMPX_TargetHeapSize;
 
+  /// Map of host pinned allocations. We track these pinned allocations so that
+  /// memory transfers involving these allocations can be optimized.
+  std::map<const void *, size_t> HostAllocations;
+  mutable std::shared_mutex HostAllocationsMutex;
+
 protected:
+  /// Check whether a buffer has been registered as host pinned memory.
+  bool isHostPinnedMemoryBuffer(const void *Buffer) const {
+    std::shared_lock<std::shared_mutex> Lock(HostAllocationsMutex);
+
+    if (HostAllocations.empty())
+      return false;
+
+    // Search the first allocation with starting address that is not less than
+    // the buffer address.
+    auto It = HostAllocations.lower_bound(Buffer);
+
+    // Direct match of starting addresses.
+    if (It != HostAllocations.end() && It->first == Buffer)
+      return true;
+
+    // Not direct match but may be a previous pinned allocation in the map which
+    // contains the buffer. Return false if there is no such a previous
+    // allocation.
+    if (It == HostAllocations.begin())
+      return false;
+
+    // Move to the previous pinned allocation.
+    --It;
+
+    // Evaluate whether the buffer is contained in the pinned allocation.
+    return ((const char *)It->first + It->second > (const char *)Buffer);
+  }
+
   /// Environment variables defined by the LLVM OpenMP implementation
   /// regarding the initial number of streams and events.
   UInt32Envar OMPX_InitialNumStreams;
   UInt32Envar OMPX_InitialNumEvents;
 
   /// Array of images loaded into the device. Images are automatically
   /// deallocated by the allocator.
   llvm::SmallVector<DeviceImageTy *> LoadedImages;

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

Error GenericDeviceTy::registerKernelOffloadEntry(
   // Set the device entry address to the kernel address and store the entry on
   // the entry table.
   DeviceEntry.addr = (void *)Kernel;
   Image.getOffloadEntryTable().addEntry(DeviceEntry);
 
   return Plugin::success();
 }
 
+Error GenericDeviceTy::registerHostPinnedMemoryBuffer(const void *Buffer,
+                                                      size_t Size) {
+  std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
+
+  auto Res = HostAllocations.insert({Buffer, Size});
+  if (!Res.second)
+    return Plugin::error("Registering an already registered pinned buffer");
+
+  return Plugin::success();
+}
+
+Error GenericDeviceTy::unregisterHostPinnedMemoryBuffer(const void *Buffer) {
+  std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
+
+  size_t Erased = HostAllocations.erase(Buffer);
+  if (!Erased)
+    return Plugin::error("Cannot find a registered host pinned buffer");
+
+  return Plugin::success();
+}
+
 Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo) {
   if (!AsyncInfo || !AsyncInfo->Queue)
     return Plugin::error("Invalid async info queue");
 
   return synchronizeImpl(*AsyncInfo);
 }
 
 Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo) {
case TARGET_ALLOC_DEVICE:
     [[fallthrough]];
   case TARGET_ALLOC_HOST:
   case TARGET_ALLOC_SHARED:
     Alloc = allocate(Size, HostPtr, Kind);
     if (!Alloc)
       return Plugin::error("Failed to allocate from device allocator");
   }
 
-  // Sucessful and valid allocation.
-  if (Alloc)
-    return Alloc;
-
-  // At this point means that we did not tried to allocate from the memory
-  // manager nor the device allocator.
+  // Report error if the memory manager or the device allocator did not return
+  // any memory buffer.
+  if (!Alloc)
     return Plugin::error("Invalid target data allocation kind or requested "
                          "allocator not implemented yet");
+
+  // Register allocated buffer as pinned memory if the type is host memory.
+  if (Kind == TARGET_ALLOC_HOST)
+    if (auto Err = registerHostPinnedMemoryBuffer(Alloc, Size))
+      return Err;
+
+  return Alloc;
 }
 
 Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind) {
   int Res;
   if (MemoryManager)
     Res = MemoryManager->free(TgtPtr);
   else
     Res = free(TgtPtr, Kind);
 
   if (Res)
     return Plugin::error("Failure to deallocate device pointer %p", TgtPtr);
 
+  // Unregister deallocated pinned memory buffer if the type is host memory.
+  if (Kind == TARGET_ALLOC_HOST)
+    if (auto Err = unregisterHostPinnedMemoryBuffer(TgtPtr))
+      return std::move(Err);
+
   return Plugin::success();
 }
 
 Error GenericDeviceTy::dataSubmit(void *TgtPtr, const void *HstPtr,
                                   int64_t Size, __tgt_async_info *AsyncInfo) {
   auto Err = Plugin::success();
   AsyncInfoWrapperTy AsyncInfoWrapper(Err, *this, AsyncInfo);
   Err = dataSubmitImpl(TgtPtr, HstPtr, Size, AsyncInfoWrapper);