[mlir] introduce type constraints for operands of LLVM dialect operations

[ftynse]

Summary

Historically, the operations in the MLIR's LLVM dialect only checked that the
operand are of LLVM dialect type without more detailed constraints. This was
due to LLVM dialect types wrapping LLVM IR types and having clunky verification
methods. With the new first-class modeling, it is possible to define type
constraints similarly to other dialects and use them to enforce some
correctness rules in verifiers instead of having LLVM assert during translation
to LLVM IR. This hardening discovered several issues where MLIR was producing
LLVM dialect operations that cannot exist in LLVM IR.

Depends On D85900

[ftynse]

Rebase

[ftynse]

Minor fix in bitwidth

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

267

nit: Align with the LLVM not the ins.

(This file seems inconsistent w.r.t this, so switch to whichever one is most common.)

449

Same here and throughout.

(Also unrelated, but we should switch away from using Arguments/Results and use the fields instead).

[ftynse]

Rebase.

[ftynse]

Formatting.

mlir/include/mlir/Dialect/GPU/GPUOps.td

 include "mlir/Dialect/GPU/GPUBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Type constraint accepting standard integers, indices and wrapped LLVM integer
 // types.
 def IntLikeOrLLVMInt : TypeConstraint<
-  Or<[AnySignlessInteger.predicate, Index.predicate, LLVMInt.predicate]>,
+  Or<[AnySignlessInteger.predicate, Index.predicate,
+      LLVM_AnyInteger.predicate]>,
   "integer, index or LLVM dialect equivalent">;
 
 //===----------------------------------------------------------------------===//
 // GPU Dialect operations.
 //===----------------------------------------------------------------------===//
 
 class GPU_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<GPU_Dialect, mnemonic, traits>;

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMIR_OP_BASE
 #define LLVMIR_OP_BASE
 
 include "mlir/IR/OpBase.td"
 
+//===----------------------------------------------------------------------===//
+// LLVM Dialect.
+//===----------------------------------------------------------------------===//
+
 def LLVM_Dialect : Dialect {
   let name = "llvm";
   let cppNamespace = "LLVM";
   let hasRegionArgAttrVerify = 1;
   let extraClassDeclaration = [{
     ~LLVMDialect();
     const llvm::DataLayout &getDataLayout();
 
   private:
     friend LLVMType;
 
     // This can't be a unique_ptr because the ctor is generated inline
     // in the class definition at the moment.
     detail::LLVMDialectImpl *impl;
   }];
 }
 
-// LLVM IR type wrapped in MLIR.
+//===----------------------------------------------------------------------===//
+// LLVM dialect type constraints.
+//===----------------------------------------------------------------------===//
+
+// LLVM dialect type.
 def LLVM_Type : DialectType<LLVM_Dialect,
                             CPred<"$_self.isa<::mlir::LLVM::LLVMType>()">,
                             "LLVM dialect type">;
 
