[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)
----------------
jhuber6 wrote:

We could handle this more intentionally but I'm unsure if it's worthwhile. The standard heuristic is to over-allocate then `munmap` the pages before the alignment you needed.

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