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

Robert Imschweiler via llvm-commits llvm-commits at lists.llvm.org
Wed Aug 12 06:28:23 PDT 2026


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

Claude assisted with this patch.

>From aa8d5ec4fed93d50c65bfe81035d2ea2ebbf7ec5 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 12 Aug 2026 07:33:34 -0500
Subject: [PATCH] [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          |  15 +
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    | 288 +++++++++++++++++-
 .../api/amdgpu_managed_memory_accessible.c    |  31 ++
 .../test/api/omp_device_managed_memory_ops.c  |  68 +++++
 openmp/docs/design/Runtimes.rst               |  13 +
 7 files changed, 406 insertions(+), 11 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 5c7ec186b0ceb..eb206a37634b9 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
@@ -66,6 +66,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 f66326a7f240e..1ed9a6275c3dc 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_system_info_t;
 
 typedef enum {
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 d26f9248e27ef..a78e6d7c5cb4a 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;
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 0f0de67b20c7c..b363bbd472013 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -10,12 +10,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <algorithm>
 #include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <deque>
 #include <functional>
+#include <map>
 #include <mutex>
 #include <string>
 #include <system_error>
@@ -46,7 +49,9 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Memory.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/raw_ostream.h"
@@ -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 {
+    const uintptr_t Addr = reinterpret_cast<uintptr_t>(Ptr);
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    // Find the allocation with the largest base address not above Addr.
+    auto It = Allocations.upper_bound(Addr);
+    if (It == Allocations.begin())
+      return false;
+    --It;
+
+    const size_t Offset = Addr - It->first;
+    const size_t AllocatedSize = It->second.allocatedSize();
+    return Offset < AllocatedSize && Size <= AllocatedSize - Offset;
+  }
+
+private:
+  /// Whether the user allows serving shared allocations through SVM.
+  BoolEnvar OMPX_UseSVM =
+      BoolEnvar("LIBOMPTARGET_AMDGPU_SHARED_ALLOC_SVM", true);
+
+  /// Whether new SVM allocations can be created. Cleared if the driver rejects
+  /// one.
+  std::atomic<bool> Supported = false;
+
+  /// The granularity of the mapped memory allocations.
+  size_t PageSize = 0;
+
+  /// The live SVM allocations, indexed by their base address.
+  std::map<uintptr_t, llvm::sys::MemoryBlock> Allocations;
+
+  /// Mutex for safe access to the allocations.
+  mutable std::mutex Mutex;
+};
+
 /// Class that implements a memory manager that gets memory from a specific
 /// memory pool.
 struct AMDGPUMemoryManagerTy : public DeviceAllocatorTy {
@@ -2701,6 +2857,9 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     return AMDImage;
   }
 
+  /// Get the plugin's manager of the shared (managed) SVM allocations.
+  AMDGPUSVMManagerTy &getSVMManager();
+
   /// Allocate memory on the device or related to the device.
   Expected<void *> allocate(size_t Size, void *, TargetAllocTy Kind,
                             size_t Alignment) override;
@@ -2710,6 +2869,15 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (TgtPtr == nullptr)
       return Plugin::success();
 
+    // Shared allocations are SVM allocations, or memory pool allocations when
+    // SVM cannot serve them. The kind is not checked here; the generic layer
+    // already diagnoses deallocations with a mismatching kind.
+    auto WasSVMAllocOrErr = getSVMManager().deallocate(TgtPtr);
+    if (!WasSVMAllocOrErr)
+      return WasSVMAllocOrErr.takeError();
+    if (*WasSVMAllocOrErr)
+      return Plugin::success();
+
     AMDGPUMemoryPoolTy *MemoryPool = nullptr;
     switch (Kind) {
     case TARGET_ALLOC_DEFAULT:
@@ -3023,9 +3191,63 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     return Plugin::success();
   }
 
+  /// Fill \p TgtPtr with the given pattern on the host. Requires \p TgtPtr to
+  /// be host accessible.
+  Error dataFillHostImpl(void *TgtPtr, const void *PatternPtr,
+                         int64_t PatternSize, int64_t Size,
+                         AsyncInfoWrapperTy &AsyncInfoWrapper) {
+    struct HostFillArgsTy {
+      void *Dst;
+      llvm::SmallVector<char, 8> Pattern;
+      int64_t NumTimes;
+    };
+    auto Args = std::make_unique<HostFillArgsTy>(HostFillArgsTy{
+        TgtPtr,
+        {reinterpret_cast<const char *>(PatternPtr),
+         reinterpret_cast<const char *>(PatternPtr) + PatternSize},
+        Size / PatternSize});
+
+    auto Fill = [](void *Data) {
+      std::unique_ptr<HostFillArgsTy> Args(static_cast<HostFillArgsTy *>(Data));
+      auto *Dst = reinterpret_cast<char *>(Args->Dst);
+      if (Args->Pattern.size() == 1) {
+        std::memset(Dst, Args->Pattern.front(), Args->NumTimes);
+      } else {
+        for (int64_t I = 0; I < Args->NumTimes; ++I)
+          Dst = std::copy(Args->Pattern.begin(), Args->Pattern.end(), Dst);
+      }
+    };
+
+    // Defer the fill until all the previously enqueued operations completed.
+    auto HasPendingWorkOrErr = hasPendingWorkImpl(AsyncInfoWrapper);
+    if (!HasPendingWorkOrErr)
+      return HasPendingWorkOrErr.takeError();
+
+    if (*HasPendingWorkOrErr) {
+      AMDGPUStreamTy *Stream = nullptr;
+      if (auto Err = getStream(AsyncInfoWrapper, Stream))
+        return Err;
+
+      if (auto Err = Stream->pushHostCallback(Fill, Args.get()))
+        return Err;
+
+      Args.release();
+      return Plugin::success();
+    }
+
+    Fill(Args.release());
+    return Plugin::success();
+  }
+
   Error dataFillImpl(void *TgtPtr, const void *PatternPtr, int64_t PatternSize,
                      int64_t Size,
                      AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+    // hsa_amd_memory_fill only operates on memory allocated by the HSA
+    // runtime. SVM allocations are host accessible, so fill them on the host.
+    if (getSVMManager().contains(TgtPtr, Size))
+      return dataFillHostImpl(TgtPtr, PatternPtr, PatternSize, Size,
+                              AsyncInfoWrapper);
+
     // Fast case, where we can use the 4 byte hsa_amd_memory_fill
     if (Size % 4 == 0 &&
         (PatternSize == 4 || PatternSize == 2 || PatternSize == 1)) {
@@ -3043,7 +3265,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;
@@ -3053,14 +3279,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) {
@@ -3071,7 +3297,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);
@@ -3488,6 +3718,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   }
 
   Expected<bool> isAccessiblePtrImpl(const void *Ptr, size_t Size) override {
+    // SVM allocations are accessible by all the kernel agents but are unknown
+    // to the HSA runtime.
+    if (getSVMManager().contains(Ptr, Size))
+      return true;
+
     hsa_amd_pointer_info_t Info;
     Info.size = sizeof(hsa_amd_pointer_info_t);
 
@@ -4056,6 +4291,8 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     if (auto Err = HostDevice->init())
       return std::move(Err);
 
+    SVMManager.init();
+
     return NumDevices;
   }
 
@@ -4070,9 +4307,12 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
       if (auto Err = HostDevice->deinit())
         return Err;
 
+    Error Err = SVMManager.deinit();
+
     // Finalize the HSA runtime.
     hsa_status_t Status = hsa_shut_down();
-    return Plugin::check(Status, "error in hsa_shut_down: %s");
+    return joinErrors(std::move(Err),
+                      Plugin::check(Status, "error in hsa_shut_down: %s"));
   }
 
   /// Creates an AMDGPU device.
