[mlir][DeclarativeParser] Add support for the TypesMatchWith trait.

[rriddle]

Summary

This allows for injecting type constraints that are not direct 1-1 mappings, for example when one type is equal to the element type of another. This allows for moving over several more parsers to the declarative form.

Depends On D74647

[rriddle]

Simplify type operand printing

[rriddle]

Rebase

[rriddle]

Rebase

[antiagainst]

Nice! Can we also update the doc regarding which traits are supported? Otherwise it's quite obscure for somebody to figure out what works and what not.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1009–1010

This pair is not trivial and it appears at lots of places. Might be worth it to define a struct to improve readability? Meta-programming is always difficult to understand; string concatenation to generate code too. :)

[rriddle]

In D74648#1887128, @antiagainst wrote:

Nice! Can we also update the doc regarding which traits are supported? Otherwise it's quite obscure for somebody to figure out what works and what not.

Yep completely agree, will add all of the currently supported traits in a followup.

[rriddle]

Resolve comments

mlir/include/mlir/Dialect/StandardOps/Ops.td

let extraClassDeclaration = [{
     operand_iterator arg_operand_end() { return operand_end(); }
 
     /// Return the callee of this operation.
     CallInterfaceCallable getCallableForCallee() { return getCallee(); }
   }];
 
   let verifier = ?;
   let hasCanonicalizer = 1;
