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

This should just be an error and handled from the user's side. If you requested SVM and the agent could not fulfill it that is an error.

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