[OpenMP] Fixed the issue that memory manager calls plugin interface to release buffered memory even after the plugin has been destroyed

[tianshilei1992]

Summary

Currently on some platforms the program might crash when the program
exits. It is because the memory manager calls plugin interface to release its
buffered memory and the plugin has already been destroyed. After some search
online, it seems that loaded dynamic libraries will be destroyed in the reverse
order of loading, which means by no means, those plugin libraries loaded by
dlopen will be destroyed after libomtarget. We need something to tell the
RTLInfoTy that the hold plugin is not valid anymore so don't call those
interfaces. This is done by registering a callback function in the plugin so that
it will be called when the plugin starts to destroy itself. The callback function
for now only does one thing: release the memory manager so that all buffered
target memory will be released. We can add more features in the future if
needed.

[jdoerfert]

openmp/libomptarget/src/MemoryManager.cpp

115

Can't we do this in the destructor MemoryManagerTy::~MemoryManagerTy() ?

[tianshilei1992]

Only check the validity in MemoryManagerTy::~MemoryManagerTy

[tianshilei1992]

openmp/libomptarget/src/MemoryManager.cpp

115

Nice point! Did that.

[jdoerfert]

One nit, works for me otherwise. Anyone else?

openmp/libomptarget/src/MemoryManager.cpp

109

Early exit or braces please, former preferred.

[tianshilei1992]

Updated the coding style

[JonChesterfield]

I suppose this is better, though it seems backwards to me. Plus we still leak memory with this design.

I was hoping the plugin would expose a function that libomptarget calls to do the cleanup, in place of today's global destructor. Libomptarget can then do an ordered cleanup, after the memory manager destructor has run.

[vzakhari]

openmp/libomptarget/src/rtl.cpp

444–445

Can we just implement this instead?

[tianshilei1992]

openmp/libomptarget/src/rtl.cpp

444–445

At this point, the plugin has already been destroyed.

[vzakhari]

openmp/libomptarget/src/rtl.cpp

444–445

How is that possible, since we never call dlclose() for the plugins?
Is the problem specific to some particular execution conditions?

Should we just "deinitialize" the memory manager here instead of postponing the "deinitialization" to MemoryManagerTy desctructor, which will obviously be executed at the wrong time.

[tianshilei1992]

openmp/libomptarget/src/rtl.cpp

444–445

That's most tough part of this question. libomptarget has no idea when its opened plugin will be destroyed. Actually, you can see there is no dlclose in current code. What we can know is, plugin will be destroyed before libomptarget. You can easily verify it by adding DP(...) to the destructor of the plugin.

[vzakhari]

openmp/libomptarget/src/rtl.cpp

444–445

This sounds really strange at least for Linux, since:

The dynamic linker maintains reference counts for object handles, so a dynamically loaded shared object is not deallocated until dlclose() has been called on it as many times as dlopen() has succeeded on it.

Do you have a reproducer for this kind of an issue? I must be missing something here...

[tianshilei1992]

openmp/libomptarget/src/rtl.cpp

444–445

You can use openmp/libomptarget/test/offloading/memory_manager.cpp and turn debug print on via setting env LIBOMPTARGET_DEBUG to a large number like 10. Then you can see a lot of logs.

...
Libomptarget --> omp_target_free deallocated device ptr
PASS
(1) Target CUDA RTL --> CUDA plugin destroyed
(2) Libomptarget --> __tgt_unregister_lib
Libomptarget --> Unloading target library!
Libomptarget --> Image 0x0000000000403870 is compatible with RTL 0x0000000000e4a8b0!
Libomptarget --> Unregistered image 0x0000000000403870 from RTL 0x0000000000e4a8b0!
Libomptarget --> Done unregistering images!
Libomptarget --> Removing translation table for descriptor 0x00000000004abb60
Libomptarget --> Done unregistering library!
Libomptarget --> Deinit target library!
Target CUDA RTL --> Error returned from cuCtxSetCurrent
Target CUDA RTL --> Unresolved CUDA error code: 4
Target CUDA RTL --> Unsuccessful cuGetErrorString return status: 4
Target CUDA RTL --> Error returned from cuCtxSetCurrent
...

