[OpenMP][libomptarget] Add utility class for reference counting

The AMDGPU NextGen plugin will use this class for counting the references of some device resources.

Differential Revision: https://reviews.llvm.org/D139787

diff --git a/openmp/libomptarget/include/Utilities.h b/openmp/libomptarget/include/Utilities.h
index 6428c07b1e7e329a63d0b56f0a5dc3f003eebfa9..813ffd08aefb8d0d0d11950fe6efe1d1b5815f27 100644 (file)
--- a/openmp/libomptarget/include/Utilities.h
+++ b/openmp/libomptarget/include/Utilities.h
@@ -19,6 +19,7 @@
 #include "Debug.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cassert>
 #include <cstdint>
 #include <cstdlib>
@@ -129,6 +130,36 @@ using UInt64Envar = Envar<uint64_t>;
 using StringEnvar = Envar<std::string>;
 using BoolEnvar = Envar<bool>;
 
+/// Utility class for thread-safe reference counting. Any class that needs
+/// objects' reference counting can inherit from this entity or have it as a
+/// class data member.
+template <typename Ty = uint32_t,
+          std::memory_order MemoryOrder = std::memory_order_relaxed>
+struct RefCountTy {
+  /// Create a refcount object initialized to zero.
+  RefCountTy() : Refs(0) {}
+
+  ~RefCountTy() { assert(Refs == 0 && "Destroying with non-zero refcount"); }
+
+  /// Increase the reference count atomically.
+  void increase() { Refs.fetch_add(1, MemoryOrder); }
+
+  /// Decrease the reference count and return whether it became zero. Decreasing
+  /// the counter in more units than it was previously increased results in
+  /// undefined behavior.
+  bool decrease() {
+    Ty Prev = Refs.fetch_sub(1, MemoryOrder);
+    assert(Prev > 0 && "Invalid refcount");
+    return (Prev == 1);
+  }
+
+  Ty get() const { return Refs.load(MemoryOrder); }
+
+private:
+  /// The atomic reference counter.
+  std::atomic<Ty> Refs;
+};
+
 template <>
 inline bool StringParser::parse(const char *ValueStr, bool &Result) {
   std::string Value(ValueStr);