https://github.com/ro-i created https://github.com/llvm/llvm-project/pull/214755
Claude assisted with this patch. >From d4ce6a6138bc4af6cb85156ab8945493cb1228ef Mon Sep 17 00:00:00 2001 From: Robert Imschweiler <[email protected]> Date: Fri, 7 Aug 2026 02:43:12 -0500 Subject: [PATCH] [offload] Thread allocation kind through async info Claude assisted with this patch. --- offload/include/Shared/APITypes.h | 7 ++- .../common/include/PluginInterface.h | 4 +- .../common/src/PluginInterface.cpp | 44 ++++++++++++++++--- 3 files changed, 44 insertions(+), 11 deletions(-) diff --git a/offload/include/Shared/APITypes.h b/offload/include/Shared/APITypes.h index 47d8c49bf7ef2..67edc7c9f60b7 100644 --- a/offload/include/Shared/APITypes.h +++ b/offload/include/Shared/APITypes.h @@ -23,6 +23,9 @@ #include <cstddef> #include <cstdint> #include <mutex> +#include <utility> + +enum TargetAllocTy : int32_t; extern "C" { @@ -75,8 +78,8 @@ struct __tgt_async_info { void *Queue = nullptr; /// A collection of allocations that are associated with this stream and that - /// should be freed after finalization. - llvm::SmallVector<void *, 2> AssociatedAllocations; + /// should be freed after finalization. Also store the type of the allocation. + llvm::SmallVector<std::pair<void *, TargetAllocTy>, 2> AssociatedAllocations; /// Mutex to guard access to AssociatedAllocations and the Queue. std::mutex Mutex; diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h index 50be67886c527..0b6a2223a75ce 100644 --- a/offload/plugins-nextgen/common/include/PluginInterface.h +++ b/offload/plugins-nextgen/common/include/PluginInterface.h @@ -168,9 +168,9 @@ struct AsyncInfoWrapperTy { /// Register \p Ptr as an associated allocation that is freed after /// finalization. - void freeAllocationAfterSynchronization(void *Ptr) { + void freeAllocationAfterSynchronization(void *Ptr, TargetAllocTy Kind) { std::lock_guard<std::mutex> AllocationGuard(AsyncInfoPtr->Mutex); - AsyncInfoPtr->AssociatedAllocations.push_back(Ptr); + AsyncInfoPtr->AssociatedAllocations.push_back({Ptr, Kind}); } private: diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp index 389d2ffbf4395..f4dcff3879ccd 100644 --- a/offload/plugins-nextgen/common/src/PluginInterface.cpp +++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp @@ -137,7 +137,8 @@ GenericKernelTy::getKernelLaunchEnvironment( return AllocOrErr.takeError(); // Remember to free the memory later. - AsyncInfoWrapper.freeAllocationAfterSynchronization(*AllocOrErr); + AsyncInfoWrapper.freeAllocationAfterSynchronization( + *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE); /// Use the KLE in the __tgt_async_info to ensure a stable address for the /// async data transfer. @@ -159,7 +160,8 @@ GenericKernelTy::getKernelLaunchEnvironment( return AllocOrErr.takeError(); LocalKLE.ReductionBuffer = *AllocOrErr; // Remember to free the memory later. - AsyncInfoWrapper.freeAllocationAfterSynchronization(*AllocOrErr); + AsyncInfoWrapper.freeAllocationAfterSynchronization( + *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE); } INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(), @@ -289,7 +291,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs, DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr; if (DynBlockMemConf.FallbackPtr) AsyncInfoWrapper.freeAllocationAfterSynchronization( - DynBlockMemConf.FallbackPtr); + DynBlockMemConf.FallbackPtr, TargetAllocTy::TARGET_ALLOC_DEVICE); auto KernelLaunchEnvOrErr = getKernelLaunchEnvironment(GenericDevice, KernelArgs, DynBlockMemConf, @@ -976,7 +978,7 @@ Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo, return Plugin::error(ErrorCode::INVALID_ARGUMENT, "invalid async info queue"); - SmallVector<void *> AllocsToDelete{}; + SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{}; { std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex}; @@ -989,8 +991,8 @@ Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo, std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations); } - for (auto *Ptr : AllocsToDelete) - if (auto Err = dataDelete(Ptr, TargetAllocTy::TARGET_ALLOC_DEVICE)) + for (auto [Ptr, Kind] : AllocsToDelete) + if (auto Err = dataDelete(Ptr, Kind)) return Err; return Plugin::success(); @@ -1003,7 +1005,35 @@ Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo, return Plugin::error(ErrorCode::INVALID_ARGUMENT, "invalid async info queue"); - return queryAsyncImpl(*AsyncInfo, ReleaseQueue, IsQueueWorkCompleted); + bool WorkCompleted = false; + SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{}; + + { + // Query and collect under the mutex, as synchronize does. Querying outside + // it would let an operation issued in between have its allocations freed + // here while it is still using them. + std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex}; + if (auto Err = queryAsyncImpl(*AsyncInfo, ReleaseQueue, &WorkCompleted)) { + if (IsQueueWorkCompleted) + *IsQueueWorkCompleted = WorkCompleted; + return Err; + } + + // An async info belonging to a nowait task is never synchronized, so this + // is the only completion notification it ever gets. Without releasing its + // allocations here they are never freed at all. + if (WorkCompleted) + std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations); + } + + for (auto [Ptr, Kind] : AllocsToDelete) + if (auto Err = dataDelete(Ptr, Kind)) + return Err; + + if (IsQueueWorkCompleted) + *IsQueueWorkCompleted = WorkCompleted; + + return Plugin::success(); } Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize) { _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