@@ -4138,6 +4378,11 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     return KernelAgents;
   }
 
+  /// Get the manager of the shared (managed) SVM allocations. A single manager
+  /// serves all the devices because SVM allocations are accessible by all the
+  /// agents.
+  AMDGPUSVMManagerTy &getSVMManager() { return SVMManager; }
+
   /// 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,
@@ -4278,6 +4523,9 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 
   /// The device representing all HSA host agents.
   AMDHostDeviceTy *HostDevice;
+
+  /// The manager of the shared (managed) SVM allocations.
+  AMDGPUSVMManagerTy SVMManager;
 };
 
 Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
@@ -4486,12 +4734,31 @@ Expected<void *> AMDGPUMemoryManagerTy::allocate(size_t Size, void *HstPtr,
   return Ptr;
 }
 
+AMDGPUSVMManagerTy &AMDGPUDeviceTy::getSVMManager() {
+  return static_cast<AMDGPUPluginTy &>(Plugin).getSVMManager();
+}
+
 Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
                                           TargetAllocTy Kind,
                                           size_t Alignment) {
   if (Size == 0)
     return nullptr;
 
+  auto &AMDGPUPlugin = static_cast<AMDGPUPluginTy &>(Plugin);
+
+  // Shared (managed) allocations use the SVM interface, which lets the driver
+  // migrate the memory between the host and the devices. Requests that SVM
+  // cannot serve fall back to the fine-grained host memory pool below.
+  auto &SVMManager = AMDGPUPlugin.getSVMManager();
+  if (Kind == TARGET_ALLOC_SHARED && SVMManager.isSupported()) {
+    auto AllocOrErr =
+        SVMManager.allocate(Size, Alignment, AMDGPUPlugin.getKernelAgents());
+    if (!AllocOrErr)
+      return AllocOrErr.takeError();
+    if (*AllocOrErr)
+      return *AllocOrErr;
+  }
+
   // Find the correct memory pool.
   AMDGPUMemoryPoolTy *MemoryPool = nullptr;
   switch (Kind) {
@@ -4521,10 +4788,9 @@ Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
     // necessary for host or shared allocations Also enabled for device memory
     // to allow device to device memcpy
     llvm::SmallVector<hsa_agent_t> Agents;
-    llvm::copy_if(static_cast<AMDGPUPluginTy &>(Plugin).getKernelAgents(),
-                  std::back_inserter(Agents), [&](hsa_agent_t Agent) {
-                    return MemoryPool->canAccess(Agent);
-                  });
+    llvm::copy_if(
+        AMDGPUPlugin.getKernelAgents(), std::back_inserter(Agents),
+        [&](hsa_agent_t Agent) { return MemoryPool->canAccess(Agent); });
 
     // Enable all valid kernel agents to access the buffer.
     if (auto Err = MemoryPool->enableAccess(Alloc, Size, Agents))
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 0000000000000..9cf01aed56d83
--- /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 0000000000000..b60ef512d3a5b
--- /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.rst b/openmp/docs/design/Runtimes.rst
index 14e300a0f531f..f6fca089559d0 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/openmp/docs/design/Runtimes.rst
@@ -1398,6 +1398,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
@@ -1498,6 +1499,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, 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 llvm-commits mailing list