Remove the `conversionCallStack` from the MLIR TypeConverter

[mehdi_amini]

Summary

This vector keeps tracks of recursive types through the recursive invocations
of convertType(). However this is something only useful for some specific
cases, in which the dedicated conversion callbacks can handle this stack
privately.

This allows removing a mutable member of the type converter.

[springerm]

Looks like a good cleanup to me. (Regardless of whether this works in multi-threading or not.) Keeping the base type converter class simple and let derived classes handle recursion edge cases.

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

166

I would mention here that this currently only works for single-threaded execution.

I was looking for ways to make this work in a multi-threaded environment. The only solution that I found (without changing the API of convertTypes) is something like this. What do you think?

DenseMap<std::thread::id, std::unique_ptr<SmallVector<Type>>> conversionCallStack;
llvm::sys::SmartRWMutex<true> callStackMutex;

...

  addConversion([&](LLVM::LLVMStructType type, SmallVectorImpl<Type> &results)
                    -> std::optional<LogicalResult> {
    ...
    SmallVector<Type> *stack = nullptr;
    std::shared_lock stackReadLock(callStackMutex);
    stackReadLock.lock();
    auto it = conversionCallStack.find(std::this_thread::get_id());
    if (it != conversionCallStack.end())
      stack = it.second().get();
    stackReadLock.unlock();
    if (!stack) {
      std::unique_lock stackWriteLock(callStackMutex);
      stackWriteLock.lock();
      stack = conversionCallStack.try_emplace(std::this_thread::get_id()).first.get();
    }
     
    // `stack` can be used without further synchronization/locking. 
    stack->push_back(type);
    auto popConversionCallStack =
        llvm::make_scope_exit([this]() { conversionCallStack->pop_back(); });

  }

167

The same field is needed in flang/lib/Optimizer/CodeGen/TypeConverter.cpp.

[springerm]

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

166

I think this would also work and does not require locking:

private:
  static thread_local DenseMap<LLVMTypeConverter *, SmallVector<Type>> conversionCallStack;

~LLVMTypeConverter() {
  conversionCallStack.erase(this);
}

addConversion([&](...) {
  SmallVector<Type> &stack = conversionCallStack[this];
  // `stack` can be used without locking.
});

[mehdi_amini]

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

166

This is still "global mutable state", which I find unsatisfactory as a general approach though.

[mehdi_amini]

Add per-thread handling and locking

[mehdi_amini]

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

167

They inherit from the class here, so it should be all fine!

[mehdi_amini]

Fix flang

[mehdi_amini]

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

167

Actually I see what you mean, I fixed flang :)

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

public:
   /// Check if a memref type can be converted to a bare pointer.
   static bool canConvertToBarePtr(BaseMemRefType type);
 
 protected:
   /// Pointer to the LLVM dialect.
   LLVM::LLVMDialect *llvmDialect;
 
 private:
+  // Recursive structure detection.
+  // We store one entry per thread here, and rely on locking.
+  DenseMap<uint64_t, std::unique_ptr<SmallVector<Type>>> conversionCallStack;
+  llvm::sys::SmartRWMutex<true> callStackMutex;
+  SmallVector<Type> &getCurrentThreadRecursiveStack();
+
   /// Convert a function type.  The arguments and results are converted one by
   /// one.  Additionally, if the function returns more than one value, pack the
   /// results into an LLVM IR structure type so that the converted function type
   /// returns at most one result.
   Type convertFunctionType(FunctionType type) const;
 
   /// Convert the index type.  Uses llvmModule data layout to create an integer
   /// of the pointer bitwidth.

mlir/include/mlir/Transforms/DialectConversion.h

