[mlir][llvm] Move LLVMFunctionType to a TypeDef

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

 DEFINE_TRIVIAL_LLVM_TYPE(LLVMX86MMXType);
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMTokenType);
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMLabelType);
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMMetadataType);
 
 #undef DEFINE_TRIVIAL_LLVM_TYPE
 
 //===----------------------------------------------------------------------===//
-// LLVMFunctionType.
-//===----------------------------------------------------------------------===//
-
-/// LLVM dialect function type. It consists of a single return type (unlike MLIR
-/// which can have multiple), a list of parameter types and can optionally be
-/// variadic.
-class LLVMFunctionType : public Type::TypeBase<LLVMFunctionType, Type,
-                                               detail::LLVMFunctionTypeStorage,
-                                               SubElementTypeInterface::Trait> {
-public:
-  /// Inherit base constructors.
-  using Base::Base;
-  using Base::getChecked;
-
-  /// Checks if the given type can be used an argument in a function type.
-  static bool isValidArgumentType(Type type);
-
-  /// Checks if the given type can be used as a result in a function type.
-  static bool isValidResultType(Type type);
-
-  /// Returns whether the function is variadic.
-  bool isVarArg() const;
-
-  /// Gets or creates an instance of LLVM dialect function in the same context
-  /// as the `result` type.
-  static LLVMFunctionType get(Type result, ArrayRef<Type> arguments,
-                              bool isVarArg = false);
-  static LLVMFunctionType
-  getChecked(function_ref<InFlightDiagnostic()> emitError, Type result,
-             ArrayRef<Type> arguments, bool isVarArg = false);
-
-  /// Returns a clone of this function type with the given argument
-  /// and result types.
-  LLVMFunctionType clone(TypeRange inputs, TypeRange results) const;
-
-  /// Returns the result type of the function.
-  Type getReturnType() const;
-
-  /// Returns the result type of the function as an ArrayRef, enabling better
-  /// integration with generic MLIR utilities.
-  ArrayRef<Type> getReturnTypes() const;
-
-  /// Returns the number of arguments to the function.
-  unsigned getNumParams();
-
-  /// Returns `i`-th argument of the function. Asserts on out-of-bounds.
-  Type getParamType(unsigned i);
-
-  /// Returns a list of argument types of the function.
-  ArrayRef<Type> getParams() const;
-  ArrayRef<Type> params() { return getParams(); }
-
-  /// Verifies that the type about to be constructed is well-formed.
-  static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
-                              Type result, ArrayRef<Type> arguments, bool);
-
-  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
-                                function_ref<void(Type)> walkTypesFn) const;
-  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                   ArrayRef<Type> replTypes) const;
-};
-
-//===----------------------------------------------------------------------===//
 // LLVMPointerType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect pointer type. This type typically represents a reference to an
 /// object in memory. Pointers may be opaque or parameterized by the element
 /// type. Both opaque and non-opaque pointers are additionally parameterized by
 /// the address space.
 class LLVMPointerType

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td

