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

Robert Imschweiler via Openmp-commits openmp-commits at lists.llvm.org
Sat Oct 3 05:32:20 PDT 2026


https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/215801

>From 8673f8d8827cb895b0b39cb2c3dcb992591b3f7b Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Sat, 3 Oct 2026 02:47:31 -0500
Subject: [PATCH 1/2] [offload][AMDGPU] Add dynamic HSA version and symbol
 checks

---
 .../amdgpu/dynamic_hsa/hsa.cpp                | 41 ++++++++++++++++++-
 .../plugins-nextgen/amdgpu/dynamic_hsa/hsa.h  |  2 +
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    | 15 +++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
index 13bd2fd064002b..724a5711a6be5e 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
@@ -11,6 +11,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/DynamicLibrary.h"
 
 #include "Shared/Debug.h"
@@ -19,6 +21,7 @@
 #include "hsa.h"
 #include "hsa_ext_amd.h"
 #include <memory>
+#include <utility>
 
 using namespace llvm::offload::debug;
 
@@ -96,14 +99,48 @@ DLWRAP_FINALIZE()
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 #endif
 
+/// Stand-in for an optional function that the HSA runtime does not provide.
+template <typename FuncTy> struct UnavailableTy;
+template <typename... ArgsTy>
+struct UnavailableTy<hsa_status_t (*)(ArgsTy...)> {
+  static hsa_status_t call(ArgsTy...) { return HSA_STATUS_ERROR; }
+};
+
+#define OPTIONAL_SYMBOL(SYMBOL)                                                \
+  {#SYMBOL, reinterpret_cast<void *>(&UnavailableTy<decltype(&SYMBOL)>::call)}
+
+/// The functions that are newer than version 1.0 of the AMD extension. If the
+/// runtime lacks one, calls to it fail instead of preventing the plugin from
+/// loading.
+static const std::pair<llvm::StringRef, void *> OptionalSymbols[] = {
+    OPTIONAL_SYMBOL(hsa_amd_vmem_address_reserve),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_address_free),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_handle_create),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_handle_release),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_map),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_unmap),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_set_access),
+};
+
+#undef OPTIONAL_SYMBOL
+
 static bool resolveSymbols(llvm::sys::DynamicLibrary &Lib, const char *Name) {
   for (size_t I = 0; I < dlwrap::size(); I++) {
     const char *Sym = dlwrap::symbol(I);
 
     void *P = Lib.getAddressOfSymbol(Sym);
     if (P == nullptr) {
-      ODBG(OLDT_Init) << "Unable to find '" << Sym << "' in '" << Name << "'!";
-      return false;
+      const auto *It = llvm::find_if(
+          OptionalSymbols, [&](const auto &S) { return S.first == Sym; });
+      if (It == std::end(OptionalSymbols)) {
+        ODBG(OLDT_Init) << "Unable to find '" << Sym << "' in '" << Name
+                        << "'!";
+        return false;
+      }
+      ODBG(OLDT_Init) << "Optional '" << Sym << "' is unavailable in '" << Name
+                      << "'";
+      *dlwrap::pointer(I) = It->second;
+      continue;
     }
     ODBG(OLDT_Init) << "Implementing " << Sym << " with dlsym(" << Sym
                     << ") -> " << P;
diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
index f66326a7f240e4..c303902ca3da7f 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
@@ -123,6 +123,8 @@ typedef enum {
   HSA_SYSTEM_INFO_VERSION_MINOR = 1,
   HSA_SYSTEM_INFO_TIMESTAMP = 2,
   HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY = 3,
+  HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR = 0x207,
+  HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR = 0x208,
 } hsa_system_info_t;
 
 typedef enum {
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index de38c2a68e4ab4..5b52849af75da1 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -48,6 +48,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/raw_ostream.h"
 
 #if !defined(__BYTE_ORDER__) || !defined(__ORDER_LITTLE_ENDIAN__) ||           \
@@ -232,6 +233,18 @@ static Error getTargetTripleAndFeatures(hsa_agent_t Agent,
   });
   return Err;
 }
+
+/// Get the version of the AMD extension that the HSA runtime implements, if the
+/// runtime reports it.
+static std::optional<llvm::VersionTuple> getAMDExtensionVersion() {
+  uint16_t Major, Minor;
+  if (hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR, &Major) !=
+          HSA_STATUS_SUCCESS ||
+      hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR, &Minor) !=
+          HSA_STATUS_SUCCESS)
+    return std::nullopt;
+  return llvm::VersionTuple(Major, Minor);
+}
 } // namespace hsa_utils
 
 /// Utility class representing generic resource references to AMDGPU resources.