public:
   std::optional<Attribute> convertTypeAttribute(Type type,
                                                 Attribute attr) const;
 
 private:
   /// The signature of the callback used to convert a type. If the new set of
   /// types is empty, the type is removed and any usages of the existing value
   /// are expected to be removed during conversion.
   using ConversionCallbackFn = std::function<std::optional<LogicalResult>(
-      Type, SmallVectorImpl<Type> &, ArrayRef<Type>)>;
+      Type, SmallVectorImpl<Type> &)>;
 
   /// The signature of the callback used to materialize a conversion.
   using MaterializationCallbackFn = std::function<std::optional<Value>(
       OpBuilder &, Type, ValueRange, Location)>;
 
   /// The signature of the callback used to convert a type attribute.
   using TypeAttributeConversionCallbackFn =
       std::function<AttributeConversionResult(Type, Attribute)>;
 
   /// Attempt to materialize a conversion using one of the provided
   /// materialization functions.
   Value
   materializeConversion(ArrayRef<MaterializationCallbackFn> materializations,
                         OpBuilder &builder, Location loc, Type resultType,
                         ValueRange inputs) const;
 
   /// Generate a wrapper for the given callback. This allows for accepting
   /// different callback forms, that all compose into a single version.
   /// With callback of form: `std::optional<Type>(T)`
   template <typename T, typename FnT>
   std::enable_if_t<std::is_invocable_v<FnT, T>, ConversionCallbackFn>
   wrapCallback(FnT &&callback) const {
-    return wrapCallback<T>(
-        [callback = std::forward<FnT>(callback)](
-            T type, SmallVectorImpl<Type> &results, ArrayRef<Type>) {
-          if (std::optional<Type> resultOpt = callback(type)) {
-            bool wasSuccess = static_cast<bool>(*resultOpt);
-            if (wasSuccess)
-              results.push_back(*resultOpt);
-            return std::optional<LogicalResult>(success(wasSuccess));
-          }
-          return std::optional<LogicalResult>();
-        });
+    return wrapCallback<T>([callback = std::forward<FnT>(callback)](
+                               T type, SmallVectorImpl<Type> &results) {
   }
   /// With callback of form: `std::optional<LogicalResult>(
-  ///     T, SmallVectorImpl<Type> &)`.
-  template <typename T, typename FnT>
-  std::enable_if_t<std::is_invocable_v<FnT, T, SmallVectorImpl<Type> &>,
-                   ConversionCallbackFn>
-  wrapCallback(FnT &&callback) const {
-    return wrapCallback<T>(
-        [callback = std::forward<FnT>(callback)](
-            T type, SmallVectorImpl<Type> &results, ArrayRef<Type>) {
-          return callback(type, results);
-        });
-  }
-  /// With callback of form: `std::optional<LogicalResult>(
   ///     T, SmallVectorImpl<Type> &, ArrayRef<Type>)`.
   template <typename T, typename FnT>
-  std::enable_if_t<
-      std::is_invocable_v<FnT, T, SmallVectorImpl<Type> &, ArrayRef<Type>>,
+  std::enable_if_t<std::is_invocable_v<FnT, T, SmallVectorImpl<Type> &>,
                    ConversionCallbackFn>
   wrapCallback(FnT &&callback) const {
     return [callback = std::forward<FnT>(callback)](
-               Type type, SmallVectorImpl<Type> &results,
-               ArrayRef<Type> callStack) -> std::optional<LogicalResult> {
+               Type type,
+               SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> {
       T derivedType = dyn_cast<T>(type);
       if (!derivedType)
         return std::nullopt;
-      return callback(derivedType, results, callStack);
+      return callback(derivedType, results);
     };
   }
 
   /// Register a type conversion.
   void registerConversion(ConversionCallbackFn callback) {
     conversions.emplace_back(std::move(callback));
     cachedDirectConversions.clear();
     cachedMultiConversions.clear();
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;
-
-  /// Stores the types that are being converted in the case when convertType
-  /// is being called recursively to convert nested types.
-  mutable SmallVector<Type, 2> conversionCallStack;
 };
 
 //===----------------------------------------------------------------------===//
 // 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/Conversion/LLVMCommon/TypeConverter.cpp

 //===- TypeConverter.cpp - Convert builtin to LLVM dialect types ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "MemRefDescriptor.h"
 #include "mlir/Conversion/LLVMCommon/MemRefBuilder.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/Threading.h"
+#include <memory>
+#include <mutex>
 #include <optional>
 
 using namespace mlir;
 
+SmallVector<Type> &LLVMTypeConverter::getCurrentThreadRecursiveStack() {
+  {
+    // Most of the time, the entry already exists in the map.
+    std::shared_lock<decltype(callStackMutex)> lock(callStackMutex,
+                                                    std::defer_lock);
+    if (getContext().isMultithreadingEnabled())
+      lock.lock();
+    auto recursiveStack = conversionCallStack.find(llvm::get_threadid());
+    if (recursiveStack != conversionCallStack.end())
+      return *recursiveStack->second;
+  }
+
+  // First time this thread gets here, we have to get an exclusive access to
+  // inset in the map
+  std::unique_lock<decltype(callStackMutex)> lock(callStackMutex);
+  auto recursiveStackInserted = conversionCallStack.insert(std::make_pair(
+      llvm::get_threadid(), std::make_unique<SmallVector<Type>>()));
+  return *recursiveStackInserted.first->second.get();
+}
+
 /// Create an LLVMTypeConverter using default LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                      const DataLayoutAnalysis *analysis)
     : LLVMTypeConverter(ctx, LowerToLLVMOptions(ctx), analysis) {}
 
 /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                      const LowerToLLVMOptions &options,
LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
   // converted even when the outer type is compatible.
   addConversion([&](LLVM::LLVMPointerType type) -> std::optional<Type> {
     if (type.isOpaque())
       return type;
     if (auto pointee = convertType(type.getElementType()))
       return LLVM::LLVMPointerType::get(pointee, type.getAddressSpace());
     return std::nullopt;
   });
-  addConversion([&](LLVM::LLVMStructType type, SmallVectorImpl<Type> &results,
-                    ArrayRef<Type> callStack) -> std::optional<LogicalResult> {
+
+  addConversion([&](LLVM::LLVMStructType type, SmallVectorImpl<Type> &results)
+                    -> std::optional<LogicalResult> {
     // Fastpath for types that won't be converted by this callback anyway.
     if (LLVM::isCompatibleType(type)) {
       results.push_back(type);
       return success();
     }
+    SmallVectorImpl<Type> &recursiveStack = getCurrentThreadRecursiveStack();
+    recursiveStack.push_back(type);
+    auto popConversionCallStack = llvm::make_scope_exit(
+        [&recursiveStack]() { recursiveStack.pop_back(); });
 
     if (type.isIdentified()) {
       auto convertedType = LLVM::LLVMStructType::getIdentified(
           type.getContext(), ("_Converted_" + type.getName()).str());
       unsigned counter = 1;
       while (convertedType.isInitialized()) {
         assert(counter != UINT_MAX &&
                "about to overflow struct renaming counter in conversion");
         convertedType = LLVM::LLVMStructType::getIdentified(
             type.getContext(),
             ("_Converted_" + std::to_string(counter) + type.getName()).str());
       }
-      if (llvm::count(callStack, type) > 1) {
+      if (llvm::count(recursiveStack, type) > 1) {
         results.push_back(convertedType);
         return success();
       }
 
       SmallVector<Type> convertedElemTypes;
       convertedElemTypes.reserve(type.getBody().size());
       if (failed(convertTypes(type.getBody(), convertedElemTypes)))
         return std::nullopt;

mlir/lib/Transforms/Utils/DialectConversion.cpp

LogicalResult TypeConverter::convertType(Type 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();
-  conversionCallStack.push_back(t);
-  auto popConversionCallStack =
-      llvm::make_scope_exit([this]() { conversionCallStack.pop_back(); });
+
   for (const ConversionCallbackFn &converter : llvm::reverse(conversions)) {
-    if (std::optional<LogicalResult> result =
-            converter(t, results, conversionCallStack)) {
+    if (std::optional<LogicalResult> result = converter(t, results)) {
       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

mlir/test/lib/Dialect/Test/TestPatterns.cpp

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/ScopeExit.h"
 
 using namespace mlir;
 using namespace test;
 
 // Native function for testing NativeCodeCall
 static Value chooseOperand(Value input1, Value input2, BoolAttr choice) {
   return choice.getValue() ? input1 : input2;
 }
StringRef getArgument() const final {
     return "test-legalize-type-conversion";
   }
   StringRef getDescription() const final {
     return "Test various type conversion functionalities in DialectConversion";
   }
 
   void runOnOperation() override {
     // Initialize the type converter.
+    SmallVector<Type, 2> conversionCallStack;
     TypeConverter converter;
 
     /// Add the legal set of type conversions.
     converter.addConversion([](Type type) -> Type {
       // Treat F64 as legal.
       if (type.isF64())
         return type;
       // Allow converting BF16/F16/F32 to F64.
       if (type.isBF16() || type.isF16() || type.isF32())
         return FloatType::getF64(type.getContext());
       // Otherwise, the type is illegal.
       return nullptr;
     });
     converter.addConversion([](IntegerType type, SmallVectorImpl<Type> &) {
       // Drop all integer types.
       return success();
     });
     converter.addConversion(
         // Convert a recursive self-referring type into a non-self-referring
         // type named "outer_converted_type" that contains a SimpleAType.
-        [&](test::TestRecursiveType type, SmallVectorImpl<Type> &results,
-            ArrayRef<Type> callStack) -> std::optional<LogicalResult> {
+        [&](test::TestRecursiveType type,
+            SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> {
           // If the type is already converted, return it to indicate that it is
           // legal.
           if (type.getName() == "outer_converted_type") {
             results.push_back(type);
             return success();
           }
 
+          conversionCallStack.push_back(type);
+          auto popConversionCallStack = llvm::make_scope_exit(
+              [&conversionCallStack]() { conversionCallStack.pop_back(); });
+
           // If the type is on the call stack more than once (it is there at
           // least once because of the _current_ call, which is always the last
           // element on the stack), we've hit the recursive case. Just return
           // SimpleAType here to create a non-recursive type as a result.
-          if (llvm::is_contained(callStack.drop_back(), type)) {
+          if (llvm::is_contained(ArrayRef(conversionCallStack).drop_back(),
+                                 type)) {
             results.push_back(test::SimpleAType::get(type.getContext()));
             return success();
           }
 
           // Convert the body recursively.
           auto result = test::TestRecursiveType::get(type.getContext(),
                                                      "outer_converted_type");
           if (failed(result.setBody(converter.convertType(type.getBody()))))