Lock the MLIR TypeConverter caches management to make it thread-safe (NFC)

mlir/include/mlir/Transforms/DialectConversion.h

 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 /// Type conversion class. Specific conversions and materializations can be
 /// registered using addConversion and addMaterialization, respectively.
 class TypeConverter {
 public:
   virtual ~TypeConverter() = default;
+  TypeConverter() = default;
+  // Copy the registered conversions, but not the caches
+  TypeConverter(const TypeConverter &other)
+      : conversions(other.conversions),
+        argumentMaterializations(other.argumentMaterializations),
+        sourceMaterializations(other.sourceMaterializations),
+        targetMaterializations(other.targetMaterializations),
+        typeAttributeConversions(other.typeAttributeConversions) {}
+  TypeConverter &operator=(const TypeConverter &other) {
+    conversions = other.conversions;
+    argumentMaterializations = other.argumentMaterializations;
+    sourceMaterializations = other.sourceMaterializations;
+    targetMaterializations = other.targetMaterializations;
+    typeAttributeConversions = other.typeAttributeConversions;
+    return *this;
+  }
 
   /// This class provides all of the information necessary to convert a type
   /// signature.
   class SignatureConversion {
   public:
     SignatureConversion(unsigned numOrigInputs)
         : remappedInputs(numOrigInputs) {}
 
private:
   SmallVector<TypeAttributeConversionCallbackFn, 2> typeAttributeConversions;
 
   /// A set of cached conversions to avoid recomputing in the common case.
   /// Direct 1-1 conversions are the most common, so this cache stores the
   /// successful 1-1 conversions as well as all failed conversions.
   mutable DenseMap<Type, Type> cachedDirectConversions;
   /// This cache stores the successful 1->N conversions, where N != 1.
   mutable DenseMap<Type, SmallVector<Type, 2>> cachedMultiConversions;
+  /// A mutex used for cache access
+  mutable llvm::sys::SmartRWMutex<true> cacheMutex;
 };
 
 //===----------------------------------------------------------------------===//
 // Conversion Patterns
 //===----------------------------------------------------------------------===//
 
 /// Base class for the conversion patterns. This pattern class enables type
 /// conversions, and other uses specific to the conversion framework. As such,

mlir/lib/Transforms/Utils/DialectConversion.cpp

void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
                                                     Value replacementValue) {
   assert(!remappedInputs[origInputNo] && "input has already been remapped");
   remappedInputs[origInputNo] =
       InputMapping{origInputNo, /*size=*/0, replacementValue};
 }
 
 LogicalResult TypeConverter::convertType(Type t,
                                          SmallVectorImpl<Type> &results) const {
+  {
+    std::shared_lock<decltype(cacheMutex)> cacheReadLock(cacheMutex,
+                                                         std::defer_lock);
+    if (t.getContext()->isMultithreadingEnabled())
+      cacheReadLock.lock();
     auto existingIt = cachedDirectConversions.find(t);
     if (existingIt != cachedDirectConversions.end()) {
       if (existingIt->second)
         results.push_back(existingIt->second);
       return success(existingIt->second != nullptr);
     }
     auto multiIt = cachedMultiConversions.find(t);
     if (multiIt != cachedMultiConversions.end()) {
       results.append(multiIt->second.begin(), multiIt->second.end());
       return success();
     }
-
+  }
   // Walk the added converters in reverse order to apply the most recently
   // registered first.
   size_t currentCount = results.size();
 
+  std::unique_lock<decltype(cacheMutex)> cacheWriteLock(cacheMutex,
+                                                        std::defer_lock);
+
   for (const ConversionCallbackFn &converter : llvm::reverse(conversions)) {
     if (std::optional<LogicalResult> result = converter(t, results)) {
+      if (t.getContext()->isMultithreadingEnabled())
+        cacheWriteLock.lock();
       if (!succeeded(*result)) {
         cachedDirectConversions.try_emplace(t, nullptr);
         return failure();
       }
       auto newTypes = ArrayRef<Type>(results).drop_front(currentCount);
       if (newTypes.size() == 1)
         cachedDirectConversions.try_emplace(t, newTypes.front());
       else