[OpenMP][CUDA] Fix potential program crash caused by double free resources

[tianshilei1992]

Summary

As we mentioned in the code comments for function ResourcePoolTy::release,
at some point there could be two identical resources on the two sides of Next
mark. It is usually not an issue, unless the following case:

1. Some resources are not returned.
2. We need to iterate the pool and free the element.

That will cause double free, which is the case for event pool. Since we don't release
events hold by the data map, it can happen that the Next mark is not reset, and
we have two identical items in the pool. When the pool is destroyed, we will call
cuEventDestroy twice on the same event. In the best case, we can only observe
CUDA errors. In the worst case, it can cause internal failures in CUDART and further
crash.

This patch fixes the issue by tracking all resources that have been given using
an unordered_set. We don't remove it when a resource is returned. When the pool
is destroyed, we merge the pool (a vector) and the set. In this way, we can make
sure that the set contains all resources allocated from the device. We just need
to iterate the set and free the resource accordingly.

For now, only event pool is set to use it. Stream pool is not because we can make
sure all streams are returned when the plugin is destroyed.

Someone might be wondering, why don't we release all events hold in the data map.
That is because, plugins are determined to be destroyed *before* libomptarget.
If we can somehow make the plugin outlast libomptarget, life will be much
easier.

[tianshilei1992]

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

302

This is the key part. Unluckily if the resources are not returned, the left one will never be overwritten.

[jdoerfert]

Why don't we need the same for Streams?

[ye-luo]

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

237

How about check in every allocated resource from runtime in ResourcesGiven and then use Resources only for pool management. In this way, ResourcesGiven is fully responsible for ownership.

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

In D122014#3393031, @jdoerfert wrote:

Why don't we need the same for Streams?

Because we always sync the stream at the end of a target task, so it is guaranteed that streams are all released in user's code. However, this is not 100% correct. If we get a "dangling" nowait target task, the stream might not be released before the end of user code. Actually, even with this patch, that could still be a problem because we only have an implicit mandatory task wait when libomp is shut down. However, that function is called via __attribute__((destructor)), which means plugins have already been destroyed. That can cause a problem. I don't have a clear clue on how to solve it.

[jdoerfert]

LG, I think this is clean.

[jdoerfert]

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

22

Nit: not needed.

[tianshilei1992]

remove unnecessary header

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

 #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"

 /// A generic pool of resources where \p T is the resource type.
 /// \p T should be copyable as the object is stored in \p std::vector .
 template <typename AllocTy> class ResourcePoolTy {
   using ElementTy = typename AllocTy::ElementTy;
   /// Index of the next available resource.
   size_t Next = 0;
   /// Mutex to guard the pool.
   std::mutex Mutex;
-  /// Pool of resources.
+  /// Pool of resources. The difference between \p Resources and \p Pool is,
+  /// when a resource is acquired and released, it is all on \p Resources. When
+  /// a batch of new resources are needed, they are both added to \p Resources
+  /// and \p Pool. The reason for this setting is, \p Resources could contain
+  /// redundant elements because resources are not released, which can cause
+  /// double free. This setting makes sure that \p Pool always has every
+  /// resource allocated from the device.
   std::vector<ElementTy> Resources;
+  std::vector<ElementTy> Pool;
   /// A reference to the corresponding allocator.
   AllocTy Allocator;
 
   /// If `Resources` is used up, we will fill in more resources. It assumes that
   /// the new size `Size` should be always larger than the current size.
   bool resize(size_t Size) {
+    assert(Resources.size() == Pool.size() && "size mismatch");
     auto CurSize = Resources.size();
     assert(Size > CurSize && "Unexpected smaller size");
+    Pool.reserve(Size);
     Resources.reserve(Size);
     for (auto I = CurSize; I < Size; ++I) {
       ElementTy NewItem;
       int Ret = Allocator.create(NewItem);
       if (Ret != OFFLOAD_SUCCESS)
         return false;
+      Pool.push_back(NewItem);
       Resources.push_back(NewItem);
     }
     return true;
   }
 
 public:
   ResourcePoolTy(AllocTy &&A, size_t Size = 0) noexcept
       : Allocator(std::move(A)) {
public:
   /// to first decrease `Next`, and then copy the resource back. It is worth
   /// noting that, the order of resources return might be different from that
   /// they're assigned, that saying, at some point, there might be two identical
   /// resources.
   /// xxax+a+++++
   ///     ^
   ///     Next
   /// However, it doesn't matter, because they're always on the two sides of
   /// `Next`. The left one will in the end be overwritten by another resource.
   /// Therefore, after several execution, the order of pool might be different
   /// from its initial state.
   void release(ElementTy R) noexcept {
     std::lock_guard<std::mutex> LG(Mutex);
     Resources[--Next] = R;
   }
 
   /// Released all stored resources and clear the pool.
   /// Note: This function is not thread safe. Be sure to guard it if necessary.
   void clear() noexcept {
-    for (auto &R : Resources)
+    for (auto &R : Pool)
       (void)Allocator.destroy(R);
+    Pool.clear();
     Resources.clear();
   }
 };
 
 } // namespace
 
 class DeviceRTLTy {
   int NumberOfDevices;