def LLVMArrayType : LLVMType<"LLVMArray", "array", [
   ];
 
   let extraClassDeclaration = [{
     /// Checks if the given type can be used inside an array type.
     static bool isValidElementType(Type type);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// LLVMFunctionType
+//===----------------------------------------------------------------------===//
+
+def LLVMFunctionType : LLVMType<"LLVMFunction", "func", [
+    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+  let summary = "LLVM function type";
+  let description = [{
+    The `!llvm.func` is a function type. It consists of a single return type
+    (unlike MLIR which can have multiple), a list of parameter types and can
+    optionally be variadic.
+
+    Example:
+
+    ```mlir
+    !llvm.func<i32 (i32)>
+    ```
+  }];
+
+  let parameters = (ins "Type":$returnType, ArrayRefParameter<"Type">:$params,
+                        "bool":$varArg);
+  let assemblyFormat = [{
+    `<` custom<PrettyLLVMType>($returnType) ` ` `(`
+    custom<FunctionTypes>($params, $varArg) `>`
+  }];
+
+  let genVerifyDecl = 1;
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins
+      "Type":$result, "ArrayRef<Type>":$arguments,
+      CArg<"bool", "false">:$isVarArg)>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Checks if the given type can be used an argument in a function type.
+    static bool isValidArgumentType(Type type);
+
+    /// Checks if the given type can be used as a result in a function type.
+    static bool isValidResultType(Type type);
+
+    /// Returns whether the function is variadic.
+    bool isVarArg() const { return getVarArg(); }
+
+    /// Returns a clone of this function type with the given argument
+    /// and result types.
+    LLVMFunctionType clone(TypeRange inputs, TypeRange results) const;
+
+    /// Returns the result type of the function as an ArrayRef, enabling better
+    /// integration with generic MLIR utilities.
+    ArrayRef<Type> getReturnTypes() const;
+
+    /// Returns the number of arguments to the function.
+    unsigned getNumParams() const { return getParams().size(); }
+
+    /// Returns `i`-th argument of the function. Asserts on out-of-bounds.
+    Type getParamType(unsigned i) { return getParams()[i]; }
+  }];
+}
+
 #endif // LLVMTYPES_TD

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

void LLVMDialect::initialize() {
 
   // clang-format off
   addTypes<LLVMVoidType,
            LLVMPPCFP128Type,
            LLVMX86MMXType,
            LLVMTokenType,
            LLVMLabelType,
            LLVMMetadataType,
-           LLVMFunctionType,
            LLVMPointerType,
            LLVMFixedVectorType,
            LLVMScalableVectorType,
            LLVMStructType>();
   // clang-format on
   registerTypes();
 
   addOperations<

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

 /// Prints a type containing a fixed number of elements.
 template <typename TypeTy>
 static void printVectorType(AsmPrinter &printer, TypeTy type) {
   printer << '<' << type.getNumElements() << " x ";
   dispatchPrint(printer, type.getElementType());
   printer << '>';
 }
 
-/// Prints a function type.
-static void printFunctionType(AsmPrinter &printer, LLVMFunctionType funcType) {
-  printer << '<';
-  dispatchPrint(printer, funcType.getReturnType());
-  printer << " (";
-  llvm::interleaveComma(
-      funcType.getParams(), printer.getStream(),
-      [&printer](Type subtype) { dispatchPrint(printer, subtype); });
-  if (funcType.isVarArg()) {
-    if (funcType.getNumParams() != 0)
-      printer << ", ";
-    printer << "...";
-  }
-  printer << ")>";
-}
-
 /// Prints the given LLVM dialect type recursively. This leverages closedness of
 /// the LLVM dialect type system to avoid printing the dialect prefix
 /// repeatedly. For recursive structures, only prints the name of the structure
 /// when printing a self-reference. Note that this does not apply to sibling
 /// references. For example,
 ///   struct<"a", (ptr<struct<"a">>)>
 ///   struct<"c", (ptr<struct<"b", (ptr<struct<"c">>)>>,
 ///                ptr<struct<"b", (ptr<struct<"c">>)>>)>
if (auto vectorType = type.dyn_cast<LLVMScalableVectorType>()) {
     printer << '>';
     return;
   }
 
   if (auto structType = type.dyn_cast<LLVMStructType>())
     return printStructType(printer, structType);
 
   if (auto funcType = type.dyn_cast<LLVMFunctionType>())
-    return printFunctionType(printer, funcType);
+    return funcType.print(printer);
 }
 
 //===----------------------------------------------------------------------===//
 // Parsing.
 //===----------------------------------------------------------------------===//
 
 static ParseResult dispatchParse(AsmParser &parser, Type &type);
 
