Fix tsan problem where the per-thread shared_ptr() can be locked right before the cache is destroyed causing a race where it tries to remove an entry from a destroyed cache.

[pschuh]

Summary

This is a rollforward with fixes of https://reviews.llvm.org/rGbcc10817d5569172ee065015747e226280e9b698 (originally https://reviews.llvm.org/D142394). The original patch exposed an asan problem on aarch64, which is fixed by simply calling the context destructors properly.

mlir/include/mlir/Support/ThreadLocalCache.h

 #include "llvm/Support/Mutex.h"
 
 namespace mlir {
 /// This class provides support for defining a thread local object with non
 /// static storage duration. This is very useful for situations in which a data
 /// cache has very large lock contention.
 template <typename ValueT>
 class ThreadLocalCache {
+  // Keep a separate shared_ptr protected state that can be acquired atomically
+  // instead of using shared_ptr's for each value. This avoids a problem
+  // where the instance shared_ptr is locked() successfully, and then the
+  // ThreadLocalCache gets destroyed before remove() can be called successfully.
+  struct PerInstanceState {
+    /// Remove the given value entry. This is generally called when a thread
+    /// local cache is destructing.
+    void remove(ValueT *value) {
+      // Erase the found value directly, because it is guaranteed to be in the
+      // list.
+      llvm::sys::SmartScopedLock<true> threadInstanceLock(instanceMutex);
+      auto it =
+          llvm::find_if(instances, [&](std::unique_ptr<ValueT> &instance) {
+            return instance.get() == value;
+          });
+      assert(it != instances.end() && "expected value to exist in cache");
+      instances.erase(it);
+    }
+
+    /// Owning pointers to all of the values that have been constructed for this
+    /// object in the static cache.
+    SmallVector<std::unique_ptr<ValueT>, 1> instances;
+
+    /// A mutex used when a new thread instance has been added to the cache for
+    /// this object.
+    llvm::sys::SmartMutex<true> instanceMutex;
+  };
+
   /// The type used for the static thread_local cache. This is a map between an
   /// instance of the non-static cache and a weak reference to an instance of
   /// 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<ThreadLocalCache<ValueT> *,
-                                                std::weak_ptr<ValueT>> {
+  struct CacheType
+      : 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.lock())
           it.first->remove(value.get());
     }
 
     /// Clear out any unused entries within the map. This method is not
~ThreadLocalCache() {
     // No cleanup is necessary here as the shared_pointer memory will go out of
     // scope and invalidate the weak pointers held by the thread_local caches.
   }
 
   /// Return an instance of the value type for the current thread.
   ValueT &get() {
     // Check for an already existing instance for this thread.
     CacheType &staticCache = getStaticCache();
-    std::weak_ptr<ValueT> &threadInstance = staticCache[this];
+    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(instanceMutex);
-    instances.push_back(std::make_shared<ValueT>());
-    std::shared_ptr<ValueT> &instance = instances.back();
-    threadInstance = instance;
+    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 *instance;
   }
   ValueT &operator*() { return get(); }
   ValueT *operator->() { return &get(); }
 
 private:
   ThreadLocalCache(ThreadLocalCache &&) = delete;
   ThreadLocalCache(const ThreadLocalCache &) = delete;
   ThreadLocalCache &operator=(const ThreadLocalCache &) = delete;
 
   /// Return the static thread local instance of the cache type.
   static CacheType &getStaticCache() {
     static LLVM_THREAD_LOCAL CacheType cache;
     return cache;
   }
 
-  /// Remove the given value entry. This is generally called when a thread local
-  /// cache is destructing.
-  void remove(ValueT *value) {
-    // Erase the found value directly, because it is guaranteed to be in the
-    // list.
-    llvm::sys::SmartScopedLock<true> threadInstanceLock(instanceMutex);
-    auto it = llvm::find_if(instances, [&](std::shared_ptr<ValueT> &instance) {
-      return instance.get() == value;
-    });
-    assert(it != instances.end() && "expected value to exist in cache");
-    instances.erase(it);
-  }
-
-  /// Owning pointers to all of the values that have been constructed for this
-  /// object in the static cache.
-  SmallVector<std::shared_ptr<ValueT>, 1> instances;
-
-  /// A mutex used when a new thread instance has been added to the cache for
-  /// this object.
-  llvm::sys::SmartMutex<true> instanceMutex;
+  std::shared_ptr<PerInstanceState> perInstanceState =
+      std::make_shared<PerInstanceState>();
 };
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_THREADLOCALCACHE_H

mlir/test/CAPI/pdl.c

 int main(void) {
   MlirContext ctx = mlirContextCreate();
   mlirDialectHandleRegisterDialect(mlirGetDialectHandle__pdl__(), ctx);
   testAttributeType(ctx);
   testOperationType(ctx);
   testRangeType(ctx);
   testTypeType(ctx);
   testValueType(ctx);
+  mlirContextDestroy(ctx);
   return EXIT_SUCCESS;
 }

mlir/test/CAPI/transform.c

void testOperationType(MlirContext ctx) {
   fprintf(stderr, "\n\n");
 }
 
 int main(void) {
   MlirContext ctx = mlirContextCreate();
   mlirDialectHandleRegisterDialect(mlirGetDialectHandle__transform__(), ctx);
   testAnyOpType(ctx);
   testOperationType(ctx);
+  mlirContextDestroy(ctx);
   return EXIT_SUCCESS;
 }