-// Type constraint accepting only wrapped LLVM integer types.
-def LLVMInt : TypeConstraint<
-  And<[LLVM_Type.predicate,
-       CPred<"$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy()">]>,
-  "LLVM dialect integer">;
-
-def LLVMIntBase : TypeConstraint<
-  And<[LLVM_Type.predicate,
-       CPred<"$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy()">]>,
-  "LLVM dialect integer">;
-
-// Integer type of a specific width.
-class LLVMI<int width>
-    : Type<And<[
-        LLVM_Type.predicate,
-        CPred<
-         "$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy(" # width # ")">]>,
-         "LLVM dialect " # width # "-bit integer">,
-      BuildableType<
-        "::mlir::LLVM::LLVMType::getIntNTy($_builder.getContext(),"
-         # width # ")">;
+// Type constraint accepting LLVM integer types.
+def LLVM_AnyInteger : Type<
+  CPred<"$_self.isa<::mlir::LLVM::LLVMIntegerType>()">,
+  "LLVM integer type">;
+
+// Type constraints accepting LLVM integer type of a specific width.
+class LLVM_IntBase<int width> :
+    Type<And<[
+        LLVM_AnyInteger.predicate,
+        CPred<"$_self.cast<::mlir::LLVM::LLVMIntegerType>().getBitWidth() == "
+              # width>]>,
+        "LLVM " # width # "-bit integer type">,
+        "::mlir::LLVM::LLVMIntegerType::get($_builder.getContext(), "
 
-def LLVMI1 : LLVMI<1>;
+def LLVM_i1 : LLVM_IntBase<1>;
+def LLVM_i8 : LLVM_IntBase<1>;
+def LLVM_i32 : LLVM_IntBase<32>;
+
+// Type constraint accepting LLVM primitive types, i.e. all types except void
+// and function.
+def LLVM_PrimitiveType : Type<
+  And<[LLVM_Type.predicate,
+       CPred<"!$_self.isa<::mlir::LLVM::LLVMVoidType, "
+                         "::mlir::LLVM::LLVMFunctionType>()">]>,
+  "primitive LLVM type">;
+
+// Type constraint accepting any LLVM floating point type.
+def LLVM_AnyFloat : Type<
+  CPred<"$_self.isa<::mlir::LLVM::LLVMBFloatType, "
+                   "::mlir::LLVM::LLVMHalfType, "
+                   "::mlir::LLVM::LLVMFloatType, "
+                   "::mlir::LLVM::LLVMDoubleType>()">,
+  "floating point LLVM type">;
+
+// Type constraint accepting any LLVM pointer type.
+def LLVM_AnyPointer : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMPointerType>()">,
+                          "LLVM pointer type">;
+
+// Type constraint accepting LLVM pointer type with an additional constraint
+// on the element type.
+class LLVM_PointerTo<Type pointee> : Type<
+  And<[LLVM_AnyPointer.predicate,
+       SubstLeaves<
+         "$_self",
+         "$_self.cast<::mlir::LLVM::LLVMPointerType>().getElementType()",
+         pointee.predicate>]>,
+  "LLVM pointer to " # pointee.description>;
+
+// Type constraint accepting any LLVM structure type.
+def LLVM_AnyStruct : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMStructType>()">,
+                         "LLVM structure type">;
+
+// Type constraint accepting opaque LLVM structure type.
+def LLVM_OpaqueStruct : Type<
+  And<[LLVM_AnyStruct.predicate,
+       CPred<"$_self.cast<::mlir::LLVM::LLVMStructType>().isOpaque()">]>>;
+
+// Type constraint accepting any LLVM type that can be loaded or stored, i.e. a
+// type that has size (not void, function or opaque struct type).
+def LLVM_LoadableType : Type<
+  And<[LLVM_PrimitiveType.predicate, Neg<LLVM_OpaqueStruct.predicate>]>,
+  "LLVM type with size">;
+
+// Type constraint accepting any LLVM aggregate type, i.e. structure or array.
+def LLVM_AnyAggregate : Type<
+  CPred<"$_self.isa<::mlir::LLVM::LLVMStructType, "
+                   "::mlir::LLVM::LLVMArrayType>()">,
+  "LLVM aggregate type">;
+
+// Type constraint accepting any LLVM non-aggregate type, i.e. not structure or
+// array.
+def LLVM_AnyNonAggregate : Type<Neg<LLVM_AnyAggregate.predicate>,
+                               "LLVM non-aggregate type">;
+
+// Type constraint accepting any LLVM vector type.
+def LLVM_AnyVector : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMVectorType>()">,
+                         "LLVM vector type">;
+
+// Type constraint accepting an LLVM vector type with an additional constraint
+// on the vector element type.
+class LLVM_VectorOf<Type element> : Type<
+  And<[LLVM_AnyVector.predicate,
+       SubstLeaves<
+         "$_self",
+         "$_self.cast<::mlir::LLVM::LLVMVectorType>().getElementType()",
+         element.predicate>]>,
+  "LLVM vector of " # element.description>;
+
+// Type constraint accepting a constrained type, or a vector of such types.
+class LLVM_ScalarOrVectorOf<Type element> :
+    AnyTypeOf<[element, LLVM_VectorOf<element>]>;
 
 // Base class for LLVM operations. Defines the interface to the llvm::IRBuilder
 // used to translate to LLVM IR proper.
 class LLVM_OpBase<Dialect dialect, string mnemonic, list<OpTrait> traits = []> :
     Op<dialect, mnemonic, traits> {
   // A pattern for constructing the LLVM IR Instruction (or other Value) that
   // corresponds to this op.  This pattern can use `builder` to refer to an
   // `llvm::IRBuilder<>` instance, $-names of arguments and results and the
   // following special variable names:
   //   - $_resultType - substituted with the LLVM IR type of the result;
   //   - $_numOperands - substituted with the number of operands (including
   //                     the variadic ones);
   //   - $_hasResult - substituted with a check that a variadic-result op does
   //                   have a result (LLVM ops can have 0 or 1 result);
   //   - $_location - mlir::Location object of the instruction.
   // Additionally, `$$` can be used to produce the dollar character.
   string llvmBuilder = "";
 }
 
+//===----------------------------------------------------------------------===//
+// Base classes for LLVM dialect operations.
+//===----------------------------------------------------------------------===//
+
 // Base class for LLVM operations. All operations get an "llvm." prefix in
 // their name automatically. LLVM operations have either zero or one result,
 // this class is specialized below for both cases and should not be used
 // directly.
 class LLVM_Op<string mnemonic, list<OpTrait> traits = []> :
     LLVM_OpBase<LLVM_Dialect, mnemonic, traits>;
 
 // Case of the LLVM enum attribute backed by I64Attr with customized string

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_TerminatorPassthroughOpBuilder : OpBuilder<
   }]>;
 
 // Base class for LLVM terminator operations.  All terminator operations have
 // zero results and an optional list of successors.
 class LLVM_TerminatorOp<string mnemonic, list<OpTrait> traits = []> :
     LLVM_Op<mnemonic, !listconcat(traits, [Terminator])>;
 
 // Class for arithmetic binary operations.