Line (1) and (2) were printed because I added DP to the destructor of CUDA plugin and function __tgt_unregister_lib. The last tens of lines are all error messages because CUDA plugin was destroyed before libomptarget.

[vzakhari]

openmp/libomptarget/src/rtl.cpp

444–445

Ah, I see. The global destructor for DeviceRTL is called before __tgt_unregister_lib is called. The plugin (the shared object) is not actually deallocated yet. Thank you for the explanation.

I am not sure libomptarget has to be involved here, since this may be done exclusively at the plugin level, i.e. DeviceRTL may be dynamically allocated/deallocated using global init/fini functions, and then all plugin APIs may return early, if they are invoked after DeviceRTL has been deallocated.

Just to be clear, I am not against this change-set :)

Anyway, libomptarget has to be redesigned a bit because the plugins loading initiated in a global constructor (which calls RTLsTy::RegisterLib) is violating Microsoft conventions for DLL initialization (e.g. you cannot create a DLL with offload and then LoadLibrary this DLL, because this will result in LoadLibrary->...->LoadLibrary call stack, which is not allowed). We should probably load the plugins on the first OpenMP construct/API encountered during the execution. I assume we will have to define clear entry/exit points for plugins then, and this will solve the issue that you are addressing here.

[tianshilei1992]

openmp/libomptarget/src/rtl.cpp

444–445

libomptarget is involved because we have a memory manager which will releases all its buffered memory in its destructor. It is an object in libomptarget. So its destructor is invoked when libomptarget is being destroyed. Now we know at this moment the plugin has already been destroyed, and the memory manager still calls plugin interfaces, causing program crash on some platforms. However, usually it will not crash on CUDA platforms, just having some errors because the plugin is using some CUDA related data which have been explicitly destroyed when the plugin is destroyed.

Redesign of libomptarget is beyond the scope of this patch. It might be done when LLVM OpenMP decides to support Windows. For now this patch is like a WA.

[tianshilei1992]

In D94256#2485287, @JonChesterfield wrote:

I suppose this is better, though it seems backwards to me. Plus we still leak memory with this design.

The memory will be eventually released by the driver if the process is terminated, right?

I was hoping the plugin would expose a function that libomptarget calls to do the cleanup, in place of today's global destructor. Libomptarget can then do an ordered cleanup, after the memory manager destructor has run.

I have some new ideas to make the cleanup possible. Will update the patch accordingly.

[JonChesterfield]

In D94256#2485965, @tianshilei1992 wrote:

In D94256#2485287, @JonChesterfield wrote:

I suppose this is better, though it seems backwards to me. Plus we still leak memory with this design.

The memory will be eventually released by the driver if the process is terminated, right?

Maybe. Probably platform-specific. I definitely have to power cycle GPUs to get them working again occasionally so the kernel doesn't (always?) leave the stack in a just-booted state after process exit.

Leaking memory by design thwarts using memory tracking tools to find accidental leaks. In this case, it's indicative that we've failed to establish static lifetimes for the components. Too much global state.

[tianshilei1992]

Fixed the issue that target memory are not released when the plugin is destroyed

[JonChesterfield]

The complexity of this approach is very high relative to instantiating MemoryManagerTy in the cuda plugin, and optionally in any others that want to use it.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

282

It looks like this stuff needs to be implemented in every plugin, since they all use the memory manager at present

[tianshilei1992]

In D94256#2489154, @JonChesterfield wrote:

The complexity of this approach is very high relative to instantiating MemoryManagerTy in the cuda plugin, and optionally in any others that want to use it.

