[mlir] Initial support for type constraints in the declarative assembly format

[rriddle]

Summary

This revision add support for accepting a few type constraints, e.g. AllTypesMatch, when inferring types for operands and results. This is used to remove the c++ parsers for several additional operations.

[merge_guards_bot]

Unit tests: fail. 62354 tests passed, 1 failed and 839 were skipped.

failed: libc++.std/containers/sequences/array/array_creation/to_array.fail.cpp

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

Looks great, let's land this!

[rriddle]

Rebase

[merge_guards_bot]

Unit tests: pass. 62405 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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

let verifier = [{
     return success();
   }];
   let parser = [{ return parseShuffleVectorOp(parser, result); }];
   let printer = [{ printShuffleVectorOp(p, *this); }];
 }
 
 // Misc operations.
 def LLVM_SelectOp
-    : LLVM_OneResultOp<"select", [NoSideEffect]>,
+    : LLVM_OneResultOp<"select",
+          [NoSideEffect, AllTypesMatch<["trueValue", "falseValue", "res"]>]>,
       Arguments<(ins LLVM_Type:$condition, LLVM_Type:$trueValue,
                  LLVM_Type:$falseValue)>,
       LLVM_Builder<
           "$res = builder.CreateSelect($condition, $trueValue, $falseValue);"> {
   let builders = [OpBuilder<
     "Builder *b, OperationState &result, Value condition, Value lhs, "
     "Value rhs", [{
       build(b, result, lhs.getType(), condition, lhs, rhs);
     }]>];
-  let parser = [{ return parseSelectOp(parser, result); }];
-  let printer = [{ printSelectOp(p, *this); }];
+  let assemblyFormat = "operands attr-dict `:` type($condition) `,` type($res)";
 }
 
 // Terminators.
 def LLVM_BrOp : LLVM_TerminatorOp<"br", []> {
   let parser = [{ return parseBrOp(parser, result); }];
   let printer = [{ printBrOp(p, *this); }];
 }
 def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br", []> {

mlir/include/mlir/Dialect/SPIRV/SPIRVBitOps.td

let availability = [
     MaxVersion<SPV_V_1_5>,
     Extension<[]>,
     Capability<[]>
   ];
 }
 
 // -----
 
