[OpenMP] Explicitly deinitialize device resources

[jdoerfert]

Summary

There are two problems this patch tries to address:
1) We currently free resources in a random order wrt. plugin and
   libomptarget destruction. This patch should ensure the CUDA plugin
   is less fragile if something during the deinitialization goes wrong.
2) We need to support (hard) pause runtime calls eventually. This patch
   allows us to free all associated resources, though we cannot
   reinitialize the device yet.

Follow up patch will associate one event pool per device/context.

[tianshilei1992]

I don't think this patch can fix the first issue because we don't know whether the plugin has already been destroyed when objects from libomptarget are destroyed. I had experience that calling plugin functions in destructor of one object from libomptarget caused crash because plugin has already been released.

[jdoerfert]

Add unregister plugin call for CUDA plugin, deinitialize all devices then

[jdoerfert]

Do not deinit the plugin device when the libomptarget device is destroyed.

[jdoerfert]

In D120089#3330880, @tianshilei1992 wrote:

I don't think this patch can fix the first issue because we don't know whether the plugin has already been destroyed when objects from libomptarget are destroyed. I had experience that calling plugin functions in destructor of one object from libomptarget caused crash because plugin has already been released.

OK, that seems to be a separate issue. I think with the __tgt_rtl_unregister_lib we might have that one under control.

So this patch does (1) introduce per device "deinit" capabilities (from libomptarget this is unused for now).
And, (2) deinit all devices as part of the user code destructor (via __tgt_rtl_unregister_lib)

[jdoerfert]

Updates

[jdoerfert]

