diff --git a/openmp/libomptarget/include/omptargetplugin.h b/openmp/libomptarget/include/omptargetplugin.h --- a/openmp/libomptarget/include/omptargetplugin.h +++ b/openmp/libomptarget/include/omptargetplugin.h @@ -133,6 +133,11 @@ // error code. int32_t __tgt_rtl_synchronize(int32_t ID, __tgt_async_info *AsyncInfoPtr); +// Register a callback function which will be invoked when the plugin is +// destructed. It accepts an argument of type void * which might be used in the +// callback function. +void __tgt_rtl_register_destruction_cb(void (*cb)(void *), void *data); + #ifdef __cplusplus } #endif diff --git a/openmp/libomptarget/plugins/cuda/src/rtl.cpp b/openmp/libomptarget/plugins/cuda/src/rtl.cpp --- a/openmp/libomptarget/plugins/cuda/src/rtl.cpp +++ b/openmp/libomptarget/plugins/cuda/src/rtl.cpp @@ -279,6 +279,15 @@ std::vector DeviceData; std::vector Modules; + // Struct to store call back inforation that will be used the plugin is + // destructed. + struct DestructionCBInfoTy { + void (*CB)(void *); + void *Data; + }; + + DestructionCBInfoTy DestructionCBInfo; + // Record entry point associated with device void addOffloadEntry(const int DeviceId, const __tgt_offload_entry entry) { FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back(); @@ -371,6 +380,10 @@ } ~DeviceRTLTy() { + // Call the callback function if it exists + if (DestructionCBInfo.CB) + DestructionCBInfo.CB(DestructionCBInfo.Data); + // First destruct stream manager in case of Contexts is destructed before it StreamManager = nullptr; @@ -982,6 +995,11 @@ return OFFLOAD_SUCCESS; } + + void registerDestructionCB(void (*CB)(void *), void *Data) { + DestructionCBInfo.CB = CB; + DestructionCBInfo.Data = Data; + } }; DeviceRTLTy DeviceRTL; @@ -1176,6 +1194,10 @@ return DeviceRTL.synchronize(device_id, async_info_ptr); } +void __tgt_rtl_register_destruction_cb(void (*cb)(void *), void *data) { + DeviceRTL.registerDestructionCB(cb, data); +} + #ifdef __cplusplus } #endif diff --git a/openmp/libomptarget/plugins/exports b/openmp/libomptarget/plugins/exports --- a/openmp/libomptarget/plugins/exports +++ b/openmp/libomptarget/plugins/exports @@ -19,6 +19,7 @@ __tgt_rtl_run_target_region; __tgt_rtl_run_target_region_async; __tgt_rtl_synchronize; + __tgt_rtl_register_destruction_cb; local: *; }; diff --git a/openmp/libomptarget/src/MemoryManager.cpp b/openmp/libomptarget/src/MemoryManager.cpp --- a/openmp/libomptarget/src/MemoryManager.cpp +++ b/openmp/libomptarget/src/MemoryManager.cpp @@ -28,6 +28,8 @@ #include "private.h" #include "rtl.h" +#define UNLIKELY(x) __builtin_expect(!!(x), 0) + namespace { constexpr const size_t BucketSize[] = { 0, 1U << 2, 1U << 3, 1U << 4, 1U << 5, 1U << 6, 1U << 7, @@ -106,6 +108,12 @@ } int MemoryManagerTy::deleteOnDevice(void *Ptr) const { + if (UNLIKELY(!Device.RTL->IsAlive)) { + DP("[MemoryManagerTy::deleteOnDevice] Device RTL has already been " + "destroyed. Return OFFLOAD_SUCCESS directly.\n"); + return OFFLOAD_SUCCESS; + } + return Device.RTL->data_delete(Device.RTLDeviceID, Ptr); } diff --git a/openmp/libomptarget/src/rtl.h b/openmp/libomptarget/src/rtl.h --- a/openmp/libomptarget/src/rtl.h +++ b/openmp/libomptarget/src/rtl.h @@ -24,6 +24,12 @@ struct DeviceTy; struct __tgt_bin_desc; +extern "C" { +// Invalidate the RTLInfoTy because the plugin dylib has been destructed so that +// any later use of plugin should be avoided. +void invalidateRTL(void *); +} + struct RTLInfoTy { typedef int32_t(is_valid_binary_ty)(void *); typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t); @@ -53,6 +59,7 @@ __tgt_async_info *); typedef int64_t(init_requires_ty)(int64_t); typedef int64_t(synchronize_ty)(int32_t, __tgt_async_info *); + typedef void(register_destruction_cb_ty)(void (*)(void *), void *); int32_t Idx = -1; // RTL index, index is the number of devices // of other RTLs that were registered before, @@ -86,10 +93,14 @@ run_team_region_async_ty *run_team_region_async = nullptr; init_requires_ty *init_requires = nullptr; synchronize_ty *synchronize = nullptr; + register_destruction_cb_ty *register_destruction_cb = nullptr; // Are there images associated with this RTL. bool isUsed = false; + // Whether the opened plugin dylib still alive (not destructed). + bool IsAlive = true; + // Mutex for thread-safety when calling RTL interface functions. // It is easier to enforce thread-safety at the libomptarget level, // so that developers of new RTLs do not have to worry about it. @@ -126,6 +137,11 @@ init_requires = r.init_requires; isUsed = r.isUsed; synchronize = r.synchronize; + register_destruction_cb = r.register_destruction_cb; + IsAlive = r.IsAlive; + + if (register_destruction_cb) + register_destruction_cb(invalidateRTL, this); } }; diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp --- a/openmp/libomptarget/src/rtl.cpp +++ b/openmp/libomptarget/src/rtl.cpp @@ -147,6 +147,8 @@ dlsym(dynlib_handle, "__tgt_rtl_data_exchange_async"); *((void **)&R.is_data_exchangable) = dlsym(dynlib_handle, "__tgt_rtl_is_data_exchangable"); + *((void **)&R.register_destruction_cb) = + dlsym(dynlib_handle, "__tgt_rtl_register_destruction_cb"); // No devices are supported by this RTL? if (!(R.NumberOfDevices = R.number_of_devices())) { @@ -154,6 +156,11 @@ continue; } + // Maybe we don't need it at all because this function will always be + // reinvoked afterwards. + if (R.register_destruction_cb) + R.register_destruction_cb(invalidateRTL, &R); + DP("Registering RTL %s supporting %d devices!\n", R.RTLName.c_str(), R.NumberOfDevices); @@ -439,3 +446,10 @@ DP("Done unregistering library!\n"); } + +extern "C" { +void invalidateRTL(void *Data) { + DP("Invalidating plugin with handler " DPxMOD ".\n", DPxPTR(Data)); + reinterpret_cast(Data)->IsAlive = false; +} +}