[mlir] First-party modeling of LLVM types

[ftynse]

Summary

The current modeling of LLVM IR types in MLIR is based on the LLVMType class
that wraps a raw llvm::Type * and delegates uniquing, printing and parsing to
LLVM itself. This model makes thread-safe type manipulation hard and is being
progressively replaced with a cleaner MLIR model that replicates the type
system. Introduce a set of classes reflecting the LLVM IR type system in MLIR
instead of wrapping the existing types. These are currently introduced as
separate classes without affecting the dialect flow, and are exercised through
a test dialect. Once feature parity is reached, the old implementation will be
gradually substituted with the new one.

Depends On D84171

[jdd]

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

2225

Why are you passing around the output stream instead of the DialogAsmPrinter? I would prefer the stack be in DialogAsmPrinter, be a stack of mlir::Type, and be automatically managed by DialogAsmPrinter so you're not relying on the closed nature of the LLVM type system. @rriddle seemed to be supportive of that idea or at least confirmed it would work. Do you consider this to be out-of-scope for this change?

[rriddle]

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

246

Please fix lint check.

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

1774

I would move these to an enum for KeyFlags/MutableFlags. Right now it is slightly confusing.

1860

Can we restructure this storage a bit? We could move the identified body Type[] to be a trailing object, and move that size here(as unsigned). We could then steal two bits from each of the sizes for the flags.

1892

You could pack this into the resultType with PointerIntPair, you just need to define PointerLikeTypeTraits for LLVMTypeNew(or drop it down to a Type and just case when accessing it).

[ftynse]

Ready for review

[ftynse]

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

1860

We could move the identified body Type[] to be a trailing object, and move that size here(as unsigned).

This would mean that we need to know the size of the storage, i.e. the number of elements in the body of an identified struct, when we construct the storage, but we don't. It may remain unknown until we mutate the storage, at which point we cannot reallocate it anymore.

I considered having TrailingObjects<Storage, char, Type> for the key, but it creates a wrong impression that both the name and the (literal) type list can be present at the same time, and increases code complexity.

1892

I've been looking for the place where alignment is indicated, but failed to find it... Can you point me to it?

2225

Well, because I only need a stream in these functions and because some functions are reused with other streams (error reporting in parsing also calls these before the type is fully constructed). LLVM types do and likely will keep relying on closedness, otherwise, e.g., !llvm.ptr<i32> is ambiguous in pointing to standard or LLVM type.

Though I am happy to use a DialectAsmPrinter/Parser with a stack instead of rolling my own, we are already looking at 1300 LoC without comments, and I'm not even done here. When somebody gets to implement those things, I can consider moving this code to use them.

[rriddle]

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

1860

SGTM.

[ftynse]

more cmake

[rriddle]

Looks really great Alex! Made a first pass through, nothing really big mostly nits.

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

56 ↗

