[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

835

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

835

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 ↗

(On Diff #437766)

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

70 ↗

(On Diff #437766)

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.

774–775

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.

820

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

820

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

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> &validTypes) {
+  if (validTypes.count(type) || !callstack.insert(type))
     return true;
 
-  callstack.insert(type);
   auto stackPopper = llvm::make_scope_exit([&] { callstack.pop_back(); });
 
   auto isCompatible = [&](Type type) {
-    return isCompatibleImpl(type, callstack);
+    return isCompatibleImpl(type, callstack, validTypes);
   };
 
-  return llvm::TypeSwitch<Type, bool>(type)
+  bool isValid =
+      llvm::TypeSwitch<Type, bool>(type)
           .Case<LLVMStructType>([&](auto structType) {
             return llvm::all_of(structType.getBody(), isCompatible);
           })
           .Case<LLVMFunctionType>([&](auto funcType) {
             return isCompatible(funcType.getReturnType()) &&
                    llvm::all_of(funcType.getParams(), isCompatible);
           })
           .Case<IntegerType>([](auto intType) { return intType.isSignless(); })
           .Case<VectorType>([&](auto vecType) {
-        return vecType.getRank() == 1 && isCompatible(vecType.getElementType());
+            return vecType.getRank() == 1 &&
+                   isCompatible(vecType.getElementType());
           })
           .Case<LLVMPointerType>([&](auto pointerType) {
             if (pointerType.isOpaque())
               return true;
             return isCompatible(pointerType.getElementType());
           })
           // clang-format off
           .Case<
               LLVMFixedVectorType,
               LLVMScalableVectorType,
               LLVMArrayType
           >([&](auto containerType) {
             return isCompatible(containerType.getElementType());
           })
           .Case<
             BFloat16Type,
             Float16Type,
             Float32Type,
             Float64Type,
             Float80Type,
             Float128Type,
             LLVMLabelType,
             LLVMMetadataType,
             LLVMPPCFP128Type,
             LLVMTokenType,
             LLVMVoidType,
             LLVMX86MMXType
           >([](Type) { return true; })
           // clang-format on
           .Default([](Type) { return false; });
+
+  if (isValid)
+    validTypes.insert(type);
+
+  return isValid;
 }
 
 bool mlir::LLVM::isCompatibleType(Type type) {
+  // Ideally, we can use a single set to memorize both the visited types and
+  // valid types -- we simply don't erase the type from set after its visiting
+  // and mark the set thread_local. Because if the type is invalid, most of the
+  // current usages of `isCompatibleType` simply bails out. However, there are
+  // still some places dipatch logics based on the result of this function and
+  // in the cases of mixing LLVM dialect with other dialects, we might end up
+  // memorizing invalid types.
+  // Thus, we're using two separate sets here.
   SetVector<Type> callstack;
-  return isCompatibleImpl(type, callstack);
+  thread_local DenseSet<Type> validTypes;
+  return isCompatibleImpl(type, callstack, validTypes);
 }
 
 bool mlir::LLVM::isCompatibleFloatingPointType(Type type) {
   return type.isa<BFloat16Type, Float16Type, Float32Type, Float64Type,
                   Float80Type, Float128Type, LLVMPPCFP128Type>();
 }
 
 bool mlir::LLVM::isCompatibleVectorType(Type type) {