This approach does not only fix the issue, but also open a door for future development. RTLInfoTy is used by DeviceTy, so each one has pointers to the other. It is a weak dependence so not a problem. We have a lot of similar things in LLVM, such as use and user.

The callback function can also be used later as it is the only chance for libomptarget to know when the plugin is destroyed and do something needed.

In the future we will add something to let a plugin exchange information with libomptarget. For example, a plugin can tell libomptarget what features it can support and libomptarget can use that information accordingly. Moving the memory manager to plugin and taking it as a common part is also feasible.

[tianshilei1992]

Proposed to move the memory manager to plugin. See D94379 for more details.

openmp/libomptarget/include/omptargetplugin.h

int32_t __tgt_rtl_run_target_team_region_async(
     int32_t ID, void *Entry, void **Args, ptrdiff_t *Offsets, int32_t NumArgs,
     int32_t NumTeams, int32_t ThreadLimit, uint64_t loop_tripcount,
     __tgt_async_info *AsyncInfoPtr);
 
 // Device synchronization. In case of success, return zero. Otherwise, return an
 // 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
 
 #endif // _OMPTARGETPLUGIN_H_

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <vector>
 
 #include "Debug.h"
 #include "omptargetplugin.h"
class DeviceRTLTy {
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 
   std::unique_ptr<StreamManagerTy> StreamManager;
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> 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();
     E.Entries.push_back(entry);
   }
 
   // Return a pointer to the entry associated with the pointer
   const __tgt_offload_entry *getOffloadEntry(const int DeviceId,
if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     }
 
     StreamManager =
         std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
   }
 
   ~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;
 
     for (CUmodule &M : Modules)
       // Close module
       if (M)
         checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
 
int synchronize(const int DeviceId, __tgt_async_info *AsyncInfoPtr) const {
     // async_info. This is to make sure the synchronization only works for its
     // own tasks.
     StreamManager->returnStream(
         DeviceId, reinterpret_cast<CUstream>(AsyncInfoPtr->Queue));
     AsyncInfoPtr->Queue = nullptr;
 
     return OFFLOAD_SUCCESS;
   }
+
+  void registerDestructionCB(void (*CB)(void *), void *Data) {
+    DestructionCBInfo.CB = CB;
+    DestructionCBInfo.Data = Data;
+  }
 };
 
 DeviceRTLTy DeviceRTL;
 } // namespace
 
 // Exposed library API function
 #ifdef __cplusplus
 extern "C" {
int32_t __tgt_rtl_synchronize(int32_t device_id,
                               __tgt_async_info *async_info_ptr) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
   assert(async_info_ptr && "async_info_ptr is nullptr");
   assert(async_info_ptr->Queue && "async_info_ptr->Queue is nullptr");
 
   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

openmp/libomptarget/plugins/exports

global:
     __tgt_rtl_data_exchange;
     __tgt_rtl_data_exchange_async;
     __tgt_rtl_data_delete;
     __tgt_rtl_run_target_team_region;
     __tgt_rtl_run_target_team_region_async;
     __tgt_rtl_run_target_region;
     __tgt_rtl_run_target_region_async;
     __tgt_rtl_synchronize;
+    __tgt_rtl_register_destruction_cb;
   local:
     *;
 };

openmp/libomptarget/src/MemoryManager.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "MemoryManager.h"
 #include "device.h"
 #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,
     1U << 8, 1U << 9, 1U << 10, 1U << 11, 1U << 12, 1U << 13};
 
 constexpr const int NumBuckets = sizeof(BucketSize) / sizeof(BucketSize[0]);
 
 /// The threshold to manage memory using memory manager. If the request size is