(On Diff #280517)

nit: The other enums in MLIR use CamelCase for elements.

158 ↗

(On Diff #280517)

Can you separate these class definitions with comment blocks?

//===-...-===//
//
//===-...-===//

189 ↗

(On Diff #280517)

nit: Checks -> Returns

194 ↗

(On Diff #280517)

Why have the duplicate method?

197 ↗

(On Diff #280517)

nit: Maybe add signless integer type

213 ↗

(On Diff #280517)

nit: of this integer type?

255 ↗

(On Diff #280517)

typos: prodiving and lits

257 ↗

(On Diff #280517)

nit: without a body

288 ↗

(On Diff #280517)

nit: Such a struct

289 ↗

(On Diff #280517)

nit: with a mutable body

296 ↗

(On Diff #280517)

nit: Set the body

432 ↗

(On Diff #280517)

nit: Can you set this to the same value as PointerLikeTypeTraits<mlir::Type>::NumLowBitsAvailable?

mlir/include/mlir/IR/DialectImplementation.h

329

nit: Can you add a one line comment here?

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

1757

Why was this removed?

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

49 ↗

(On Diff #280856)

Could we use llvm/ADT/BitFields here instead?

105 ↗

(On Diff #280856)

Can you move these to a TypeDetail header in this directory?

115 ↗

(On Diff #280856)

nit: whether the

133 ↗

(On Diff #280856)

nit: I think you could state - mutable (Identified structs only) which would avoid the need to state it for each point? Feel free to ignore, given that the current description matches the style immutable part.

193 ↗

(On Diff #280856)

Can you use the enum names instead of numbers here?

203 ↗

(On Diff #280856)

if (isIdentified())
  return other.isIdentified() && identifier.equals(...);
return !other.isIdentified() && other.isPacked() == ...;

?

227 ↗

(On Diff #280856)

nit: Prefer getFoo style getters and setters. Same applies below.

259 ↗

(On Diff #280856)

typo: unintilized

289 ↗

(On Diff #280856)

nit: Prefer early return here.

353 ↗

(On Diff #280856)

BumpPtrAllocator::Allocate<T> (used by copyInto) ensures address alignment for a given T. Are you wanting the char * key to have the same alignment as the Type *?

Also, what would be within the pointer union?

462 ↗

(On Diff #280856)

Please remove the ::mlir:: in this file.

629 ↗

(On Diff #280856)

Could you also move the printing/parsing to a different file?

894 ↗

(On Diff #280856)

Please add a comment here.

941 ↗

(On Diff #280856)

nit: Can you just capture the diagnostic as a variable, and then attach the note?

auto diag = parser.emitError(subtypesLoc) << ...;
diag.attachNote() << ...;

969 ↗

(On Diff #280856)

nit: Just return directly.

[ftynse]

Address review

[ftynse]

Try to appease MSVC

[ftynse]

Even more MSVC

[ftynse]

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

194 ↗

(On Diff #280517)

LLVM has it as params but that contradicts the usual MLIR style. I am fine either way.

257 ↗

(On Diff #280517)

Intentionality matters here. A struct can be created without a body with the intention of setting the body later, or with the intention of keeping it opaque. I referred to the latter as intentionally opaque, do you have a better name?

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

1757

code motion accident

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

49 ↗

(On Diff #280856)

Sure, thanks for mentioning this!

193 ↗

(On Diff #280856)

The "packed" flag is present in two enums and while it currently has the same value in both, it can change leading to subtle problems with hashing here, e.g., if it has the same value as identified. I can add named variables here instead.

353 ↗

(On Diff #280856)

BumpPtrAllocator::Allocate<T> (used by copyInto) ensures address alignment for a given T. Are you wanting the char * key to have the same alignment as the Type *?

I want at least one bit for the pointer union itself. If I get three bits, I can use the remaining bits for flags instead of stealing them from the size. But that may consume extra space, so I'm also fine with writing the casts myself.

Also, what would be within the pointer union?

The same thing as in keyPtr with type erasure: first element of the key.

969 ↗

(On Diff #280856)

debug leftover, sorry

[ftynse]

Use LLVMTypeNew inside struct storage for consistency with other storages

[rriddle]

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

1 ↗

(On Diff #281543)

Should we name this TypeParser instead? This would match the naming that the Quant dialect used.

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

193 ↗

(On Diff #280856)

I can add named variables here instead.

Ah, thanks for explaining. Named variables works for me.

353 ↗

(On Diff #280856)

I want at least one bit for the pointer union itself. If I get three bits, I can use the
remaining bits for flags instead of stealing them from the size. But that may consume
extra space, so I'm also fine with writing the casts myself.

Ah, we could add an alignment parameter to StorageAllocator::copyInto. We'd just need to reimplement what BumpPtrAllocator<T>::Allocate does(effectively just uses the type-erased Allocate and casts back to a T*) and allow for a new desired alignment. We'd just need to assert that the desired alignment is still valid for T.

[ftynse]

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

1 ↗

(On Diff #281543)

I considered that, but the file contains both parsing and printing so it felt incorrect to emphasize the parser. I can rename if you think it makes more sense.

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

353 ↗

(On Diff #280856)

Ah, we could add an alignment parameter to StorageAllocator::copyInto. We'd just need to reimplement what BumpPtrAllocator<T>::Allocate does(effectively just uses the type-erased Allocate and casts back to a T*) and allow for a new desired alignment. We'd just need to assert that the desired alignment is still valid for T.

I'll keep it as a more detailed TODO, it is unclear to me if we want to potentially waste additional memory for alignment just to remove a couple of casts when we pack bits into sizes in the same object.

[ftynse]

Rebase

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

 
 namespace mlir {
 namespace LLVM {
 class LLVMDialect;
 
 namespace detail {
 struct LLVMTypeStorage;
 struct LLVMDialectImpl;
+
+struct LLVMFunctionTypeStorage;
+struct LLVMIntegerTypeStorage;
+struct LLVMPointerTypeStorage;
+struct LLVMStructTypeStorage;
+struct LLVMTypeAndSizeStorage;
 } // namespace detail
 
 class LLVMType : public mlir::Type::TypeBase<LLVMType, mlir::Type,
                                              detail::LLVMTypeStorage> {
 public:
   enum Kind {
     LLVM_TYPE = FIRST_LLVM_TYPE,
   };
private:
   static LLVMType get(MLIRContext *context, llvm::Type *llvmType);
 
   /// Get an LLVMType with an llvm type that may cause changes to the underlying
   /// llvm context when constructed.
   static LLVMType getLocked(LLVMDialect *dialect,
                             function_ref<llvm::Type *()> typeBuilder);
 };
 
+class LLVMTypeNew : public Type {
+public:
+  enum Kind {
+    // Keep non-parametric types contiguous in the enum.
+    VOID_TYPE = FIRST_LLVM_TYPE + 1,
+    HALF_TYPE,
+    BFLOAT_TYPE,
+    FLOAT_TYPE,
+    DOUBLE_TYPE,
+    FP128_TYPE,
+    X86_FP80_TYPE,
+    PPC_FP128_TYPE,
+    X86_MMX_TYPE,
+    LABEL_TYPE,
+    TOKEN_TYPE,
+    METADATA_TYPE,
+    // End of non-parametric types.
+    FUNCTION_TYPE,
+    INTEGER_TYPE,
+    POINTER_TYPE,
+    FIXED_VECTOR_TYPE,
+    SCALABLE_VECTOR_TYPE,
+    ARRAY_TYPE,
+    STRUCT_TYPE,
+    FIRST_NEW_LLVM_TYPE = VOID_TYPE,
+    LAST_NEW_LLVM_TYPE = STRUCT_TYPE,
+    FIRST_TRIVIAL_TYPE = VOID_TYPE,
+    LAST_TRIVIAL_TYPE = METADATA_TYPE
+  };
+
+  using Type::Type;
+
+  static bool kindof(unsigned kind) {
+    return FIRST_NEW_LLVM_TYPE <= kind <= LAST_NEW_LLVM_TYPE;
+  }
+};
+
+template <typename Derived, unsigned Kind>
+class LLVMTrivialType
+    : public Type::TypeBase<Derived, LLVMTypeNew, TypeStorage> {
+public:
+  using Type::TypeBase<Derived, LLVMTypeNew, TypeStorage>::TypeBase;
+  using Base = LLVMTrivialType<Derived, Kind>;
+
+  static bool kindof(unsigned kind) { return kind == Kind; }
+
+  static Derived get(MLIRContext *context) {
+    return Type::TypeBase<Derived, LLVMTypeNew, TypeStorage>::get(context,
+                                                                  Kind);
+  }
+};
+
+#define DEFINE_TRIVIAL_LLVM_TYPE(ClassName, Kind)                              \
+  class ClassName : public LLVMTrivialType<ClassName, Kind> {                  \
+  public:                                                                      \
+    using Base::Base;                                                          \
+  }
+
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMVoidType, LLVMTypeNew::VOID_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMHalfType, LLVMTypeNew::HALF_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMBFloatType, LLVMTypeNew::BFLOAT_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMFloatType, LLVMTypeNew::FLOAT_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMDoubleType, LLVMTypeNew::DOUBLE_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMFP128Type, LLVMTypeNew::FP128_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMX86FP80Type, LLVMTypeNew::X86_FP80_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMPPCFP128Type, LLVMTypeNew::PPC_FP128_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMX86MMXType, LLVMTypeNew::X86_MMX_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMTokenType, LLVMTypeNew::TOKEN_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMLabelType, LLVMTypeNew::LABEL_TYPE);
+DEFINE_TRIVIAL_LLVM_TYPE(LLVMMetadataType, LLVMTypeNew::METADATA_TYPE);
+
+#undef DEFINE_TRIVIAL_LLVM_TYPE
+
+class LLVMArrayType : public Type::TypeBase<LLVMArrayType, LLVMTypeNew,
+                                            detail::LLVMTypeAndSizeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) { return kind == LLVMTypeNew::ARRAY_TYPE; }
+
+  static LLVMArrayType get(LLVMTypeNew elementType, unsigned numElements);
+  LLVMTypeNew getElementType();
+  unsigned getNumElements();
+};
+
+class LLVMFunctionType
+    : public Type::TypeBase<LLVMFunctionType, LLVMTypeNew,
+                            detail::LLVMFunctionTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::FUNCTION_TYPE;
+  }
+
+  static LLVMFunctionType get(LLVMTypeNew result,
+                              ArrayRef<LLVMTypeNew> arguments,
+                              bool isVarArg = false);
+
+  LLVMTypeNew getReturnType();
+  unsigned getNumParams();
+  LLVMTypeNew getParamType(unsigned i);
+  bool isVarArg();
+  ArrayRef<LLVMTypeNew> getParams();
+  ArrayRef<LLVMTypeNew> params() { return getParams(); }
+};
+
+class LLVMIntegerType : public Type::TypeBase<LLVMIntegerType, LLVMTypeNew,
+                                              detail::LLVMIntegerTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::INTEGER_TYPE;
+  }
+
+  static LLVMIntegerType get(MLIRContext *ctx, unsigned bitwidth);
+  unsigned getBitWidth();
+};
+
+class LLVMPointerType : public Type::TypeBase<LLVMPointerType, LLVMTypeNew,
+                                              detail::LLVMPointerTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::POINTER_TYPE;
+  }
+
+  static LLVMPointerType get(LLVMTypeNew pointee, unsigned addressSpace = 0);
+  LLVMTypeNew getElementType();
+  unsigned getAddressSpace();
+};
+
+class LLVMStructType : public Type::TypeBase<LLVMStructType, LLVMTypeNew,
+                                             detail::LLVMStructTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) { return kind == LLVMTypeNew::STRUCT_TYPE; }
+
+  static LLVMStructType getIdentified(MLIRContext *context, StringRef name);
+  static LLVMStructType getLiteral(MLIRContext *context, ArrayRef<Type> types,
+                                   bool isPacked = false);
+  static LLVMStructType getOpaque(StringRef name, MLIRContext *context);
+  LogicalResult setBody(ArrayRef<Type> types, bool isPacked);
+
+  bool isPacked();
+  bool isIdentified();
+  bool isOpaque();
+  StringRef identifier();
+  ArrayRef<Type> getBody();
+};
+
+class LLVMVectorType : public LLVMTypeNew {
+public:
+  using LLVMTypeNew::LLVMTypeNew;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::FIXED_VECTOR_TYPE ||
+           kind == LLVMTypeNew::SCALABLE_VECTOR_TYPE;
+  }
+
+  LLVMTypeNew getElementType();
+  llvm::ElementCount getElementCount();
+};
+
+class LLVMFixedVectorType
+    : public Type::TypeBase<LLVMFixedVectorType, LLVMVectorType,
+                            detail::LLVMTypeAndSizeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::FIXED_VECTOR_TYPE;
+  }
+
+  static LLVMFixedVectorType get(LLVMTypeNew elementType, unsigned numElements);
+
+  unsigned getNumElements();
+};
+
+class LLVMScalableVectorType
+    : public Type::TypeBase<LLVMScalableVectorType, LLVMVectorType,
+                            detail::LLVMTypeAndSizeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == LLVMTypeNew::SCALABLE_VECTOR_TYPE;
+  }
+
+  static LLVMScalableVectorType get(LLVMTypeNew elementType,
+                                    unsigned minNumElements);
+
+  unsigned getMinNumElements();
+};
+
 ///// Ops /////
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.h.inc"
 
 #include "mlir/Dialect/LLVMIR/LLVMOpsDialect.h.inc"
 
 /// Create an LLVM global containing the string "value" at the module containing
 /// surrounding the insertion point of builder. Obtain the address of that
 /// global and use it to compute the address of the first character in the
 /// string (operations inserted at the builder insertion point).
 Value createGlobalString(Location loc, OpBuilder &builder, StringRef name,
                          StringRef value, LLVM::Linkage linkage,
                          LLVM::LLVMDialect *llvmDialect);
 
 /// LLVM requires some operations to be inside of a Module operation. This
 /// function confirms that the Operation has the desired properties.
 bool satisfiesLLVMModule(Operation *op);
 
 /// Clones the given module into the provided context. This is implemented by
 /// transforming the module into bitcode and then reparsing the bitcode in the
 /// provided context.
 std::unique_ptr<llvm::Module>
 cloneModuleIntoNewContext(llvm::LLVMContext *context, llvm::Module *module);
 
+namespace detail {
+LLVMTypeNew parseType(DialectAsmParser &parser);
+void printType(LLVMTypeNew type, DialectAsmPrinter &printer);
+} // end namespace detail
+
+class LLVMDialectNewTypes : public Dialect {
+public:
+  LLVMDialectNewTypes(MLIRContext *ctx) : Dialect(getDialectNamespace(), ctx) {
+    // clang-format off
+    addTypes<LLVMVoidType,
+             LLVMHalfType,
+             LLVMBFloatType,
+             LLVMFloatType,
+             LLVMDoubleType,
+             LLVMFP128Type,
+             LLVMX86FP80Type,
+             LLVMPPCFP128Type,
+             LLVMX86MMXType,
+             LLVMTokenType,
+             LLVMLabelType,
+             LLVMMetadataType,
+             LLVMFunctionType,
+             LLVMIntegerType,
+             LLVMPointerType,
+             LLVMFixedVectorType,
+             LLVMScalableVectorType,
+             LLVMArrayType,
+             LLVMStructType>();
+    // clang-format on
+  }
+  static StringRef getDialectNamespace() { return "llvm2"; }
+
+  Type parseType(DialectAsmParser &parser) const override {
+    return detail::parseType(parser);
+  }
+  void printType(Type type, DialectAsmPrinter &printer) const override {
+    detail::printType(type.cast<LLVMTypeNew>(), printer);
+  }
+};
+
 } // end namespace LLVM
 } // end namespace mlir
 
+namespace llvm {
+
+// Type hash just like pointers.
+template <>
+struct DenseMapInfo<::mlir::LLVM::LLVMTypeNew> {
+  static ::mlir::LLVM::LLVMTypeNew getEmptyKey() {
+    void *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
+    return ::mlir::LLVM::LLVMTypeNew(
+        static_cast<::mlir::LLVM::LLVMTypeNew::ImplType *>(pointer));
+  }
+  static mlir::Type getTombstoneKey() {
+    void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
+    return ::mlir::LLVM::LLVMTypeNew(
+        static_cast<::mlir::LLVM::LLVMTypeNew::ImplType *>(pointer));
+  }
+  static unsigned getHashValue(mlir::Type val) { return mlir::hash_value(val); }
+  static bool isEqual(mlir::Type LHS, mlir::Type RHS) { return LHS == RHS; }
+};
+
+} // namespace llvm
+
 #endif // MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_

mlir/include/mlir/IR/DialectImplementation.h

public:
   virtual ParseResult parseOptionalComma() = 0;
 
   /// Parse a `=` token.
   virtual ParseResult parseEqual() = 0;
 
   /// Parse a `=` token if present.
   virtual ParseResult parseOptionalEqual() = 0;
 
+  /// Parse a quoted string token if present.
+  virtual ParseResult parseOptionalString(StringRef *string) = 0;
+
   /// Parse a given keyword.
   ParseResult parseKeyword(StringRef keyword, const Twine &msg = "") {
     auto loc = getCurrentLocation();
     if (parseOptionalKeyword(keyword))
       return emitError(loc, "expected '") << keyword << "'" << msg;
     return success();
   }
 
template <typename TypeType> ParseResult parseType(TypeType &result) {
 
     // Check for the right kind of attribute.
     result = type.dyn_cast<TypeType>();
     if (!result)
       return emitError(loc, "invalid kind of type specified");
     return success();
   }
 
+  virtual OptionalParseResult parseOptionalType(Type &result) = 0;
+
   /// Parse a 'x' separated dimension list. This populates the dimension list,
   /// using -1 for the `?` dimensions if `allowDynamic` is set and errors out on
   /// `?` otherwise.
   ///
   ///   dimension-list ::= (dimension `x`)*
   ///   dimension ::= `?` | integer
   ///
   /// When `allowDynamic` is not set, this is used to parse:

mlir/include/mlir/InitAllDialects.h

 // all the possible dialects to be made available to the context automatically.
 inline void registerAllDialects() {
   static bool init_once = []() {
     registerDialect<AffineDialect>();
     registerDialect<avx512::AVX512Dialect>();
     registerDialect<gpu::GPUDialect>();
     registerDialect<LLVM::LLVMAVX512Dialect>();
     registerDialect<LLVM::LLVMDialect>();
+    registerDialect<LLVM::LLVMDialectNewTypes>();
     registerDialect<linalg::LinalgDialect>();
     registerDialect<scf::SCFDialect>();
     registerDialect<omp::OpenMPDialect>();
     registerDialect<quant::QuantizationDialect>();
     registerDialect<spirv::SPIRVDialect>();
     registerDialect<StandardOpsDialect>();
     registerDialect<vector::VectorDialect>();
     registerDialect<NVVM::NVVMDialect>();

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

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
 
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Mutex.h"

 }
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace LLVM {
+
 namespace detail {
 struct LLVMDialectImpl {
   LLVMDialectImpl() : module("LLVMDialectModule", llvmContext) {}
 
   llvm::LLVMContext llvmContext;
   llvm::Module module;
 
   /// A set of LLVMTypes that are cached on construction to avoid any lookups or
Type LLVMDialect::parseType(DialectAsmParser &parser) const {
   llvm::SMDiagnostic errorMessage;
   llvm::Type *type = llvm::parseType(tyData, errorMessage, impl->module);
   if (!type)
     return (parser.emitError(parser.getNameLoc(), errorMessage.getMessage()),
             nullptr);
   return LLVMType::get(getContext(), type);
 }
 
+void printLLVMTypeImpl(llvm::raw_ostream &os, LLVMTypeNew type,
+                       llvm::SetVector<StringRef> &stack);
+
 /// Print a type registered to this dialect.
 void LLVMDialect::printType(Type type, DialectAsmPrinter &os) const {
   auto llvmType = type.dyn_cast<LLVMType>();
   assert(llvmType && "printing wrong type");
   assert(llvmType.getUnderlyingType() && "no underlying LLVM type");
   llvmType.getUnderlyingType()->print(os.getStream());
 }
 
 /// Verify LLVMIR function argument attributes.
 LogicalResult LLVMDialect::verifyRegionArgAttribute(Operation *op,
                                                     unsigned regionIdx,
                                                     unsigned argIdx,
                                                     NamedAttribute argAttr) {
   // Check that llvm.noalias is a boolean attribute.
   if (argAttr.first == "llvm.noalias" && !argAttr.second.isa<BoolAttr>())
     return op->emitError()
            << "llvm.noalias argument attribute of non boolean type";
   // Check that llvm.align is an integer attribute.
   if (argAttr.first == "llvm.align" && !argAttr.second.isa<IntegerAttr>())
     return op->emitError()
            << "llvm.align argument attribute of non integer type";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMType.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace LLVM {
 namespace detail {
+struct LLVMStructTypeStorage : public ::mlir::TypeStorage {
+  // Key: flags(identified, opaqueOrPacked), name-or-type-list
+  //            for identified -> opaque, for non-identified -> packed
+  // Body: initialized, type-list and packed for identified
+
+  constexpr static uint8_t IdentifiedFlag = 0x1;
+  constexpr static uint8_t InitializedFlag = 0x1;
+  constexpr static uint8_t PackedFlag = 0x2;
+  constexpr static uint8_t OpaqueFlag = 0x2;
+
+  using KeyTy = std::tuple<const void *, unsigned, uint8_t>;
+
+  LLVMStructTypeStorage(const void *ptr, unsigned size, uint8_t flags)
+      : keyPtr(ptr), keySize(size), keyFlags(flags), mutableFlags(0) {}
+
+  static bool isIdentified(const KeyTy &key) {
+    return std::get<2>(key) & IdentifiedFlag;
+  }
+
+  bool isIdentified() const { return keyFlags & IdentifiedFlag; }
+
+  bool isPacked() const {
+    return isIdentified() ? (mutableFlags & PackedFlag)
+                          : (keyFlags & PackedFlag);
+  }
+
+  bool isOpaque() const { return keyFlags & OpaqueFlag; }
+
+  bool isInitialized() const { return mutableFlags & InitializedFlag; }
+
+  static StringRef identifier(const KeyTy &key) {
+    return StringRef(static_cast<const char *>(std::get<0>(key)),
+                     std::get<1>(key));
+  }
+  StringRef identifier() const {
+    return StringRef(static_cast<const char *>(keyPtr), keySize);
+  }
+
+  static ArrayRef<Type> typelist(const KeyTy &key) {
+    return ArrayRef<Type>(static_cast<const Type *>(std::get<0>(key)),
+                          std::get<1>(key));
+  }
+  ArrayRef<Type> typelist() const {
+    return ArrayRef<Type>(static_cast<const Type *>(keyPtr), keySize);
+  }
+
+  bool operator==(const KeyTy &other) const {
+    uint8_t otherFlags = std::get<2>(other);
+    if (keyFlags != otherFlags)
+      return false;
+
+    if (isIdentified(other))
+      return identifier(other).equals(identifier());
+    return typelist(other) == typelist();
+  };
+
+  static llvm::hash_code hashKey(const KeyTy &key) {
+    if (isIdentified(key))
+      return llvm::hash_combine(identifier(key), std::get<2>(key));
+    return llvm::hash_combine(typelist(key), std::get<2>(key));
+  }
+
+  static LLVMStructTypeStorage *construct(TypeStorageAllocator &allocator,
+                                          const KeyTy &key) {
+    const void *ptr = isIdentified(key)
+                          ? static_cast<const void *>(
+                                allocator.copyInto(identifier(key)).data())
+                          : allocator.copyInto(typelist(key)).data();
+
+    return new (allocator.allocate<LLVMStructTypeStorage>())
+        LLVMStructTypeStorage(ptr, std::get<1>(key), std::get<2>(key));
+  }
+
+  LogicalResult mutate(TypeStorageAllocator &allocator, ArrayRef<Type> body,
+                       bool packed) {
+    if (!isIdentified())
+      return failure();
+    if (isInitialized())
+      return success(body == identifiedStructBody && packed == isPacked());
+
+    mutableFlags |= InitializedFlag;
+    if (packed)
+      mutableFlags |= PackedFlag;
+    identifiedStructBody = allocator.copyInto(body);
+    return success();
+  }
+
+  // Note: cannot use PointerUnion because bump-ptr allocator does not guarantee
+  // address alignment.
+  const void *keyPtr;
+  unsigned keySize;
+  ArrayRef<Type> identifiedStructBody;
+  uint8_t keyFlags : 2;
+  uint8_t mutableFlags : 2;
+};
+
+struct LLVMFunctionTypeStorage : public ::mlir::TypeStorage {
+  using KeyTy = std::tuple<LLVMTypeNew, ArrayRef<LLVMTypeNew>, bool>;
+
+  LLVMFunctionTypeStorage(LLVMTypeNew result, ArrayRef<LLVMTypeNew> arguments,
+                          bool var)
+      : resultType(result), argumentTypes(arguments), variadic(var) {}
+
+  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),
+                              std::get<2>(key));
+  }
+
+  bool operator==(const KeyTy &key) const {
+    return std::make_tuple(std::get<0>(key), argumentTypes, variadic) == key;
+  }
+
+  LLVMTypeNew resultType;
+  ArrayRef<LLVMTypeNew> argumentTypes;
+  int8_t variadic : 1;
+};
+
+struct LLVMIntegerTypeStorage : public ::mlir::TypeStorage {
+  using KeyTy = unsigned;
+
+  LLVMIntegerTypeStorage(unsigned width) : bitwidth(width) {}
+
+  static LLVMIntegerTypeStorage *construct(TypeStorageAllocator &allocator,
+                                           const KeyTy &key) {
+    return new (allocator.allocate<LLVMIntegerTypeStorage>())
+        LLVMIntegerTypeStorage(key);
+  }
+
+  bool operator==(const KeyTy &key) const { return key == bitwidth; }
+
+  unsigned bitwidth;
+};
+
+struct LLVMPointerTypeStorage : public ::mlir::TypeStorage {
+  using KeyTy = std::tuple<LLVMTypeNew, unsigned>;
+
+  LLVMPointerTypeStorage(const KeyTy &key)
+      : pointeeType(std::get<0>(key)), addressSpace(std::get<1>(key)) {}
+
+  static LLVMPointerTypeStorage *construct(TypeStorageAllocator &allocator,
+                                           const KeyTy &key) {
+    return new (allocator.allocate<LLVMPointerTypeStorage>())
+        LLVMPointerTypeStorage(key);
+  }
+
+  bool operator==(const KeyTy &key) const {
+    return std::make_tuple(pointeeType, addressSpace) == key;
+  }
+
+  LLVMTypeNew pointeeType;
+  unsigned addressSpace;
+};
+
+/// Common storage used for LLVM dialect types that need an element type and a
+/// number: arrays, fixed and scalable vectors. The actual semantics of the
+/// type is defined by its kind.
+struct LLVMTypeAndSizeStorage : public ::mlir::TypeStorage {
+  using KeyTy = std::tuple<LLVMTypeNew, unsigned>;
+
+  LLVMTypeAndSizeStorage(const KeyTy &key)
+      : elementType(std::get<0>(key)), numElements(std::get<1>(key)) {}
+
+  static LLVMTypeAndSizeStorage *construct(TypeStorageAllocator &allocator,
+                                           const KeyTy &key) {
+    return new (allocator.allocate<LLVMTypeAndSizeStorage>())
+        LLVMTypeAndSizeStorage(key);
+  }
+
+  bool operator==(const KeyTy &key) const {
+    return std::make_tuple(elementType, numElements) == key;
+  }
+
+  LLVMTypeNew elementType;
+  unsigned numElements;
+};
+
+} // end namespace detail
+} // end namespace LLVM
+} // end namespace mlir
+
+// Array type.
+
+LLVMArrayType LLVMArrayType::get(LLVMTypeNew elementType,
+                                 unsigned numElements) {
+  assert(elementType && "expected non-null subtype");
+  return Base::get(elementType.getContext(), LLVMTypeNew::ARRAY_TYPE,
+                   elementType, numElements);
+}
+
+LLVMTypeNew LLVMArrayType::getElementType() { return getImpl()->elementType; }
+
+unsigned LLVMArrayType::getNumElements() { return getImpl()->numElements; }
+
+//
+// Function type.
+//
+
+LLVMFunctionType LLVMFunctionType::get(LLVMTypeNew result,
+                                       ArrayRef<LLVMTypeNew> arguments,
+                                       bool isVarArg) {
+  assert(result && "expected non-null result");
+  return Base::get(result.getContext(), LLVMTypeNew::FUNCTION_TYPE, result,
+                   arguments, isVarArg);
+}
+
+LLVMTypeNew LLVMFunctionType::getReturnType() {
+  return getImpl()->resultType.cast<LLVMTypeNew>();
+}
+
+unsigned LLVMFunctionType::getNumParams() {
+  return getImpl()->argumentTypes.size();
+}
+
+LLVMTypeNew LLVMFunctionType::getParamType(unsigned i) {
+  return getImpl()->argumentTypes[i].cast<LLVMTypeNew>();
+}
+
+bool LLVMFunctionType::isVarArg() { return getImpl()->variadic; }
+
+ArrayRef<LLVMTypeNew> LLVMFunctionType::getParams() {
+  return getImpl()->argumentTypes;
+}
+
+//
+// Integer type.
+//
+
+LLVMIntegerType LLVMIntegerType::get(MLIRContext *ctx, unsigned bitwidth) {
+  return Base::get(ctx, LLVMTypeNew::INTEGER_TYPE, bitwidth);
+}
+
+unsigned LLVMIntegerType::getBitWidth() { return getImpl()->bitwidth; }
+
+//
+// Pointer type.
+//
+
+LLVMPointerType LLVMPointerType::get(LLVMTypeNew pointee,
+                                     unsigned addressSpace) {
+  assert(pointee && "expected non-null subtype");
+  return Base::get(pointee.getContext(), LLVMTypeNew::POINTER_TYPE, pointee,
+                   addressSpace);
+}
+
+LLVMTypeNew LLVMPointerType::getElementType() { return getImpl()->pointeeType; }
+
+unsigned LLVMPointerType::getAddressSpace() { return getImpl()->addressSpace; }
+
+//
+// Struct type.
+//
+
+LLVMStructType LLVMStructType::getIdentified(MLIRContext *context,
+                                             StringRef name) {
+  return Base::get(context, LLVMTypeNew::STRUCT_TYPE, name.data(), name.size(),
+                   detail::LLVMStructTypeStorage::IdentifiedFlag);
+}
+
+LLVMStructType LLVMStructType::getLiteral(MLIRContext *context,
+                                          ArrayRef<Type> types, bool isPacked) {
+  uint8_t flags = 0;
+  if (isPacked)
+    flags |= detail::LLVMStructTypeStorage::PackedFlag;
+  return Base::get(context, LLVMTypeNew::STRUCT_TYPE, types.data(),
+                   types.size(), flags);
+}
+
+LLVMStructType LLVMStructType::getOpaque(StringRef name, MLIRContext *context) {
+  auto flags = detail::LLVMStructTypeStorage::OpaqueFlag |
+               detail::LLVMStructTypeStorage::IdentifiedFlag;
+  return Base::get(context, LLVMTypeNew::STRUCT_TYPE, name.data(), name.size(),
+                   flags);
+}
+
+LogicalResult LLVMStructType::setBody(ArrayRef<Type> types, bool isPacked) {
+  assert(isIdentified() && "can only set bodies of identified structs");
+  return Base::mutate(types, isPacked);
+}
+
+bool LLVMStructType::isPacked() { return getImpl()->isPacked(); }
+bool LLVMStructType::isIdentified() { return getImpl()->isIdentified(); }
+bool LLVMStructType::isOpaque() {
+  return getImpl()->isOpaque() || !getImpl()->isInitialized();
+}
+StringRef LLVMStructType::identifier() { return getImpl()->identifier(); }
+ArrayRef<Type> LLVMStructType::getBody() {
+  return isIdentified() ? getImpl()->identifiedStructBody
+                        : getImpl()->typelist();
+}
+
+static StringRef getTypeKeyword(LLVMTypeNew type) {
+  switch (type.getKind()) {
+  case LLVMTypeNew::VOID_TYPE:
+    return "void";
+  case LLVMTypeNew::HALF_TYPE:
+    return "half";
+  case LLVMTypeNew::BFLOAT_TYPE:
+    return "bfloat";
+  case LLVMTypeNew::FLOAT_TYPE:
+    return "float";
+  case LLVMTypeNew::DOUBLE_TYPE:
+    return "double";
+  case LLVMTypeNew::FP128_TYPE:
+    return "fp128";
+  case LLVMTypeNew::X86_FP80_TYPE:
+    return "x86_fp80";
+  case LLVMTypeNew::PPC_FP128_TYPE:
+    return "ppc_fp128";
+  case LLVMTypeNew::X86_MMX_TYPE:
+    return "x86_mmx";
+  case LLVMTypeNew::TOKEN_TYPE:
+    return "token";
+  case LLVMTypeNew::LABEL_TYPE:
+    return "label";
+  case LLVMTypeNew::METADATA_TYPE:
+    return "metadata";
+  case LLVMTypeNew::FUNCTION_TYPE:
+    return "func";
+  case LLVMTypeNew::INTEGER_TYPE:
+    return "i";
+  case LLVMTypeNew::POINTER_TYPE:
+    return "ptr";
+  case LLVMTypeNew::FIXED_VECTOR_TYPE:
+  case LLVMTypeNew::SCALABLE_VECTOR_TYPE:
+    return "vec";
+  case LLVMTypeNew::ARRAY_TYPE:
+    return "array";
+  case LLVMTypeNew::STRUCT_TYPE:
+    return "struct";
+  }
+  llvm_unreachable("unhandled type kind");
+}
+
+static void printLLVMStructTypeBody(llvm::raw_ostream &os, LLVMStructType type,
+                                    llvm::SetVector<StringRef> &stack) {
+  if (type.isPacked())
+    os << "packed ";
+
+  os << '(';
+  if (type.isIdentified())
+    stack.insert(type.identifier());
+  llvm::interleaveComma(type.getBody(), os, [&](Type subtype) {
+    printLLVMTypeImpl(os, subtype.cast<LLVMTypeNew>(), stack);
+  });
+  if (type.isIdentified())
+    stack.pop_back();
+  os << ')';
+}
+
+void printLLVMStructTypeImpl(llvm::raw_ostream &os, LLVMStructType type,
+                             llvm::SetVector<StringRef> &stack) {
+  os << "<";
+  if (type.isIdentified()) {
+    os << '"' << type.identifier() << '"';
+    // If we are printing a reference to one of the enclosing structs, just
+    // print the name and stop to avoid infinitely long output.
+    if (stack.count(type.identifier())) {
+      os << '>';
+      return;
+    }
+    os << ", ";
+
+    if (type.isOpaque()) {
+      os << "opaque>";
+      return;
+    }
+  }
+
+  printLLVMStructTypeBody(os, type, stack);
+  os << '>';
+}
+
+template <typename TypeTy>
+static void printArrayOrVectorType(llvm::raw_ostream &os, TypeTy type,
+                                   llvm::SetVector<StringRef> &stack) {
+  os << '<' << type.getNumElements() << " x ";
+  printLLVMTypeImpl(os, type.getElementType(), stack);
+  os << '>';
+}
+
+static void printFunctionType(llvm::raw_ostream &os, LLVMFunctionType funcType,
+                              llvm::SetVector<StringRef> &stack) {
+  os << '<';
+  printLLVMTypeImpl(os, funcType.getReturnType(), stack);
+  os << " (";
+  llvm::interleaveComma(funcType.getParams(), os,
+                        [&os, &stack](LLVMTypeNew subtype) {
+                          printLLVMTypeImpl(os, subtype, stack);
+                        });
+  if (funcType.isVarArg()) {
+    if (funcType.getNumParams() != 0)
+      os << ", ";
+    os << "...";
+  }
+  os << ")>";
+}
+
+void printLLVMTypeImpl(llvm::raw_ostream &os, LLVMTypeNew type,
+                       llvm::SetVector<StringRef> &stack) {
+  if (!type) {
+    os << "<<NULL-TYPE>>";
+    return;
+  }
+
+  unsigned kind = type.getKind();
+  os << getTypeKeyword(type);
+
+  // Trivial types only consist of their keyword.
+  if (LLVMTypeNew::FIRST_TRIVIAL_TYPE <= kind &&
+      kind <= LLVMTypeNew::LAST_TRIVIAL_TYPE)
+    return;
+
+  if (auto intType = type.dyn_cast<LLVMIntegerType>()) {
+    os << intType.getBitWidth();
+    return;
+  }
+
+  if (auto ptrType = type.dyn_cast<LLVMPointerType>()) {
+    os << '<';
+    printLLVMTypeImpl(os, ptrType.getElementType(), stack);
+    if (ptrType.getAddressSpace() != 0)
+      os << ", " << ptrType.getAddressSpace();
+    os << '>';
+    return;
+  }
+
+  if (auto arrayType = type.dyn_cast<LLVMArrayType>())
+    return printArrayOrVectorType(os, arrayType, stack);
+  if (auto vectorType = type.dyn_cast<LLVMFixedVectorType>())
+    return printArrayOrVectorType(os, vectorType, stack);
+
+  if (auto vectorType = type.dyn_cast<LLVMScalableVectorType>()) {
+    os << "<? x " << vectorType.getMinNumElements() << " x ";
+    printLLVMTypeImpl(os, vectorType.getElementType(), stack);
+    os << '>';
+    return;
+  }
+
+  if (auto structType = type.dyn_cast<LLVMStructType>())
+    return printLLVMStructTypeImpl(os, structType, stack);
+
+  printFunctionType(os, type.cast<LLVMFunctionType>(), stack);
+}
+
+void mlir::LLVM::detail::printType(LLVMTypeNew type,
+                                   DialectAsmPrinter &printer) {
+  llvm::SetVector<StringRef> stack;
+  return printLLVMTypeImpl(printer.getStream(), type, stack);
+}
+
+//
+// Parsing.
+//
+
+static LLVMTypeNew parseLLVMTypeImpl(DialectAsmParser &parser,
+                                     llvm::SetVector<StringRef> &stack);
+
+static ParseResult parseLLVMTypeImpl(DialectAsmParser &parser,
+                                     llvm::SetVector<StringRef> &stack,
+                                     LLVMTypeNew &result) {
+  result = parseLLVMTypeImpl(parser, stack);
+  return success(result != nullptr);
+}
+
+static LLVMFunctionType parseFunctionType(DialectAsmParser &parser,
+                                          llvm::SetVector<StringRef> &stack) {
+  LLVMTypeNew returnType;
+  if (parser.parseLess() || parseLLVMTypeImpl(parser, stack, returnType) ||
+      parser.parseLParen())
+    return LLVMFunctionType();
+
+  // Function type without arguments.
+  if (succeeded(parser.parseOptionalRParen())) {
+    if (succeeded(parser.parseGreater()))
+      return LLVMFunctionType::get(returnType, {}, /*isVarArg=*/false);
+    return LLVMFunctionType();
+  }
+
+  SmallVector<LLVMTypeNew, 8> argTypes;
+  do {
+    if (succeeded(parser.parseOptionalEllipsis())) {
+      if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
+        return LLVMFunctionType();
+      return LLVMFunctionType::get(returnType, argTypes, /*isVarArg=*/true);
+    }
+
+    argTypes.push_back(parseLLVMTypeImpl(parser, stack));
+    if (!argTypes.back())
+      return LLVMFunctionType();
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
+    return LLVMFunctionType();
+  return LLVMFunctionType::get(returnType, argTypes, /*isVarArg=*/false);
+}
+
+static LLVMPointerType parsePointerType(DialectAsmParser &parser,
+                                        llvm::SetVector<StringRef> &stack) {
+  LLVMTypeNew elementType;
+  if (parser.parseLess() || parseLLVMTypeImpl(parser, stack, elementType))
+    return LLVMPointerType();
+
+  unsigned addressSpace = 0;
+  if (succeeded(parser.parseOptionalComma()) &&
+      failed(parser.parseInteger(addressSpace)))
+    return LLVMPointerType();
+  if (failed(parser.parseGreater()))
+    return LLVMPointerType();
+  return LLVMPointerType::get(elementType, addressSpace);
+}
+
+static LLVMVectorType parseVectorType(DialectAsmParser &parser,
+                                      llvm::SetVector<StringRef> &stack) {
+  SmallVector<int64_t, 2> dims;
+  llvm::SMLoc dimPos;
+  LLVMTypeNew elementType;
+  if (parser.parseLess() || parser.getCurrentLocation(&dimPos) ||
+      parser.parseDimensionList(dims, /*allowDynamic=*/true) ||
+      parseLLVMTypeImpl(parser, stack, elementType) || parser.parseGreater())
+    return LLVMVectorType();
+
+  if (dims.empty() || dims.size() > 2 || (dims.size() == 2 ^ dims[0] == -1)) {
+    parser.emitError(dimPos)
+        << "expected '? x <integer> x <type>' or '<integer> x <type>'";
+    return LLVMVectorType();
+  }
+
+  bool isScalable = dims.size() == 2;
+  return isScalable ? static_cast<LLVMVectorType>(
+                          LLVMScalableVectorType::get(elementType, dims[1]))
+                    : LLVMFixedVectorType::get(elementType, dims[0]);
+}
+
+static LLVMArrayType parseArrayType(DialectAsmParser &parser,
+                                    llvm::SetVector<StringRef> &stack) {
+  SmallVector<int64_t, 1> dims;
+  llvm::SMLoc sizePos;
+  LLVMTypeNew elementType;
+  if (parser.parseLess() || parser.getCurrentLocation(&sizePos) ||
+      parser.parseDimensionList(dims, /*allowDynamic=*/false) ||
+      parseLLVMTypeImpl(parser, stack, elementType) || parser.parseGreater())
+    return LLVMArrayType();
+
+  if (dims.size() != 1) {
+    parser.emitError(sizePos) << "expected ? x <type>";
+    return LLVMArrayType();
+  }
+
+  return LLVMArrayType::get(elementType, dims[0]);
+}
+
+static LLVMStructType trySetStructBody(LLVMStructType type,
+                                       ArrayRef<Type> subtypes, bool isPacked,
+                                       DialectAsmParser &parser,
+                                       llvm::SMLoc subtypesLoc,
+                                       llvm::SetVector<StringRef> &stack) {
+  if (succeeded(type.setBody(subtypes, isPacked)))
+    return type;
+
+  std::string currentBody;
+  llvm::raw_string_ostream currentBodyStream(currentBody);
+  printLLVMStructTypeBody(currentBodyStream, type, stack);
+  (parser.emitError(subtypesLoc)
+   << "identified type already used with a different body")
+          .attachNote()
+      << "existing body: " << currentBodyStream.str();
+  return LLVMStructType();
+}
+
+static LLVMStructType parseStructType(DialectAsmParser &parser,
+                                      llvm::SetVector<StringRef> &stack) {
+  MLIRContext *ctx = parser.getBuilder().getContext();
+
+  if (failed(parser.parseLess()))
+    return LLVMStructType();
+
+  StringRef name;
+  bool isIdentified = succeeded(parser.parseOptionalString(&name));
+  if (isIdentified) {
+    if (stack.count(name)) {
+      if (failed(parser.parseGreater()))
+        return LLVMStructType();
+      return LLVMStructType::getIdentified(ctx, name);
+    }
+    if (failed(parser.parseComma()))
+      return LLVMStructType();
+  }
+
+  llvm::SMLoc kwLoc = parser.getCurrentLocation();
+  if (succeeded(parser.parseOptionalKeyword("opaque"))) {
+    if (!isIdentified)
+      return parser.emitError(kwLoc, "only identified structs can be opaque"),
+             LLVMStructType();
+    if (failed(parser.parseGreater()))
+      return LLVMStructType();
+    return LLVMStructType::getOpaque(name, ctx);
+  }
+
+  bool isPacked = succeeded(parser.parseOptionalKeyword("packed"));
+  if (failed(parser.parseLParen()))
+    return LLVMStructType();
+
+  if (succeeded(parser.parseOptionalRParen())) {
+    if (failed(parser.parseGreater()))
+      return LLVMStructType();
+    if (!isIdentified)
+      return LLVMStructType::getLiteral(ctx, {}, isPacked);
+    auto type = LLVMStructType::getIdentified(ctx, name);
+    return trySetStructBody(type, {}, isPacked, parser, kwLoc, stack);
+  }
+
+  SmallVector<Type, 4> subtypes;
+  llvm::SMLoc subtypesLoc = parser.getCurrentLocation();
+  do {
+    if (isIdentified)
+      stack.insert(name);
+    Type type = parseLLVMTypeImpl(parser, stack);
+    if (!type)
+      return LLVMStructType();
+    subtypes.push_back(type);
+    if (isIdentified)
+      stack.pop_back();
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (parser.parseRParen() || parser.parseGreater())
+    return LLVMStructType();
+
+  if (!isIdentified)
+    return LLVMStructType::getLiteral(ctx, subtypes, isPacked);
+
+  auto type = LLVMStructType::getIdentified(ctx, name);
+  return trySetStructBody(type, subtypes, isPacked, parser, subtypesLoc, stack);
+}
+
+static LLVMTypeNew parseLLVMTypeImpl(DialectAsmParser &parser,
+                                     llvm::SetVector<StringRef> &stack) {
+  // Special case for integers (i[1-9][0-9]*) that are literals rather than
+  // keywords for the parser, so they are not caught by the main dispatch below.
+  // Try parsing it a built-in integer type instead.
+  Type maybeIntegerType;
+  MLIRContext *ctx = parser.getBuilder().getContext();
+  llvm::SMLoc keyLoc = parser.getCurrentLocation();
+  OptionalParseResult result = parser.parseOptionalType(maybeIntegerType);
+  if (result.hasValue()) {
+    if (failed(*result))
+      return LLVMTypeNew();
+
+    if (!maybeIntegerType.isSignlessInteger()) {
+      parser.emitError(keyLoc) << "unexpected type, expected i* or keyword";
+      return LLVMTypeNew();
+    }
+    return LLVMIntegerType::get(ctx, maybeIntegerType.getIntOrFloatBitWidth());
+  }
+
+  StringRef key;
+  if (failed(parser.parseKeyword(&key)))
+    return LLVMTypeNew();
+
+  return llvm::StringSwitch<function_ref<LLVMTypeNew()>>(key)
+      .Case("void", [&] { return LLVMVoidType::get(ctx); })
+      .Case("half", [&] { return LLVMHalfType::get(ctx); })
+      .Case("bfloat", [&] { return LLVMBFloatType::get(ctx); })
+      .Case("float", [&] { return LLVMFloatType::get(ctx); })
+      .Case("double", [&] { return LLVMDoubleType::get(ctx); })
+      .Case("fp128", [&] { return LLVMFP128Type::get(ctx); })
+      .Case("x86_fp80", [&] { return LLVMX86FP80Type::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, stack); })
+      .Case("ptr", [&] { return parsePointerType(parser, stack); })
+      .Case("vec", [&] { return parseVectorType(parser, stack); })
+      .Case("array", [&] { return parseArrayType(parser, stack); })
+      .Case("struct", [&] { return parseStructType(parser, stack); })
+      .Default([&] {
+        parser.emitError(keyLoc) << "unknown LLVM type: " << key;
+        return LLVMTypeNew();
+      })();
+}
+
+LLVMTypeNew mlir::LLVM::detail::parseType(DialectAsmParser &parser) {
+  llvm::SetVector<StringRef> stack;
+  return parseLLVMTypeImpl(parser, stack);
+}
+
+//
+// Vector types.
+//
+
+LLVMTypeNew LLVMVectorType::getElementType() {
+  // Both derived classes share the implementation type.
+  return static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->elementType;
+}
+
+llvm::ElementCount LLVMVectorType::getElementCount() {
+  // Both derived classes share the implementation type.
+  return llvm::ElementCount(
+      static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->numElements,
+      this->isa<LLVMScalableVectorType>());
+}
+
+LLVMFixedVectorType LLVMFixedVectorType::get(LLVMTypeNew elementType,
+                                             unsigned numElements) {
+  assert(elementType && "expected non-null subtype");
+  return Base::get(elementType.getContext(), LLVMTypeNew::FIXED_VECTOR_TYPE,
+                   elementType, numElements)
+      .cast<LLVMFixedVectorType>();
+}
+
+unsigned LLVMFixedVectorType::getNumElements() {
+  return getImpl()->numElements;
+}
+
+LLVMScalableVectorType LLVMScalableVectorType::get(LLVMTypeNew elementType,
+                                                   unsigned minNumElements) {
+  assert(elementType && "expected non-null subtype");
+  return Base::get(elementType.getContext(), LLVMTypeNew::SCALABLE_VECTOR_TYPE,
+                   elementType, minNumElements)
+      .cast<LLVMScalableVectorType>();
+}
+
+unsigned LLVMScalableVectorType::getMinNumElements() {
+  return getImpl()->numElements;
+}
+
+namespace mlir {
+namespace LLVM {
+namespace detail {
 struct LLVMTypeStorage : public ::mlir::TypeStorage {
   LLVMTypeStorage(llvm::Type *ty) : underlyingType(ty) {}
 
   // LLVM types are pointer-unique.
   using KeyTy = llvm::Type *;
   bool operator==(const KeyTy &key) const { return key == underlyingType; }
 
   static LLVMTypeStorage *construct(TypeStorageAllocator &allocator,

mlir/lib/Parser/DialectSymbolParser.cpp

ParseResult parseOptionalQuestion() override {
     return success(parser.consumeIf(Token::question));
   }
 
   /// Parses a '*' if present.
   ParseResult parseOptionalStar() override {
     return success(parser.consumeIf(Token::star));
   }
 
+  /// Parses a quoted string token if present.
+  ParseResult parseOptionalString(StringRef *string) override {
+    if (!parser.getToken().is(Token::string))
+      return failure();
+
+    if (string)
+      *string = parser.getTokenSpelling().drop_front().drop_back();
+    parser.consumeToken();
+    return success();
+  }
+
   /// Returns if the current token corresponds to a keyword.
   bool isCurrentTokenAKeyword() const {
     return parser.getToken().is(Token::bare_identifier) ||
            parser.getToken().isKeyword();
   }
 
   /// Parse the given keyword if present.
   ParseResult parseOptionalKeyword(StringRef keyword) override {
ParseResult parseType(Type &result) override {
     return success(static_cast<bool>(result));
   }
 
   ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
                                  bool allowDynamic) override {
     return parser.parseDimensionListRanked(dimensions, allowDynamic);
   }
 
+  OptionalParseResult parseOptionalType(Type &result) override {
+    return parser.parseOptionalType(result);
+  }
+
 private:
   /// The full symbol specification.
   StringRef fullSpec;
 
   /// The source location of the dialect symbol.
   SMLoc nameLoc;
 
   /// The main parser.

mlir/test/Dialect/LLVMIR/types-invalid.mlir

+// RUN: mlir-opt --allow-unregistered-dialect -split-input-file -verify-diagnostics %s
+
+func @repeated_struct_name() {
+  "some.op"() : () -> !llvm2.struct<"a", (ptr<struct<"a">>)>
+  // expected-error @+2 {{identified type already used with a different body}}
+  // expected-note @+1 {{existing body: (ptr<struct<"a">>)}}
+  "some.op"() : () -> !llvm2.struct<"a", (i32)>
+}
+
+// -----
+
+func @repeated_struct_name_packed() {
+  "some.op"() : () -> !llvm2.struct<"a", packed (i32)>
+  // expected-error @+2 {{identified type already used with a different body}}
+  // expected-note @+1 {{existing body: packed (i32)}}
+  "some.op"() : () -> !llvm2.struct<"a", (i32)>
+}
+
+// -----
+
+func @unexpected_type() {
+  // expected-error @+1 {{unexpected type, expected i* or keyword}}
+  "some.op"() : () -> !llvm2.f32
+}
+
+// -----
+
+func @unexpected_type() {
+  // expected-error @+1 {{unknown LLVM type}}
+  "some.op"() : () -> !llvm2.ifoo
+}

mlir/test/Dialect/LLVMIR/types.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file | mlir-opt -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: @primitive
+func @primitive() {
+  // CHECK: !llvm2.void
+  "some.op"() : () -> !llvm2.void
+  // CHECK: !llvm2.half
+  "some.op"() : () -> !llvm2.half
+  // CHECK: !llvm2.bfloat
+  "some.op"() : () -> !llvm2.bfloat
+  // CHECK: !llvm2.float
+  "some.op"() : () -> !llvm2.float
+  // CHECK: !llvm2.double
+  "some.op"() : () -> !llvm2.double
+  // CHECK: !llvm2.fp128
+  "some.op"() : () -> !llvm2.fp128
+  // CHECK: !llvm2.x86_fp80
+  "some.op"() : () -> !llvm2.x86_fp80
+  // CHECK: !llvm2.ppc_fp128
+  "some.op"() : () -> !llvm2.ppc_fp128
+  // CHECK: !llvm2.x86_mmx
+  "some.op"() : () -> !llvm2.x86_mmx
+  // CHECK: !llvm2.token
+  "some.op"() : () -> !llvm2.token
+  // CHECK: !llvm2.label
+  "some.op"() : () -> !llvm2.label
+  // CHECK: !llvm2.metadata
+  "some.op"() : () -> !llvm2.metadata
+  return
+}
+
+// CHECK-LABEL: @func
+func @func() {
+  // CHECK: !llvm2.func<void ()>
+  "some.op"() : () -> !llvm2.func<void ()>
+  // CHECK: !llvm2.func<void (i32)>
+  "some.op"() : () -> !llvm2.func<void (i32)>
+  // CHECK: !llvm2.func<i32 ()>
+  "some.op"() : () -> !llvm2.func<i32 ()>
+  // CHECK: !llvm2.func<i32 (half, bfloat, float, double)>
+  "some.op"() : () -> !llvm2.func<i32 (half, bfloat, float, double)>
+  // CHECK: !llvm2.func<i32 (i32, i32)>
+  "some.op"() : () -> !llvm2.func<i32 (i32, i32)>
+  // CHECK: !llvm2.func<void (...)>
+  "some.op"() : () -> !llvm2.func<void (...)>
+  // CHECK: !llvm2.func<void (i32, i32, ...)>
+  "some.op"() : () -> !llvm2.func<void (i32, i32, ...)>
+  return
+}
+
+// CHECK-LABEL: @integer
+func @integer() {
+  // CHECK: !llvm2.i1
+  "some.op"() : () -> !llvm2.i1
+  // CHECK: !llvm2.i8
+  "some.op"() : () -> !llvm2.i8
+  // CHECK: !llvm2.i16
+  "some.op"() : () -> !llvm2.i16
+  // CHECK: !llvm2.i32
+  "some.op"() : () -> !llvm2.i32
+  // CHECK: !llvm2.i64
+  "some.op"() : () -> !llvm2.i64
+  // CHECK: !llvm2.i57
+  "some.op"() : () -> !llvm2.i57
+  // CHECK: !llvm2.i129
+  "some.op"() : () -> !llvm2.i129
+  return
+}
+
+// CHECK-LABEL: @ptr
+func @ptr() {
+  // CHECK: !llvm2.ptr<i8>
+  "some.op"() : () -> !llvm2.ptr<i8>
+  // CHECK: !llvm2.ptr<float>
+  "some.op"() : () -> !llvm2.ptr<float>
+  // CHECK: !llvm2.ptr<ptr<i8>>
+  "some.op"() : () -> !llvm2.ptr<ptr<i8>>
+  // CHECK: !llvm2.ptr<ptr<ptr<ptr<ptr<i8>>>>>
+  "some.op"() : () -> !llvm2.ptr<ptr<ptr<ptr<ptr<i8>>>>>
+  // CHECK: !llvm2.ptr<i8>
+  "some.op"() : () -> !llvm2.ptr<i8, 0>
+  // CHECK: !llvm2.ptr<i8, 1>
+  "some.op"() : () -> !llvm2.ptr<i8, 1>
+  // CHECK: !llvm2.ptr<i8, 42>
+  "some.op"() : () -> !llvm2.ptr<i8, 42>
+  // CHECK: !llvm2.ptr<ptr<i8, 42>, 9>
+  "some.op"() : () -> !llvm2.ptr<ptr<i8, 42>, 9>
+  return
+}
+
+// CHECK-LABEL: @vec
+func @vec() {
+  // CHECK: !llvm2.vec<4 x i32>
+  "some.op"() : () -> !llvm2.vec<4 x i32>
+  // CHECK: !llvm2.vec<4 x float>
+  "some.op"() : () -> !llvm2.vec<4 x float>
+  // CHECK: !llvm2.vec<? x 4 x i32>
+  "some.op"() : () -> !llvm2.vec<? x 4 x i32>
+  // CHECK: !llvm2.vec<? x 8 x half>
+  "some.op"() : () -> !llvm2.vec<? x 8 x half>
+  // CHECK: !llvm2.vec<4 x ptr<i8>>
+  "some.op"() : () -> !llvm2.vec<4 x ptr<i8>>
+  return
+}
+
+// CHECK-LABEL: @array
+func @array() {
+  // CHECK: !llvm2.array<10 x i32>
+  "some.op"() : () -> !llvm2.array<10 x i32>
+  // CHECK: !llvm2.array<8 x float>
+  "some.op"() : () -> !llvm2.array<8 x float>
+  // CHECK: !llvm2.array<10 x ptr<i32, 4>>
+  "some.op"() : () -> !llvm2.array<10 x ptr<i32, 4>>
+  // CHECK: !llvm2.array<10 x array<4 x float>>
+  "some.op"() : () -> !llvm2.array<10 x array<4 x float>>
+  return
+}
+
+// CHECK-LABEL: @literal_struct
+func @literal_struct() {
+  // CHECK: !llvm2.struct<()>
+  "some.op"() : () -> !llvm2.struct<()>
+  // CHECK: !llvm2.struct<(i32)>
+  "some.op"() : () -> !llvm2.struct<(i32)>
+  // CHECK: !llvm2.struct<(float, i32)>
+  "some.op"() : () -> !llvm2.struct<(float, i32)>
+  // CHECK: !llvm2.struct<(struct<(i32)>)>
+  "some.op"() : () -> !llvm2.struct<(struct<(i32)>)>
+  // CHECK: !llvm2.struct<(i32, struct<(i32)>, float)>
+  "some.op"() : () -> !llvm2.struct<(i32, struct<(i32)>, float)>
+
+  // CHECK: !llvm2.struct<packed ()>
+  "some.op"() : () -> !llvm2.struct<packed ()>
+  // CHECK: !llvm2.struct<packed (i32)>
+  "some.op"() : () -> !llvm2.struct<packed (i32)>
+  // CHECK: !llvm2.struct<packed (float, i32)>
+  "some.op"() : () -> !llvm2.struct<packed (float, i32)>
+  // CHECK: !llvm2.struct<packed (float, i32)>
+  "some.op"() : () -> !llvm2.struct<packed (float, i32)>
+  // CHECK: !llvm2.struct<packed (struct<(i32)>)>
+  "some.op"() : () -> !llvm2.struct<packed (struct<(i32)>)>
+  // CHECK: !llvm2.struct<packed (i32, struct<(i32, i1)>, float)>
+  "some.op"() : () -> !llvm2.struct<packed (i32, struct<(i32, i1)>, float)>
+
+  // CHECK: !llvm2.struct<(struct<packed (i32)>)>
+  "some.op"() : () -> !llvm2.struct<(struct<packed (i32)>)>
+  // CHECK: !llvm2.struct<packed (struct<(i32)>)>
+  "some.op"() : () -> !llvm2.struct<packed (struct<(i32)>)>
+  return
+}
+
+// CHECK-LABEL: @identified_struct
+func @identified_struct() {
+  // CHECK: !llvm2.struct<"empty", ()>
+  "some.op"() : () -> !llvm2.struct<"empty", ()>
+  // CHECK: !llvm2.struct<"opaque", opaque>
+  "some.op"() : () -> !llvm2.struct<"opaque", opaque>
+  // CHECK: !llvm2.struct<"long", (i32, struct<(i32, i1)>, float, ptr<func<void ()>>)>
+  "some.op"() : () -> !llvm2.struct<"long", (i32, struct<(i32, i1)>, float, ptr<func<void ()>>)>
+  // CHECK: !llvm2.struct<"self-recursive", (ptr<struct<"self-recursive">>)>
+  "some.op"() : () -> !llvm2.struct<"self-recursive", (ptr<struct<"self-recursive">>)>
+  // CHECK: !llvm2.struct<"unpacked", (i32)>
+  "some.op"() : () -> !llvm2.struct<"unpacked", (i32)>
+  // CHECK: !llvm2.struct<"packed", packed (i32)>
+  "some.op"() : () -> !llvm2.struct<"packed", packed (i32)>
+  // CHECK: !llvm2.struct<"name with spaces and !^$@$#", packed (i32)>
+  "some.op"() : () -> !llvm2.struct<"name with spaces and !^$@$#", packed (i32)>
+
+  // CHECK: !llvm2.struct<"mutually-a", (ptr<struct<"mutually-b", (ptr<struct<"mutually-a">, 3>)>>)>
+  "some.op"() : () -> !llvm2.struct<"mutually-a", (ptr<struct<"mutually-b", (ptr<struct<"mutually-a">, 3>)>>)>
+  // CHECK: !llvm2.struct<"mutually-b", (ptr<struct<"mutually-a", (ptr<struct<"mutually-b">>)>, 3>)>
+  "some.op"() : () -> !llvm2.struct<"mutually-b", (ptr<struct<"mutually-a", (ptr<struct<"mutually-b">>)>, 3>)>
+  // CHECK: !llvm2.struct<"referring-another", (ptr<struct<"unpacked", (i32)>>)>
+  "some.op"() : () -> !llvm2.struct<"referring-another", (ptr<struct<"unpacked", (i32)>>)>
+
+  // CHECK: !llvm2.struct<"struct-of-arrays", (array<10 x i32>)>
+  "some.op"() : () -> !llvm2.struct<"struct-of-arrays", (array<10 x i32>)>
+  // CHECK: !llvm2.array<10 x struct<"array-of-structs", (i32)>>
+  "some.op"() : () -> !llvm2.array<10 x struct<"array-of-structs", (i32)>>
+  // CHECK: !llvm2.ptr<struct<"ptr-to-struct", (i8)>>
+  "some.op"() : () -> !llvm2.ptr<struct<"ptr-to-struct", (i8)>>
+  return
+}
+