[mlir][LLVMIR] Memorize compatible LLVM types

[myhsu]

Summary

This patch memorizes compatible LLVM types in LLVM::isCompatibleType in order to avoid redundant works.

This is especially useful when the size of program is big and there are multiple occurrences of some deeply nested LLVM struct types, in which case we can gain quite some speedups with this patch.

Existing tests should be sufficient to cover the changes here.

[Mogball]

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

836–844

Definitely do not want to be using thread_local or static or anything whose lifetime could exceed that of an MLIRContext. I believe the problem is with nested LLVM types and not nested builtin types like vector<*xvector<*xvector<...>>>. You could make all LLVM types (or just the problematic one) store a verified bit instead of using a set. Alternatively, you can put the set on the LLVMDialect (and make sure its access is thread-safe) and move this function to be a member function of LLVMDialect.

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

836–844

Definitely do not want to be using thread_local or static or anything whose lifetime could exceed that of an MLIRContext.

+1 here

[myhsu]

This is a major revision. I choose to put the type cache (wrapped with ThreadLocalCache) inside LLVM dialect.
The reason I didn't put a verified flag inside aggregate types was because (i) it effectively makes the affected type mutable. Though it's fine for LLVMStructType, which is already mutable, we also want to cache other LLVM types like LLVMPointer as well. Making them mutable solely based on this purpose didn't really justify IMHO. (ii) The LLVM::isCompatibleType function is highly likely to be used in different threads, so we need to lock on this verified flag in every affected types.

[Mogball]

On the matter of using bits: yes, it would require making more types mutable. On locking the verified bit: mutate already locks the storage uniquer (if threading is enabled). Reading the bit, however, doesn't need to be locked.

I assume ThreadLocalCache is faster than a shared DenseSet with a mutex?

Mostly LG otherwise

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

61

I wouldn't really want these to be publicly exposed. Can you make the isCompatibleType a static function of LLVMDialect?

70

drop name spaces (like in the rest of the file)

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

775

type.getContext()->getLoadedDialect<LLVMDialect>()?

This would allow you to cache any type (if the LLVM dialect is loaded, which it should if there are LLVM types floating around in the context).

[myhsu]

• Put isCompatibleType as a static member function of LLVMDialect.
• Cache non-LLVM aggregate types as well.

[myhsu]

In D127918#3590877, @Mogball wrote:

On the matter of using bits: yes, it would require making more types mutable. On locking the verified bit: mutate already locks the storage uniquer (if threading is enabled). Reading the bit, however, doesn't need to be locked.

Good to know!

I assume ThreadLocalCache is faster than a shared DenseSet with a mutex?

Correct. It's using thread_local underlying but avoid the lifetime issue you mentioned in earlier comments. Thus, it only locks when a new thread is created.

Mostly LG otherwise

[Mogball]

LGTM but please take a look at the comments

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

766–772

I think this should speed up the case of a type already being verified.

I wonder if it would not be possible to remove callstack. I remember you had a comment about it but I don't know where it went. Perhaps something like: optimistically insert the type being verified into compatibleTypes to break cycles but then remove the type if it turns out to be incompatible. I don't think a SetVector is necessary.

And if a cached compatibleTypes set can't be retrieved from the LLVMDialect, just use a local one.

784

I would outline the call to cacheCompatibleType:

bool result = llvm::TypeSwitch<Type, bool>(type)...
if (result && compatibleTypes)
  compatibleTypes->insert(type);
return result;

[Mogball]

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

766–772

(this can be a follow-up patch, since I think it would improve performance).

[myhsu]

Use a single DenseSet for both dependency breaking and cache for compatible types. Please take another look.

[myhsu]

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

766–772

I wonder if it would not be possible to remove callstack. I remember you had a comment about it but I don't know where it went.

Yes I still remember that (I forgot why I removed it though): originally I only used a single DenseSet to break the cycle and memorize the compatible types, and I simply didn't remove the type after its visiting, based on an (not so solid) assumption that if isCompatibleType returns false, most of the caller will just bail out. It turned out that assumption worked well if we only use a single dialect, if we loaded multiple dialects, isCompatibleType will be used to dispatch some of the parsing / printing logics and the DenseSet would eventually memorize some incompatible types.

optimistically insert the type being verified into compatibleTypes to break cycles but then remove the type if it turns out to be incompatible. I don't think a SetVector is necessary.

sounds like a way. I think I can do it in this patch.

876

FYI: SetVector::pop_back is not faster than DenseSet::erase since the former actually uses DenseSet::erase underlying.

[Mogball]

Nice! LGTM

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

876