for (auto Itr = PtrToNodeTable.begin(); Itr != PtrToNodeTable.end(); ++Itr) {
     assert(Itr->second.Ptr && "nullptr in map table");
     deleteOnDevice(Itr->second.Ptr);
   }
 }
 
 void *MemoryManagerTy::allocateOnDevice(size_t Size, void *HstPtr) const {
   return Device.RTL->data_alloc(Device.RTLDeviceID, Size, HstPtr);
 }
 
 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);
 }
 
 void *MemoryManagerTy::freeAndAllocate(size_t Size, void *HstPtr) {
   std::vector<void *> RemoveList;
 
   // Deallocate all memory in FreeList
   for (int I = 0; I < NumBuckets; ++I) {

openmp/libomptarget/src/rtl.h

 #include <mutex>
 #include <string>
 #include <vector>
 
 // Forward declarations.
 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);
   typedef int32_t(number_of_devices_ty)();
   typedef int32_t(init_device_ty)(int32_t);
   typedef __tgt_target_table *(load_binary_ty)(int32_t, void *);
   typedef void *(data_alloc_ty)(int32_t, int64_t, void *);
   typedef int32_t(data_submit_ty)(int32_t, void *, void *, int64_t);
struct RTLInfoTy {
   typedef int32_t(run_team_region_ty)(int32_t, void *, void **, ptrdiff_t *,
                                       int32_t, int32_t, int32_t, uint64_t);
   typedef int32_t(run_team_region_async_ty)(int32_t, void *, void **,
                                             ptrdiff_t *, int32_t, int32_t,
                                             int32_t, uint64_t,
                                             __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,
                                 // i.e. the OpenMP index of the first device
                                 // to be registered with this RTL.
   int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
 
   void *LibraryHandler = nullptr;
#endif
   data_exchange_async_ty *data_exchange_async = nullptr;
   data_delete_ty *data_delete = nullptr;
   run_region_ty *run_region = nullptr;
   run_region_async_ty *run_region_async = nullptr;
   run_team_region_ty *run_team_region = nullptr;
   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.
   std::mutex Mtx;
 
   // The existence of the mutex above makes RTLInfoTy non-copyable.
   // We need to provide a copy constructor explicitly.
   RTLInfoTy() = default;
#endif
     data_delete = r.data_delete;
     run_region = r.run_region;
     run_region_async = r.run_region_async;
     run_team_region = r.run_team_region;
     run_team_region_async = r.run_team_region_async;
     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);
   }
 };
 
 /// RTLs identified in the system.
 class RTLsTy {
 private:
   // Mutex-like object to guarantee thread-safety and unique initialization
   // (i.e. the library attempts to load the RTLs (plugins) only once).

openmp/libomptarget/src/rtl.cpp

*((void **)&R.run_team_region_async) =
         dlsym(dynlib_handle, "__tgt_rtl_run_target_team_region_async");
     *((void **)&R.synchronize) = dlsym(dynlib_handle, "__tgt_rtl_synchronize");
     *((void **)&R.data_exchange) =
         dlsym(dynlib_handle, "__tgt_rtl_data_exchange");
     *((void **)&R.data_exchange_async) =
         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())) {
       DP("No devices supported in this RTL\n");
       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);
 
     // The RTL is valid! Will save the information in the RTLs list.
     AllRTLs.push_back(R);
   }
 
   DP("RTLs loaded!\n");
if (TransTable != PM->HostEntriesBeginToTransTable.end()) {
     PM->HostEntriesBeginToTransTable.erase(TransTable);
   } else {
     DP("Translation table for descriptor " DPxMOD " cannot be found, probably "
         "it has been already removed.\n", DPxPTR(desc->HostEntriesBegin));
   }
 
   PM->TblMapMtx.unlock();
 
   // TODO: Remove RTL and the devices it manages if it's not used anymore?
   // TODO: Write some RTL->unload_image(...) function?
 
   DP("Done unregistering library!\n");
 }
+
+extern "C" {
+void invalidateRTL(void *Data) {
+  DP("Invalidating plugin with handler " DPxMOD ".\n", DPxPTR(Data));
+  reinterpret_cast<RTLInfoTy *>(Data)->IsAlive = false;
+}
+}