[mlir] Avoid needlessly converting LLVM named structs with compatible elements

[ftynse]

Summary

Conversion of LLVM named structs leads to them being renamed since we cannot
modify the body of the struct type once it is set. Previously, this applied to
all named struct types, even if their element types were not affected by the
conversion. Make this behvaior only applicable when element types are changed.
This requires making the LLVM dialect type-compatibility check recursively look
at the element types (arguably, it should have been doing than since the moment
the LLVM dialect type system stopped being closed). In addition, have a more
lax check for outer types only to avoid repeated check when necessary (e.g.,
parser, verifiers that are going to also look at the inner type).

[wsmoses]

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

1

Why this deletion?

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

 
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
 /// Returns `true` if the given type is compatible with the LLVM dialect.
 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.
 bool isCompatibleFloatingPointType(Type type);
 
 /// Returns `true` if the given type is a vector type compatible with the LLVM
 /// dialect. Compatible types include 1D built-in vector types of built-in
 /// integers and floating-point values, LLVM dialect fixed vector types of LLVM
 /// dialect pointers and LLVM dialect scalable vector types.

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
   // converted even when the outer type is compatible.
   addConversion([&](LLVM::LLVMPointerType type) -> llvm::Optional<Type> {
     if (auto pointee = convertType(type.getElementType()))
       return LLVM::LLVMPointerType::get(pointee, type.getAddressSpace());
     return llvm::None;
   });
   addConversion([&](LLVM::LLVMStructType type, SmallVectorImpl<Type> &results,
                     ArrayRef<Type> callStack) -> llvm::Optional<LogicalResult> {
+    // Fastpath for types that won't be converted by this callback anyway.
+    if (LLVM::isCompatibleType(type)) {
+      results.push_back(type);
+      return success();
+    }
+
     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(

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

 }
 
 /// Parses one of the LLVM dialect types.
 Type mlir::LLVM::detail::parseType(DialectAsmParser &parser) {
   llvm::SMLoc loc = parser.getCurrentLocation();
   Type type = dispatchParse(parser, /*allowAny=*/false);
   if (!type)
     return type;
-  if (!isCompatibleType(type)) {
+  if (!isCompatibleOuterType(type)) {
     parser.emitError(loc) << "unexpected type, expected keyword";
     return nullptr;
   }
   return type;
 }

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

-//===- LLVMTypes.cpp - MLIR 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the types for the LLVM dialect in MLIR. These MLIR types
 // correspond to the LLVM IR type system.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TypeDetail.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/TypeSize.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 //===----------------------------------------------------------------------===//
 // Array type.
LLVMFunctionType::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Pointer type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMPointerType::isValidElementType(Type type) {
-  return isCompatibleType(type) ? !type.isa<LLVMVoidType, LLVMTokenType,
-                                            LLVMMetadataType, LLVMLabelType>()
-                                : type.isa<PointerElementTypeInterface>();
+  return isCompatibleOuterType(type)
+             ? !type.isa<LLVMVoidType, LLVMTokenType, LLVMMetadataType,
+                         LLVMLabelType>()
 }
 
 LLVMPointerType LLVMPointerType::get(Type pointee, unsigned addressSpace) {
   assert(pointee && "expected non-null subtype");
   return Base::get(pointee.getContext(), pointee, addressSpace);
 }
 
 LLVMPointerType
LLVMScalableVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
   return verifyVectorConstructionInvariants<LLVMScalableVectorType>(
       emitError, elementType, numElements);
 }
 
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
-bool mlir::LLVM::isCompatibleType(Type type) {
-  // Only signless integers are compatible.
-  if (auto intType = type.dyn_cast<IntegerType>())
-    return intType.isSignless();
-
-  // 1D vector types are compatible if their element types are.
-  if (auto vecType = type.dyn_cast<VectorType>())
-    return vecType.getRank() == 1 && isCompatibleType(vecType.getElementType());
-
+bool mlir::LLVM::isCompatibleOuterType(Type type) {
   // clang-format off
-  return type.isa<
+  if (type.isa<
       BFloat16Type,
       Float16Type,
       Float32Type,
       Float64Type,
       Float80Type,
       Float128Type,
       LLVMArrayType,
       LLVMFunctionType,
       LLVMLabelType,
       LLVMMetadataType,
       LLVMPPCFP128Type,
       LLVMPointerType,
       LLVMStructType,
       LLVMTokenType,
       LLVMFixedVectorType,
       LLVMScalableVectorType,
       LLVMVoidType,
       LLVMX86MMXType
-  >();
+    >()) {
     // clang-format on
+    return true;
+  }
+
+  // Only signless integers are compatible.
+  if (auto intType = type.dyn_cast<IntegerType>())
+    return intType.isSignless();
+
+  // 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))
+    return true;
+
+  callstack.insert(type);
+  auto stackPopper = llvm::make_scope_exit([&] { callstack.pop_back(); });
+
+  auto isCompatible = [&](Type type) {
+    return isCompatibleImpl(type, callstack);
+  };
+
+  return 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());
+      })
+      // clang-format off
+      .Case<
+          LLVMPointerType,
+          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; });
+}
+
+bool mlir::LLVM::isCompatibleType(Type type) {
+  SetVector<Type> callstack;
+  return isCompatibleImpl(type, callstack);
 }
 
 bool mlir::LLVM::isCompatibleFloatingPointType(Type type) {
   return type.isa<BFloat16Type, Float16Type, Float32Type, Float64Type,
                   Float80Type, Float128Type, LLVMPPCFP128Type>();
 }
 
 bool mlir::LLVM::isCompatibleVectorType(Type type) {

mlir/test/Conversion/StandardToLLVM/convert-types.mlir

 // CHECK-LABEL: @struct_ptr()
 // CHECK: !llvm.struct<(ptr<i42>)> 
 func private @struct_ptr() -> !llvm.struct<(ptr<!test.smpla>)>
 
 // CHECK-LABEL: @named_struct_ptr()
 // CHECK: !llvm.struct<"_Converted_named", (ptr<i42>)>
 func private @named_struct_ptr() -> !llvm.struct<"named", (ptr<!test.smpla>)>
 
+// CHECK-LABEL: @named_no_convert
+// CHECK: !llvm.struct<"no_convert", (ptr<struct<"no_convert">>)>
+func private @named_no_convert() -> !llvm.struct<"no_convert", (ptr<struct<"no_convert">>)>
+
 // CHECK-LABEL: @array_ptr()
 // CHECK: !llvm.array<10 x ptr<i42>> 
 func private @array_ptr() -> !llvm.array<10 x ptr<!test.smpla>>
 
 // CHECK-LABEL: @func()
 // CHECK: !llvm.ptr<func<i42 (i42)>>
 func private @func() -> !llvm.ptr<!llvm.func<!test.smpla (!test.smpla)>>
 
 // TODO: support conversion of recursive types in the conversion infra.
 // CHECK-LABEL: @named_recursive()
 // CHECK: !llvm.struct<"_Converted_recursive", (ptr<i42>, ptr<struct<"_Converted_recursive">>)>
 func private @named_recursive() -> !llvm.struct<"recursive", (ptr<!test.smpla>, ptr<struct<"recursive">>)>