`

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

781

The conditionals here might help us avoid the crashes we've seen before.

[tianshilei1992]

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

566

Why do we need it?

[jdoerfert]

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

566

Because we need to know the state of a device. Is it initialized or not. This flag tells us.

[tianshilei1992]

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

566

Here they are assigned all at once, but we only initialize the device on demand.

[jdoerfert]

Addressed comments

[tianshilei1992]

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

590

Actually this could have data race. Say if multiple threads are trying to initialize it at the same time. I don't remember we have the guard in libomptarget. We could potentially use std::vector<std::once_flag> instead.

[jdoerfert]

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

590

This is something libomptarget should handle. There is little reason N plugins should do it and also little reason to assume plugins should allow concurrent initialization of the same device. That said, DeviceTy::init() in libomptarget is guarded by a std::call_once which should do the trick just fine.

[tianshilei1992]

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

26

Why do we want omptarget.h here?

590

I'm fine with that, as long as we really do it.

767

This none_of is useless here because InitializedFlags[DeviceId] is reset beforehand.

openmp/libomptarget/src/device.cpp

471

Who calls it?

[tianshilei1992]

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

767

Oh, none_of, nvm.

[jdoerfert]

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

26

Because my clang complete things we use it (probably for some types).

590

If we wouldn't, how is this patch making it worse? All the assignments below would be racy as well.

767

This is going to be replaced in the follow up when we have one event pool per device (which we need).

openmp/libomptarget/src/device.cpp

471

Nobody anymore, probably. Was called but I removed it. We would call it for hard_shutdown though.

[tianshilei1992]

LG w/ some nits.

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

26

omptargetplugin.h includes omptarget.h already, but it's a good practice to include what we need. Good to keep it.

590

I agree. No problem. I mean, we need to do guard it. If not in this patch, then we do it in another patch.

openmp/libomptarget/src/device.cpp

471

We can call it in the deinit function emitted by clang in another patch once we land it.

[tianshilei1992]

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

760

Modules are not indexed by device ID. We need std::vector<std::vector<CUmodule>> instead.

openmp/libomptarget/include/device.h

struct DeviceTy {
 
   /// Destroy the event.
   int32_t destroyEvent(void *Event);
   /// }
 
 private:
   // Call to RTL
   void init(); // To be called only via DeviceTy::initOnce()
+
+  /// Deinitialize the device (and plugin).
+  void deinit();
 };
 
 extern bool device_is_ready(int device_num);
 
 /// Struct for the data required to handle plugins
 struct PluginManager {
   PluginManager(bool UseEventsForAtomicTransfers)
       : UseEventsForAtomicTransfers(UseEventsForAtomicTransfers) {}

openmp/libomptarget/include/omptargetplugin.h

 
 // Initialize the requires flags for the device.
 int64_t __tgt_rtl_init_requires(int64_t RequiresFlags);
 
 // Initialize the specified device. In case of success return 0; otherwise
 // return an error code.
 int32_t __tgt_rtl_init_device(int32_t ID);
 
+// Deinitialize the specified device. In case of success return 0; otherwise
+// return an error code.
+int32_t __tgt_rtl_deinit_device(int32_t ID);
+
 // Pass an executable image section described by image to the specified
 // device and prepare an address table of target entities. In case of error,
 // return NULL. Otherwise, return a pointer to the built address table.
 // Individual entries in the table may also be NULL, when the corresponding
 // offload region is not supported on the target device.
 __tgt_target_table *__tgt_rtl_load_binary(int32_t ID,
                                           __tgt_device_image *Image);
 

openmp/libomptarget/include/rtl.h

 struct DeviceTy;
 struct __tgt_bin_desc;
 
 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 int32_t(deinit_device_ty)(int32_t);
   typedef __tgt_target_table *(load_binary_ty)(int32_t, void *);
   typedef void *(data_alloc_ty)(int32_t, int64_t, void *, int32_t);
   typedef int32_t(data_submit_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_submit_async_ty)(int32_t, void *, void *, int64_t,
                                         __tgt_async_info *);
   typedef int32_t(data_retrieve_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_retrieve_async_ty)(int32_t, void *, void *, int64_t,
                                           __tgt_async_info *);
#ifdef OMPTARGET_DEBUG
   std::string RTLName;
 #endif
 
   // Functions implemented in the RTL.
   is_valid_binary_ty *is_valid_binary = nullptr;
   is_data_exchangable_ty *is_data_exchangable = nullptr;
   number_of_devices_ty *number_of_devices = nullptr;
   init_device_ty *init_device = nullptr;
+  deinit_device_ty *deinit_device = nullptr;
   load_binary_ty *load_binary = nullptr;
   data_alloc_ty *data_alloc = nullptr;
   data_submit_ty *data_submit = nullptr;
   data_submit_async_ty *data_submit_async = nullptr;
   data_retrieve_ty *data_retrieve = nullptr;
   data_retrieve_async_ty *data_retrieve_async = nullptr;
   data_exchange_ty *data_exchange = nullptr;
   data_exchange_async_ty *data_exchange_async = nullptr;

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 <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
+#include "omptarget.h"
 #include "omptargetplugin.h"
 
 #define TARGET_NAME CUDA
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 
 #include "MemoryManager.h"
 
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
class DeviceRTLTy {
   using StreamPoolTy = ResourcePoolTy<CUstream>;
   std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
 
   ResourcePoolTy<CUevent> EventPool;
 
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> Modules;
 
+  /// Vector of flags indicating the initalization status of all associated
+  /// devices.
+  std::vector<bool> InitializedFlags;
+
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     const int DeviceId;
     const std::vector<DeviceDataTy> &DeviceData;
     std::unordered_map<void *, TargetAllocTy> HostPinnedAllocs;
 
   public:
class DeviceRTLTy {
   void clearOffloadEntriesTable(const int DeviceId) {
     DeviceData[DeviceId].FuncGblEntries.emplace_back();
     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = nullptr;
   }
 
 public:
-
   CUstream getStream(const int DeviceId, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     if (!AsyncInfo->Queue) {
       CUstream S;
       if (StreamPool[DeviceId]->acquire(S) != OFFLOAD_SUCCESS)
         return nullptr;
 
DeviceRTLTy()
     std::pair<size_t, bool> Res = MemoryManagerTy::getSizeThresholdFromEnv();
     UseMemoryManager = Res.second;
     size_t MemoryManagerThreshold = Res.first;
 
     if (UseMemoryManager)
       for (int I = 0; I < NumberOfDevices; ++I)
         MemoryManagers.emplace_back(std::make_unique<MemoryManagerTy>(
             DeviceAllocators[I], MemoryManagerThreshold));
+
+    // We lazily initialize all devices later.
+    InitializedFlags.assign(NumberOfDevices, false);
   }
 
   ~DeviceRTLTy() {
-    // We first destruct memory managers in case that its dependent data are
-    // destroyed before it.
-    for (auto &M : MemoryManagers)
-      M.release();
-
-    for (CUmodule &M : Modules)
-      // Close module
-      if (M)
-        checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
-
-    for (auto &S : StreamPool)
-      S.reset();
-
-    EventPool.clear();
-
-    for (DeviceDataTy &D : DeviceData) {
-      // Destroy context
-      if (D.Context) {
-        checkResult(cuCtxSetCurrent(D.Context),
-                    "Error returned from cuCtxSetCurrent\n");
-        CUdevice Device;
-        checkResult(cuCtxGetDevice(&Device),
-                    "Error returned from cuCtxGetDevice\n");
-        checkResult(cuDevicePrimaryCtxRelease(Device),
-                    "Error returned from cuDevicePrimaryCtxRelease\n");
-      }
-    }
+    for (int DeviceId = 0; DeviceId < NumberOfDevices; ++DeviceId)
+      deinitDevice(DeviceId);
   }
 
   // Check whether a given DeviceId is valid
   bool isValidDeviceId(const int DeviceId) const {
     return DeviceId >= 0 && DeviceId < NumberOfDevices;
   }
 
   int getNumOfDevices() const { return NumberOfDevices; }
 
   void setRequiresFlag(const int64_t Flags) { this->RequiresFlags = Flags; }
 
   int initDevice(const int DeviceId) {
     CUdevice Device;
 
     DP("Getting device %d\n", DeviceId);
     CUresult Err = cuDeviceGet(&Device, DeviceId);
     if (!checkResult(Err, "Error returned from cuDeviceGet\n"))
       return OFFLOAD_FAIL;
 
+    assert(InitializedFlags[DeviceId] == false && "Reinitializing device!");
+    InitializedFlags[DeviceId] = true;
+
     // Query the current flags of the primary context and set its flags if
     // it is inactive
     unsigned int FormerPrimaryCtxFlags = 0;
     int FormerPrimaryCtxIsActive = 0;
     Err = cuDevicePrimaryCtxGetState(Device, &FormerPrimaryCtxFlags,
                                      &FormerPrimaryCtxIsActive);
     if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxGetState\n"))
       return OFFLOAD_FAIL;
if (DeviceData[DeviceId].NumThreads >
       DP("Default number of threads exceeds device limit, capping at %d\n",
          DeviceData[DeviceId].ThreadsPerBlock);
       DeviceData[DeviceId].NumThreads = DeviceData[DeviceId].ThreadsPerBlock;
     }
 
     return OFFLOAD_SUCCESS;
   }
 
+  int deinitDevice(const int DeviceId) {
+    auto IsInitialized = InitializedFlags[DeviceId];
+    if (!IsInitialized)
+      return OFFLOAD_SUCCESS;
+    InitializedFlags[DeviceId] = false;
+
+    if (UseMemoryManager)
+      MemoryManagers[DeviceId].release();
+
+    // Close module
+    if (CUmodule &M = Modules[DeviceId])
+      checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
+
+    StreamPool[DeviceId].reset();
+
+    // The event pool is shared, we initialize it once all devices have been
+    // deinitialized.
+    if (std::none_of(InitializedFlags.begin(), InitializedFlags.end(),
+                     [](bool IsInitialized) { return IsInitialized; }))
+      EventPool.clear();
+
+    // Destroy context
+    DeviceDataTy &D = DeviceData[DeviceId];
+    if (D.Context) {
+      if (checkResult(cuCtxSetCurrent(D.Context),
+                      "Error returned from cuCtxSetCurrent\n")) {
+        CUdevice Device;
+        if (checkResult(cuCtxGetDevice(&Device),
+                        "Error returned from cuCtxGetDevice\n"))
+          checkResult(cuDevicePrimaryCtxRelease(Device),
+                      "Error returned from cuDevicePrimaryCtxRelease\n");
+      }
+    }
+    return OFFLOAD_SUCCESS;
+  }
+
   __tgt_target_table *loadBinary(const int DeviceId,
                                  const __tgt_device_image *Image) {
     // Set the context we are using
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return nullptr;
 
     // Clear the offload table as we are going to create a new one.

 }
 
 int32_t __tgt_rtl_init_device(int32_t device_id) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   return DeviceRTL.initDevice(device_id);
 }
 
+int32_t __tgt_rtl_deinit_device(int32_t device_id) {
+  assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
+
+  return DeviceRTL.deinitDevice(device_id);
+}
+
 __tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
                                           __tgt_device_image *image) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   return DeviceRTL.loadBinary(device_id, image);
 }
 
 void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *,

openmp/libomptarget/plugins/exports

 VERS1.0 {
   global:
     __tgt_rtl_is_valid_binary;
     __tgt_rtl_is_data_exchangable;
     __tgt_rtl_number_of_devices;
     __tgt_rtl_init_requires;
     __tgt_rtl_init_device;
+    __tgt_rtl_deinit_device;
     __tgt_rtl_load_binary;
     __tgt_rtl_data_alloc;
     __tgt_rtl_data_submit;
     __tgt_rtl_data_submit_async;
     __tgt_rtl_data_retrieve;
     __tgt_rtl_data_retrieve_async;
     __tgt_rtl_data_exchange;
     __tgt_rtl_data_exchange_async;

openmp/libomptarget/src/device.cpp

int32_t DeviceTy::initOnce() {
   // that some other thread already attempted to execute init() and if IsInit
   // is still false, return OFFLOAD_FAIL.
   if (IsInit)
     return OFFLOAD_SUCCESS;
   else
     return OFFLOAD_FAIL;
 }
 
+void DeviceTy::deinit() {
+  if (RTL->deinit_device)
+    RTL->deinit_device(RTLDeviceID);
+}
+
 // Load binary to device.
 __tgt_target_table *DeviceTy::load_binary(void *Img) {
   std::lock_guard<decltype(RTL->Mtx)> LG(RTL->Mtx);
   __tgt_target_table *rc = RTL->load_binary(RTLDeviceID, Img);
   return rc;
 }
 
 void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {

openmp/libomptarget/src/rtl.cpp

 #ifdef OMPTARGET_DEBUG
     R.RTLName = Name;
 #endif
 
     DP("Registering RTL %s supporting %d devices!\n", R.RTLName.c_str(),
        R.NumberOfDevices);
 
     // Optional functions
+    *((void **)&R.deinit_device) =
+        dlsym(dynlib_handle, "__tgt_rtl_deinit_device");
     *((void **)&R.init_requires) =
         dlsym(dynlib_handle, "__tgt_rtl_init_requires");
     *((void **)&R.data_submit_async) =
         dlsym(dynlib_handle, "__tgt_rtl_data_submit_async");
     *((void **)&R.data_retrieve_async) =
         dlsym(dynlib_handle, "__tgt_rtl_data_retrieve_async");
     *((void **)&R.run_region_async) =
         dlsym(dynlib_handle, "__tgt_rtl_run_target_region_async");