-class LLVM_ArithmeticOp<string mnemonic, string builderFunc,
-                        list<OpTrait> traits = []> :
+class LLVM_ArithmeticOpBase<Type type, string mnemonic,
+                            string builderFunc, list<OpTrait> traits = []> :
     LLVM_OneResultOp<mnemonic,
            !listconcat([NoSideEffect, SameOperandsAndResultType], traits)>,
-    Arguments<(ins LLVM_Type:$lhs, LLVM_Type:$rhs)>,
+    Arguments<(ins LLVM_ScalarOrVectorOf<type>:$lhs,
+               LLVM_ScalarOrVectorOf<type>:$rhs)>,
     LLVM_Builder<"$res = builder." # builderFunc # "($lhs, $rhs);"> {
-  let parser = [{ return impl::parseOneResultSameOperandTypeOp(parser, result); }];
+  let parser =
+      [{ return impl::parseOneResultSameOperandTypeOp(parser, result); }];
   let printer = [{ mlir::impl::printOneResultOp(this->getOperation(), p); }];
 }
-class LLVM_UnaryArithmeticOp<string mnemonic, string builderFunc,
+class LLVM_IntArithmeticOp<string mnemonic, string builderFunc,
+                           list<OpTrait> traits = []> :
+    LLVM_ArithmeticOpBase<LLVM_AnyInteger, mnemonic, builderFunc, traits>;
+class LLVM_FloatArithmeticOp<string mnemonic, string builderFunc,
                              list<OpTrait> traits = []> :
+    LLVM_ArithmeticOpBase<LLVM_AnyFloat, mnemonic, builderFunc, traits>;
+
+// Class for arithmetic unary operations.
+class LLVM_UnaryArithmeticOp<Type type, string mnemonic,
+                             string builderFunc, list<OpTrait> traits = []> :
     LLVM_OneResultOp<mnemonic,
            !listconcat([NoSideEffect, SameOperandsAndResultType], traits)>,
-    Arguments<(ins LLVM_Type:$operand)>,
+    Arguments<(ins type:$operand)>,
     LLVM_Builder<"$res = builder." # builderFunc # "($operand);"> {
-  let parser = [{ return impl::parseOneResultSameOperandTypeOp(parser, result); }];
+  let parser =
+      [{ return impl::parseOneResultSameOperandTypeOp(parser, result); }];
   let printer = [{ mlir::impl::printOneResultOp(this->getOperation(), p); }];
 }
 
 // Integer binary operations.
-def LLVM_AddOp : LLVM_ArithmeticOp<"add", "CreateAdd", [Commutative]>;
-def LLVM_SubOp : LLVM_ArithmeticOp<"sub", "CreateSub">;
-def LLVM_MulOp : LLVM_ArithmeticOp<"mul", "CreateMul", [Commutative]>;
-def LLVM_UDivOp : LLVM_ArithmeticOp<"udiv", "CreateUDiv">;
-def LLVM_SDivOp : LLVM_ArithmeticOp<"sdiv", "CreateSDiv">;
-def LLVM_URemOp : LLVM_ArithmeticOp<"urem", "CreateURem">;
-def LLVM_SRemOp : LLVM_ArithmeticOp<"srem", "CreateSRem">;
-def LLVM_AndOp : LLVM_ArithmeticOp<"and", "CreateAnd">;
-def LLVM_OrOp : LLVM_ArithmeticOp<"or", "CreateOr">;
-def LLVM_XOrOp : LLVM_ArithmeticOp<"xor", "CreateXor">;
-def LLVM_ShlOp : LLVM_ArithmeticOp<"shl", "CreateShl">;
-def LLVM_LShrOp : LLVM_ArithmeticOp<"lshr", "CreateLShr">;
-def LLVM_AShrOp : LLVM_ArithmeticOp<"ashr", "CreateAShr">;
+def LLVM_AddOp : LLVM_IntArithmeticOp<"add", "CreateAdd", [Commutative]>;
+def LLVM_SubOp : LLVM_IntArithmeticOp<"sub", "CreateSub">;
+def LLVM_MulOp : LLVM_IntArithmeticOp<"mul", "CreateMul", [Commutative]>;
+def LLVM_UDivOp : LLVM_IntArithmeticOp<"udiv", "CreateUDiv">;
+def LLVM_SDivOp : LLVM_IntArithmeticOp<"sdiv", "CreateSDiv">;
+def LLVM_URemOp : LLVM_IntArithmeticOp<"urem", "CreateURem">;
+def LLVM_SRemOp : LLVM_IntArithmeticOp<"srem", "CreateSRem">;
+def LLVM_AndOp : LLVM_IntArithmeticOp<"and", "CreateAnd">;
+def LLVM_OrOp : LLVM_IntArithmeticOp<"or", "CreateOr">;
+def LLVM_XOrOp : LLVM_IntArithmeticOp<"xor", "CreateXor">;
+def LLVM_ShlOp : LLVM_IntArithmeticOp<"shl", "CreateShl">;
+def LLVM_LShrOp : LLVM_IntArithmeticOp<"lshr", "CreateLShr">;
+def LLVM_AShrOp : LLVM_IntArithmeticOp<"ashr", "CreateAShr">;
 
 // Predicate for integer comparisons.
 def ICmpPredicateEQ  : I64EnumAttrCase<"eq", 0>;
 def ICmpPredicateNE  : I64EnumAttrCase<"ne", 1>;
 def ICmpPredicateSLT : I64EnumAttrCase<"slt", 2>;
 def ICmpPredicateSLE : I64EnumAttrCase<"sle", 3>;
 def ICmpPredicateSGT : I64EnumAttrCase<"sgt", 4>;
 def ICmpPredicateSGE : I64EnumAttrCase<"sge", 5>;
 def ICmpPredicateULT : I64EnumAttrCase<"ult", 6>;
 def ICmpPredicateULE : I64EnumAttrCase<"ule", 7>;
 def ICmpPredicateUGT : I64EnumAttrCase<"ugt", 8>;
 def ICmpPredicateUGE : I64EnumAttrCase<"uge", 9>;
 def ICmpPredicate : I64EnumAttr<
     "ICmpPredicate",
     "llvm.icmp comparison predicate",
     [ICmpPredicateEQ, ICmpPredicateNE, ICmpPredicateSLT, ICmpPredicateSLE,
      ICmpPredicateSGT, ICmpPredicateSGE, ICmpPredicateULT, ICmpPredicateULE,
      ICmpPredicateUGT, ICmpPredicateUGE]> {
   let cppNamespace = "::mlir::LLVM";
 }
 
 // Other integer operations.
 def LLVM_ICmpOp : LLVM_OneResultOp<"icmp", [NoSideEffect]>,