-def SPV_BitFieldInsertOp : SPV_Op<"BitFieldInsert", [NoSideEffect]> {
+def SPV_BitFieldInsertOp : SPV_Op<"BitFieldInsert",
+    [NoSideEffect, AllTypesMatch<["base", "insert", "result"]>]> {
   let summary = [{
     Make a copy of an object, with a modified bit field that comes from
     another object.
   }];
 
   let description = [{
      Results are computed per component.
 
let arguments = (ins
     SPV_ScalarOrVectorOf<SPV_Integer>:$insert,
     SPV_Integer:$offset,
     SPV_Integer:$count
   );
 
   let results = (outs
     SPV_ScalarOrVectorOf<SPV_Integer>:$result
   );
+
+  let verifier = [{ return success(); }];
+
+  let assemblyFormat = [{
+    operands attr-dict `:` type($base) `,` type($offset) `,` type($count)
+  }];
 }
 
 // -----
 
 def SPV_BitFieldSExtractOp : SPV_BitFieldExtractOp<"BitFieldSExtract", []> {
   let summary = "Extract a bit field from an object, with sign extension.";
 
   let description = [{

mlir/include/mlir/Dialect/SPIRV/SPIRVLogicalOps.td

let description = [{
     %5 = spv.SLessThanEqual %2, %3 : vector<4xi32>
 
     ```
   }];
 }
 
 // -----
 
-def SPV_SelectOp : SPV_Op<"Select", [NoSideEffect]> {
+def SPV_SelectOp : SPV_Op<"Select",
+    [NoSideEffect, AllTypesMatch<["true_value", "false_value", "result"]>]> {
   let summary = [{
     Select between two objects. Before version 1.4, results are only
     computed per component.
   }];
 
   let description = [{
     Before version 1.4, Result Type must be a pointer, scalar, or vector.
 
def SPV_SelectOp : SPV_Op<"Select",
 
   let results = (outs
     SPV_SelectType:$result
   );
 
   let builders = [OpBuilder<[{Builder *builder, OperationState &state,
                               Value cond, Value trueValue,
                               Value falseValue}]>];
+
+  let assemblyFormat = [{
+    operands attr-dict `:` type($condition) `,` type($result)
+  }];
 }
 
 // -----
 
 def SPV_UGreaterThanOp : SPV_LogicalBinaryOp<"UGreaterThan", SPV_Integer, []> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is greater than  Operand 2.
   }];

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

let extraClassDeclaration = [{
     }
   }];
 }
 
 def Vector_InsertOp :
   Vector_Op<"insert", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
-     PredOpTrait<"dest operand and result have same type",
-                 TCresIsSameAsOpBase<0, 1>>]>,
+     AllTypesMatch<["dest", "res"]>]>,
      Arguments<(ins AnyType:$source, AnyVector:$dest, I64ArrayAttr:$position)>,
-     Results<(outs AnyVector)> {
+     Results<(outs AnyVector:$res)> {
   let summary = "insert operation";
   let description = [{
     Takes an n-D source vector, an (n+k)-D destination vector and a k-D position
     and inserts the n-D source into the (n+k)-D destination at the proper
     position. Degenerates to a scalar source type when n = 0.
 
     Examples:
     ```
       %2 = vector.insert %0, %1[3]:
         vector<8x16xf32> into vector<4x8x16xf32>
       %5 = vector.insert %3, %4[3, 3, 3]:
         f32 into vector<4x8x16xf32>
     ```
   }];
+  let assemblyFormat = [{
+    $source `,` $dest $position attr-dict `:` type($source) `into` type($dest)
+  }];
+
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value source, " #
     "Value dest, ArrayRef<int64_t>">];
   let extraClassDeclaration = [{
     static StringRef getPositionAttrName() { return "position"; }
     Type getSourceType() { return source().getType(); }
     VectorType getDestVectorType() {
       return dest().getType().cast<VectorType>();
let assemblyFormat = [{
     type(results)
   }];
 }
 
 def Vector_InsertStridedSliceOp :
   Vector_Op<"insert_strided_slice", [NoSideEffect,
     PredOpTrait<"operand #0 and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
-    PredOpTrait<"dest operand and result have same type",
-                 TCresIsSameAsOpBase<0, 1>>]>,
+    AllTypesMatch<["dest", "res"]>]>,
     Arguments<(ins AnyVector:$source, AnyVector:$dest, I64ArrayAttr:$offsets,
                I64ArrayAttr:$strides)>,
-    Results<(outs AnyVector)> {
+    Results<(outs AnyVector:$res)> {
   let summary = "strided_slice operation";
   let description = [{
     Takes a k-D source vector, an n-D destination vector (n >= k), n-sized
     `offsets` integer array attribute, a k-sized `strides` integer array attribute
     and inserts the k-D source vector as a strided subvector at the proper offset
     into the n-D destination vector.
 
     At the moment strides must contain only 1s.
 
     Returns an n-D vector that is a copy of the n-D destination vector in which
     the last k-D dimensions contain the k-D source vector elements strided at
     the proper location as specified by the offsets.
 
     Examples:
     ```
       %2 = vector.insert_strided_slice %0, %1
           {offsets = [0, 0, 2], strides = [1, 1]}:
         vector<2x4xf32> into vector<16x4x8xf32>
     ```
   }];
+
+  let assemblyFormat = [{
+    $source `,` $dest attr-dict `:` type($source) `into` type($dest)
+  }];
+
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value source, Value dest, " #
     "ArrayRef<int64_t> offsets, ArrayRef<int64_t> strides">];
   let extraClassDeclaration = [{
     static StringRef getOffsetsAttrName() { return "offsets"; }
     static StringRef getStridesAttrName() { return "strides"; }
     VectorType getSourceVectorType() {
       return source().getType().cast<VectorType>();

mlir/include/mlir/IR/OpBase.td

 // TODO(b/135032064): Only works for non-variadic.
 class AllMatchSameOperatorPred<list<string> names, string operator> :
     AllMatchPred<!foreach(n, names, !subst("$_self", "$" # n, operator))>;
 
 class AllMatchSameOperatorTrait<list<string> names, string operator,
                                 string description> :
     PredOpTrait<
         "all of {" # StrJoin<names>.result # "} have same " # description,
-        AllMatchSameOperatorPred<names, operator>>;
+        AllMatchSameOperatorPred<names, operator>> {
+  list<string> values = names;
+}
 
 class AllElementCountsMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, ElementCount<"_self">.result,
                               "element count">;
 
 class AllElementTypesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, ElementType<"_self">.result,
                               "element type">;

mlir/include/mlir/TableGen/OpTrait.h

public:
 
   explicit OpTrait(Kind kind, const llvm::Record *def);
 
   // Returns an OpTrait corresponding to the init provided.
   static OpTrait create(const llvm::Init *init);
 
   Kind getKind() const { return kind; }
 
+  // Returns the Tablegen definition this operator was constructed from.
+  const llvm::Record &getDef() const { return *def; }
+
 protected:
   // The TableGen definition of this trait.
   const llvm::Record *def;
   Kind kind;
 };
 
 // OpTrait corresponding to a native C++ OpTrait.
 class NativeOpTrait : public OpTrait {

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

if (parser.resolveOperand(container, containerType, result.operands) ||
       parser.resolveOperand(value, valueType, result.operands))
     return failure();
 
   result.addTypes(containerType);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Printing/parsing for LLVM::SelectOp.
-//===----------------------------------------------------------------------===//
-
-static void printSelectOp(OpAsmPrinter &p, SelectOp &op) {
-  p << op.getOperationName() << ' ' << op.condition() << ", " << op.trueValue()
-    << ", " << op.falseValue();
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.condition().getType() << ", " << op.trueValue().getType();
-}
-
-// <operation> ::= `llvm.select` ssa-use `,` ssa-use `,` ssa-use
-//                 attribute-dict? `:` type, type
-static ParseResult parseSelectOp(OpAsmParser &parser, OperationState &result) {
-  OpAsmParser::OperandType condition, trueValue, falseValue;
-  Type conditionType, argType;
-
-  if (parser.parseOperand(condition) || parser.parseComma() ||
-      parser.parseOperand(trueValue) || parser.parseComma() ||
-      parser.parseOperand(falseValue) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(conditionType) || parser.parseComma() ||
-      parser.parseType(argType))
-    return failure();
-
-  if (parser.resolveOperand(condition, conditionType, result.operands) ||
-      parser.resolveOperand(trueValue, argType, result.operands) ||
-      parser.resolveOperand(falseValue, argType, result.operands))
-    return failure();
-
-  result.addTypes(argType);
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::BrOp.
 //===----------------------------------------------------------------------===//
 
 static void printBrOp(OpAsmPrinter &p, BrOp &op) {
   p << op.getOperationName() << ' ';
   p.printSuccessorAndUseList(op.getOperation(), 0);
   p.printOptionalAttrDict(op.getAttrs());
 }

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 }
 
 void spirv::BitcastOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ConvertChainedBitcast>(context);
 }
 
 //===----------------------------------------------------------------------===//
-// spv.BitFieldInsert
-//===----------------------------------------------------------------------===//
-
-static ParseResult parseBitFieldInsertOp(OpAsmParser &parser,
-                                         OperationState &state) {
-  SmallVector<OpAsmParser::OperandType, 4> operandInfo;
-  Type baseType;
-  Type offsetType;
-  Type countType;
-  auto loc = parser.getCurrentLocation();
-
-  if (parser.parseOperandList(operandInfo, 4) || parser.parseColon() ||
-      parser.parseType(baseType) || parser.parseComma() ||
-      parser.parseType(offsetType) || parser.parseComma() ||
-      parser.parseType(countType) ||
-      parser.resolveOperands(operandInfo,
-                             {baseType, baseType, offsetType, countType}, loc,
-                             state.operands)) {
-    return failure();
-  }
-  state.addTypes(baseType);
-  return success();
-}
-
-static void print(spirv::BitFieldInsertOp bitFieldInsertOp,
-                  OpAsmPrinter &printer) {
-  printer << spirv::BitFieldInsertOp::getOperationName() << ' '
-          << bitFieldInsertOp.getOperands() << " : "
-          << bitFieldInsertOp.base().getType() << ", "
-          << bitFieldInsertOp.offset().getType() << ", "
-          << bitFieldInsertOp.count().getType();
-}
-
-static LogicalResult verify(spirv::BitFieldInsertOp bitFieldOp) {
-  auto baseType = bitFieldOp.base().getType();
-  auto insertType = bitFieldOp.insert().getType();
-  auto resultType = bitFieldOp.getResult().getType();
-
-  if ((baseType != insertType) || (baseType != resultType)) {
-    return bitFieldOp.emitError("expected the same type for the base operand, "
-                                "insert operand, and "
-                                "result, but provided ")
-           << baseType << ", " << insertType << " and " << resultType;
-  }
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // spv.BranchOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseBranchOp(OpAsmParser &parser, OperationState &state) {
   Block *dest;
   SmallVector<Value, 4> destOperands;
   if (parser.parseSuccessorAndUseList(dest, destOperands))
     return failure();

 // spv.Select
 //===----------------------------------------------------------------------===//
 
 void spirv::SelectOp::build(Builder *builder, OperationState &state, Value cond,
                             Value trueValue, Value falseValue) {
   build(builder, state, trueValue.getType(), cond, trueValue, falseValue);
 }
 
-static ParseResult parseSelectOp(OpAsmParser &parser, OperationState &state) {
-  OpAsmParser::OperandType condition;
-  SmallVector<OpAsmParser::OperandType, 2> operands;
-  SmallVector<Type, 2> types;
-  auto loc = parser.getCurrentLocation();
-  if (parser.parseOperand(condition) || parser.parseComma() ||
-      parser.parseOperandList(operands, 2) ||
-      parser.parseColonTypeList(types)) {
-    return failure();
-  }
-  if (types.size() != 2) {
-    return parser.emitError(
-        loc, "need exactly two trailing types for select condition and object");
-  }
-  if (parser.resolveOperand(condition, types[0], state.operands) ||
-      parser.resolveOperands(operands, types[1], state.operands)) {
-    return failure();
-  }
-  return parser.addTypesToList(types[1], state.types);
-}
-
-static void print(spirv::SelectOp op, OpAsmPrinter &printer) {
-  printer << spirv::SelectOp::getOperationName() << " " << op.getOperands()
-          << " : " << op.condition().getType() << ", " << op.result().getType();
-}
-
 static LogicalResult verify(spirv::SelectOp op) {
-  auto resultTy = op.result().getType();
-  if (op.true_value().getType() != resultTy) {
-    return op.emitOpError("result type and true value type must be the same");
-  }
-  if (op.false_value().getType() != resultTy) {
-    return op.emitOpError("result type and false value type must be the same");
-  }
   if (auto conditionTy = op.condition().getType().dyn_cast<VectorType>()) {
-    auto resultVectorTy = resultTy.dyn_cast<VectorType>();
+    auto resultVectorTy = op.result().getType().dyn_cast<VectorType>();
     if (!resultVectorTy) {
       return op.emitOpError("result expected to be of vector type when "
                             "condition is of vector type");
     }
     if (resultVectorTy.getNumElements() != conditionTy.getNumElements()) {
       return op.emitOpError("result should have the same number of elements as "
                             "the condition when condition is of vector type");
     }

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 void InsertOp::build(Builder *builder, OperationState &result, Value source,
                      Value dest, ArrayRef<int64_t> position) {
   result.addOperands({source, dest});
   auto positionAttr = getVectorSubscriptAttr(*builder, position);
   result.addTypes(dest.getType());
   result.addAttribute(getPositionAttrName(), positionAttr);
 }
 
-static void print(OpAsmPrinter &p, InsertOp op) {
-  p << op.getOperationName() << " " << op.source() << ", " << op.dest()
-    << op.position();
-  p.printOptionalAttrDict(op.getAttrs(), {InsertOp::getPositionAttrName()});
-  p << " : " << op.getSourceType() << " into " << op.getDestVectorType();
-}
-
-static ParseResult parseInsertOp(OpAsmParser &parser, OperationState &result) {
-  SmallVector<NamedAttribute, 4> attrs;
-  OpAsmParser::OperandType source, dest;
-  Type sourceType;
-  VectorType destType;
-  Attribute attr;
-  return failure(parser.parseOperand(source) || parser.parseComma() ||
-                 parser.parseOperand(dest) ||
-                 parser.parseAttribute(attr, InsertOp::getPositionAttrName(),
-                                       result.attributes) ||
-                 parser.parseOptionalAttrDict(attrs) ||
-                 parser.parseColonType(sourceType) ||
-                 parser.parseKeywordType("into", destType) ||
-                 parser.resolveOperand(source, sourceType, result.operands) ||
-                 parser.resolveOperand(dest, destType, result.operands) ||
-                 parser.addTypeToList(destType, result.types));
-}
-
 static LogicalResult verify(InsertOp op) {
   auto positionAttr = op.position().getValue();
   if (positionAttr.empty())
     return op.emitOpError("expected non-empty position attribute");
   auto destVectorType = op.getDestVectorType();
   if (positionAttr.size() > static_cast<unsigned>(destVectorType.getRank()))
     return op.emitOpError(
         "expected position attribute of rank smaller than dest vector rank");
void InsertStridedSliceOp::build(Builder *builder, OperationState &result,
   result.addOperands({source, dest});
   auto offsetsAttr = getVectorSubscriptAttr(*builder, offsets);
   auto stridesAttr = getVectorSubscriptAttr(*builder, strides);
   result.addTypes(dest.getType());
   result.addAttribute(getOffsetsAttrName(), offsetsAttr);
   result.addAttribute(getStridesAttrName(), stridesAttr);
 }
 
-static void print(OpAsmPrinter &p, InsertStridedSliceOp op) {
-  p << op.getOperationName() << " " << op.source() << ", " << op.dest() << " ";
-  p.printOptionalAttrDict(op.getAttrs());
-  p << " : " << op.getSourceVectorType() << " into " << op.getDestVectorType();
-}
-
-static ParseResult parseInsertStridedSliceOp(OpAsmParser &parser,
-                                             OperationState &result) {
-  OpAsmParser::OperandType source, dest;
-  VectorType sourceVectorType, destVectorType;
-  return failure(
-      parser.parseOperand(source) || parser.parseComma() ||
-      parser.parseOperand(dest) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(sourceVectorType) ||
-      parser.parseKeywordType("into", destVectorType) ||
-      parser.resolveOperand(source, sourceVectorType, result.operands) ||
-      parser.resolveOperand(dest, destVectorType, result.operands) ||
-      parser.addTypeToList(destVectorType, result.types));
-}
-
 // TODO(ntv) Should be moved to Tablegen Confined attributes.
 template <typename OpType>
 static LogicalResult isIntegerArrayAttrSmallerThanShape(OpType op,
                                                         ArrayAttr arrayAttr,
                                                         ArrayRef<int64_t> shape,
                                                         StringRef attrName) {
   if (arrayAttr.size() > shape.size())
     return op.emitOpError("expected ")

mlir/test/Dialect/SPIRV/ops.mlir

func @bit_field_insert_vec(%base: vector<3xi32>, %insert: vector<3xi32>, %offset: i32, %count: i16) -> vector<3xi32> {
   // CHECK: {{%.*}} = spv.BitFieldInsert {{%.*}}, {{%.*}}, {{%.*}}, {{%.*}} : vector<3xi32>, i32, i16
   %0 = spv.BitFieldInsert %base, %insert, %offset, %count : vector<3xi32>, i32, i16
   spv.ReturnValue %0 : vector<3xi32>
 }
 
 // -----
 
 func @bit_field_insert_invalid_insert_type(%base: vector<3xi32>, %insert: vector<2xi32>, %offset: i32, %count: i16) -> vector<3xi32> {
-  // expected-error @+1 {{expected the same type for the base operand, insert operand, and result, but provided 'vector<3xi32>', 'vector<2xi32>' and 'vector<3xi32>'}}
+  // expected-error @+1 {{all of {base, insert, result} have same type}}
   %0 = "spv.BitFieldInsert" (%base, %insert, %offset, %count) : (vector<3xi32>, vector<2xi32>, i32, i16) -> vector<3xi32>
   spv.ReturnValue %0 : vector<3xi32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.BitFieldSExtract
func @select_op_vec_condn_vec(%arg0: vector<3xi1>) -> () {
   return
 }
 
 // -----
 
 func @select_op(%arg0: i1) -> () {
   %0 = spv.constant 2 : i32
   %1 = spv.constant 3 : i32
-  // expected-error @+1 {{need exactly two trailing types for select condition and object}}
+  // expected-error @+2 {{expected ','}}
   %2 = spv.Select %arg0, %0, %1 : i1
   return
 }
 
 // -----
 
 func @select_op(%arg1: vector<3xi1>) -> () {
   %0 = spv.constant 2 : i32
func @select_op(%arg1: vector<4xi1>) -> () {
   return
 }
 
 // -----
 
 func @select_op(%arg1: vector<4xi1>) -> () {
   %0 = spv.constant dense<[2.0, 3.0, 4.0]> : vector<3xf32>
   %1 = spv.constant dense<[5, 6, 7]> : vector<3xi32>
-  // expected-error @+1 {{op result type and true value type must be the same}}
+  // expected-error @+1 {{all of {true_value, false_value, result} have same type}}
   %2 = "spv.Select"(%arg1, %0, %1) : (vector<4xi1>, vector<3xf32>, vector<3xi32>) -> vector<3xi32>
   return
 }
 
 // -----
 
 func @select_op(%arg1: vector<4xi1>) -> () {
   %0 = spv.constant dense<[2.0, 3.0, 4.0]> : vector<3xf32>
   %1 = spv.constant dense<[5, 6, 7]> : vector<3xi32>
-  // expected-error @+1 {{op result type and false value type must be the same}}
+  // expected-error @+1 {{all of {true_value, false_value, result} have same type}}
   %2 = "spv.Select"(%arg1, %1, %0) : (vector<4xi1>, vector<3xi32>, vector<3xf32>) -> vector<3xi32>
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.ShiftLeftLogical

mlir/test/mlir-tblgen/op-format-spec.td

 // RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s 2>&1 | FileCheck %s --dump-input-on-failure
 
 // This file contains tests for the specification of the declarative op format.
 
 include "mlir/IR/OpBase.td"
 
 def TestDialect : Dialect {
   let name = "test";
 }
-class TestFormat_Op<string name, string fmt> : Op<TestDialect, name> {
+class TestFormat_Op<string name, string fmt, list<OpTrait> traits = []>
+    : Op<TestDialect, name, traits> {
   let assemblyFormat = fmt;
 }
 
 //===----------------------------------------------------------------------===//
 // Directives
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

 }]>, Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
 def ZCoverageValidB : TestFormat_Op<"variable_valid_b", [{
   $operand type(operands) type(results) attr-dict
 }]>, Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
 def ZCoverageValidC : TestFormat_Op<"variable_valid_c", [{
   operands functional-type(operands, results) attr-dict
 }]>, Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
 
+// Check that we can infer type equalities from certain traits.
+def ZCoverageValidD : TestFormat_Op<"variable_valid_d", [{
+  operands type($result) attr-dict
+}], [SameOperandsAndResultType]>, Arguments<(ins AnyMemRef:$operand)>,
+     Results<(outs AnyMemRef:$result)>;
+def ZCoverageValidE : TestFormat_Op<"variable_valid_e", [{
+  $operand type($operand) attr-dict
+}], [SameOperandsAndResultType]>, Arguments<(ins AnyMemRef:$operand)>,
+     Results<(outs AnyMemRef:$result)>;
+def ZCoverageValidF : TestFormat_Op<"variable_valid_f", [{
+  operands type($other) attr-dict
+}], [SameTypeOperands]>, Arguments<(ins AnyMemRef:$operand, AnyMemRef:$other)>;
+def ZCoverageValidG : TestFormat_Op<"variable_valid_g", [{
+  operands type($other) attr-dict
+}], [AllTypesMatch<["operand", "other"]>]>,
+     Arguments<(ins AnyMemRef:$operand, AnyMemRef:$other)>;
+def ZCoverageValidH : TestFormat_Op<"variable_valid_h", [{
+  operands type($result) attr-dict
+}], [AllTypesMatch<["operand", "result"]>]>,
+     Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
+

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // OperationFormat
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct OperationFormat {
+  /// This class represents a specific resolver for an operand or result type.
+  class TypeResolution {
+  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; }
+
+  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;
+  };
+
   OperationFormat(const Operator &op)
       : allOperandTypes(false), allResultTypes(false) {
-    buildableOperandTypes.resize(op.getNumOperands(), llvm::None);
-    buildableResultTypes.resize(op.getNumResults(), llvm::None);
+    operandTypes.resize(op.getNumOperands(), TypeResolution());
+    resultTypes.resize(op.getNumResults(), TypeResolution());
   }
 
   /// Generate the operation parser from this format.
   void genParser(Operator &op, OpClass &opClass);
   /// Generate the c++ to resolve the types of operands and results during
   /// parsing.
   void genParserTypeResolution(Operator &op, OpMethodBody &body);
 
   /// Generate the operation printer from this format.
   void genPrinter(Operator &op, OpClass &opClass);
 
   /// The various elements in this format.
   std::vector<std::unique_ptr<Element>> elements;
 
   /// A flag indicating if all operand/result types were seen. If the format
   /// contains these, it can not contain individual type resolvers.
   bool allOperandTypes, allResultTypes;
 
   /// A map of buildable types to indices.
   llvm::MapVector<StringRef, int, llvm::StringMap<int>> buildableTypes;
 
   /// The index of the buildable type, if valid, for every operand and result.
-  std::vector<Optional<int>> buildableOperandTypes, buildableResultTypes;
+  std::vector<TypeResolution> operandTypes, resultTypes;
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // Parser Gen
 
 /// The code snippet used to generate a parser call for an attribute.
 ///
body << "  Type odsBuildableType" << it.second << " = parser.getBuilder()."
          << it.first << ";\n";
 
   // 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 = buildableResultTypes[i])
+      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";
       body << ");\n";
     }
   }
 
   // Early exit if there are no operands.
   if (op.getNumOperands() == 0)
if (hasAllOperands) {
     body << "  if (parser.resolveOperands(allOperands, ";
 
     // 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, [&](int i) {
-        if (Optional<int> val = buildableOperandTypes[i])
+        if (Optional<int> val = operandTypes[i].getBuilderIdx())
           body << "ArrayRef<Type>(odsBuildableType" << *val << ")";
+        else if (Optional<StringRef> var = operandTypes[i].getVariable())
+          body << *var << "Types";
         else
           body << op.getOperand(i).name << "Types";
       });
       body << ")";
     } else {
       body << op.operand_begin()->name << "Types";
     }
 
     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 = buildableOperandTypes[i])
+    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";
   }
 }
 
 //===----------------------------------------------------------------------===//
 // PrinterGen
Token::Kind kind = llvm::StringSwitch<Token::Kind>(str)
                          .Default(Token::identifier);
   return Token(kind, str);
 }
 
 //===----------------------------------------------------------------------===//
 // FormatParser
 //===----------------------------------------------------------------------===//
 
+/// Function to find an element within the given range that has the same name as
+/// 'name'.
+template <typename RangeT> static auto findArg(RangeT &&range, StringRef name) {
+  auto it = llvm::find_if(range, [=](auto &arg) { return arg.name == name; });
+  return it != range.end() ? &*it : nullptr;
+}
+
 namespace {
 /// This class implements a parser for an instance of an operation assembly
 /// format.
 class FormatParser {
 public:
   FormatParser(llvm::SourceMgr &mgr, OperationFormat &format, Operator &op)
       : lexer(mgr), curToken(lexer.lexToken()), fmt(format), op(op),
         seenOperandTypes(op.getNumOperands()),
         seenResultTypes(op.getNumResults()) {}
 
   /// 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);
+  /// 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,
+      bool includeResults);
+
   /// 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);
private:
   // Fields
   //===--------------------------------------------------------------------===//
 
   FormatLexer lexer;
   Token curToken;
   OperationFormat &fmt;
   Operator &op;
 
-  // The following are various bits of format state used for verification during
-  // parsing.
+  // The following are various bits of format state used for verification
+  // during parsing.
   bool hasAllOperands = false, hasAttrDict = false;
   llvm::SmallBitVector seenOperandTypes, seenResultTypes;
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedAttribute *> seenAttrs;
 };
 } // end anonymous namespace
 
 LogicalResult FormatParser::parse() {
   llvm::SMLoc loc = curToken.getLoc();
 
   // Parse each of the format elements into the main format.
   while (curToken.getKind() != Token::eof) {
     std::unique_ptr<Element> element;
     if (failed(parseElement(element, /*isTopLevel=*/true)))
       return failure();
     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;
+  for (const OpTrait &trait : op.getTraits()) {
+    const llvm::Record &def = trait.getDef();
+    if (def.isSubClassOf("AllTypesMatch"))
+      handleAllTypesMatchConstraint(def.getValueAsListOfStrings("values"),
+                                    variableTyResolver);
+    else if (def.getName() == "SameTypeOperands")
+      handleSameTypesConstraint(variableTyResolver, /*includeResults=*/false);
+    else if (def.getName() == "SameOperandsAndResultType")
+      handleSameTypesConstraint(variableTyResolver, /*includeResults=*/true);
+  }
+
   // 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);
+        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()) {
         return emitError(loc, "format missing instance of result #" + Twine(i) +
                                   "('" + result.name + "') type");
       }
       // Note in the format that this result uses the custom builder.
       auto it = buildableTypes.insert({*builder, buildableTypes.size()});
-      fmt.buildableResultTypes[i] = it.first->second;
+      fmt.resultTypes[i].setBuilderIdx(it.first->second);
     }
   }
 
   // Check that all of the operands are within the format, and their types can
   // be inferred.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);
 
     // Check that the operand itself is in the format.
     if (!hasAllOperands && !seenOperands.count(&operand)) {
       return emitError(loc, "format missing instance of operand #" + Twine(i) +
                                 "('" + operand.name + "')");
     }
 
     // Check that the operand type is in the format, or that it can be inferred.
-    if (!fmt.allOperandTypes && !seenOperandTypes.test(i)) {
-      // 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.buildableOperandTypes[i] = it.first->second;
-    }
+    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);
+      continue;
+    }
+
+      return emitError(loc, "format missing instance of operand #" + Twine(i) +
+                                "('" + operand.name + "') type");
+    fmt.operandTypes[i].setBuilderIdx(it.first->second);
   }
   return success();
 }
 
+void FormatParser::handleAllTypesMatchConstraint(
+    ArrayRef<StringRef> values,
+    llvm::StringMap<const NamedTypeConstraint *> &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 {
+      continue;
+    }
+
+    // Mark this value as the type resolver for the other variables.
+    for (unsigned j = 0; j != i; ++j)
+      variableTyResolver[values[j]] = arg;
+    for (unsigned j = i + 1; j != e; ++j)
+      variableTyResolver[values[j]] = arg;
+  }
+}
+
+void FormatParser::handleSameTypesConstraint(
+    llvm::StringMap<const NamedTypeConstraint *> &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;
+  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;
+  }
+}
+
 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();
 
-  // Functor used to find an element within the given range that has the same
-  // name as 'name'.
-  auto findArg = [&](auto &&range) {
-    auto it = llvm::find_if(range, [=](auto &arg) { return arg.name == name; });
-    return it != range.end() ? &*it : nullptr;
-  };
-
   // Check that the parsed argument is something actually registered on the op.
   /// Attributes
-  if (const NamedAttribute *attr = findArg(op.getAttributes())) {
+  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
-  if (const NamedTypeConstraint *operand = findArg(op.getOperands())) {
+  if (const NamedTypeConstraint *operand = findArg(op.getOperands(), name)) {
     if (isTopLevel) {
       if (hasAllOperands || !seenOperands.insert(operand).second)
         return emitError(loc, "operand '" + name + "' is already bound");
     }
     element = std::make_unique<OperandVariable>(operand);
     return success();
   }
   /// Results.
-  if (const NamedTypeConstraint *result = findArg(op.getResults())) {
+  if (const auto *result = findArg(op.getResults(), name)) {
     if (isTopLevel)
       return emitError(loc, "results can not be used at the top level");
     element = std::make_unique<ResultVariable>(result);
     return success();
   }
   return emitError(loc, "expected variable to refer to a argument or result");
 }