@@ -3241,6 +3254,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
       Info.add("HSA Runtime Version",
                std::to_string(Major) + "." + std::to_string(Minor), "",
                DeviceInfo::DRIVER_VERSION);
+    if (auto Version = hsa_utils::getAMDExtensionVersion())
+      Info.add("HSA AMD Extension Version", Version->getAsString());
 
     Info.add("HSA OpenMP Device Number", DeviceId);
 

>From 1a2b0e29fae08d26f1aed080a4c372282775ab4e Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Sat, 3 Oct 2026 02:47:31 -0500
Subject: [PATCH 2/2] [offload] Use HSA SVM for AMDGPU shared memory

---
 .../amdgpu/dynamic_hsa/hsa.cpp                |   1 +
 .../plugins-nextgen/amdgpu/dynamic_hsa/hsa.h  |   1 +
 .../amdgpu/dynamic_hsa/hsa_ext_amd.h          |  24 ++-
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    | 183 ++++++++++++++++--
 .../api/amdgpu_managed_memory_accessible.c    |  31 +++
 .../test/api/omp_device_managed_memory_ops.c  |  68 +++++++
 openmp/docs/design/Runtimes.md                |  13 ++
 7 files changed, 299 insertions(+), 22 deletions(-)
 create mode 100644 offload/test/api/amdgpu_managed_memory_accessible.c
 create mode 100644 offload/test/api/omp_device_managed_memory_ops.c

diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
index 724a5711a6be5e..fa384233327b70 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
@@ -69,6 +69,7 @@ DLWRAP(hsa_amd_agents_allow_access, 4)
 DLWRAP(hsa_amd_memory_lock, 5)
 DLWRAP(hsa_amd_memory_unlock, 1)
 DLWRAP(hsa_amd_memory_fill, 3)
+DLWRAP(hsa_amd_svm_attributes_set, 4)
 DLWRAP(hsa_amd_register_system_event_handler, 2)
 DLWRAP(hsa_amd_signal_create, 5)
 DLWRAP(hsa_amd_signal_async_handler, 5)
diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
index c303902ca3da7f..108a8ca9635079 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
@@ -123,6 +123,7 @@ typedef enum {
   HSA_SYSTEM_INFO_VERSION_MINOR = 1,
   HSA_SYSTEM_INFO_TIMESTAMP = 2,
   HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY = 3,
+  HSA_AMD_SYSTEM_INFO_SVM_SUPPORTED = 0x201,
   HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR = 0x207,
   HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR = 0x208,
 } hsa_system_info_t;
diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
index c736ec0759841d..220b54901e67c1 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
@@ -125,6 +125,21 @@ hsa_status_t hsa_amd_memory_unlock(void* host_ptr);
 
 hsa_status_t hsa_amd_memory_fill(void *ptr, uint32_t value, size_t count);
 
+typedef enum hsa_amd_svm_attribute_s {
+  HSA_AMD_SVM_ATTRIB_GLOBAL_FLAG = 0,
+  HSA_AMD_SVM_ATTRIB_AGENT_ACCESSIBLE = 0x200,
+} hsa_amd_svm_attribute_t;
+
+typedef struct hsa_amd_svm_attribute_pair_s {
+  uint64_t attribute;
+  uint64_t value;
+} hsa_amd_svm_attribute_pair_t;
+
+hsa_status_t
+hsa_amd_svm_attributes_set(void *ptr, size_t size,
+                           hsa_amd_svm_attribute_pair_t *attribute_list,
+                           size_t attribute_count);
+
 typedef enum hsa_amd_event_type_s {
   HSA_AMD_GPU_MEMORY_FAULT_EVENT = 0,
 } hsa_amd_event_type_t;
