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

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

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D85901

Author: ftynse

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

@@ -21,7 +21,8 @@ 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">;
 
 //===----------------------------------------------------------------------===//

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

@@ -17,6 +17,10 @@
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
+//===----------------------------------------------------------------------===//
+// LLVM Dialect.
+//===----------------------------------------------------------------------===//
+
 def LLVM_Dialect : Dialect {
   let name = "llvm";
   let cppNamespace = "LLVM";
@@ -38,34 +42,108 @@ def LLVM_Dialect : Dialect {
   }];
 }
 
-// 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">;
+// 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">,
+    BuildableType<
+        "::mlir::LLVM::LLVMIntegerType::get($_builder.getContext(), "
+        # width # ")">;
+
+def LLVM_i1 : LLVM_IntBase<1>;
+def LLVM_i8 : LLVM_IntBase<8>;
+def LLVM_i32 : LLVM_IntBase<32>;
 
-def LLVMIntBase : TypeConstraint<
+// Type constraint accepting LLVM primitive types, i.e. all types except void
+// and function.
+def LLVM_PrimitiveType : Type<
   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 # ")">;
-
-def LLVMI1 : LLVMI<1>;
+       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.
@@ -85,6 +163,10 @@ class LLVM_OpBase<Dialect dialect, string mnemonic, list<OpTrait> traits = []> :
   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