[llvm-branch-commits] [llvm] [offload][omp] Use olMemRegister for memory locking (PR #221982)

Alex Duran via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Tue Sep 8 05:47:55 PDT 2026


https://github.com/adurang updated https://github.com/llvm/llvm-project/pull/221982

>From 8b2c415449b81f4e09f1951612eba72045ff41b6 Mon Sep 17 00:00:00 2001
From: "Duran, Alex" <alejandro.duran at intel.com>
Date: Tue, 8 Sep 2026 05:32:22 -0700
Subject: [PATCH] [offload][omp] Use olMemRegister for memory locking

---
 offload/include/device.h                      | 15 ++++
 offload/liboffload/exports                    |  4 -
 offload/libompaccsupport/device.cpp           | 69 ++++++++++++++--
 offload/libomptarget/omptarget.cpp            | 18 ++---
 .../common/include/PluginInterface.h          | 43 ----------
 .../common/src/PluginInterface.cpp            | 80 -------------------
 6 files changed, 87 insertions(+), 142 deletions(-)

diff --git a/offload/include/device.h b/offload/include/device.h
index 514cde4728d40..61e5b6aaecc85 100644
--- a/offload/include/device.h
+++ b/offload/include/device.h
@@ -108,6 +108,15 @@ struct DeviceTy {
   int32_t dataExchange(void *SrcPtr, DeviceTy &DstDev, void *DstPtr,
                        int64_t Size, AsyncInfoTy &AsyncInfo);
 
+  /// Register (and, if \p LockMemory, page-lock) the host buffer \p HstPtr
+  /// with \p Size bytes, returning the device-accessible pointer.
+  llvm::Expected<void *> registerMemory(void *HstPtr, int64_t Size,
+                                        bool LockMemory = true);
+
+  /// Unregister (and, if \p UnlockMemory, page-unlock) a host buffer
+  /// previously registered via registerMemory.
+  llvm::Error unregisterMemory(void *HstPtr, bool UnlockMemory = true);
+
   /// Notify the plugin about a new mapping starting at the host address
   /// \p HstPtr and \p Size bytes.
   int32_t notifyDataMapped(void *HstPtr, int64_t Size);
@@ -191,6 +200,12 @@ struct DeviceTy {
 
   /// Flag to indicate pending images (true after construction).
   bool HasPendingImages = true;
+
+  /// Indicate whether mapped host buffers should be locked automatically.
+  bool LockMappedBuffers = false;
+
+  /// Indicate whether failures when locking mapped buffers should be ignored.
+  bool IgnoreLockMappedFailures = true;
 };
 
 /// Resolve the device address of the global variable \p Name in \p Program,
diff --git a/offload/liboffload/exports b/offload/liboffload/exports
index d50ef17941228..b7d234f84e643 100644
--- a/offload/liboffload/exports
+++ b/offload/liboffload/exports
@@ -11,10 +11,6 @@ global:
     "llvm::omp::target::plugin::GenericPluginTy::create_interop(int, int, interop_spec_t*)";
     "llvm::omp::target::plugin::GenericPluginTy::data_alloc(int, long, void*, int)";
     "llvm::omp::target::plugin::GenericPluginTy::data_delete(int, void*, int)";
-    "llvm::omp::target::plugin::GenericPluginTy::data_lock(int, void*, long, void**)";
-    "llvm::omp::target::plugin::GenericPluginTy::data_notify_mapped(int, void*, long)";
-    "llvm::omp::target::plugin::GenericPluginTy::data_notify_unmapped(int, void*)";
-    "llvm::omp::target::plugin::GenericPluginTy::data_unlock(int, void*)";
     "llvm::omp::target::plugin::GenericPluginTy::flush_queue(omp_interop_val_t*)";
     "llvm::omp::target::plugin::GenericPluginTy::initialize_record_replay(int, long, void*, bool, bool, bool, bool, char const*, char const*)";
     "llvm::omp::target::plugin::GenericPluginTy::is_initialized() const";
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 872160ac2a7d0..2856925a93f19 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -81,6 +81,32 @@ llvm::Error DeviceTy::init() {
                                      DeviceID);
   }
 
+  // Envar that indicates whether mapped host buffers should be locked
+  // automatically. The possible values are boolean (on/off) and a special:
+  //   off:       Mapped host buffers are not locked.
+  //   on:        Mapped host buffers are locked in a best-effort approach.
+  //              Failure to lock the buffers are silent.
+  //   mandatory: Mapped host buffers are always locked and failures to lock
+  //              a buffer results in a fatal error.
+  StringEnvar OMPX_LockMappedBuffers("LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS",
+                                     "off");
+  bool Enabled;
+  if (StringParser::parse(OMPX_LockMappedBuffers.get().data(), Enabled)) {
+    // Parsed as a boolean value. Enable the feature if necessary.
+    LockMappedBuffers = Enabled;
+    IgnoreLockMappedFailures = true;
+  } else if (OMPX_LockMappedBuffers.get() == "mandatory") {
+    // Enable the feature and failures are fatal.
+    LockMappedBuffers = true;
+    IgnoreLockMappedFailures = false;
+  } else {
+    // Disable by default.
+    ODBG(ODT_Alloc) << "Invalid value LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS="
+                    << OMPX_LockMappedBuffers.get();
+    LockMappedBuffers = false;
+    IgnoreLockMappedFailures = true;
+  }
+
   // Enables recording kernels if set.
   BoolEnvar OMPX_RecordKernel("LIBOMPTARGET_RECORD", false);
   if (OMPX_RecordKernel) {
@@ -387,13 +413,40 @@ int32_t DeviceTy::dataExchange(void *SrcPtr, DeviceTy &DstDev, void *DstPtr,
   return OFFLOAD_SUCCESS;
 }
 
+llvm::Expected<void *> DeviceTy::registerMemory(void *HstPtr, int64_t Size,
+                                                bool LockMemory) {
+  void *LockedPtr = nullptr;
+  ol_memory_register_flags_t Flags =
+      LockMemory ? OL_MEMORY_REGISTER_FLAG_LOCK_MEMORY : 0;
+  if (auto Res = olMemRegister(DeviceHandle, HstPtr, Size, Flags, &LockedPtr))
+    return error::createOffloadError(error::ErrorCode::UNKNOWN,
+                                     "failed to lock memory %p: %s", HstPtr,
+                                     Res->Details);
+  return LockedPtr;
+}
+
+llvm::Error DeviceTy::unregisterMemory(void *HstPtr, bool UnlockMemory) {
+  ol_memory_register_flags_t Flags =
+      UnlockMemory ? OL_MEMORY_REGISTER_FLAG_UNLOCK_MEMORY : 0;
+  if (auto Res = olMemUnregister(DeviceHandle, HstPtr, Flags))
+    return error::createOffloadError(error::ErrorCode::UNKNOWN,
+                                     "failed to unlock memory %p: %s", HstPtr,
+                                     Res->Details);
+  return llvm::Error::success();
+}
+
 int32_t DeviceTy::notifyDataMapped(void *HstPtr, int64_t Size) {
   ODBG(ODT_Mapping) << "Notifying about new mapping: HstPtr=" << HstPtr
                     << ", Size=" << Size;
 
-  if (RTL->data_notify_mapped(RTLDeviceID, HstPtr, Size)) {
-    REPORT() << "Notifying about data mapping failed.";
-    return OFFLOAD_FAIL;
+  auto LockedPtrOrErr = registerMemory(HstPtr, Size, LockMappedBuffers);
+  if (!LockedPtrOrErr) {
+    if (!IgnoreLockMappedFailures) {
+      REPORT() << "Notifying about data mapping failed: "
+               << llvm::toString(LockedPtrOrErr.takeError());
+      return OFFLOAD_FAIL;
+    }
+    llvm::consumeError(LockedPtrOrErr.takeError());
   }
   return OFFLOAD_SUCCESS;
 }
@@ -401,9 +454,13 @@ int32_t DeviceTy::notifyDataMapped(void *HstPtr, int64_t Size) {
 int32_t DeviceTy::notifyDataUnmapped(void *HstPtr) {
   ODBG(ODT_Mapping) << "Notifying about an unmapping: HstPtr=" << HstPtr;
 
-  if (RTL->data_notify_unmapped(RTLDeviceID, HstPtr)) {
-    REPORT() << "Notifying about data unmapping failed.";
-    return OFFLOAD_FAIL;
+  if (auto Err = unregisterMemory(HstPtr, LockMappedBuffers)) {
+    if (!IgnoreLockMappedFailures) {
+      REPORT() << "Notifying about data unmapping failed: "
+               << llvm::toString(std::move(Err));
+      return OFFLOAD_FAIL;
+    }
+    llvm::consumeError(std::move(Err));
   }
   return OFFLOAD_SUCCESS;
 }
diff --git a/offload/libomptarget/omptarget.cpp b/offload/libomptarget/omptarget.cpp
index 24a9b93d1a0ea..70330e6a0da04 100644
--- a/offload/libomptarget/omptarget.cpp
+++ b/offload/libomptarget/omptarget.cpp
@@ -288,20 +288,18 @@ void *targetLockExplicit(void *HostPtr, size_t Size, int DeviceNum,
     return NULL;
   }
 
-  void *RC = NULL;
-
   auto DeviceOrErr = PM->getDevice(DeviceNum);
   if (!DeviceOrErr)
     FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str());
 
-  int32_t Err = 0;
-  Err = DeviceOrErr->RTL->data_lock(DeviceNum, HostPtr, Size, &RC);
-  if (Err) {
-    ODBG(ODT_Interface) << "Could not lock ptr " << HostPtr;
+  auto LockedPtrOrErr = DeviceOrErr->registerMemory(HostPtr, Size);
+  if (!LockedPtrOrErr) {
+    ODBG(ODT_Interface) << "Could not lock ptr " << HostPtr << ": "
+                        << toString(LockedPtrOrErr.takeError());
     return nullptr;
   }
-  ODBG(ODT_Interface) << Name << " returns device ptr " << RC;
-  return RC;
+  ODBG(ODT_Interface) << Name << " returns device ptr " << *LockedPtrOrErr;
+  return *LockedPtrOrErr;
 }
 
 void targetUnlockExplicit(void *HostPtr, int DeviceNum, const char *Name) {
@@ -312,7 +310,9 @@ void targetUnlockExplicit(void *HostPtr, int DeviceNum, const char *Name) {
   if (!DeviceOrErr)
     FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str());
 
-  DeviceOrErr->RTL->data_unlock(DeviceNum, HostPtr);
+  if (auto Err = DeviceOrErr->unregisterMemory(HostPtr))
+    ODBG(ODT_Interface) << "Could not unlock ptr " << HostPtr << ": "
+                        << toString(std::move(Err));
   ODBG(ODT_Interface) << Name << " returns";
 }
 
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index aa9799e2b4a80..8eecbd4e19e1c 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1074,30 +1074,6 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Unlock a previously locked host buffer starting at \p HstPtr.
   virtual Error dataUnlockImpl(void *HstPtr) = 0;
 
-  /// Mark the host buffer with address \p HstPtr and \p Size bytes as a mapped
-  /// buffer. This means that libomptarget created a new mapping of that host
-  /// buffer (e.g., because a user OpenMP target map) and the buffer may be used
-  /// as source/destination of memory transfers. We can use this information to
-  /// lock the host buffer and optimize its memory transfers.
-  Error notifyDataMapped(void *HstPtr, int64_t Size) {
-    auto Err = PinnedAllocs.registerMemory(HstPtr, Size, LockMappedBuffers);
-    if (!Err && !IgnoreLockMappedFailures)
-      return Err.takeError();
-    return Plugin::success();
-  }
-
-  /// Mark the host buffer with address \p HstPtr as unmapped. This means that
-  /// libomptarget removed an existing mapping. If the plugin locked the buffer
-  /// in notifyDataMapped, this function should unlock it.
-  Error notifyDataUnmapped(void *HstPtr) {
-    auto Err = PinnedAllocs.unregisterMemory(HstPtr, LockMappedBuffers);
-    if (IgnoreLockMappedFailures) {
-      consumeError(std::move(Err));
-      return Plugin::success();
-    }
-    return Err;
-  }
-
   /// Check whether the host buffer with address \p HstPtr is pinned by the
   /// underlying vendor-specific runtime (if any). Retrieve the host pointer,
   /// the device accessible pointer and the size of the original pinned buffer.
@@ -1439,12 +1415,6 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   BoolEnvar OMPX_ReuseBlocksForHighTripCount =
       BoolEnvar("LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT", true);
 
-  /// Indicate whether mapped host buffers should be locked automatically.
-  bool LockMappedBuffers;
-
-  /// Indicate whether failures when locking mapped buffers should be ignored.
-  bool IgnoreLockMappedFailures;
-
   /// Record and replay manager.
   RecordReplayTy *RecordReplay = nullptr;
 
@@ -1712,19 +1682,6 @@ struct GenericPluginTy {
   /// Deallocates memory on the given device.
   int32_t data_delete(int32_t DeviceId, void *TgtPtr, int32_t Kind);
 
-  /// Locks / pins host memory using the plugin runtime.
-  int32_t data_lock(int32_t DeviceId, void *Ptr, int64_t Size,
-                    void **LockedPtr);
-
-  /// Unlocks / unpins host memory using the plugin runtime.
-  int32_t data_unlock(int32_t DeviceId, void *Ptr);
-
-  /// Notify the runtime about a new mapping that has been created outside.
-  int32_t data_notify_mapped(int32_t DeviceId, void *HstPtr, int64_t Size);
-
-  /// Notify t he runtime about a mapping that has been deleted.
-  int32_t data_notify_unmapped(int32_t DeviceId, void *HstPtr);
-
   /// Begin executing a kernel on the given device.
   int32_t launch_kernel(int32_t DeviceId, void *TgtEntryPtr,
                         KernelLaunchArgsTy &LaunchArgs,
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 5de0c9302c573..c303cbf204880 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -482,32 +482,6 @@ GenericDeviceTy::GenericDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId,
 #undef bindOmptCallback
 
 #endif
-
-  // Envar that indicates whether mapped host buffers should be locked
-  // automatically. The possible values are boolean (on/off) and a special:
-  //   off:       Mapped host buffers are not locked.
-  //   on:        Mapped host buffers are locked in a best-effort approach.
-  //              Failure to lock the buffers are silent.
-  //   mandatory: Mapped host buffers are always locked and failures to lock
-  //              a buffer results in a fatal error.
-  StringEnvar OMPX_LockMappedBuffers("LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS",
-                                     "off");
-
-  bool Enabled;
-  if (StringParser::parse(OMPX_LockMappedBuffers.get().data(), Enabled)) {
-    // Parsed as a boolean value. Enable the feature if necessary.
-    LockMappedBuffers = Enabled;
-    IgnoreLockMappedFailures = true;
-  } else if (OMPX_LockMappedBuffers.get() == "mandatory") {
-    // Enable the feature and failures are fatal.
-    LockMappedBuffers = true;
-    IgnoreLockMappedFailures = false;
-  } else {
-    // Disable by default.
-    ODBG(OLDT_Alloc) << "Invalid value LIBOMPTARGET_LOCK_MAPPED_HOST_BUFFERS="
-                     << OMPX_LockMappedBuffers.get();
-    LockMappedBuffers = false;
-  }
 }
 
 Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
@@ -1530,60 +1504,6 @@ int32_t GenericPluginTy::data_delete(int32_t DeviceId, void *TgtPtr,
   return OFFLOAD_SUCCESS;
 }
 
-int32_t GenericPluginTy::data_lock(int32_t DeviceId, void *Ptr, int64_t Size,
-                                   void **LockedPtr) {
-  auto LockedPtrOrErr = getDevice(DeviceId).registerMemory(Ptr, Size);
-  if (!LockedPtrOrErr) {
-    auto Err = LockedPtrOrErr.takeError();
-    REPORT() << "Failure to lock memory " << Ptr << ": "
-             << toString(std::move(Err));
-    return OFFLOAD_FAIL;
-  }
-
-  if (!(*LockedPtrOrErr)) {
-    REPORT() << "Failure to lock memory " << Ptr
-             << ": obtained a null locked pointer";
-    return OFFLOAD_FAIL;
-  }
-  *LockedPtr = *LockedPtrOrErr;
-
-  return OFFLOAD_SUCCESS;
-}
-
-int32_t GenericPluginTy::data_unlock(int32_t DeviceId, void *Ptr) {
-  auto Err = getDevice(DeviceId).unregisterMemory(Ptr);
-  if (Err) {
-    REPORT() << "Failure to unlock memory " << Ptr << ": "
-             << toString(std::move(Err));
-    return OFFLOAD_FAIL;
-  }
-
-  return OFFLOAD_SUCCESS;
-}
-
-int32_t GenericPluginTy::data_notify_mapped(int32_t DeviceId, void *HstPtr,
-                                            int64_t Size) {
-  auto Err = getDevice(DeviceId).notifyDataMapped(HstPtr, Size);
-  if (Err) {
-    REPORT() << "Failure to notify data mapped " << HstPtr << ": "
-             << toString(std::move(Err));
-    return OFFLOAD_FAIL;
-  }
-
-  return OFFLOAD_SUCCESS;
-}
-
-int32_t GenericPluginTy::data_notify_unmapped(int32_t DeviceId, void *HstPtr) {
-  auto Err = getDevice(DeviceId).notifyDataUnmapped(HstPtr);
-  if (Err) {
-    REPORT() << "Failure to notify data unmapped " << HstPtr << ": "
-             << toString(std::move(Err));
-    return OFFLOAD_FAIL;
-  }
-
-  return OFFLOAD_SUCCESS;
-}
-
 int32_t GenericPluginTy::launch_kernel(int32_t DeviceId, void *TgtEntryPtr,
                                        KernelLaunchArgsTy &LaunchArgs,
                                        __tgt_async_info *AsyncInfoPtr) {



More information about the llvm-branch-commits mailing list