[llvm-branch-commits] [mlir] 50cb216 - Revert "[mlir] Optimize ThreadLocalCache by removing atomic bottleneck (#93270)"

[via llvm-branch-commits]

Author: Jacques Pienaar
Date: 2024-05-24T07:45:37-07:00
New Revision: 50cb2160464ef916146c199ab8c9fe3ea9f054d6

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

LOG: Revert "[mlir] Optimize ThreadLocalCache by removing atomic bottleneck (#93270)"

This reverts commit 1b803fe53dda11ca63f008f5ab8e206f3954ce48.

Added: 
    

Modified: 
    mlir/include/mlir/Support/ThreadLocalCache.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Support/ThreadLocalCache.h b/mlir/include/mlir/Support/ThreadLocalCache.h
index d19257bf6e25e..1be94ca14bcfa 100644
--- a/mlir/include/mlir/Support/ThreadLocalCache.h
+++ b/mlir/include/mlir/Support/ThreadLocalCache.h
@@ -58,12 +58,11 @@ class ThreadLocalCache {
   /// ValueT. We use a weak reference here so that the object can be destroyed
   /// without needing to lock access to the cache itself.
   struct CacheType
-      : public llvm::SmallDenseMap<PerInstanceState *,
-                                   std::pair<std::weak_ptr<ValueT>, ValueT *>> {
+      : public llvm::SmallDenseMap<PerInstanceState *, std::weak_ptr<ValueT>> {
     ~CacheType() {
       // Remove the values of this cache that haven't already expired.
       for (auto &it : *this)
-        if (std::shared_ptr<ValueT> value = it.second.first.lock())
+        if (std::shared_ptr<ValueT> value = it.second.lock())
           it.first->remove(value.get());
     }
 
@@ -72,7 +71,7 @@ class ThreadLocalCache {
     void clearExpiredEntries() {
       for (auto it = this->begin(), e = this->end(); it != e;) {
         auto curIt = it++;
-        if (curIt->second.first.expired())
+        if (curIt->second.expired())
           this->erase(curIt);
       }
     }
@@ -89,27 +88,22 @@ class ThreadLocalCache {
   ValueT &get() {
     // Check for an already existing instance for this thread.
     CacheType &staticCache = getStaticCache();
-    std::pair<std::weak_ptr<ValueT>, ValueT *> &threadInstance =
-        staticCache[perInstanceState.get()];
-    if (ValueT *value = threadInstance.second)
+    std::weak_ptr<ValueT> &threadInstance = staticCache[perInstanceState.get()];
+    if (std::shared_ptr<ValueT> value = threadInstance.lock())
       return *value;
 
     // Otherwise, create a new instance for this thread.
-    {
-      llvm::sys::SmartScopedLock<true> threadInstanceLock(
-          perInstanceState->instanceMutex);
-      threadInstance.second =
-          perInstanceState->instances.emplace_back(std::make_unique<ValueT>())
-              .get();
-    }
-    threadInstance.first =
-        std::shared_ptr<ValueT>(perInstanceState, threadInstance.second);
+    llvm::sys::SmartScopedLock<true> threadInstanceLock(
+        perInstanceState->instanceMutex);
+    perInstanceState->instances.push_back(std::make_unique<ValueT>());
+    ValueT *instance = perInstanceState->instances.back().get();
+    threadInstance = std::shared_ptr<ValueT>(perInstanceState, instance);
 
     // Before returning the new instance, take the chance to clear out any used
     // entries in the static map. The cache is only cleared within the same
     // thread to remove the need to lock the cache itself.
     staticCache.clearExpiredEntries();
-    return *threadInstance.second;
+    return *instance;
   }
   ValueT &operator*() { return get(); }
   ValueT *operator->() { return &get(); }