-                  Arguments<(ins ICmpPredicate:$predicate, LLVM_Type:$lhs,
-                             LLVM_Type:$rhs)> {
+                  Arguments<(ins ICmpPredicate:$predicate,
+                             LLVM_ScalarOrVectorOf<LLVM_AnyInteger>:$lhs,
+                             LLVM_ScalarOrVectorOf<LLVM_AnyInteger>:$rhs)> {
   let llvmBuilder = [{
     $res = builder.CreateICmp(getLLVMCmpPredicate($predicate), $lhs, $rhs);
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, ICmpPredicate predicate, Value lhs, "
     "Value rhs", [{
       build(b, result, LLVMType::getInt1Ty(lhs.getType().getContext()),
             b.getI64IntegerAttr(static_cast<int64_t>(predicate)), lhs, rhs);
def FCmpPredicate : I64EnumAttr<
      FCmpPredicateUEQ, FCmpPredicateUGT, FCmpPredicateUGE, FCmpPredicateULT,
      FCmpPredicateULE, FCmpPredicateUNE, FCmpPredicateUNO, FCmpPredicateTRUE
     ]> {
   let cppNamespace = "::mlir::LLVM";
 }
 
 // Other integer operations.
 def LLVM_FCmpOp : LLVM_OneResultOp<"fcmp", [NoSideEffect]>,
-                  Arguments<(ins FCmpPredicate:$predicate, LLVM_Type:$lhs,
-                             LLVM_Type:$rhs)> {
+                  Arguments<(ins FCmpPredicate:$predicate,
+                             LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$lhs,
+                             LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$rhs)> {
   let llvmBuilder = [{
     $res = builder.CreateFCmp(getLLVMCmpPredicate($predicate), $lhs, $rhs);
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, FCmpPredicate predicate, Value lhs, "
     "Value rhs", [{
       build(b, result, LLVMType::getInt1Ty(lhs.getType().getContext()),
             b.getI64IntegerAttr(static_cast<int64_t>(predicate)), lhs, rhs);
     }]>];
   let parser = [{ return parseCmpOp<FCmpPredicate>(parser, result); }];
   let printer = [{ printFCmpOp(p, *this); }];
 }
 
 // Floating point binary operations.
-def LLVM_FAddOp : LLVM_ArithmeticOp<"fadd", "CreateFAdd">;
-def LLVM_FSubOp : LLVM_ArithmeticOp<"fsub", "CreateFSub">;
-def LLVM_FMulOp : LLVM_ArithmeticOp<"fmul", "CreateFMul">;
-def LLVM_FDivOp : LLVM_ArithmeticOp<"fdiv", "CreateFDiv">;
-def LLVM_FRemOp : LLVM_ArithmeticOp<"frem", "CreateFRem">;
-def LLVM_FNegOp : LLVM_UnaryArithmeticOp<"fneg", "CreateFNeg">;
+def LLVM_FAddOp : LLVM_FloatArithmeticOp<"fadd", "CreateFAdd">;
+def LLVM_FSubOp : LLVM_FloatArithmeticOp<"fsub", "CreateFSub">;
+def LLVM_FMulOp : LLVM_FloatArithmeticOp<"fmul", "CreateFMul">;
+def LLVM_FDivOp : LLVM_FloatArithmeticOp<"fdiv", "CreateFDiv">;
+def LLVM_FRemOp : LLVM_FloatArithmeticOp<"frem", "CreateFRem">;
+def LLVM_FNegOp : LLVM_UnaryArithmeticOp<LLVM_ScalarOrVectorOf<LLVM_AnyFloat>,
+                                         "fneg", "CreateFNeg">;
 
 // Common code definition that is used to verify and set the alignment attribute
 // of LLVM ops that accept such an attribute.
 class MemoryOpWithAlignmentBase {
   code alignmentVerifierCode = [{
     if (alignment().hasValue()) {
       auto align = alignment().getValue().getSExtValue();
       if (align < 0)
code setNonTemporalMetadataCode = [{
     }
   }];
 }
 
 // Memory-related operations.
 def LLVM_AllocaOp :
     MemoryOpWithAlignmentBase,
     LLVM_OneResultOp<"alloca">,
-    Arguments<(ins LLVM_Type:$arraySize, OptionalAttr<I64Attr>:$alignment)> {
+    Arguments<(ins LLVM_AnyInteger:$arraySize,
+               OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
     auto *inst = builder.CreateAlloca(
       $_resultType->getPointerElementType(), $arraySize);
     }] # setAlignmentCode # [{
     $res = inst;
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Type resultType, Value arraySize, "
     "unsigned alignment",
     [{
       if (alignment == 0)
         return build(b, result, resultType, arraySize, IntegerAttr());
       build(b, result, resultType, arraySize, b.getI64IntegerAttr(alignment));
   }]>];
   let parser = [{ return parseAllocaOp(parser, result); }];
   let printer = [{ printAllocaOp(p, *this); }];
   let verifier = alignmentVerifierCode;
 }
+
 def LLVM_GEPOp : LLVM_OneResultOp<"getelementptr", [NoSideEffect]>,
-                 Arguments<(ins LLVM_Type:$base, Variadic<LLVM_Type>:$indices)>,
+                 Arguments<(ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
+                            Variadic<LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>:$indices)>,
                  LLVM_Builder<"$res = builder.CreateGEP($base, $indices);"> {
   let assemblyFormat = [{
     $base `[` $indices `]` attr-dict `:` functional-type(operands, results)
   }];
 }
 def LLVM_LoadOp :
     MemoryOpWithAlignmentAndAttributes,
     LLVM_OneResultOp<"load">,
-    Arguments<(ins LLVM_Type:$addr,
+    Arguments<(ins LLVM_PointerTo<LLVM_LoadableType>:$addr,
                    OptionalAttr<I64Attr>:$alignment,
                    UnitAttr:$volatile_,
                    UnitAttr:$nontemporal)> {
   string llvmBuilder = [{
     auto *inst = builder.CreateLoad($addr, $volatile_);
   }] # setAlignmentCode # setNonTemporalMetadataCode # [{
     $res = inst;
   }];
let builders = [OpBuilder<
     "bool isNonTemporal = false">];
   let parser = [{ return parseLoadOp(parser, result); }];
   let printer = [{ printLoadOp(p, *this); }];
   let verifier = alignmentVerifierCode;
 }
 def LLVM_StoreOp :
     MemoryOpWithAlignmentAndAttributes,
     LLVM_ZeroResultOp<"store">,
-    Arguments<(ins LLVM_Type:$value,
-                   LLVM_Type:$addr,
+    Arguments<(ins LLVM_LoadableType:$value,
+                   LLVM_PointerTo<LLVM_LoadableType>:$addr,
                    OptionalAttr<I64Attr>:$alignment,
                    UnitAttr:$volatile_,
                    UnitAttr:$nontemporal)> {
   string llvmBuilder = [{
     auto *inst = builder.CreateStore($value, $addr, $volatile_);
   }] # setAlignmentCode # setNonTemporalMetadataCode;
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Value value, Value addr, "
     "unsigned alignment = 0, bool isVolatile = false, "
     "bool isNonTemporal = false">
     ];
   let parser = [{ return parseStoreOp(parser, result); }];
   let printer = [{ printStoreOp(p, *this); }];
   let verifier = alignmentVerifierCode;
 }
 
 // Casts.
-class LLVM_CastOp<string mnemonic, string builderFunc,
+class LLVM_CastOp<string mnemonic, string builderFunc, Type type,
                   list<OpTrait> traits = []> :
     LLVM_OneResultOp<mnemonic,
            !listconcat([NoSideEffect], traits)>,
-    Arguments<(ins LLVM_Type:$arg)>,
+    Arguments<(ins type:$arg)>,
     LLVM_Builder<"$res = builder." # builderFunc # "($arg, $_resultType);"> {
   let parser = [{ return mlir::impl::parseCastOp(parser, result); }];
   let printer = [{ mlir::impl::printCastOp(this->getOperation(), p); }];
 }
-def LLVM_BitcastOp : LLVM_CastOp<"bitcast", "CreateBitCast">;
-def LLVM_AddrSpaceCastOp : LLVM_CastOp<"addrspacecast", "CreateAddrSpaceCast">;
-def LLVM_IntToPtrOp : LLVM_CastOp<"inttoptr", "CreateIntToPtr">;
-def LLVM_PtrToIntOp : LLVM_CastOp<"ptrtoint", "CreatePtrToInt">;
-def LLVM_SExtOp : LLVM_CastOp<"sext", "CreateSExt">;
-def LLVM_ZExtOp : LLVM_CastOp<"zext", "CreateZExt">;
-def LLVM_TruncOp : LLVM_CastOp<"trunc", "CreateTrunc">;
-def LLVM_SIToFPOp : LLVM_CastOp<"sitofp", "CreateSIToFP">;
-def LLVM_UIToFPOp : LLVM_CastOp<"uitofp", "CreateUIToFP">;
-def LLVM_FPToSIOp : LLVM_CastOp<"fptosi", "CreateFPToSI">;
-def LLVM_FPToUIOp : LLVM_CastOp<"fptoui", "CreateFPToUI">;
-def LLVM_FPExtOp : LLVM_CastOp<"fpext", "CreateFPExt">;
-def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "CreateFPTrunc">;
+def LLVM_BitcastOp : LLVM_CastOp<"bitcast", "CreateBitCast",
+                                 LLVM_AnyNonAggregate>;
+def LLVM_AddrSpaceCastOp : LLVM_CastOp<"addrspacecast", "CreateAddrSpaceCast",
+                                       LLVM_ScalarOrVectorOf<LLVM_AnyPointer>>;
+def LLVM_IntToPtrOp : LLVM_CastOp<"inttoptr", "CreateIntToPtr",
+                                  LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_PtrToIntOp : LLVM_CastOp<"ptrtoint", "CreatePtrToInt",
+                                  LLVM_ScalarOrVectorOf<LLVM_AnyPointer>>;
+def LLVM_SExtOp : LLVM_CastOp<"sext", "CreateSExt",
+                              LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_ZExtOp : LLVM_CastOp<"zext", "CreateZExt",
+                              LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_TruncOp : LLVM_CastOp<"trunc", "CreateTrunc",
+                               LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_SIToFPOp : LLVM_CastOp<"sitofp", "CreateSIToFP",
+                                LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_UIToFPOp : LLVM_CastOp<"uitofp", "CreateUIToFP",
+                                LLVM_ScalarOrVectorOf<LLVM_AnyInteger>>;
+def LLVM_FPToSIOp : LLVM_CastOp<"fptosi", "CreateFPToSI",
+                                LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
+def LLVM_FPToUIOp : LLVM_CastOp<"fptoui", "CreateFPToUI",
+                                LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
+def LLVM_FPExtOp : LLVM_CastOp<"fpext", "CreateFPExt",
+                                LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
+def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "CreateFPTrunc",
+                                 LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
 
 // Call-related operations.
 def LLVM_InvokeOp : LLVM_Op<"invoke", [
                       AttrSizedOperandSegments,
                       DeclareOpInterfaceMethods<BranchOpInterface>,
                       Terminator
                     ]>,
                     Arguments<(ins OptionalAttr<FlatSymbolRefAttr>:$callee,
let verifier = [{
     if (getNumResults() > 1)
       return emitOpError("must have 0 or 1 result");
     return success();
   }];
   let parser = [{ return parseCallOp(parser, result); }];
   let printer = [{ printCallOp(p, *this); }];
 }
 def LLVM_ExtractElementOp : LLVM_OneResultOp<"extractelement", [NoSideEffect]>,
-                          Arguments<(ins LLVM_Type:$vector,
-                                     LLVM_Type:$position)> {
+                          Arguments<(ins LLVM_AnyVector:$vector,
+                                     LLVM_AnyInteger:$position)> {
   string llvmBuilder = [{
     $res = builder.CreateExtractElement($vector, $position);
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Value vector, Value position,"
     "ArrayRef<NamedAttribute> attrs = {}">];
   let parser = [{ return parseExtractElementOp(parser, result); }];
   let printer = [{ printExtractElementOp(p, *this); }];
 }
 def LLVM_ExtractValueOp : LLVM_OneResultOp<"extractvalue", [NoSideEffect]>,
-                          Arguments<(ins LLVM_Type:$container,
+                          Arguments<(ins LLVM_AnyAggregate:$container,
                                      ArrayAttr:$position)> {
   string llvmBuilder = [{
     $res = builder.CreateExtractValue($container, extractPosition($position));
   }];
   let parser = [{ return parseExtractValueOp(parser, result); }];
   let printer = [{ printExtractValueOp(p, *this); }];
 }
 def LLVM_InsertElementOp : LLVM_OneResultOp<"insertelement", [NoSideEffect]>,
-                         Arguments<(ins LLVM_Type:$vector, LLVM_Type:$value,
-                                    LLVM_Type:$position)> {
+                         Arguments<(ins LLVM_AnyVector:$vector,
+                                    LLVM_PrimitiveType:$value,
+                                    LLVM_AnyInteger:$position)> {
   string llvmBuilder = [{
     $res = builder.CreateInsertElement($vector, $value, $position);
   }];
   let parser = [{ return parseInsertElementOp(parser, result); }];
   let printer = [{ printInsertElementOp(p, *this); }];
 }
 def LLVM_InsertValueOp : LLVM_OneResultOp<"insertvalue", [NoSideEffect]>,
-                         Arguments<(ins LLVM_Type:$container, LLVM_Type:$value,
+                         Arguments<(ins LLVM_AnyAggregate:$container,
+                                    LLVM_PrimitiveType:$value,
                                     ArrayAttr:$position)> {
   string llvmBuilder = [{
     $res = builder.CreateInsertValue($container, $value,
                                      extractPosition($position));
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Value container, Value value, "
     "ArrayAttr position",
     [{
       build(b, result, container.getType(), container, value, position);
     }]>];
   let parser = [{ return parseInsertValueOp(parser, result); }];
   let printer = [{ printInsertValueOp(p, *this); }];
 }
 def LLVM_ShuffleVectorOp
     : LLVM_OneResultOp<"shufflevector", [NoSideEffect]>,
-      Arguments<(ins LLVM_Type:$v1, LLVM_Type:$v2, ArrayAttr:$mask)> {
+      Arguments<(ins LLVM_AnyVector:$v1, LLVM_AnyVector:$v2, ArrayAttr:$mask)> {
   string llvmBuilder = [{
       SmallVector<unsigned, 4> position = extractPosition($mask);
       SmallVector<int, 4> mask(position.begin(), position.end());
       $res = builder.CreateShuffleVector($v1, $v2, mask);
   }];
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Value v1, Value v2, "
     "ArrayAttr mask, ArrayRef<NamedAttribute> attrs = {}">];
def LLVM_ShuffleVectorOp
   let parser = [{ return parseShuffleVectorOp(parser, result); }];
   let printer = [{ printShuffleVectorOp(p, *this); }];
 }
 
 // Misc operations.
 def LLVM_SelectOp
     : LLVM_OneResultOp<"select",
           [NoSideEffect, AllTypesMatch<["trueValue", "falseValue", "res"]>]>,
-      Arguments<(ins LLVM_Type:$condition, LLVM_Type:$trueValue,
+      Arguments<(ins LLVM_ScalarOrVectorOf<LLVM_i1>:$condition,
+                 LLVM_Type:$trueValue,
                  LLVM_Type:$falseValue)>,
       LLVM_Builder<
           "$res = builder.CreateSelect($condition, $trueValue, $falseValue);"> {
   let builders = [OpBuilder<
     "OpBuilder &b, OperationState &result, Value condition, Value lhs, "
     "Value rhs", [{
       build(b, result, lhs.getType(), condition, lhs, rhs);
     }]>];
def LLVM_BrOp : LLVM_TerminatorOp<"br",
   let assemblyFormat = [{
     $dest (`(` $destOperands^ `:` type($destOperands) `)`)? attr-dict
   }];
   let builders = [LLVM_TerminatorPassthroughOpBuilder];
 }
 def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br",
     [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
      NoSideEffect]> {
-  let arguments = (ins LLVMI1:$condition,
+  let arguments = (ins LLVM_i1:$condition,
                    Variadic<LLVM_Type>:$trueDestOperands,
                    Variadic<LLVM_Type>:$falseDestOperands,
                    OptionalAttr<ElementsAttr>:$branch_weights);
   let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
   let assemblyFormat = [{
     $condition ( `weights` `(` $branch_weights^ `)` )? `,`
     $trueDest (`(` $trueDestOperands^ `:` type($trueDestOperands) `)`)? `,`
     $falseDest (`(` $falseDestOperands^ `:` type($falseDestOperands) `)`)?
def AtomicOrdering : I64EnumAttr<
     "AtomicOrdering",
     "Atomic ordering for LLVM's memory model",
     [AtomicOrderingNotAtomic, AtomicOrderingUnordered, AtomicOrderingMonotonic,
      AtomicOrderingAcquire, AtomicOrderingRelease, AtomicOrderingAcquireRelease,
      AtomicOrderingSequentiallyConsistent]> {
   let cppNamespace = "::mlir::LLVM";
 }
 
+def LLVM_AtomicRMWType : AnyTypeOf<[LLVM_AnyFloat, LLVM_AnyInteger]>;
+
 def LLVM_AtomicRMWOp : LLVM_Op<"atomicrmw">,
-    Arguments<(ins AtomicBinOp:$bin_op, LLVM_Type:$ptr, LLVM_Type:$val,
-                   AtomicOrdering:$ordering)>,
+    Arguments<(ins AtomicBinOp:$bin_op, LLVM_PointerTo<LLVM_AtomicRMWType>:$ptr,
+               LLVM_AtomicRMWType:$val, AtomicOrdering:$ordering)>,
     Results<(outs LLVM_Type:$res)> {
   let llvmBuilder = [{
     $res = builder.CreateAtomicRMW(getLLVMAtomicBinOp($bin_op), $ptr, $val,
                                    getLLVMAtomicOrdering($ordering));
   }];
   let parser = [{ return parseAtomicRMWOp(parser, result); }];
   let printer = [{ printAtomicRMWOp(p, *this); }];
   let verifier = "return ::verify(*this);";
 }
 
+def LLVM_AtomicCmpXchgType : AnyTypeOf<[LLVM_AnyInteger, LLVM_AnyPointer]>;
+
 def LLVM_AtomicCmpXchgOp : LLVM_Op<"cmpxchg">,
-    Arguments<(ins LLVM_Type:$ptr, LLVM_Type:$cmp, LLVM_Type:$val,
+    Arguments<(ins LLVM_PointerTo<LLVM_AtomicCmpXchgType>:$ptr,
+                   LLVM_AtomicCmpXchgType:$cmp, LLVM_AtomicCmpXchgType:$val,
                    AtomicOrdering:$success_ordering,
                    AtomicOrdering:$failure_ordering)>,
     Results<(outs LLVM_Type:$res)> {
   let llvmBuilder = [{
     $res = builder.CreateAtomicCmpXchg($ptr, $cmp, $val,
                    getLLVMAtomicOrdering($success_ordering),
                    getLLVMAtomicOrdering($failure_ordering));
   }];

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

if (parseAtomicBinOp(parser, result, "bin_op") || parser.parseOperand(ptr) ||
     return failure();
 
   result.addTypes(type);
   return success();
 }
 
 static LogicalResult verify(AtomicRMWOp op) {
   auto ptrType = op.ptr().getType().cast<LLVM::LLVMType>();
-  if (!ptrType.isPointerTy())
-    return op.emitOpError("expected LLVM IR pointer type for operand #0");
   auto valType = op.val().getType().cast<LLVM::LLVMType>();
   if (valType != ptrType.getPointerElementTy())
     return op.emitOpError("expected LLVM IR element type for operand #0 to "
                           "match type for operand #1");
   auto resType = op.res().getType().cast<LLVM::LLVMType>();
   if (resType != valType)
     return op.emitOpError(
         "expected LLVM IR result type to match type for operand #1");

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

LogicalResult LLVMStructType::setBody(ArrayRef<LLVMType> types, bool isPacked) {
   assert(llvm::all_of(types, LLVMStructType::isValidElementType) &&
          "expected valid body types");
   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();
+  return getImpl()->isIdentified() &&
+         (getImpl()->isOpaque() || !getImpl()->isInitialized());
 }
 bool LLVMStructType::isInitialized() { return getImpl()->isInitialized(); }
 StringRef LLVMStructType::getName() { return getImpl()->getIdentifier(); }
 ArrayRef<LLVMType> LLVMStructType::getBody() {
   return isIdentified() ? getImpl()->getIdentifiedStructBody()
                         : getImpl()->getTypeList();
 }
 

mlir/test/Dialect/LLVMIR/invalid.mlir

func @nvvm_invalid_mma_7(%a0 : !llvm.vec<2 x half>, %a1 : !llvm.vec<2 x half>,
   %0:2 = nvvm.mma.sync %a0, %a1, %b0, %b1, %c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7 {alayout="col", blayout="row"} : (!llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float) -> (!llvm.struct<(float, float, float, float, float, float, float, float)>, !llvm.i32)
   llvm.return %0#0 : !llvm.struct<(float, float, float, float, float, float, float, float)>
 }
 
 // -----
 
 // CHECK-LABEL: @atomicrmw_expected_ptr
 func @atomicrmw_expected_ptr(%f32 : !llvm.float) {
-  // expected-error@+1 {{expected LLVM IR pointer type for operand #0}}
+  // expected-error@+1 {{operand #0 must be LLVM pointer to floating point LLVM type or LLVM integer type}}
   %0 = "llvm.atomicrmw"(%f32, %f32) {bin_op=11, ordering=1} : (!llvm.float, !llvm.float) -> !llvm.float
   llvm.return
 }
 
 // -----
 
 // CHECK-LABEL: @atomicrmw_mismatched_operands
 func @atomicrmw_mismatched_operands(%f32_ptr : !llvm.ptr<float>, %i32 : !llvm.i32) {
func @atomicrmw_expected_int(%f32_ptr : !llvm.ptr<float>, %f32 : !llvm.float) {
   %0 = llvm.atomicrmw max %f32_ptr, %f32 unordered : !llvm.float
   llvm.return
 }
 
 // -----
 
 // CHECK-LABEL: @cmpxchg_expected_ptr
 func @cmpxchg_expected_ptr(%f32_ptr : !llvm.ptr<float>, %f32 : !llvm.float) {
-  // expected-error@+1 {{expected LLVM IR pointer type for operand #0}}
+  // expected-error@+1 {{op operand #0 must be LLVM pointer to LLVM integer type or LLVM pointer type}}
   %0 = "llvm.cmpxchg"(%f32, %f32, %f32) {success_ordering=2,failure_ordering=2} : (!llvm.float, !llvm.float, !llvm.float) -> !llvm.struct<(float, i1)>
   llvm.return
 }
 
 // -----
 
 // CHECK-LABEL: @cmpxchg_mismatched_operands
 func @cmpxchg_mismatched_operands(%i64_ptr : !llvm.ptr<i64>, %i32 : !llvm.i32) {