-/// Parses an LLVM dialect function type.
-///   llvm-type :: = `func<` llvm-type `(` llvm-type-list `...`? `)>`
-static LLVMFunctionType parseFunctionType(AsmParser &parser) {
-  SMLoc loc = parser.getCurrentLocation();
-  Type returnType;
-  if (parser.parseLess() || dispatchParse(parser, returnType) ||
-      parser.parseLParen())
-    return LLVMFunctionType();
-
-  // Function type without arguments.
-  if (succeeded(parser.parseOptionalRParen())) {
-    if (succeeded(parser.parseGreater()))
-      return parser.getChecked<LLVMFunctionType>(loc, returnType, llvm::None,
-                                                 /*isVarArg=*/false);
-    return LLVMFunctionType();
-  }
-
-  // Parse arguments.
-  SmallVector<Type, 8> argTypes;
-  do {
-    if (succeeded(parser.parseOptionalEllipsis())) {
-      if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
-        return LLVMFunctionType();
-      return parser.getChecked<LLVMFunctionType>(loc, returnType, argTypes,
-                                                 /*isVarArg=*/true);
-    }
-
-    Type arg;
-    if (dispatchParse(parser, arg))
-      return LLVMFunctionType();
-    argTypes.push_back(arg);
-  } while (succeeded(parser.parseOptionalComma()));
-
-  if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
-    return LLVMFunctionType();
-  return parser.getChecked<LLVMFunctionType>(loc, returnType, argTypes,
-                                             /*isVarArg=*/false);
-}
-
 /// Parses an LLVM dialect pointer type.
 ///   llvm-type ::= `ptr<` llvm-type (`,` integer)? `>`
 ///               | `ptr` (`<` integer `>`)?
 static LLVMPointerType parsePointerType(AsmParser &parser) {
   SMLoc loc = parser.getCurrentLocation();
   Type elementType;
   if (parser.parseOptionalLess()) {
     return parser.getChecked<LLVMPointerType>(loc, parser.getContext(),
static Type dispatchParse(AsmParser &parser, bool allowAny = true) {
   MLIRContext *ctx = parser.getContext();
   return StringSwitch<function_ref<Type()>>(key)
       .Case("void", [&] { return LLVMVoidType::get(ctx); })
       .Case("ppc_fp128", [&] { return LLVMPPCFP128Type::get(ctx); })
       .Case("x86_mmx", [&] { return LLVMX86MMXType::get(ctx); })
       .Case("token", [&] { return LLVMTokenType::get(ctx); })
       .Case("label", [&] { return LLVMLabelType::get(ctx); })
       .Case("metadata", [&] { return LLVMMetadataType::get(ctx); })
-      .Case("func", [&] { return parseFunctionType(parser); })
+      .Case("func", [&] { return LLVMFunctionType::parse(parser); })
       .Case("ptr", [&] { return parsePointerType(parser); })
       .Case("vec", [&] { return parseVectorType(parser); })
       .Case("array", [&] { return LLVMArrayType::parse(parser); })
       .Case("struct", [&] { return parseStructType(parser); })
       .Default([&] {
         parser.emitError(keyLoc) << "unknown LLVM type: " << key;
         return Type();
       })();

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

 #include "llvm/Support/TypeSize.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 constexpr const static unsigned kBitsInByte = 8;
 
 //===----------------------------------------------------------------------===//
+// custom<FunctionTypes>
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseFunctionTypes(AsmParser &p,
+                                      FailureOr<SmallVector<Type>> &params,
+                                      FailureOr<bool> &isVarArg) {
+  params.emplace();
+  isVarArg = false;
+  // `(` `)`
+  if (succeeded(p.parseOptionalRParen()))
+    return success();
+
+  // `(` `...` `)`
+  if (succeeded(p.parseOptionalEllipsis())) {
+    isVarArg = true;
+    return p.parseRParen();
+  }
+
+  // type (`,` type)* (`,` `...`)?
+  FailureOr<Type> type;
+  if (parsePrettyLLVMType(p, type))
+    return failure();
+  params->push_back(*type);
+  while (succeeded(p.parseOptionalComma())) {
+    if (succeeded(p.parseOptionalEllipsis())) {
+      isVarArg = true;
+      return p.parseRParen();
+    }
+    if (parsePrettyLLVMType(p, type))
+      return failure();
+    params->push_back(*type);
+  }
+  return p.parseRParen();
+}
+
+static void printFunctionTypes(AsmPrinter &p, ArrayRef<Type> params,
+                               bool isVarArg) {
+  llvm::interleaveComma(params, p,
+                        [&](Type type) { printPrettyLLVMType(p, type); });
+  if (isVarArg) {
+    if (!params.empty())
+      p << ", ";
+    p << "...";
+  }
+  p << ')';
+}
+
+//===----------------------------------------------------------------------===//
+// ODS-Generated Definitions
+//===----------------------------------------------------------------------===//
+
+/// These are unused for now.
+/// TODO: Move over to these once more types have been migrated to TypeDef.
+LLVM_ATTRIBUTE_UNUSED static OptionalParseResult
+generatedTypeParser(AsmParser &parser, StringRef *mnemonic, Type &value);
+LLVM_ATTRIBUTE_UNUSED static LogicalResult
+generatedTypePrinter(Type def, AsmPrinter &printer);
+
+#include "mlir/Dialect/LLVMIR/LLVMTypeInterfaces.cpp.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/LLVMIR/LLVMTypes.cpp.inc"
+
+//===----------------------------------------------------------------------===//
 // LLVMArrayType
 //===----------------------------------------------------------------------===//
 
 bool LLVMArrayType::isValidElementType(Type type) {
   return !type.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                    LLVMFunctionType, LLVMTokenType, LLVMScalableVectorType>();
 }
 

 }
 
 LLVMFunctionType LLVMFunctionType::clone(TypeRange inputs,
                                          TypeRange results) const {
   assert(results.size() == 1 && "expected a single result type");
   return get(results[0], llvm::to_vector(inputs), isVarArg());
 }
 
-Type LLVMFunctionType::getReturnType() const {
-  return getImpl()->getReturnType();
-}
 ArrayRef<Type> LLVMFunctionType::getReturnTypes() const {
-  return getImpl()->getReturnType();
-}
-
-unsigned LLVMFunctionType::getNumParams() {
-  return getImpl()->getArgumentTypes().size();
-}
-
-Type LLVMFunctionType::getParamType(unsigned i) {
-  return getImpl()->getArgumentTypes()[i];
-}
-
-bool LLVMFunctionType::isVarArg() const { return getImpl()->isVariadic(); }
-
-ArrayRef<Type> LLVMFunctionType::getParams() const {
-  return getImpl()->getArgumentTypes();
+  return static_cast<detail::LLVMFunctionTypeStorage *>(getImpl())->returnType;
 }
 
 LogicalResult
 LLVMFunctionType::verify(function_ref<InFlightDiagnostic()> emitError,
                          Type result, ArrayRef<Type> arguments, bool) {
   if (!isValidResultType(result))
     return emitError() << "invalid function result type: " << result;
 
   for (Type arg : arguments)
     if (!isValidArgumentType(arg))
       return emitError() << "invalid function argument type: " << arg;
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// SubElementTypeInterface
+
 void LLVMFunctionType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
-  for (Type type : llvm::concat<const Type>(getReturnTypes(), getParams()))
+  walkTypesFn(getReturnType());
+  for (Type type : getParams())
     walkTypesFn(type);
 }
 
 Type LLVMFunctionType::replaceImmediateSubElements(
     ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
   return get(replTypes.front(), replTypes.drop_front(), isVarArg());
 }
 
return llvm::TypeSwitch<Type, llvm::TypeSize>(type)
                        LLVMTokenType, LLVMStructType, LLVMArrayType,
                        LLVMPointerType, LLVMFunctionType>()) &&
                "unexpected missing support for primitive type");
         return llvm::TypeSize::Fixed(0);
       });
 }
 
 //===----------------------------------------------------------------------===//