Yeah, that's exactly what I mean when I said SetVector is not needed...

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMOpsEnums.h.inc"
 
 namespace mlir {

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

let extraClassDeclaration = [{
 
     /// Verifies if the given string is a well-formed data layout descriptor.
     /// Uses `reportError` to report errors.
     static LogicalResult verifyDataLayoutString(
         StringRef descr, llvm::function_ref<void (const Twine &)> reportError);
 
     /// Name of the target triple attribute.
     static StringRef getTargetTripleAttrName() { return "llvm.target_triple"; }
+
+    /// Returns `true` if the given type is compatible with the LLVM dialect.
+    static bool isCompatibleType(Type);
+
+  private:
+    /// A cache storing compatible LLVM types that have been verified. This
+    /// can save us lots of verification time if there are many occurrences
+    /// of some deeply-nested aggregate types in the program.
+    ThreadLocalCache<DenseSet<Type>> compatibleTypes;
   }];
 
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM dialect type constraints.
 //===----------------------------------------------------------------------===//
 
 // LLVM dialect type.
 def LLVM_Type : DialectType<LLVM_Dialect,

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 /// Prints an LLVM Dialect type.
 void printType(Type type, AsmPrinter &printer);
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
-/// Returns `true` if the given type is compatible with the LLVM dialect.
+/// Returns `true` if the given type is compatible with the LLVM dialect. This
+/// is an alias to `LLVMDialect::isCompatibleType`.
 bool isCompatibleType(Type type);
 
 /// Returns `true` if the given outer type is compatible with the LLVM dialect
 /// without checking its potential nested types such as struct elements.
 bool isCompatibleOuterType(Type type);
 
 /// Returns `true` if the given type is a floating-point type compatible with
 /// the LLVM dialect.

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

bool mlir::LLVM::isCompatibleOuterType(Type type) {
 
   // 1D vector types are compatible.
   if (auto vecType = type.dyn_cast<VectorType>())
     return vecType.getRank() == 1;
 
   return false;
 }
 
-static bool isCompatibleImpl(Type type, SetVector<Type> &callstack) {
-  if (callstack.contains(type))
+static bool isCompatibleImpl(Type type, SetVector<Type> &callstack,
+                             DenseSet<Type> *compatibleTypes) {
+  if (!callstack.insert(type))
     return true;
 
-  callstack.insert(type);
   auto stackPopper = llvm::make_scope_exit([&] { callstack.pop_back(); });
 
+  if (compatibleTypes && compatibleTypes->count(type))
+    return true;
+
   auto isCompatible = [&](Type type) {
-    return isCompatibleImpl(type, callstack);
+    return isCompatibleImpl(type, callstack, compatibleTypes);
+  };
+
+  auto cacheCompatibleType = [&](bool result, Type typeToCache) -> bool {
+    if (result && compatibleTypes)
+      compatibleTypes->insert(typeToCache);
+    return result;
   };
 
   return llvm::TypeSwitch<Type, bool>(type)
       .Case<LLVMStructType>([&](auto structType) {
-        return llvm::all_of(structType.getBody(), isCompatible);
+        return cacheCompatibleType(
+            llvm::all_of(structType.getBody(), isCompatible), structType);
       })
       .Case<LLVMFunctionType>([&](auto funcType) {
-        return isCompatible(funcType.getReturnType()) &&
-               llvm::all_of(funcType.getParams(), isCompatible);
+        return cacheCompatibleType(
+            isCompatible(funcType.getReturnType()) &&
+                llvm::all_of(funcType.getParams(), isCompatible),
+            funcType);
       })
       .Case<IntegerType>([](auto intType) { return intType.isSignless(); })
       .Case<VectorType>([&](auto vecType) {
-        return vecType.getRank() == 1 && isCompatible(vecType.getElementType());
+        return cacheCompatibleType(vecType.getRank() == 1 &&
+                                       isCompatible(vecType.getElementType()),
+                                   vecType);
       })
       .Case<LLVMPointerType>([&](auto pointerType) {
         if (pointerType.isOpaque())
           return true;
-        return isCompatible(pointerType.getElementType());
+        return cacheCompatibleType(isCompatible(pointerType.getElementType()),
+                                   pointerType);
       })
       // clang-format off
       .Case<
           LLVMFixedVectorType,
           LLVMScalableVectorType,
           LLVMArrayType
       >([&](auto containerType) {
-        return isCompatible(containerType.getElementType());
+        return cacheCompatibleType(isCompatible(containerType.getElementType()),
+                                   containerType);
       })
       .Case<
         BFloat16Type,
         Float16Type,
         Float32Type,
         Float64Type,
         Float80Type,
         Float128Type,
         LLVMLabelType,
         LLVMMetadataType,
         LLVMPPCFP128Type,
         LLVMTokenType,
         LLVMVoidType,
         LLVMX86MMXType
       >([](Type) { return true; })
       // clang-format on
       .Default([](Type) { return false; });
 }
 
-bool mlir::LLVM::isCompatibleType(Type type) {
+bool LLVMDialect::isCompatibleType(Type type) {
   SetVector<Type> callstack;
-  return isCompatibleImpl(type, callstack);
+  DenseSet<Type> *compatibleTypes = nullptr;
+  if (auto *llvmDialect =
+          type.getContext()->getLoadedDialect<LLVM::LLVMDialect>())
+    compatibleTypes = &llvmDialect->compatibleTypes.get();
+  return isCompatibleImpl(type, callstack, compatibleTypes);
+}
+
+bool mlir::LLVM::isCompatibleType(Type type) {
+  return LLVMDialect::isCompatibleType(type);
 }
 
 bool mlir::LLVM::isCompatibleFloatingPointType(Type type) {
   return type.isa<BFloat16Type, Float16Type, Float32Type, Float64Type,
                   Float80Type, Float128Type, LLVMPPCFP128Type>();
 }
 
 bool mlir::LLVM::isCompatibleVectorType(Type type) {

 Type mlir::LLVM::getVectorElementType(Type type) {
   return llvm::TypeSwitch<Type, Type>(type)
       .Case<LLVMFixedVectorType, LLVMScalableVectorType, VectorType>(
           [](auto ty) { return ty.getElementType(); })
       .Default([](Type) -> Type {
         llvm_unreachable("incompatible with LLVM vector type");
       });
 }
 
 llvm::ElementCount mlir::LLVM::getVectorNumElements(Type type) {
   return llvm::TypeSwitch<Type, llvm::ElementCount>(type)
       .Case([](VectorType ty) {
         if (ty.isScalable())
           return llvm::ElementCount::getScalable(ty.getNumElements());
         return llvm::ElementCount::getFixed(ty.getNumElements());
       })
       .Case([](LLVMFixedVectorType ty) {