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" <[email protected]> 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) { _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