@@ -163,7 +178,8 @@ typedef enum {
 typedef enum {
   HSA_EXT_POINTER_TYPE_UNKNOWN = 0,
   HSA_EXT_POINTER_TYPE_HSA = 1,
-  HSA_EXT_POINTER_TYPE_LOCKED = 2
+  HSA_EXT_POINTER_TYPE_LOCKED = 2,
+  HSA_EXT_POINTER_TYPE_RESERVED_ADDR = 5
 } hsa_amd_pointer_type_t;
 
 typedef struct hsa_amd_pointer_info_s {
@@ -174,6 +190,8 @@ typedef struct hsa_amd_pointer_info_s {
   size_t sizeInBytes;
   void *userData;
   hsa_agent_t agentOwner;
+  uint32_t global_flags;
+  bool registered;
 } hsa_amd_pointer_info_t;
 
 typedef enum {
@@ -211,6 +229,10 @@ hsa_status_t hsa_amd_profiling_set_profiler_enabled(hsa_queue_t *queue,
 hsa_status_t hsa_amd_vmem_address_reserve(void **va, size_t size,
                                           uint64_t address, uint64_t flags);
 
+typedef enum hsa_amd_vmem_address_reserve_flag_s {
+  HSA_AMD_VMEM_ADDRESS_NO_REGISTER = (1UL << 0),
+} hsa_amd_vmem_address_reserve_flag_t;
+
 hsa_status_t hsa_amd_vmem_address_free(void *va, size_t size);
 
 hsa_status_t
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 5b52849af75da1..a5313c0bf329ac 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -46,6 +46,7 @@
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/VersionTuple.h"
@@ -245,6 +246,28 @@ static std::optional<llvm::VersionTuple> getAMDExtensionVersion() {
     return std::nullopt;
   return llvm::VersionTuple(Major, Minor);
 }
+
+/// Get the information of the allocation containing \p Ptr.
+static Expected<hsa_amd_pointer_info_t> getPointerInfo(const void *Ptr) {
+  hsa_amd_pointer_info_t Info{};
+  Info.size = sizeof(hsa_amd_pointer_info_t);
+  hsa_status_t Status = hsa_amd_pointer_info(Ptr, &Info, /*Allocator=*/nullptr,
+                                             /*num_agents_accessible=*/nullptr,
+                                             /*accessible=*/nullptr);
+  if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
+    return std::move(Err);
+  return Info;
+}
+
+/// Whether \p Info describes shared virtual memory (SVM), i.e., addresses that
+/// are reserved without being registered in the driver. Runtimes before version
+/// 1.13 of the AMD extension do not report the registration.
+static bool isSVM(const hsa_amd_pointer_info_t &Info) {
+  return Info.type == HSA_EXT_POINTER_TYPE_RESERVED_ADDR &&
+         Info.size >= offsetof(hsa_amd_pointer_info_t, registered) +
+                          sizeof(Info.registered) &&
+         !Info.registered;
+}
 } // namespace hsa_utils
 
 /// Utility class representing generic resource references to AMDGPU resources.
@@ -2723,6 +2746,12 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (TgtPtr == nullptr)
       return Plugin::success();
 
+    // See allocate() for the registration of host / shared memory as pinned
+    // memory.
+    if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
+      if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
+        return Err;
+
     AMDGPUMemoryPoolTy *MemoryPool = nullptr;
     switch (Kind) {
     case TARGET_ALLOC_DEFAULT:
@@ -2732,21 +2761,27 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     case TARGET_ALLOC_HOST:
       MemoryPool = &HostDevice.getFineGrainedMemoryPool();
       break;
-    case TARGET_ALLOC_SHARED:
+    case TARGET_ALLOC_SHARED: {
+      // Shared allocations are SVM allocations unless SVM is unsupported.
+      auto InfoOrErr = hsa_utils::getPointerInfo(TgtPtr);
+      if (!InfoOrErr)
+        return InfoOrErr.takeError();
+
+      if (hsa_utils::isSVM(*InfoOrErr)) {
+        hsa_status_t Status =
+            hsa_amd_vmem_address_free(TgtPtr, InfoOrErr->sizeInBytes);
+        return Plugin::check(Status, "error in hsa_amd_vmem_address_free: %s");
+      }
+
       MemoryPool = &HostDevice.getFineGrainedMemoryPool();
       break;
     }
+    }
 
     if (!MemoryPool)
       return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
                            "no memory pool for the specified allocation kind");
 
-    // See allocate() for the registration of host / shared memory as pinned
-    // memory.
-    if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
-      if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
-        return Err;
-
     if (auto Err = MemoryPool->deallocate(TgtPtr))
       return Err;
 
@@ -2835,9 +2870,13 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
       return std::move(Err);
 
     // The buffer may be locked or allocated through HSA allocators. Assume that
-    // the buffer is host pinned if the runtime reports a HSA type.
-    if (Info.type != HSA_EXT_POINTER_TYPE_LOCKED &&
-        Info.type != HSA_EXT_POINTER_TYPE_HSA)
+    // the buffer is host pinned if the runtime reports a HSA type. SVM is not
+    // pinned, but all the agents access it through its host address, so it
+    // must not be locked either.
+    if (hsa_utils::isSVM(Info))
+      Info.agentBaseAddress = Info.hostBaseAddress;
+    else if (Info.type != HSA_EXT_POINTER_TYPE_LOCKED &&
+             Info.type != HSA_EXT_POINTER_TYPE_HSA)
       return false;
 
     assert(Info.hostBaseAddress && "Invalid host pinned address");
@@ -3045,8 +3084,15 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   Error dataFillImpl(void *TgtPtr, const void *PatternPtr, int64_t PatternSize,
                      int64_t Size,
                      AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+    // hsa_amd_memory_fill fills SVM on the host, which runtimes before ROCr
+    // commit 9ab39639dd ("fix CPU side hsa_amd_memory_fill", May 2026) get
+    // wrong for patterns wider than a byte. SVM thus takes the slow case below.
+    auto InfoOrErr = hsa_utils::getPointerInfo(TgtPtr);
+    if (!InfoOrErr)
+      return InfoOrErr.takeError();
+
     // Fast case, where we can use the 4 byte hsa_amd_memory_fill
-    if (Size % 4 == 0 &&
+    if (!hsa_utils::isSVM(*InfoOrErr) && Size % 4 == 0 &&
         (PatternSize == 4 || PatternSize == 2 || PatternSize == 1)) {
       uint32_t Pattern;
       if (PatternSize == 1) {
@@ -3062,7 +3108,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
         llvm_unreachable("Invalid pattern size");
       }
 
-      if (hasPendingWorkImpl(AsyncInfoWrapper)) {
+      auto HasPendingWorkOrErr = hasPendingWorkImpl(AsyncInfoWrapper);
+      if (!HasPendingWorkOrErr)
+        return HasPendingWorkOrErr.takeError();
+
+      if (*HasPendingWorkOrErr) {
         AMDGPUStreamTy *Stream = nullptr;
         if (auto Err = getStream(AsyncInfoWrapper, Stream))
           return Err;
@@ -3072,14 +3122,14 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
           uint32_t Pattern;
           int64_t Size;
         };
-        auto *Args = new MemFillArgsTy{TgtPtr, Pattern, Size / 4};
+        auto Args = std::make_unique<MemFillArgsTy>(
+            MemFillArgsTy{TgtPtr, Pattern, Size / 4});
         auto Fill = [](void *Data) {
-          MemFillArgsTy *Args = reinterpret_cast<MemFillArgsTy *>(Data);
-          assert(Args && "Invalid arguments");
+          std::unique_ptr<MemFillArgsTy> Args(
+              static_cast<MemFillArgsTy *>(Data));
 
           auto Status =
               hsa_amd_memory_fill(Args->Dst, Args->Pattern, Args->Size);
-          delete Args;
           auto Err =
               Plugin::check(Status, "error in hsa_amd_memory_fill: %s\n");
           if (Err) {
@@ -3090,7 +3140,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
 
         // hsa_amd_memory_fill doesn't signal completion using a signal, so use
         // the existing host callback logic to handle that instead
-        return Stream->pushHostCallback(Fill, Args);
+        if (auto Err = Stream->pushHostCallback(Fill, Args.get()))
+          return Err;
+
+        Args.release();
+        return Plugin::success();
       }
       // If there is no pending work, do the fill synchronously
       auto Status = hsa_amd_memory_fill(TgtPtr, Pattern, Size / 4);
@@ -3514,6 +3568,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
       return std::move(Err);
 
+    // SVM is accessible by all the agents, see initSVM().
+    if (hsa_utils::isSVM(Info))
+      return Size <=
+             Info.sizeInBytes - utils::getPtrDiff(Ptr, Info.hostBaseAddress);
+
     // Checks if the pointer is known by HSA and accessible by the device
     for (uint32_t i = 0; i < Count; i++) {
       if (Agents[i].handle == getAgent().handle)
@@ -4098,6 +4157,8 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     if (auto Err = HostDevice->init())
       return std::move(Err);
 
+    initSVM(HostAgents);
+
     return NumDevices;
   }
 
@@ -4185,6 +4246,59 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     return KernelAgents;
   }
 
+  /// Enable SVM if the user, the driver and the runtime support it.
+  void initSVM(ArrayRef<hsa_agent_t> HostAgents) {
+    bool Supported = false;
+    if (!OMPX_UseSVM ||
+        hsa_system_get_info(HSA_AMD_SYSTEM_INFO_SVM_SUPPORTED, &Supported) !=
+            HSA_STATUS_SUCCESS ||
+        !Supported)
+      return;
+
+    // SVM is recognized through the registration of its addresses, which the
+    // runtime reports since version 1.13 of the AMD extension.
+    auto Version = hsa_utils::getAMDExtensionVersion();
+    if (!Version || *Version < llvm::VersionTuple(1, 13))
+      return;
+
+    // Accesses may incur a page fault and the migration of the memory to the
+    // accessing agent.
+    for (hsa_agent_t Agent :
+         llvm::concat<const hsa_agent_t>(HostAgents, KernelAgents))
+      SVMAttrs.push_back({HSA_AMD_SVM_ATTRIB_AGENT_ACCESSIBLE, Agent.handle});
+  }
+
+  /// Allocate shared virtual memory (SVM), which is accessible by all the
+  /// agents and which the driver migrates between them on demand. Returns null
+  /// if SVM is unsupported.
+  Expected<void *> allocateSVM(size_t Size, size_t Alignment) {
+    // The reserved addresses are page aligned.
+    if (SVMAttrs.empty() ||
+        Alignment > llvm::sys::Process::getPageSizeEstimate())
+      return nullptr;
+
+    // Reserve addresses that are not registered in the driver, so that the SVM
+    // attributes govern the memory backing them.
+    void *Ptr = nullptr;
+    hsa_status_t Status = hsa_amd_vmem_address_reserve(
+        &Ptr, Size, /*address=*/0, HSA_AMD_VMEM_ADDRESS_NO_REGISTER);
+    if (auto Err =
+            Plugin::check(Status, "error in hsa_amd_vmem_address_reserve: %s"))
+      return std::move(Err);
+
+    Status =
+        hsa_amd_svm_attributes_set(Ptr, Size, SVMAttrs.data(), SVMAttrs.size());
+    if (auto Err =
+            Plugin::check(Status, "error in hsa_amd_svm_attributes_set: %s")) {
+      Status = hsa_amd_vmem_address_free(Ptr, Size);
+      return joinErrors(
+          std::move(Err),
+          Plugin::check(Status, "error in hsa_amd_vmem_address_free: %s"));
+    }
+
+    return Ptr;
+  }
+
   /// Create an HSA signal for the RPC doorbell and return the fields needed
   /// for the GPU to fire interrupts that wake the server thread.
   Error initRPCDoorbell(uint64_t *&Value, uint64_t *&Mailbox,
@@ -4325,6 +4439,14 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 
   /// The device representing all HSA host agents.
   AMDHostDeviceTy *HostDevice;
+
+  /// Whether the user allows serving shared allocations through SVM.
+  BoolEnvar OMPX_UseSVM =
+      BoolEnvar("LIBOMPTARGET_AMDGPU_SHARED_ALLOC_SVM", true);
+
+  /// The attributes making SVM accessible by all the agents. Empty if SVM is
+  /// unsupported.
+  llvm::SmallVector<hsa_amd_svm_attribute_pair_t> SVMAttrs;
 };
 
 Expected<void *> AMDGPUPluginContextTy::allocate(GenericDeviceTy &Device,
@@ -4356,7 +4478,8 @@ Error AMDGPUPluginContextTy::deallocate(GenericDeviceTy &Device, void *Ptr,
 Expected<PluginAllocInfoTy>
 AMDGPUPluginContextTy::getAllocInfo(const void *Ptr) {
   // HSA gives the base of the region containing Ptr, so interior pointers
-  // resolve to the same tracker entry.
+  // resolve to the same tracker entry. SVM only reports its host base address,
+  // which all the agents share.
   hsa_amd_pointer_info_t HsaInfo{};
   HsaInfo.size = sizeof(hsa_amd_pointer_info_t);
   hsa_status_t Status = hsa_amd_pointer_info(
@@ -4364,19 +4487,20 @@ AMDGPUPluginContextTy::getAllocInfo(const void *Ptr) {
       /*num_agents_accessible=*/nullptr, /*accessible=*/nullptr);
   if (auto Err = Plugin::check(Status, "error in hsa_amd_pointer_info: %s"))
     return std::move(Err);
+  void *Base = hsa_utils::isSVM(HsaInfo) ? HsaInfo.hostBaseAddress
+                                         : HsaInfo.agentBaseAddress;
 
   AllocInfo Info;
   {
     std::lock_guard<std::mutex> Lock(AllocationsMutex);
-    auto It = Allocations.find(HsaInfo.agentBaseAddress);
+    auto It = Allocations.find(Base);
     if (It == Allocations.end())
       return Plugin::error(ErrorCode::NOT_FOUND,
                            "pointer is not a known allocation in this context");
     Info = It->second;
   }
 
-  return PluginAllocInfoTy{Info.Device, Info.Kind, HsaInfo.agentBaseAddress,
-                           HsaInfo.sizeInBytes};
+  return PluginAllocInfoTy{Info.Device, Info.Kind, Base, HsaInfo.sizeInBytes};
 }
 
 Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
@@ -4588,6 +4712,23 @@ Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
   if (Size == 0)
     return nullptr;
 
+  // Shared (managed) allocations use SVM, which the driver migrates between
+  // the host and the devices on demand. They fall back to the fine-grained host
+  // memory pool below if SVM is unsupported or cannot serve the alignment.
+  if (Kind == TARGET_ALLOC_SHARED) {
+    auto AllocOrErr =
+        static_cast<AMDGPUPluginTy &>(Plugin).allocateSVM(Size, Alignment);
+    if (!AllocOrErr)
+      return AllocOrErr.takeError();
+
+    // Register as pinned memory like the memory pool allocations below.
+    if (void *Alloc = *AllocOrErr) {
+      if (auto Err = PinnedAllocs.registerHostBuffer(Alloc, Alloc, Size))
+        return std::move(Err);
+      return Alloc;
+    }
+  }
+
   // Find the correct memory pool.
   AMDGPUMemoryPoolTy *MemoryPool = nullptr;
   switch (Kind) {
diff --git a/offload/test/api/amdgpu_managed_memory_accessible.c b/offload/test/api/amdgpu_managed_memory_accessible.c
new file mode 100644
index 00000000000000..9cf01aed56d83c
--- /dev/null
+++ b/offload/test/api/amdgpu_managed_memory_accessible.c
@@ -0,0 +1,31 @@
+// RUN: %libomptarget-compile-run-and-check-amdgcn-amd-amdhsa
+
+// REQUIRES: amdgcn-amd-amdhsa
+
+// Device managed memory is accessible by the device it was allocated for.
+
+#include <omp.h>
+#include <stdio.h>
+
+void *llvm_omp_target_alloc_shared(size_t, int);
+void llvm_omp_target_free_shared(void *, int);
+
+int main() {
+  const size_t Size = 1024;
+  const int Device = omp_get_default_device();
+
+  char *Shared = llvm_omp_target_alloc_shared(Size, Device);
+  if (!Shared) {
+    printf("FAIL: allocation\n");
+    return 1;
+  }
+
+  int Accessible =
+      omp_target_is_accessible(Shared, Size, Device) &&
+      omp_target_is_accessible(&Shared[Size / 2], Size / 2, Device);
+
+  llvm_omp_target_free_shared(Shared, Device);
+
+  // CHECK: PASS
+  printf(Accessible ? "PASS\n" : "FAIL\n");
+}
diff --git a/offload/test/api/omp_device_managed_memory_ops.c b/offload/test/api/omp_device_managed_memory_ops.c
new file mode 100644
index 00000000000000..b60ef512d3a5b4
--- /dev/null
+++ b/offload/test/api/omp_device_managed_memory_ops.c
@@ -0,0 +1,68 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+
+// Check the data movement API on device managed memory.
+
+#include <omp.h>
+#include <stdio.h>
+
+void *llvm_omp_target_alloc_shared(size_t, int);
+void llvm_omp_target_free_shared(void *, int);
+
+int main() {
+  const int N = 128;
+  const int Device = omp_get_default_device();
+  const int Host = omp_get_initial_device();
+
+  int *Shared = llvm_omp_target_alloc_shared(N * sizeof(int), Device);
+  int Buffer[N];
+  int Failures = 0;
+
+  if (!Shared) {
+    printf("FAIL: allocation\n");
+    return 1;
+  }
+
+  // The host can access the allocation directly.
+  for (int I = 0; I < N; ++I)
+    Shared[I] = I;
+
+  // The device can access the allocation directly.
+#pragma omp target teams distribute parallel for device(Device)                \
+    is_device_ptr(Shared)
+  for (int I = 0; I < N; ++I)
+    Shared[I] += 1;
+
+  for (int I = 0; I < N; ++I)
+    Failures += (Shared[I] != I + 1);
+
+  // Filling the whole allocation.
+  if (omp_target_memset(Shared, 0, N * sizeof(int), Device) != Shared)
+    ++Failures;
+  for (int I = 0; I < N; ++I)
+    Failures += (Shared[I] != 0);
+
+  // Filling a subrange of the allocation.
+  if (omp_target_memset(&Shared[1], 0xFF, sizeof(int), Device) != &Shared[1])
+    ++Failures;
+  Failures += (Shared[0] != 0) + (Shared[1] != -1) + (Shared[2] != 0);
+
+  // Copying into and out of the allocation.
+  for (int I = 0; I < N; ++I)
+    Buffer[I] = 2 * I;
+  Failures += (omp_target_memcpy(Shared, Buffer, N * sizeof(int), 0, 0, Device,
+                                 Host) != 0);
+  for (int I = 0; I < N; ++I)
+    Buffer[I] = 0;
+  Failures += (omp_target_memcpy(Buffer, Shared, N * sizeof(int), 0, 0, Host,
+                                 Device) != 0);
+  for (int I = 0; I < N; ++I)
+    Failures += (Buffer[I] != 2 * I);
+
+  llvm_omp_target_free_shared(Shared, Device);
+
+  // CHECK: PASS
+  if (!Failures)
+    printf("PASS\n");
+  else
+    printf("FAIL: %d\n", Failures);
+}
diff --git a/openmp/docs/design/Runtimes.md b/openmp/docs/design/Runtimes.md
index 6fae6a84a549a6..bb97a2fba7c568 100644
--- a/openmp/docs/design/Runtimes.md
+++ b/openmp/docs/design/Runtimes.md
@@ -1335,6 +1335,7 @@ There are several environment variables to change the behavior of the plugins:
 - `LIBOMPTARGET_AMDGPU_MAX_ASYNC_COPY_BYTES`
 - `LIBOMPTARGET_AMDGPU_NUM_INITIAL_HSA_SIGNALS`
 - `LIBOMPTARGET_AMDGPU_STREAM_BUSYWAIT`
+- `LIBOMPTARGET_AMDGPU_SHARED_ALLOC_SVM`
 
 The environment variables `LIBOMPTARGET_STACK_SIZE` and
 `LIBOMPTARGET_HEAP_SIZE` are described in
@@ -1425,6 +1426,18 @@ HSA wait state within the AMDGPU plugin. For the duration of this value
 the HSA runtime may busy wait. This can reduce overall latency.
 The default value is `2000000`.
 
+#### LIBOMPTARGET_AMDGPU_SHARED_ALLOC_SVM
+
+This environment variable controls whether shared (managed) allocations, e.g.,
+the ones performed by `omp_target_alloc_shared`, are served through the HSA
+shared virtual memory (SVM) interface in the AMDGPU plugin. SVM allocations are
+backed by ordinary system memory that the driver can migrate between the host
+and the devices on demand, which greatly improves the device access bandwidth on
+systems with XNACK enabled. If SVM is disabled, or if it is not supported by the
+driver or the HSA runtime, shared allocations are served from a fine-grained host
+memory pool and therefore remain resident on the host. The default value is
+`true`.
+
 (remote_offloading_plugin)=
 
 ### Remote Offloading Plugin:



More information about the Openmp-commits mailing list