[Openmp-commits] [llvm] [openmp] [offload] Use HSA SVM for AMDGPU shared memory (PR #215801)

[Joseph Huber via Openmp-commits]

================
@@ -450,6 +455,157 @@ struct AMDGPUMemoryPoolTy {
   size_t PoolAllocationAlignment;
 };
 
+/// Class that implements shared (managed) allocations on top of the HSA shared
+/// virtual memory (SVM) interface.
+///
+/// SVM allocations are backed by ordinary system memory that is made
+/// accessible to all the kernel agents. With XNACK enabled, the driver
+/// migrates the pages between the host and the devices on demand; otherwise
+/// they remain resident in system memory.
+struct AMDGPUSVMManagerTy {
+  /// Determine whether the SVM interface can be used for shared allocations.
+  void init() {
+    if (!OMPX_UseSVM)
+      return;
+
+    bool SVMSupported = false;
+    if (hsa_system_get_info(HSA_AMD_SYSTEM_INFO_SVM_SUPPORTED, &SVMSupported) !=
+        HSA_STATUS_SUCCESS)
+      return;
+
+    Supported = SVMSupported;
+    PageSize = llvm::sys::Process::getPageSizeEstimate();
+  }
+
+  /// Release the allocations that are still live.
+  Error deinit() {
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    Error Err = Plugin::success();
+    for (auto &Allocation : Allocations)
+      if (std::error_code EC =
+              llvm::sys::Memory::releaseMappedMemory(Allocation.second))
+        Err = joinErrors(std::move(Err),
+                         Plugin::error(ErrorCode::UNKNOWN,
+                                       "error releasing SVM memory: %s",
+                                       EC.message().c_str()));
+
+    Allocations.clear();
+    return Err;
+  }
+
+  /// Whether new SVM allocations can be created.
+  bool isSupported() const { return Supported; }
+
+  /// Allocate system memory and make it accessible to all the \p Agents.
+  /// Returns a null pointer if the request cannot be served, either because
+  /// the requested alignment is too large or because the driver rejected the
+  /// allocation.
+  Expected<void *> allocate(size_t Size, size_t Alignment,
+                            ArrayRef<hsa_agent_t> Agents) {
+    // A concurrent allocation may have disabled the SVM interface.
+    if (!Supported)
+      return nullptr;
+
+    // Mapped memory is only page aligned.
+    if (Alignment > PageSize)
+      return nullptr;
+
+    std::error_code EC;
+    llvm::sys::MemoryBlock Block = llvm::sys::Memory::allocateMappedMemory(
+        Size, /*NearBlock=*/nullptr,
+        llvm::sys::Memory::MF_READ | llvm::sys::Memory::MF_WRITE, EC);
+    if (EC)
+      return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
+                           "error allocating SVM memory: %s",
+                           EC.message().c_str());
+
+    // Give all the kernel agents access to the allocation. Accesses may incur
+    // a page fault and the migration of the memory to the accessing agent.
+    llvm::SmallVector<hsa_amd_svm_attribute_pair_t> Attrs;
+    for (hsa_agent_t Agent : Agents)
+      Attrs.push_back({HSA_AMD_SVM_ATTRIB_AGENT_ACCESSIBLE, Agent.handle});
+
+    hsa_status_t Status = hsa_amd_svm_attributes_set(
+        Block.base(), Block.allocatedSize(), Attrs.data(), Attrs.size());
+    if (auto Err =
+            Plugin::check(Status, "error in hsa_amd_svm_attributes_set: %s")) {
+      // The driver rejected the allocation, e.g., because one of the agents
+      // does not support SVM. Stop creating SVM allocations, reporting the
+      // rejection once.
+      if (Supported.exchange(false))
+        REPORT() << "Serving shared allocations from host memory: "
+                 << toString(std::move(Err));
+      else
+        consumeError(std::move(Err));
+
+      // Ignore any error from undoing the allocation.
+      consumeError(llvm::errorCodeToError(
+          llvm::sys::Memory::releaseMappedMemory(Block)));
+      return nullptr;
+    }
+
+    std::lock_guard<std::mutex> Lock(Mutex);
+    Allocations[reinterpret_cast<uintptr_t>(Block.base())] = Block;
+    return Block.base();
+  }
+
+  /// Release the SVM allocation starting at \p Ptr. Returns whether \p Ptr is
+  /// an SVM allocation.
+  Expected<bool> deallocate(void *Ptr) {
+    llvm::sys::MemoryBlock Block;
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      auto It = Allocations.find(reinterpret_cast<uintptr_t>(Ptr));
+      if (It == Allocations.end())
+        return false;
+      Block = It->second;
+      Allocations.erase(It);
+    }
+
+    if (std::error_code EC = llvm::sys::Memory::releaseMappedMemory(Block))
+      return Plugin::error(ErrorCode::UNKNOWN, "error releasing SVM memory: %s",
+                           EC.message().c_str());
+    return true;
+  }
+
+  /// Whether the \p Size bytes starting at \p Ptr are within an SVM
+  /// allocation. Allocations stay live after the SVM interface is disabled, so
+  /// this must not check Supported.
+  bool contains(const void *Ptr, size_t Size = 1) const {
----------------
jhuber6 wrote:

Do we need a full region check? The user shouldn't corrupt the pointers returned by the allocator. I'd only expect to need this if we were caching pages. i.e. several allocations share a 4KiB page and before freeing it all others need to be freed.

https://github.com/llvm/llvm-project/pull/215801