-// ODS-Generated Definitions
-//===----------------------------------------------------------------------===//
-
-/// These are unused for now.
-/// TODO: Move over to these once more types have been migrated to TypeDef.
-LLVM_ATTRIBUTE_UNUSED static OptionalParseResult
-generatedTypeParser(AsmParser &parser, StringRef *mnemonic, Type &value);
-LLVM_ATTRIBUTE_UNUSED static LogicalResult
-generatedTypePrinter(Type def, AsmPrinter &printer);
-
-#include "mlir/Dialect/LLVMIR/LLVMTypeInterfaces.cpp.inc"
-
-#define GET_TYPEDEF_CLASSES
-#include "mlir/Dialect/LLVMIR/LLVMTypes.cpp.inc"
-
-//===----------------------------------------------------------------------===//
 // LLVMDialect
 //===----------------------------------------------------------------------===//
 
 void LLVMDialect::registerTypes() {
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/LLVMIR/LLVMTypes.cpp.inc"
       >();

mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h

private:
   unsigned keySizeAndFlags = 0;
 
   /// Number of the types contained in an identified struct combined with
   /// mutable flags. Must only be used through the Mutable* bitfields.
   unsigned identifiedBodySizeAndFlags = 0;
 };
 
 //===----------------------------------------------------------------------===//
