From bbcffb08f0fdc0be8c8cba48410f9cb556ea661d Mon Sep 17 00:00:00 2001
From: Kevin Sala <kevin.sala@bsc.es>
Date: Sun, 11 Dec 2022 16:51:05 +0100
Subject: [PATCH] [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
---
 openmp/libomptarget/include/Utilities.h | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/openmp/libomptarget/include/Utilities.h b/openmp/libomptarget/include/Utilities.h
index 6428c07b..813ffd0 100644
--- 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);
-- 
2.7.4