+
+  let assemblyFormat = "$callee `(` $operands `)` attr-dict `:` type($callee)";
 }
 
 def CeilFOp : FloatUnaryOp<"ceilf"> {
   let summary = "ceiling of the specified value";
   let description = [{
     The `ceilf` operation computes the ceiling of a given value. It takes one
     operand and returns one result of the same type. This type may be a float
     scalar type, a vector whose element type is float, or a tensor of floats.
CmpIPredicate getPredicate() {
       return (CmpIPredicate)getAttrOfType<IntegerAttr>(getPredicateAttrName())
           .getInt();
     }
   }];
 
   let verifier = [{ return success(); }];
 
   let hasFolder = 1;
+
+  let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
 def CondBranchOp : Std_Op<"cond_br", [Terminator]> {
   let summary = "conditional branch operation";
   let description = [{
     The "cond_br" operation represents a conditional branch operation in a
     function. The operation takes variable number of operands and produces
     no results. The operand number and types for each successor must match the
let extraClassDeclaration = [{
     Value getAggregate() { return getOperand(0); }
 
     operand_range getIndices() {
       return {operand_begin() + 1, operand_end()};
     }
   }];
 
   let hasFolder = 1;
+
+  let assemblyFormat = [{
+    $aggregate `[` $indices `]` attr-dict `:` type($aggregate)
+  }];
 }
 
 def IndexCastOp : CastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
   let summary = "cast between index and integer types";
   let description = [{
     Casts between integer scalars and 'index' scalars.  Index is an integer of
     platform-specific bit width.  If casting to a wider integer, the value is
     sign-extended.  If casting to a narrower integer, the value is truncated.
let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return getMemRef().getType().cast<MemRefType>();
     }
 
     operand_range getIndices() { return {operand_begin() + 1, operand_end()}; }
   }];
 
   let hasFolder = 1;
+
+  let assemblyFormat = "$memref `[` $indices `]` attr-dict `:` type($memref)";
 }
 
 def LogOp : FloatUnaryOp<"log"> {
   let summary = "base-e logarithm of the specified value";
 }
 
 def Log10Op : FloatUnaryOp<"log10"> {
   let summary = "base-10 logarithm of the specified value";
def SelectOp : Std_Op<"select", [NoSideEffect, SameOperandsAndResultShape,
 
   let extraClassDeclaration = [{
       Value getCondition() { return condition(); }
       Value getTrueValue() { return true_value(); }
       Value getFalseValue() { return false_value(); }
   }];
 
   let hasFolder = 1;
+
+  let assemblyFormat = [{
+    $condition `,` $true_value `,` $false_value attr-dict `:` type($result)
+  }];
 }
 
 def SignExtendIOp : Std_Op<"sexti",
     [NoSideEffect, SameOperandsAndResultShape]> {
   let summary = "integer sign extension operation";
   let description = [{
     The integer sign extension operation takes an integer input of
     width M and an integer destination type of width N. The destination
def SplatOp : Std_Op<"splat", [NoSideEffect,
   let results = (outs AnyTypeOf<[AnyVector, AnyStaticShapeTensor]>:$aggregate);
 
   let builders =
       [OpBuilder<"Builder *builder, OperationState &result, Value element, "
                   "Type aggregateType",
                   [{ build(builder, result, aggregateType, element); }]>];
 
   let hasFolder = 1;
+
+  let assemblyFormat = "$input attr-dict `:` type($aggregate)";
 }
 
 def StoreOp : Std_Op<"store",
      [TypesMatchWith<"value type matches element type of memref",
                      "memref", "value",
                      "$_self.cast<MemRefType>().getElementType()">]> {
   let summary = "store operation";
   let description = [{
let extraClassDeclaration = [{
       }
 
       operand_range getIndices() {
         return {operand_begin() + 2, operand_end()};
       }
   }];
 
   let hasFolder = 1;
+
+  let assemblyFormat = [{
+    $value `,` $memref `[` $indices `]` attr-dict `:` type($memref)
+  }];
 }
 
 def SubFOp : FloatArithmeticOp<"subf"> {
   let summary = "floating point subtraction operation";
   let hasFolder = 1;
 }
 
 def SubIOp : IntArithmeticOp<"subi"> {
let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value memref", [{
       auto memrefType = memref.getType().cast<MemRefType>();
       auto resultType = RankedTensorType::get(memrefType.getShape(),
                                               memrefType.getElementType());
       result.addOperands(memref);
       result.addTypes(resultType);
   }]>];
 
-
   let extraClassDeclaration = [{
     /// The result of a tensor_load is always a tensor.
     TensorType getType() { return getResult().getType().cast<TensorType>(); }
   }];
+
+  let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 def TensorStoreOp : Std_Op<"tensor_store",
     [SameOperandsShape, SameOperandsElementType,
      TypesMatchWith<"value type matches tensor equivalent of memref",
                     "memref", "tensor",
                     "getTensorTypeFromMemRefType($_self)">]> {
   let summary = "tensor store operation";
   let description = [{
     The "tensor_store" operation stores the contents of a tensor into a memref.
     The first operand is a value of tensor type, the second operand is a value
     of memref type. The shapes and element types of these must match, and are
     specified by the memref type.
 
     Example:
        %9 = dim %8, 1 : tensor<4x?xf32>
        %10 = alloc(%9) : memref<4x?xf32, #layout, memspace0>
        tensor_store %8, %10 : memref<4x?xf32, #layout, memspace0>
   }];
 
   let arguments = (ins AnyTensor:$tensor, AnyMemRef:$memref);
   // TensorStoreOp is fully verified by traits.
   let verifier = ?;
+
+  let assemblyFormat = "$tensor `,` $memref attr-dict `:` type($memref)";
 }
 
 def TruncateIOp : Std_Op<"trunci", [NoSideEffect, SameOperandsAndResultShape]> {
   let summary = "integer truncation operation";
   let description = [{
     The integer truncation operation takes an integer input of
     width M and an integer destination type of width N. The destination
     bit-width must be smaller than the input bit-width (N < M).

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

```
       %1 = vector.extractelement %0[%c : i32]: vector<16xf32>
     ```
   }];
   let extraClassDeclaration = [{
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
   }];
+
+  let assemblyFormat = [{
+    $vector `[` $position `:` type($position) `]` attr-dict `:` type($vector)
+  }];
 }
 
 def Vector_ExtractOp :
   Vector_Op<"extract", [NoSideEffect,
      PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
     Arguments<(ins AnyVector:$vector, I64ArrayAttr:$position)>,
     Results<(outs AnyType)> {
let description = [{
     ```
   }];
   let extraClassDeclaration = [{
     Type getSourceType() { return source().getType(); }
     VectorType getDestVectorType() {
       return dest().getType().cast<VectorType>();
     }
   }];
+
+  let assemblyFormat = [{
+    $source `,` $dest `[` $position `:` type($position) `]` attr-dict `:`
+    type($result)
+  }];
 }
 
 def Vector_InsertOp :
   Vector_Op<"insert", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      AllTypesMatch<["dest", "res"]>]>,
      Arguments<(ins AnyType:$source, AnyVector:$dest, I64ArrayAttr:$position)>,

mlir/lib/Dialect/StandardOps/Ops.cpp

PatternMatchResult matchAndRewrite(CallIndirectOp indirectCall,
     rewriter.replaceOpWithNewOp<CallOp>(indirectCall, calledFn,
                                         indirectCall.getResultTypes(),
                                         indirectCall.getArgOperands());
     return matchSuccess();
   }
 };
 } // end anonymous namespace.
 
-static ParseResult parseCallIndirectOp(OpAsmParser &parser,
-                                       OperationState &result) {
-  FunctionType calleeType;
-  OpAsmParser::OperandType callee;
-  llvm::SMLoc operandsLoc;
-  SmallVector<OpAsmParser::OperandType, 4> operands;
-  return failure(
-      parser.parseOperand(callee) || parser.getCurrentLocation(&operandsLoc) ||
-      parser.parseOperandList(operands, OpAsmParser::Delimiter::Paren) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(calleeType) ||
-      parser.resolveOperand(callee, calleeType, result.operands) ||
-      parser.resolveOperands(operands, calleeType.getInputs(), operandsLoc,
-                             result.operands) ||
-      parser.addTypesToList(calleeType.getResults(), result.types));
-}
-
-static void print(OpAsmPrinter &p, CallIndirectOp op) {
-  p << "call_indirect " << op.getCallee() << '(' << op.getArgOperands() << ')';
-  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"callee"});
-  p << " : " << op.getCallee().getType();
-}
-
 void CallIndirectOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyIndirectCallWithKnownCallee>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // General helpers for comparison ops
 //===----------------------------------------------------------------------===//
static void buildCmpIOp(Builder *build, OperationState &result,
                         CmpIPredicate predicate, Value lhs, Value rhs) {
   result.addOperands({lhs, rhs});
   result.types.push_back(getI1SameShape(lhs.getType()));
   result.addAttribute(
       CmpIOp::getPredicateAttrName(),
       build->getI64IntegerAttr(static_cast<int64_t>(predicate)));
 }
 
-static ParseResult parseCmpIOp(OpAsmParser &parser, OperationState &result) {
-  SmallVector<OpAsmParser::OperandType, 2> ops;
-  SmallVector<NamedAttribute, 4> attrs;
-  Attribute predicateNameAttr;
-  Type type;
-  if (parser.parseAttribute(predicateNameAttr, CmpIOp::getPredicateAttrName(),
-                            attrs) ||
-      parser.parseComma() || parser.parseOperandList(ops, 2) ||
-      parser.parseOptionalAttrDict(attrs) || parser.parseColonType(type) ||
-      parser.resolveOperands(ops, type, result.operands))
-    return failure();
-
-  if (!predicateNameAttr.isa<StringAttr>())
-    return parser.emitError(parser.getNameLoc(),
-                            "expected string comparison predicate attribute");
-
-  // Rewrite string attribute to an enum value.
-  StringRef predicateName = predicateNameAttr.cast<StringAttr>().getValue();
-  Optional<CmpIPredicate> predicate = symbolizeCmpIPredicate(predicateName);
-  if (!predicate.hasValue())
-    return parser.emitError(parser.getNameLoc())
-           << "unknown comparison predicate \"" << predicateName << "\"";
-
-  auto builder = parser.getBuilder();
-  Type i1Type = getCheckedI1SameShape(type);
-  if (!i1Type)
-    return parser.emitError(parser.getNameLoc(),
-                            "expected type with valid i1 shape");
-
-  attrs[0].second = builder.getI64IntegerAttr(static_cast<int64_t>(*predicate));
-  result.attributes = attrs;
-
-  result.addTypes({i1Type});
-  return success();
-}
-
-static void print(OpAsmPrinter &p, CmpIOp op) {
-  p << "cmpi ";
-
-  Builder b(op.getContext());
-  auto predicateValue =
-      op.getAttrOfType<IntegerAttr>(CmpIOp::getPredicateAttrName()).getInt();
-  p << '"' << stringifyCmpIPredicate(static_cast<CmpIPredicate>(predicateValue))
-    << '"' << ", " << op.lhs() << ", " << op.rhs();
-  p.printOptionalAttrDict(op.getAttrs(),
-                          /*elidedAttrs=*/{CmpIOp::getPredicateAttrName()});
-  p << " : " << op.lhs().getType();
-}
-
 // Compute `lhs` `pred` `rhs`, where `pred` is one of the known integer
 // comparison predicates.
 static bool applyCmpPredicate(CmpIPredicate predicate, const APInt &lhs,
                               const APInt &rhs) {
   switch (predicate) {
   case CmpIPredicate::eq:
     return lhs.eq(rhs);
   case CmpIPredicate::ne:
LogicalResult DmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
   /// dma_wait(memrefcast) -> dma_wait
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, ExtractElementOp op) {
-  p << "extract_element " << op.getAggregate() << '[' << op.getIndices();
-  p << ']';
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getAggregate().getType();
-}
-
-static ParseResult parseExtractElementOp(OpAsmParser &parser,
-                                         OperationState &result) {
-  OpAsmParser::OperandType aggregateInfo;
-  SmallVector<OpAsmParser::OperandType, 4> indexInfo;
-  ShapedType type;
-
-  auto indexTy = parser.getBuilder().getIndexType();
-  return failure(
-      parser.parseOperand(aggregateInfo) ||
-      parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type) ||
-      parser.resolveOperand(aggregateInfo, type, result.operands) ||
-      parser.resolveOperands(indexInfo, indexTy, result.operands) ||
-      parser.addTypeToList(type.getElementType(), result.types));
-}
-
 static LogicalResult verify(ExtractElementOp op) {
   // Verify the # indices match if we have a ranked type.
   auto aggregateType = op.getAggregate().getType().cast<ShapedType>();
   if (aggregateType.hasRank() &&
       aggregateType.getRank() != op.getNumOperands() - 1)
     return op.emitOpError("incorrect number of indices for extract_element");
 
   return success();
OpFoldResult IndexCastOp::fold(ArrayRef<Attribute> cstOperands) {
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, LoadOp op) {
-  p << "load " << op.getMemRef() << '[' << op.getIndices() << ']';
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getMemRefType();
-}
-
-static ParseResult parseLoadOp(OpAsmParser &parser, OperationState &result) {
-  OpAsmParser::OperandType memrefInfo;
-  SmallVector<OpAsmParser::OperandType, 4> indexInfo;
-  MemRefType type;
-
-  auto indexTy = parser.getBuilder().getIndexType();
-  return failure(
-      parser.parseOperand(memrefInfo) ||
-      parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type) ||
-      parser.resolveOperand(memrefInfo, type, result.operands) ||
-      parser.resolveOperands(indexInfo, indexTy, result.operands) ||
-      parser.addTypeToList(type.getElementType(), result.types));
-}
-
 static LogicalResult verify(LoadOp op) {
   if (op.getNumOperands() != 1 + op.getMemRefType().getRank())
     return op.emitOpError("incorrect number of indices for load");
   return success();
 }
 
 OpFoldResult LoadOp::fold(ArrayRef<Attribute> cstOperands) {
   /// load(memrefcast) -> load

 bool SIToFPOp::areCastCompatible(Type a, Type b) {
   return a.isa<IntegerType>() && b.isa<FloatType>();
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
-static ParseResult parseSelectOp(OpAsmParser &parser, OperationState &result) {
-  SmallVector<OpAsmParser::OperandType, 3> ops;
-  SmallVector<NamedAttribute, 4> attrs;
-  Type type;
-  if (parser.parseOperandList(ops, 3) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type))
-    return failure();
-
-  auto i1Type = getCheckedI1SameShape(type);
-  if (!i1Type)
-    return parser.emitError(parser.getNameLoc(),
-                            "expected type with valid i1 shape");
-
-  SmallVector<Type, 3> types = {i1Type, type, type};
-  return failure(parser.resolveOperands(ops, types, parser.getNameLoc(),
-                                        result.operands) ||
-                 parser.addTypeToList(type, result.types));
-}
-
-static void print(OpAsmPrinter &p, SelectOp op) {
-  p << "select " << op.getOperands() << " : " << op.getTrueValue().getType();
-  p.printOptionalAttrDict(op.getAttrs());
-}
-
 OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
   auto condition = getCondition();
 
   // select true, %0, %1 => %0
   if (matchPattern(condition, m_One()))
     return getTrueValue();
 
   // select false, %0, %1 => %1
static LogicalResult verify(SignExtendIOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SplatOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, SplatOp op) {
-  p << "splat " << op.getOperand();
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getType();
-}
-
-static ParseResult parseSplatOp(OpAsmParser &parser, OperationState &result) {
-  OpAsmParser::OperandType splatValueInfo;
-  ShapedType shapedType;
-
-  return failure(parser.parseOperand(splatValueInfo) ||
-                 parser.parseOptionalAttrDict(result.attributes) ||
-                 parser.parseColonType(shapedType) ||
-                 parser.resolveOperand(splatValueInfo,
-                                       shapedType.getElementType(),
-                                       result.operands) ||
-                 parser.addTypeToList(shapedType, result.types));
-}
-
 static LogicalResult verify(SplatOp op) {
   // TODO: we could replace this by a trait.
   if (op.getOperand().getType() !=
       op.getType().cast<ShapedType>().getElementType())
     return op.emitError("operand should be of elemental type of result type");
 
   return success();
 }
OpFoldResult SplatOp::fold(ArrayRef<Attribute> operands) {
   // SplatElementsAttr::get treats single value for second arg as being a splat.
   return SplatElementsAttr::get(shapedType, {constOperand});
 }
 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, StoreOp op) {
-  p << "store " << op.getValueToStore();
-  p << ", " << op.getMemRef() << '[' << op.getIndices() << ']';
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getMemRefType();
-}
-
-static ParseResult parseStoreOp(OpAsmParser &parser, OperationState &result) {
-  OpAsmParser::OperandType storeValueInfo;
-  OpAsmParser::OperandType memrefInfo;
-  SmallVector<OpAsmParser::OperandType, 4> indexInfo;
-  MemRefType memrefType;
-
-  auto indexTy = parser.getBuilder().getIndexType();
-  return failure(
-      parser.parseOperand(storeValueInfo) || parser.parseComma() ||
-      parser.parseOperand(memrefInfo) ||
-      parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(memrefType) ||
-      parser.resolveOperand(storeValueInfo, memrefType.getElementType(),
-                            result.operands) ||
-      parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
-      parser.resolveOperands(indexInfo, indexTy, result.operands));
-}
-
 static LogicalResult verify(StoreOp op) {
   if (op.getNumOperands() != 2 + op.getMemRefType().getRank())
     return op.emitOpError("store index operand count not equal to memref rank");
 
   return success();
 }
 
 LogicalResult StoreOp::fold(ArrayRef<Attribute> cstOperands,

 
 static Type getTensorTypeFromMemRefType(Type type) {
   if (auto memref = type.dyn_cast<MemRefType>())
     return RankedTensorType::get(memref.getShape(), memref.getElementType());
   return NoneType::get(type.getContext());
 }
 
 //===----------------------------------------------------------------------===//
-// TensorLoadOp
-//===----------------------------------------------------------------------===//
-
-static void print(OpAsmPrinter &p, TensorLoadOp op) {
-  p << "tensor_load " << op.getOperand();
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getOperand().getType();
-}
-
-static ParseResult parseTensorLoadOp(OpAsmParser &parser,
-                                     OperationState &result) {
-  OpAsmParser::OperandType op;
-  Type type;
-  return failure(
-      parser.parseOperand(op) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type) ||
-      parser.resolveOperand(op, type, result.operands) ||
-      parser.addTypeToList(getTensorTypeFromMemRefType(type), result.types));
-}
-
-//===----------------------------------------------------------------------===//
-// TensorStoreOp
-//===----------------------------------------------------------------------===//
-
-static void print(OpAsmPrinter &p, TensorStoreOp op) {
-  p << "tensor_store " << op.tensor() << ", " << op.memref();
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.memref().getType();
-}
-
-static ParseResult parseTensorStoreOp(OpAsmParser &parser,
-                                      OperationState &result) {
-  SmallVector<OpAsmParser::OperandType, 2> ops;
-  Type type;
-  llvm::SMLoc loc = parser.getCurrentLocation();
-  return failure(
-      parser.parseOperandList(ops, /*requiredOperandCount=*/2) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type) ||
-      parser.resolveOperands(ops, {getTensorTypeFromMemRefType(type), type},
-                             loc, result.operands));
-}
-
-//===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(TruncateIOp op) {
   auto srcType = getElementTypeOrSelf(op.getOperand().getType());
   auto dstType = getElementTypeOrSelf(op.getType());
 
   if (srcType.isa<IndexType>())

mlir/lib/Dialect/VectorOps/VectorOps.cpp

for (auto mapAttr : mapAttrs)
     res.push_back(mapAttr.cast<AffineMapAttr>().getValue());
   return res;
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, vector::ExtractElementOp op) {
-  p << op.getOperationName() << " " << op.vector() << "[" << op.position()
-    << " : " << op.position().getType() << "]";
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.vector().getType();
-}
-
-static ParseResult parseExtractElementOp(OpAsmParser &parser,
-                                         OperationState &result) {
-  OpAsmParser::OperandType vector, position;
-  Type positionType;
-  VectorType vectorType;
-  if (parser.parseOperand(vector) || parser.parseLSquare() ||
-      parser.parseOperand(position) || parser.parseColonType(positionType) ||
-      parser.parseRSquare() ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(vectorType))
-    return failure();
-  Type resultType = vectorType.getElementType();
-  return failure(
-      parser.resolveOperand(vector, vectorType, result.operands) ||
-      parser.resolveOperand(position, positionType, result.operands) ||
-      parser.addTypeToList(resultType, result.types));
-}
-
 static LogicalResult verify(vector::ExtractElementOp op) {
   VectorType vectorType = op.getVectorType();
   if (vectorType.getRank() != 1)
     return op.emitOpError("expected 1-D vector");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &result) {
   parser.addTypeToList(resType, result.types);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // InsertElementOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, InsertElementOp op) {
-  p << op.getOperationName() << " " << op.source() << ", " << op.dest() << "["
-    << op.position() << " : " << op.position().getType() << "]";
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.dest().getType();
-}
-
-static ParseResult parseInsertElementOp(OpAsmParser &parser,
-                                        OperationState &result) {
-  OpAsmParser::OperandType source, dest, position;
-  Type positionType;
-  VectorType destType;
-  if (parser.parseOperand(source) || parser.parseComma() ||
-      parser.parseOperand(dest) || parser.parseLSquare() ||
-      parser.parseOperand(position) || parser.parseColonType(positionType) ||
-      parser.parseRSquare() ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(destType))
-    return failure();
-  Type sourceType = destType.getElementType();
-  return failure(
-      parser.resolveOperand(source, sourceType, result.operands) ||
-      parser.resolveOperand(dest, destType, result.operands) ||
-      parser.resolveOperand(position, positionType, result.operands) ||
-      parser.addTypeToList(destType, result.types));
-}
-
 static LogicalResult verify(InsertElementOp op) {
   auto dstVectorType = op.getDestVectorType();
   if (dstVectorType.getRank() != 1)
     return op.emitOpError("expected 1-D vector");
   return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/IR/invalid-ops.mlir

^bb0(%a : i32, %b : i32):
   %r = cmpi %a, %b : i32 // expected-error {{expected non-function type}}
 }
 
 // -----
 
 // Integer comparisons are not recognized for float types.
 func @func_with_ops(f32, f32) {
 ^bb0(%a : f32, %b : f32):
-  %r = cmpi "eq", %a, %b : f32 // expected-error {{operand #0 must be integer-like}}
+  %r = cmpi "eq", %a, %b : f32 // expected-error {{'lhs' must be integer-like, but got 'f32'}}
 }
 
 // -----
 
 // Result type must be boolean like.
 func @func_with_ops(i32, i32) {
 ^bb0(%a : i32, %b : i32):
   %r = "std.cmpi"(%a, %b) {predicate = 0} : (i32, i32) -> i32 // expected-error {{op result #0 must be bool-like}}

 ^bb0(%cond : i1, %t : tensor<42xi32>, %f : tensor<?xi32>):
   // expected-error@+1 {{ op requires the same shape for all operands and results}}
   %r = "std.select"(%cond, %t, %f) : (i1, tensor<42xi32>, tensor<?xi32>) -> tensor<42xi32>
 }
 
 // -----
 
 func @invalid_select_shape(%cond : i1, %idx : () -> ()) {
-  // expected-error@+1 {{expected type with valid i1 shape}}
+  // expected-error@+1 {{'result' must be integer-like or floating-point-like, but got '() -> ()'}}
   %sel = select %cond, %idx, %idx : () -> ()
 
 // -----
 
 func @invalid_cmp_shape(%idx : () -> ()) {
-  // expected-error@+1 {{expected type with valid i1 shape}}
+  // expected-error@+1 {{'lhs' must be integer-like, but got '() -> ()'}}
   %cmp = cmpi "eq", %idx, %idx : () -> ()
 
 // -----
 
 func @dma_no_src_memref(%m : f32, %tag : f32, %c0 : index) {
   // expected-error@+1 {{expected source to be of memref type}}
   dma_start %m[%c0], %m[%c0], %c0, %tag[%c0] : f32, f32, f32
 }

 func @dma_wait_no_tag_memref(%tag : f32, %c0 : index) {
   // expected-error@+1 {{expected tag to be of memref type}}
   dma_wait %tag[%c0], %arg0 : f32
 }
 
 // -----
 
 func @invalid_cmp_attr(%idx : i32) {
-  // expected-error@+1 {{expected string comparison predicate attribute}}
+  // expected-error@+1 {{invalid kind of attribute specified}}
   %cmp = cmpi i1, %idx, %idx : i32
 
 // -----
 
 func @cmpf_generic_invalid_predicate_value(%a : f32) {
   // expected-error@+1 {{'predicate' attribute value out of range}}
   %r = "std.cmpf"(%a, %a) {predicate = 42} : (f32, f32) -> i1
 }

mlir/tools/mlir-tblgen/OpFormatGen.cpp

struct OperationFormat {
   public:
     TypeResolution() = default;
 
     /// Get the index into the buildable types for this type, or None.
     Optional<int> getBuilderIdx() const { return builderIdx; }
     void setBuilderIdx(int idx) { builderIdx = idx; }
 
     /// Get the variable this type is resolved to, or None.
-    Optional<StringRef> getVariable() const { return variableName; }
-    void setVariable(StringRef variable) { variableName = variable; }
+    const NamedTypeConstraint *getVariable() const { return variable; }
+    Optional<StringRef> getVarTransformer() const {
+      return variableTransformer;
+    }
+    void setVariable(const NamedTypeConstraint *var,
+                     Optional<StringRef> transformer) {
+      variable = var;
+      variableTransformer = transformer;
+    }
 
   private:
     /// If the type is resolved with a buildable type, this is the index into
     /// 'buildableTypes' in the parent format.
     Optional<int> builderIdx;
     /// If the type is resolved based upon another operand or result, this is
-    /// the name of the variable that this type is resolved to.
-    Optional<StringRef> variableName;
+    /// the variable that this type is resolved to.
+    const NamedTypeConstraint *variable;
+    /// If the type is resolved based upon another operand or result, this is
+    /// a transformer to apply to the variable when resolving.
+    Optional<StringRef> variableTransformer;
   };
 
   OperationFormat(const Operator &op)
       : allOperandTypes(false), allResultTypes(false) {
     operandTypes.resize(op.getNumOperands(), TypeResolution());
     resultTypes.resize(op.getNumResults(), TypeResolution());
   }
 
void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   // Generate the code to resolve the operand and result types now that they
   // have been parsed.
   genParserTypeResolution(op, body);
   body << "  return success();\n";
 }
 
 void OperationFormat::genParserTypeResolution(Operator &op,
                                               OpMethodBody &body) {
+  // If any of type resolutions use transformed variables, make sure that the
+  // types of those variables are resolved.
+  SmallPtrSet<const NamedTypeConstraint *, 8> verifiedVariables;
+  FmtContext verifierFCtx;
+  for (TypeResolution &resolver :
+       llvm::concat<TypeResolution>(resultTypes, operandTypes)) {
+    Optional<StringRef> transformer = resolver.getVarTransformer();
+    if (!transformer)
+      continue;
+    // Ensure that we don't verify the same variables twice.
+    const NamedTypeConstraint *variable = resolver.getVariable();
+    if (!verifiedVariables.insert(variable).second)
+      continue;
+
+    auto constraint = variable->constraint;
+    body << "  for (Type type : " << variable->name << "Types) {\n"
+         << "    (void)type;\n"
+         << "    if (!("
+         << tgfmt(constraint.getConditionTemplate(),
+                  &verifierFCtx.withSelf("type"))
+         << ")) {\n"
+         << formatv("      return parser.emitError(parser.getNameLoc()) << "
+                    "\"'{0}' must be {1}, but got \" << type;\n",
+                    variable->name, constraint.getDescription())
+         << "    }\n"
+         << "  }\n";
+  }
+
   // Initialize the set of buildable types.
   if (!buildableTypes.empty()) {
     body << "  Builder &builder = parser.getBuilder();\n";
 
     FmtContext typeBuilderCtx;
     typeBuilderCtx.withBuilder("builder");
     for (auto &it : buildableTypes)
       body << "  Type odsBuildableType" << it.second << " = "
            << tgfmt(it.first, &typeBuilderCtx) << ";\n";
   }
 
+  // Emit the code necessary for a type resolver. If `forceRange` is true,
+  // we always generate a range of types.
+  auto emitTypeResolver = [&](TypeResolution &resolver, StringRef curVar,
+                              bool forceRange) {
+    if (Optional<int> val = resolver.getBuilderIdx()) {
+      if (forceRange)
+        body << "ArrayRef<Type>(odsBuildableType" << *val << ")";
+      else
+        body << "odsBuildableType" << *val;
+    } else if (const NamedTypeConstraint *var = resolver.getVariable()) {
+      if (Optional<StringRef> tform = resolver.getVarTransformer())
+        body << "ArrayRef<Type>("
+             << tgfmt(*tform, &FmtContext().withSelf(var->name + "Types[0]"))
+             << ")";
+      else
+        body << var->name << "Types";
+    } else {
+      body << curVar << "Types";
+    }
+  };
+
   // Resolve each of the result types.
   if (allResultTypes) {
     body << "  result.addTypes(allResultTypes);\n";
   } else {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
       body << "  result.addTypes(";
-      if (Optional<int> val = resultTypes[i].getBuilderIdx())
-        body << "odsBuildableType" << *val;
-      else if (Optional<StringRef> var = resultTypes[i].getVariable())
-        body << *var << "Types";
-      else
-        body << op.getResultName(i) << "Types";
+      emitTypeResolver(resultTypes[i], op.getResultName(i),
+                       /*forceRange=*/false);
       body << ");\n";
     }
   }
 
   // Early exit if there are no operands.
   if (op.getNumOperands() == 0)
     return;
 
if (allOperandTypes) {
     body << ", allOperandTypes, parser.getNameLoc(), result.operands))\n"
          << "    return failure();\n";
     return;
   }
   // Handle the case where all of the operands were grouped together.
   if (hasAllOperands) {
     body << "  if (parser.resolveOperands(allOperands, ";
 
-    auto emitOperandType = [&](int idx) {
-      if (Optional<int> val = operandTypes[idx].getBuilderIdx())
-        body << "ArrayRef<Type>(odsBuildableType" << *val << ")";
-      else if (Optional<StringRef> var = operandTypes[idx].getVariable())
-        body << *var << "Types";
-      else
-        body << op.getOperand(idx).name << "Types";
-    };
-
     // Group all of the operand types together to perform the resolution all at
     // once. Use llvm::concat to perform the merge. llvm::concat does not allow
     // the case of a single range, so guard it here.
     if (op.getNumOperands() > 1) {
       body << "llvm::concat<const Type>(";
-      interleaveComma(llvm::seq<int>(0, op.getNumOperands()), body,
-                      emitOperandType);
+      interleaveComma(llvm::seq<int>(0, op.getNumOperands()), body, [&](int i) {
+        emitTypeResolver(operandTypes[i], op.getOperand(i).name,
+                         /*forceRange=*/true);
+      });
       body << ")";
     } else {
-      emitOperandType(/*idx=*/0);
+      emitTypeResolver(operandTypes.front(), op.getOperand(0).name,
+                       /*forceRange=*/true);
     }
 
     body << ", allOperandLoc, result.operands))\n"
          << "    return failure();\n";
     return;
   }
 
   // The final case is the one where each of the operands types are resolved
   // separately.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);
     body << "  if (parser.resolveOperands(" << operand.name << "Operands, ";
-    if (Optional<int> val = operandTypes[i].getBuilderIdx())
-      body << "odsBuildableType" << *val << ", ";
-    else if (Optional<StringRef> var = operandTypes[i].getVariable())
-      body << *var << "Types, " << operand.name << "OperandsLoc, ";
-    else
-      body << operand.name << "Types, " << operand.name << "OperandsLoc, ";
-    body << "result.operands))\n    return failure();\n";
+    emitTypeResolver(operandTypes[i], operand.name, /*forceRange=*/false);
+
+    // If this isn't a buildable type, verify the sizes match by adding the loc.
+    if (!operandTypes[i].getBuilderIdx())
+      body << ", " << operand.name << "OperandsLoc";
+    body << ", result.operands))\n    return failure();\n";
   }
 }
 
 //===----------------------------------------------------------------------===//
 // PrinterGen
 
 /// Generate the printer for the 'attr-dict' directive.
 static void genAttrDictPrinter(OperationFormat &fmt, OpMethodBody &body) {
public:
   /// Parse the operation assembly format.
   LogicalResult parse();
 
 private:
   /// Given the values of an `AllTypesMatch` trait, check for inferrable type
   /// resolution.
   void handleAllTypesMatchConstraint(
       ArrayRef<StringRef> values,
-      llvm::StringMap<const NamedTypeConstraint *> &variableTyResolver);
+      llvm::StringMap<std::pair<const NamedTypeConstraint *,
+                                Optional<StringRef>>> &variableTyResolver);
   /// Check for inferrable type resolution given all operands, and or results,
   /// have the same type. If 'includeResults' is true, the results also have the
   /// same type as all of the operands.
   void handleSameTypesConstraint(
-      llvm::StringMap<const NamedTypeConstraint *> &variableTyResolver,
+      llvm::StringMap<std::pair<const NamedTypeConstraint *,
+                                Optional<StringRef>>> &variableTyResolver,
       bool includeResults);
 
+  /// Returns an argument with the given name that has been seen within the
+  /// format.
+  const NamedTypeConstraint *findSeenArg(StringRef name);
+
   /// Parse a specific element.
   LogicalResult parseElement(std::unique_ptr<Element> &element,
                              bool isTopLevel);
   LogicalResult parseVariable(std::unique_ptr<Element> &element,
                               bool isTopLevel);
   LogicalResult parseDirective(std::unique_ptr<Element> &element,
                                bool isTopLevel);
   LogicalResult parseLiteral(std::unique_ptr<Element> &element);
while (curToken.getKind() != Token::eof) {
     fmt.elements.push_back(std::move(element));
   }
 
   // Check that the attribute dictionary is in the format.
   if (!hasAttrDict)
     return emitError(loc, "format missing 'attr-dict' directive");
 
   // Check for any type traits that we can use for inferring types.
-  llvm::StringMap<const NamedTypeConstraint *> variableTyResolver;
+  llvm::StringMap<std::pair<const NamedTypeConstraint *, Optional<StringRef>>>
+      variableTyResolver;
   for (const OpTrait &trait : op.getTraits()) {
     const llvm::Record &def = trait.getDef();
-    if (def.isSubClassOf("AllTypesMatch"))
+    if (def.isSubClassOf("AllTypesMatch")) {
       handleAllTypesMatchConstraint(def.getValueAsListOfStrings("values"),
                                     variableTyResolver);
-    else if (def.getName() == "SameTypeOperands")
+    } else if (def.getName() == "SameTypeOperands") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/false);
-    else if (def.getName() == "SameOperandsAndResultType")
+    } else if (def.getName() == "SameOperandsAndResultType") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/true);
+    } else if (def.isSubClassOf("TypesMatchWith")) {
+      if (const auto *lhsArg = findSeenArg(def.getValueAsString("lhs")))
+        variableTyResolver[def.getValueAsString("rhs")] = {
+            lhsArg, def.getValueAsString("transformer")};
+    }
   }
 
   // Check that all of the result types can be inferred.
   auto &buildableTypes = fmt.buildableTypes;
   if (!fmt.allResultTypes) {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
       if (seenResultTypes.test(i))
         continue;
 
       // Check to see if we can infer this type from another variable.
       auto varResolverIt = variableTyResolver.find(op.getResultName(i));
       if (varResolverIt != variableTyResolver.end()) {
-        fmt.resultTypes[i].setVariable(varResolverIt->second->name);
+        fmt.resultTypes[i].setVariable(varResolverIt->second.first,
+                                       varResolverIt->second.second);
         continue;
       }
 
       // If the result is not variadic, allow for the case where the type has a
       // builder that we can use.
       NamedTypeConstraint &result = op.getResult(i);
       Optional<StringRef> builder = result.constraint.getBuilderCall();
       if (!builder || result.constraint.isVariadic()) {
for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
 
     // Check that the operand type is in the format, or that it can be inferred.
     if (fmt.allOperandTypes || seenOperandTypes.test(i))
       continue;
 
     // Check to see if we can infer this type from another variable.
     auto varResolverIt = variableTyResolver.find(op.getOperand(i).name);
     if (varResolverIt != variableTyResolver.end()) {
-      fmt.operandTypes[i].setVariable(varResolverIt->second->name);
+      fmt.operandTypes[i].setVariable(varResolverIt->second.first,
+                                      varResolverIt->second.second);
       continue;
     }
 
     // Similarly to results, allow a custom builder for resolving the type if
     // we aren't using the 'operands' directive.
     Optional<StringRef> builder = operand.constraint.getBuilderCall();
     if (!builder || (hasAllOperands && operand.isVariadic())) {
       return emitError(loc, "format missing instance of operand #" + Twine(i) +
                                 "('" + operand.name + "') type");
     }
     auto it = buildableTypes.insert({*builder, buildableTypes.size()});
     fmt.operandTypes[i].setBuilderIdx(it.first->second);
   }
   return success();
 }
 
 void FormatParser::handleAllTypesMatchConstraint(
     ArrayRef<StringRef> values,
-    llvm::StringMap<const NamedTypeConstraint *> &variableTyResolver) {
+    llvm::StringMap<std::pair<const NamedTypeConstraint *, Optional<StringRef>>>
+        &variableTyResolver) {
   for (unsigned i = 0, e = values.size(); i != e; ++i) {
     // Check to see if this value matches a resolved operand or result type.
-    const NamedTypeConstraint *arg = nullptr;
-    if ((arg = findArg(op.getOperands(), values[i]))) {
-      if (!seenOperandTypes.test(arg - op.operand_begin()))
-        continue;
-    } else if ((arg = findArg(op.getResults(), values[i]))) {
-      if (!seenResultTypes.test(arg - op.result_begin()))
-        continue;
-    } else {
+    const NamedTypeConstraint *arg = findSeenArg(values[i]);
+    if (!arg)
       continue;
-    }
 
     // Mark this value as the type resolver for the other variables.
     for (unsigned j = 0; j != i; ++j)
-      variableTyResolver[values[j]] = arg;
+      variableTyResolver[values[j]] = {arg, llvm::None};
     for (unsigned j = i + 1; j != e; ++j)
-      variableTyResolver[values[j]] = arg;
+      variableTyResolver[values[j]] = {arg, llvm::None};
   }
 }
 
 void FormatParser::handleSameTypesConstraint(
-    llvm::StringMap<const NamedTypeConstraint *> &variableTyResolver,
+    llvm::StringMap<std::pair<const NamedTypeConstraint *, Optional<StringRef>>>
+        &variableTyResolver,
     bool includeResults) {
   const NamedTypeConstraint *resolver = nullptr;
   int resolvedIt = -1;
 
   // Check to see if there is an operand or result to use for the resolution.
   if ((resolvedIt = seenOperandTypes.find_first()) != -1)
     resolver = &op.getOperand(resolvedIt);
   else if (includeResults && (resolvedIt = seenResultTypes.find_first()) != -1)
     resolver = &op.getResult(resolvedIt);
   else
     return;
 
   // Set the resolvers for each operand and result.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i)
     if (!seenOperandTypes.test(i) && !op.getOperand(i).name.empty())
-      variableTyResolver[op.getOperand(i).name] = resolver;
+      variableTyResolver[op.getOperand(i).name] = {resolver, llvm::None};
   if (includeResults) {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i)
       if (!seenResultTypes.test(i) && !op.getResultName(i).empty())
-        variableTyResolver[op.getResultName(i)] = resolver;
+        variableTyResolver[op.getResultName(i)] = {resolver, llvm::None};
   }
 }
 
+const NamedTypeConstraint *FormatParser::findSeenArg(StringRef name) {
+  if (auto *arg = findArg(op.getOperands(), name))
+    return seenOperandTypes.test(arg - op.operand_begin()) ? arg : nullptr;
+  if (auto *arg = findArg(op.getResults(), name))
+    return seenResultTypes.test(arg - op.result_begin()) ? arg : nullptr;
+  return nullptr;
+}
+
 LogicalResult FormatParser::parseElement(std::unique_ptr<Element> &element,
                                          bool isTopLevel) {
   // Directives.
   if (curToken.isKeyword())
     return parseDirective(element, isTopLevel);
   // Literals.
   if (curToken.getKind() == Token::literal)
     return parseLiteral(element);
   // Variables.
   if (curToken.getKind() == Token::variable)
     return parseVariable(element, isTopLevel);
   return emitError(curToken.getLoc(),
                    "expected directive, literal, or variable");
 }
 
 LogicalResult FormatParser::parseVariable(std::unique_ptr<Element> &element,
                                           bool isTopLevel) {
   Token varTok = curToken;
   consumeToken();
 
   StringRef name = varTok.getSpelling().drop_front();
   llvm::SMLoc loc = varTok.getLoc();
 
-  // Check that the parsed argument is something actually registered on the op.
+  // Check that the parsed argument is something actually registered on the
+  // op.
   /// Attributes
   if (const NamedAttribute *attr = findArg(op.getAttributes(), name)) {
     if (isTopLevel && !seenAttrs.insert(attr).second)
       return emitError(loc, "attribute '" + name + "' is already bound");
     element = std::make_unique<AttributeVariable>(attr);
     return success();
   }
   /// Operands