-// LLVMFunctionTypeStorage.
-//===----------------------------------------------------------------------===//
-
-/// Type storage for LLVM dialect function types. These are uniqued using the
-/// list of types they contain and the vararg bit.
-struct LLVMFunctionTypeStorage : public TypeStorage {
-  using KeyTy = std::tuple<Type, ArrayRef<Type>, bool>;
-
-  /// Construct a storage from the given components. The list is expected to be
-  /// allocated in the context.
-  LLVMFunctionTypeStorage(Type result, ArrayRef<Type> arguments, bool variadic)
-      : resultType(result), isVariadicFlag(variadic),
-        numArguments(arguments.size()), argumentTypes(arguments.data()) {}
-
-  /// Hook into the type uniquing infrastructure.
-  static LLVMFunctionTypeStorage *construct(TypeStorageAllocator &allocator,
-                                            const KeyTy &key) {
-    return new (allocator.allocate<LLVMFunctionTypeStorage>())
-        LLVMFunctionTypeStorage(std::get<0>(key),
-                                allocator.copyInto(std::get<1>(key)),
-                                std::get<2>(key));
-  }
-
-  static unsigned hashKey(const KeyTy &key) {
-    // LLVM doesn't like hashing bools in tuples.
-    return llvm::hash_combine(std::get<0>(key), std::get<1>(key),
-                              static_cast<int>(std::get<2>(key)));
-  }
-
-  bool operator==(const KeyTy &key) const {
-    return std::make_tuple(getReturnType(), getArgumentTypes(), isVariadic()) ==
-           key;
-  }
-
-  /// Returns the list of function argument types.
-  ArrayRef<Type> getArgumentTypes() const {
-    return ArrayRef<Type>(argumentTypes, numArguments);
-  }
-
-  /// Checks whether the function type is variadic.
-  bool isVariadic() const { return isVariadicFlag; }
-
-  /// Returns the function result type.
-  const Type &getReturnType() const { return resultType; }
-
-private:
-  /// The result type of the function.
-  Type resultType;
-  /// Flag indicating if the function is variadic.
-  bool isVariadicFlag;
-  /// The argument types of the function.
-  unsigned numArguments;
-  const Type *argumentTypes;
-};
-
-//===----------------------------------------------------------------------===//
 // LLVMPointerTypeStorage.
 //===----------------------------------------------------------------------===//
 
 /// Storage type for LLVM dialect pointer types. These are uniqued by a pair of
 /// element type and address space. The element type may be null indicating that
 /// the pointer is opaque.
 struct LLVMPointerTypeStorage : public TypeStorage {
   using KeyTy = std::tuple<Type, unsigned>;