[Openmp-commits] [openmp] bbcffb0 - [OpenMP][libomptarget] Add utility class for reference counting

[Kevin Sala via Openmp-commits]

Author: Kevin Sala
Date: 2022-12-11T21:39:25+01:00
New Revision: bbcffb08f0fdc0be8c8cba48410f9cb556ea661d

URL: https://github.com/llvm/llvm-project/commit/bbcffb08f0fdc0be8c8cba48410f9cb556ea661d
DIFF: https://github.com/llvm/llvm-project/commit/bbcffb08f0fdc0be8c8cba48410f9cb556ea661d.diff

LOG: [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

Added: 
    

Modified: 
    openmp/libomptarget/include/Utilities.h

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/include/Utilities.h b/openmp/libomptarget/include/Utilities.h
index 6428c07b1e7e..813ffd08aefb 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);