[llvm] [OFFLOAD] Add support for indexed per-thread containers (PR #164263)

[Kevin Sala Penades via llvm-commits]

================
@@ -91,24 +177,114 @@ template <typename ContainerType, typename ObjectType> struct PerThreadTable {
   // Iterators to traverse objects owned by
   // the current thread
   iterator begin() {
-    auto &Entry = getThreadEntry();
+    ContainerType &Entry = getThreadEntry();
     return Entry.begin();
   }
   iterator end() {
-    auto &Entry = getThreadEntry();
+    ContainerType &Entry = getThreadEntry();
     return Entry.end();
   }
 
-  template <class F> void clear(F f) {
-    std::lock_guard<std::mutex> Lock(Mtx);
-    for (auto ThData : ThreadDataList) {
-      if (!ThData->ThEntry || ThData->NElements == 0)
+  template <class ClearFuncTy> void clear(ClearFuncTy ClearFunc) {
+    assert(Mutex.try_lock() && (Mutex.unlock(), true) &&
+           "Clear cannot be called while other threads are adding entries");
+    for (std::shared_ptr<PerThreadData> ThreadData : ThreadDataList) {
+      if (!ThreadData->ThreadEntry || ThreadData->NElements == 0)
         continue;
-      ThData->ThEntry->clear(f);
-      ThData->NElements = 0;
+      if constexpr (has_clearAll<ContainerType>::value) {
+        ThreadData->ThreadEntry->clearAll(ClearFunc);
+      } else if constexpr (has_iterator<ContainerType>::value &&
+                           has_clear<ContainerType>::value) {
+        for (auto &Obj : *ThreadData->ThreadEntry) {
+          if constexpr (is_associative<ContainerType>::value) {
+            ClearFunc(Obj.second);
+          } else {
+            ClearFunc(Obj);
+          }
+        }
+        ThreadData->ThreadEntry->clear();
+      } else {
+        static_assert(true, "Container type not supported");
+      }
+      ThreadData->NElements = 0;
     }
     ThreadDataList.clear();
   }
+
+  template <class DeinitFuncTy> llvm::Error deinit(DeinitFuncTy DeinitFunc) {
+    assert(Mutex.try_lock() && (Mutex.unlock(), true) &&
+           "Deinit cannot be called while other threads are adding entries");
+    for (std::shared_ptr<PerThreadData> ThreadData : ThreadDataList) {
+      if (!ThreadData->ThreadEntry || ThreadData->NElements == 0)
+        continue;
+      for (auto &Obj : *ThreadData->ThreadEntry) {
+        if constexpr (is_associative<ContainerType>::value) {
+          if (auto Err = DeinitFunc(Obj.second))
+            return Err;
+        } else {
+          if (auto Err = DeinitFunc(Obj))
+            return Err;
+        }
+      }
+    }
+    return llvm::Error::success();
+  }
+};
+
+template <typename T, typename = std::void_t<>> struct ContainerValueType {
+  using type = typename T::value_type;
+};
+template <typename T>
+struct ContainerValueType<T, std::void_t<typename T::mapped_type>> {
+  using type = typename T::mapped_type;
+};
+
+template <typename ContainerType, size_t reserveSize = 0>
+struct PerThreadContainer
+    : public PerThreadTable<ContainerType,
+                            typename ContainerValueType<ContainerType>::type> {
+
+  // helpers
+  template <typename T, typename = std::void_t<>> struct indexType {
+    using type = typename T::size_type;
+  };
+  template <typename T> struct indexType<T, std::void_t<typename T::key_type>> {
+    using type = typename T::key_type;
+  };
+  template <typename T, typename = std::void_t<>>
+  struct has_resize : std::false_type {};
+  template <typename T>
+  struct has_resize<T, std::void_t<decltype(std::declval<T>().resize(1))>>
+      : std::true_type {};
+
+  template <typename T, typename = std::void_t<>>
+  struct has_reserve : std::false_type {};
+  template <typename T>
+  struct has_reserve<T, std::void_t<decltype(std::declval<T>().reserve(1))>>
+      : std::true_type {};
+
+  using IndexType = typename indexType<ContainerType>::type;
+  using ObjectType = typename ContainerValueType<ContainerType>::type;
+
+  // Get the object for the given index in the current thread
+  ObjectType &get(IndexType Index) {
+    ContainerType &Entry = this->getThreadEntry();
+
+    // specialized code for vector-like containers
+    if constexpr (has_resize<ContainerType>::value) {
+      if (Index >= Entry.size()) {
----------------
kevinsala wrote:

I understand you're delaying the action of reserving until you are sure the reserved size will be used. That's right? You don't make an initial reservation to avoid consuming memory that may not be used?

https://github.com/llvm/llvm-project/pull/164263