[mlir][ods] Generate inferReturnTypes for ops with TypesMatchWith

[Mogball]

Summary

Ops that use TypesMatchWith to constrain result types for verification
and to infer result types during parser generation should also be able
to have the inferReturnTypes method auto generated. This patch
upgrades the logic for generating inferReturnTypes to handle the
TypesMatchWith trait by building a type inference graph where each edge
corresponds to "type of A can be inferred from type of B", supporting
transformers other than "$_self".

[jpienaar]

I like all the red :)

mlir/lib/TableGen/Operator.cpp

394

Comments? (esp what does each element of the vector correspond to if any). Esp. given the name o fthe struct is quite general.

405–415

This comment feels out of place here.

436

I recall there was some uncertainty with variadics, but I also seem to recall there was a proposal to address it. What does this do with variadic operands & results?

475

known type vs already inferred confuses me here. E.g., an inferred type is known. Could we use other phrases?

509

Why do we need this here in condition slot vs in increment one? (also not sure why post-incrementing here)

512

Can we add a little static helper that would be more self-documenting than checking for negative or not?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

2525–2530

Topological order for result types feel weird :) This is in order to create valid computation structure so that inferred types can be functions of other inferred types correct?

2526

Document convention used here (I think it is just sentinel or known)

[Mogball]

mlir/lib/TableGen/Operator.cpp

436

If the op has variadic results, type inference is not generated (this is checked at the start of the function). It's not generally valid to have TypesMatchWith or AllTypesMatch mix variadic and non-variadic types anyways, so I didn't bother checking that.

509

This is an llvm::make_early_inc_range except I still need the iterator object itself to call std::list::erase.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

2525–2530

Yes. I'll clarify what this means: topological wrt the type inference (directed acyclic) graph.

[lattner]

This is so great Jeff, thank you for driving this forward, I love the boilerplate destruction!

[rriddle]

mlir/lib/TableGen/Operator.cpp

436

Why not? CallIndirectOp is an easy example of when this happens.

453

Missing continue; after this?

[Mogball]

mlir/lib/TableGen/Operator.cpp

436

This is a pre-existing restriction. I can make another pass and see if I can add variadic result support.

[Mogball]

review comments

[Mogball]

fix build

[rriddle]

mlir/include/mlir/TableGen/Operator.h

49

Can you document this? This is referring to the index we use to transform right?

[Mogball]

mlir/include/mlir/TableGen/Operator.h

49

Yes. I'll add a doc

[Mogball]

rebase+review comments

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

%x = "arith.cmpi"(%lhs, %rhs) {predicate = 0 : i64}
         : (vector<4xi64>, vector<4xi64>) -> vector<4xi1>
     ```
   }];
 
   let arguments = (ins Arith_CmpIPredicateAttr:$predicate,
                        SignlessIntegerLikeOfAnyRank:$lhs,
                        SignlessIntegerLikeOfAnyRank:$rhs);
 
-  let builders = [
-    OpBuilder<(ins "CmpIPredicate":$predicate, "Value":$lhs, "Value":$rhs), [{
-      build($_builder, $_state, ::getI1SameShape(lhs.getType()),
-            predicate, lhs, rhs);
-    }]>
-  ];
-
   let extraClassDeclaration = [{
     static arith::CmpIPredicate getPredicateByName(StringRef name);
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
let description = [{
     %r3 = "arith.cmpf"(%0, %1) {predicate: 0} : (f8, f8) -> i1
     ```
   }];
 
   let arguments = (ins Arith_CmpFPredicateAttr:$predicate,
                        FloatLike:$lhs,
                        FloatLike:$rhs);
 
-  let builders = [
-    OpBuilder<(ins "CmpFPredicate":$predicate, "Value":$lhs, "Value":$rhs), [{
-      build($_builder, $_state, ::getI1SameShape(lhs.getType()),
-            predicate, lhs, rhs);
-    }]>
-  ];
-
   let extraClassDeclaration = [{
     static arith::CmpFPredicate getPredicateByName(StringRef name);
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let description = [{
     If tensor load is used in the bufferization steps, mutating the source
     buffer after loading leads to undefined behavior.
   }];
 
   let arguments = (ins Arg<AnyRankedOrUnrankedMemRef,
                        "the reference to load from", [MemRead]>:$memref);
   let results = (outs AnyTensor:$result);
 
-  let builders = [
-    OpBuilder<(ins "Value":$memref), [{
-      $_state.addOperands(memref);
-      $_state.addTypes(memref::getTensorTypeFromMemRefType(memref.getType()));
-    }]>];
-
   let extraClassDeclaration = [{
     /// The result of a to_tensor is always a tensor.
     TensorType getType() {
       Type resultType = getResult().getType();
       if (resultType.isa<TensorType>())
         return resultType.cast<TensorType>();
       return {};
     }

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

 
 // Other integer operations.
 def LLVM_ICmpOp : LLVM_ArithmeticCmpOp<"icmp", [Pure]> {
   let arguments = (ins ICmpPredicate:$predicate,
                    AnyTypeOf<[LLVM_ScalarOrVectorOf<AnyInteger>,
                               LLVM_ScalarOrVectorOf<LLVM_AnyPointer>]>:$lhs,
                    AnyTypeOf<[LLVM_ScalarOrVectorOf<AnyInteger>,
                               LLVM_ScalarOrVectorOf<LLVM_AnyPointer>]>:$rhs);
-  let builders = [
-    OpBuilder<(ins "ICmpPredicate":$predicate, "Value":$lhs, "Value":$rhs)>
-  ];
   let hasCustomAssemblyFormat = 1;
   string llvmInstName = "ICmp";
   string llvmBuilder = [{
     $res = builder.CreateICmp(getLLVMCmpPredicate($predicate), $lhs, $rhs);
   }];
   string mlirBuilder = [{
     auto *iCmpInst = cast<llvm::ICmpInst>(inst);
     $res = $_builder.create<$_qualCppClassName>(
       $_location, getICmpPredicate(iCmpInst->getPredicate()), $lhs, $rhs);
   }];
   // Set the $predicate index to -1 to indicate there is no matching operand
   // and decrement the following indices.
   list<int> llvmArgIndices = [-1, 0, 1];
 }
 
 // Other floating-point operations.
 def LLVM_FCmpOp : LLVM_ArithmeticCmpOp<"fcmp", [
     Pure, DeclareOpInterfaceMethods<FastmathFlagsInterface>]> {
   let arguments = (ins FCmpPredicate:$predicate,
                    LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$lhs,
                    LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$rhs,
                    DefaultValuedAttr<LLVM_FastmathFlagsAttr,
                                      "{}">:$fastmathFlags);
-  let builders = [
-    OpBuilder<(ins "FCmpPredicate":$predicate, "Value":$lhs, "Value":$rhs)>
-  ];
   let hasCustomAssemblyFormat = 1;
   string llvmInstName = "FCmp";
   string llvmBuilder = [{
     $res = builder.CreateFCmp(getLLVMCmpPredicate($predicate), $lhs, $rhs);
   }];
   string mlirBuilder = [{
     auto *fCmpInst = cast<llvm::FCmpInst>(inst);
     auto op = $_builder.create<$_qualCppClassName>(

 def LLVM_ExtractElementOp : LLVM_Op<"extractelement", [Pure,
     TypesMatchWith<"result type matches vector element type", "vector", "res",
                    "LLVM::getVectorElementType($_self)">]> {
   let summary = "Extract an element from an LLVM vector.";
 
   let arguments = (ins LLVM_AnyVector:$vector, AnyInteger:$position);
   let results = (outs LLVM_Type:$res);
 
-  let builders = [
-    OpBuilder<(ins "Value":$vector, "Value":$position,
-                   CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
-  ];
-
   let assemblyFormat = [{
     $vector `[` $position `:` type($position) `]` attr-dict `:` type($vector)
   }];
 
   string llvmInstName = "ExtractElement";
   string llvmBuilder = [{
     $res = builder.CreateExtractElement($vector, $position);
   }];

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     in these contexts.
   }];
 
   let arguments = (ins Arg<AnyMemRef, "the reference to load from",
                            [MemRead]>:$memref,
                        Variadic<Index>:$indices);
   let results = (outs AnyType:$result);
 
-  let builders = [
-    OpBuilder<(ins "Value":$memref, CArg<"ValueRange", "{}">:$indices), [{
-      auto memrefType = memref.getType().cast<MemRefType>();
-      $_state.addOperands(memref);
-      $_state.addOperands(indices);
-      $_state.types.push_back(memrefType.getElementType());
-    }]>];
-
   let extraClassDeclaration = [{
     Value getMemRef() { return getOperand(0); }
     void setMemRef(Value value) { setOperand(0, value); }
     MemRefType getMemRefType() {
       return getMemRef().getType().cast<MemRefType>();
     }
   }];
 

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

class SPIRV_LogicalBinaryOp<string mnemonic, Type operandsType,
                      Pure, SameTypeOperands,
                      SameOperandsAndResultShape,
                      TypesMatchWith<"type of result to correspond to the `i1` "
                                     "equivalent of the operand",
                                     "operand1", "result",
                                     "getUnaryOpResultType($_self)"
                      >])> {
   let assemblyFormat = "$operand1 `,` $operand2 `:` type($operand1) attr-dict";
-
-  let builders = [
-    OpBuilder<(ins "Value":$lhs, "Value":$rhs),
-    [{::buildLogicalBinaryOp($_builder, $_state, lhs, rhs);}]>
-  ];
 }
 
 class SPIRV_LogicalUnaryOp<string mnemonic, Type operandType,
                          list<Trait> traits = []> :
       // Result type is SPIRV_Bool.
       SPIRV_UnaryOp<mnemonic, SPIRV_Bool, operandType,
                   !listconcat(traits, [
                      Pure, SameTypeOperands, SameOperandsAndResultShape,
                      TypesMatchWith<"type of result to correspond to the `i1` "
                                     "equivalent of the operand",
                                     "operand", "result",
                                     "getUnaryOpResultType($_self)"
                      >])> {
   let assemblyFormat = "$operand `:` type($operand) attr-dict";
-
-  let builders = [
-    OpBuilder<(ins "Value":$value),
-    [{::buildLogicalUnaryOp($_builder, $_state, value);}]>
-  ];
 }
 
 // -----
 
 def SPIRV_FOrdEqualOp : SPIRV_LogicalBinaryOp<"FOrdEqual", SPIRV_Float, [Commutative]> {
   let summary = "Floating-point comparison for being ordered and equal.";
 
   let description = [{

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let description = [{
     %5 = tensor.extract %rt[%1, %2] : tensor<?x?xi32>
     ```
   }];
 
   let arguments = (ins AnyRankedTensor:$tensor, Variadic<Index>:$indices);
   let results = (outs AnyType:$result);
   let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)";
 
-  let builders = [
-    OpBuilder<(ins "Value":$tensor, CArg<"ValueRange", "{}">:$indices), [{
-      auto resType = tensor.getType().cast<ShapedType>().getElementType();
-      build($_builder, $_state, resType, tensor, indices);
-    }]>];
-
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 
 //===----------------------------------------------------------------------===//
 // ExtractSliceOp
let description = [{
 
     An extract_slice operation may additionally reduce the rank of the resulting
     tensor by removing dimensions that are statically known to be of size 1.
     This rank-reduction behavior is not required by the op semantics: this
     flexibility allows to progressively drop unit dimensions while lowering
     between different flavors of ops on that operate on tensors.
 
     #### Verification vs Inference in the rank-reduced case
-    
     Note that there may be multiple ways to infer a resulting rank-reduced type.
       e.g. 1x6x1 could potentially rank-reduce to either 1x6 or 6x1 2-D shapes.
 
     To disambiguate, the inference helpers `inferCanonicalRankReducedResultType`
     only drop the first unit dimensions, in order:
       e.g. 1x6x1 rank-reduced to 2-D will infer the 6x1 2-D shape, but not 1x6.
 
     Verification however has access to result type and does not need to infer.
def Tensor_InsertOp : Tensor_Op<"insert", [
   let arguments = (ins AnyType:$scalar,
                        AnyRankedTensor:$dest,
                        Variadic<Index>:$indices);
   let results = (outs AnyRankedTensor:$result);
   let assemblyFormat = [{
     $scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest)
   }];
 
-  let builders = [
-    OpBuilder<(ins "Value":$scalar, "Value":$dest,
-      CArg<"ValueRange", "{}">:$indices), [{
-      auto resType = dest.getType();
-      build($_builder, $_state, resType, scalar, dest, indices);
-    }]>];
-
   let extraClassDeclaration = [{
     std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
       return {1, 2};  // `dest` operand
     }
   }];
 
   let hasFolder = 1;
   let hasVerifier = 1;
let description = [{
     known to be of size 1.
     This rank-altering behavior is not required by the op semantics: this
     flexibility allows to progressively drop unit dimensions while lowering
     between different flavors of ops on that operate on tensors.
     The rank-altering behavior of tensor.insert_slice matches the rank-reducing
     behavior of tensor.extract_slice.
 
     #### Verification in the rank-reduced case
-    
     The same verification discussion and mechanisms apply as for ExtractSliceOp.
     Unlike ExtractSliceOp however, there is no need for a specific inference.
 
     Example:
 
     ```mlir
     // Rank-altering insert_slice.
     %1 = tensor.insert_slice %t into %0[0, 0, 0][1, 16, 4][1, 1, 1] :
let description = [{
     statically known to be of size 1.
     This rank-altering behavior is not required by the op semantics: this
     flexibility allows to progressively drop unit dimensions while lowering
     between different flavors of ops on that operate on tensors.
     The rank-altering behavior of tensor.parallel_insert_slice matches the
     rank-reducing behavior of tensor.insert_slice and tensor.extract_slice.
 
     #### Verification in the rank-reduced case
-    
     The same verification discussion and mechanisms apply as for ExtractSliceOp.
     Unlike ExtractSliceOp however, there is no need for a specific inference.
   }];
 
   let arguments = (ins
     AnyRankedTensor:$source,
     AnyRankedTensor:$dest,
     Variadic<Index>:$offsets,

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

def Vector_ExtractElementOp :
   }];
   let assemblyFormat = [{
     $vector `[` ($position^ `:` type($position))? `]` attr-dict `:` type($vector)
   }];
 
   let builders = [
     // 0-D builder.
     OpBuilder<(ins "Value":$source)>,
-    // 1-D + position builder.
-    OpBuilder<(ins "Value":$source, "Value":$position)>,
   ];
   let extraClassDeclaration = [{
     VectorType getVectorType() {
       return getVector().getType().cast<VectorType>();
     }
   }];
   let hasVerifier = 1;
   let hasFolder = 1;

mlir/include/mlir/TableGen/Operator.h

 class DefInit;
 class Record;
 class StringInit;
 } // namespace llvm
 
 namespace mlir {
 namespace tblgen {
 
+/// This class represents an inferred result type. The result type can be
+/// inferred from an argument or result type. If it is inferred from another
+/// result type, that type must be buildable or inferred from yet another type.
+class InferredResultType {
+public:
+  InferredResultType(int index, std::string transformer)
+      : index(index), transformer(std::move(transformer)) {}
+
+  /// Returns true if result type is inferred from an argument type.
+  bool isArg() const { return isArgIndex(index); }
+  /// Return the mapped argument or result index.
+  int getIndex() const { return index; }
+  /// If the type is inferred from a result, return the result index.
+  int getResultIndex() const { return unmapResultIndex(index); }
+
+  // Mapping from result index to combined argument and result index.
+  // Arguments are indexed to match getArg index, while the result indexes are
+  // mapped to avoid overlap.
+  static int mapResultIndex(int i) { return -1 - i; }
+  static int unmapResultIndex(int i) { return -i - 1; }
+  static bool isResultIndex(int i) { return i < 0; }
+  static bool isArgIndex(int i) { return i >= 0; }
+
+  StringRef getTransformer() const { return transformer; }
+
+private:
+  /// The index of the source argument or result.
+  int index;
+
+  /// The transfer to apply to the type to obtain the inferred type.
+  std::string transformer;
+};
+
 /// Wrapper class that contains a MLIR op's information (e.g., operands,
 /// attributes) defined in TableGen and provides helper methods for
 /// accessing them.
 class Operator {
 public:
   explicit Operator(const llvm::Record &def);
   explicit Operator(const llvm::Record *def) : Operator(*def) {}
 
public:
 
   /// Prints the contents in this operator to the given `os`. This is used for
   /// debugging purposes.
   void print(llvm::raw_ostream &os) const;
 
   /// Return whether all the result types are known.
   bool allResultTypesKnown() const { return allResultsHaveKnownTypes; };
 
-  /// Pair representing either a index to an argument or a type constraint. Only
-  /// one of these entries should have the non-default value.
-  struct ArgOrType {
-    explicit ArgOrType(int index) : index(index), constraint(std::nullopt) {}
-    explicit ArgOrType(TypeConstraint constraint)
-        : index(std::nullopt), constraint(constraint) {}
-    bool isArg() const {
-      assert(constraint.has_value() ^ index.has_value());
-      return index.has_value();
-    }
-    bool isType() const {
-      assert(constraint.has_value() ^ index.has_value());
-      return constraint.has_value();
-    }
-
-    int getArg() const { return *index; }
-    TypeConstraint getType() const { return *constraint; }
-
-  private:
-    std::optional<int> index;
-    std::optional<TypeConstraint> constraint;
-  };
-
-  /// Return all arguments or type constraints with same type as result[index].
+  ///  Return all arguments or type constraints with same type as result[index].
   /// Requires: all result types are known.
-  ArrayRef<ArgOrType> getSameTypeAsResult(int index) const;
+  const InferredResultType &getInferredResultType(int index) const;
 
   /// Pair consisting kind of argument and index into operands or attributes.
   struct OperandOrAttribute {
     enum class Kind { Operand, Attribute };
     OperandOrAttribute(Kind kind, int index) {
       packed = (index << 1) | (kind == Kind::Attribute);
     }
     int operandOrAttributeIndex() const { return (packed >> 1); }
private:
 
   /// The traits of the op.
   SmallVector<Trait, 4> traits;
 
   /// The regions of this op.
   SmallVector<NamedRegion, 1> regions;
 
   /// The argument with the same type as the result.
-  SmallVector<SmallVector<ArgOrType, 2>, 4> resultTypeMapping;
+  SmallVector<InferredResultType> resultTypeMapping;
 
   /// Map from argument to attribute or operand number.
   SmallVector<OperandOrAttribute, 4> attrOrOperandMapping;
 
   /// The builders of this operator.
   SmallVector<Builder> builders;
 
   /// The number of native attributes stored in the leading positions of

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

static BufferAllocs allocBuffers(OpBuilder &b, OpTy xferOp) {
   auto bufferType = MemRefType::get({}, xferOp.getVectorType());
   result.dataBuffer = b.create<memref::AllocaOp>(loc, bufferType);
 
   if (xferOp.getMask()) {
     auto maskType = MemRefType::get({}, xferOp.getMask().getType());
     auto maskBuffer = b.create<memref::AllocaOp>(loc, maskType);
     b.setInsertionPoint(xferOp);
     b.create<memref::StoreOp>(loc, xferOp.getMask(), maskBuffer);
-    result.maskBuffer = b.create<memref::LoadOp>(loc, maskBuffer);
+    result.maskBuffer = b.create<memref::LoadOp>(loc, maskBuffer, ValueRange());
   }
 
   return result;
 }
 
 /// Given a MemRefType with VectorType element type, unpack one dimension from
 /// the VectorType into the MemRefType.
 ///

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

if (LLVM::isCompatibleVectorType(type))
     return LLVM::getVectorType(i1Type, LLVM::getVectorNumElements(type));
   return i1Type;
 }
 
 //===----------------------------------------------------------------------===//
 // Printing, parsing and builder for LLVM::CmpOp.
 //===----------------------------------------------------------------------===//
 
-void ICmpOp::build(OpBuilder &builder, OperationState &result,
-                   ICmpPredicate predicate, Value lhs, Value rhs) {
-  build(builder, result, getI1SameShape(lhs.getType()), predicate, lhs, rhs);
-}
-
-void FCmpOp::build(OpBuilder &builder, OperationState &result,
-                   FCmpPredicate predicate, Value lhs, Value rhs) {
-  build(builder, result, getI1SameShape(lhs.getType()), predicate, lhs, rhs);
-}
-
 void ICmpOp::print(OpAsmPrinter &p) {
   p << " \"" << stringifyICmpPredicate(getPredicate()) << "\" " << getOperand(0)
     << ", " << getOperand(1);
   p.printOptionalAttrDict((*this)->getAttrs(), {"predicate"});
   p << " : " << getLhs().getType();
 }
 
 void FCmpOp::print(OpAsmPrinter &p) {
if (isDirect) {
     if (!llvmResultType.isa<LLVM::LLVMVoidType>())
       result.addTypes(llvmResultType);
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// ExtractElementOp
-//===----------------------------------------------------------------------===//
-
-/// Expects vector to be an LLVM vector type and position to be an integer type.
-void LLVM::ExtractElementOp::build(OpBuilder &b, OperationState &result,
-                                   Value vector, Value position,
-                                   ArrayRef<NamedAttribute> attrs) {
-  auto vectorType = vector.getType();
-  auto llvmType = LLVM::getVectorElementType(vectorType);
-  build(b, result, llvmType, vector, position);
-  result.addAttributes(attrs);
-}
-
-//===----------------------------------------------------------------------===//
 // ExtractValueOp
 //===----------------------------------------------------------------------===//
 
 /// Extract the type at `position` in the LLVM IR aggregate type
 /// `containerType`. Each element of `position` is an index into a nested
 /// aggregate type. Return the resulting type or emit an error.
 static Type getInsertExtractValueElementType(
     function_ref<InFlightDiagnostic(StringRef)> emitError, Type containerType,

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

if (op->getOperand(0).getType() != op->getResult(0).getType()) {
     return op->emitError("expected the same type for the first operand and "
                          "result, but provided ")
            << op->getOperand(0).getType() << " and "
            << op->getResult(0).getType();
   }
   return success();
 }
 
-static void buildLogicalBinaryOp(OpBuilder &builder, OperationState &state,
-                                 Value lhs, Value rhs) {
-  assert(lhs.getType() == rhs.getType());
-
-  Type boolType = builder.getI1Type();
-  if (auto vecType = lhs.getType().dyn_cast<VectorType>())
-    boolType = VectorType::get(vecType.getShape(), boolType);
-  state.addTypes(boolType);
-
-  state.addOperands({lhs, rhs});
-}
-
-static void buildLogicalUnaryOp(OpBuilder &builder, OperationState &state,
-                                Value value) {
-  Type boolType = builder.getI1Type();
-  if (auto vecType = value.getType().dyn_cast<VectorType>())
-    boolType = VectorType::get(vecType.getShape(), boolType);
-  state.addTypes(boolType);
-
-  state.addOperands(value);
-}
-
 //===----------------------------------------------------------------------===//
 // spirv.AccessChainOp
 //===----------------------------------------------------------------------===//
 
 static Type getElementPtrType(Type type, ValueRange indices, Location baseLoc) {
   auto ptrType = type.dyn_cast<spirv::PointerType>();
   if (!ptrType) {
     emitError(baseLoc, "'spirv.AccessChain' op expected a pointer "

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 //===----------------------------------------------------------------------===//
 
 void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
                                      Value source) {
   result.addOperands({source});
   result.addTypes(source.getType().cast<VectorType>().getElementType());
 }
 
-void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
-                                     Value source, Value position) {
-  result.addOperands({source, position});
-  result.addTypes(source.getType().cast<VectorType>().getElementType());
-}
-
 LogicalResult vector::ExtractElementOp::verify() {
   VectorType vectorType = getVectorType();
   if (vectorType.getRank() == 0) {
     if (getPosition())
       return emitOpError("expected position to be empty with 0-D vector");
     return success();
   }
   if (vectorType.getRank() != 1)

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

static void createFullPartialVectorTransferWrite(RewriterBase &b,
   Location loc = xferOp.getLoc();
   auto notInBounds = b.create<arith::XOrIOp>(
       loc, inBoundsCond, b.create<arith::ConstantIntOp>(loc, true, 1));
   b.create<scf::IfOp>(loc, notInBounds, [&](OpBuilder &b, Location loc) {
     IRMapping mapping;
     Value load = b.create<memref::LoadOp>(
         loc,
         b.create<vector::TypeCastOp>(
-            loc, MemRefType::get({}, xferOp.getVector().getType()), alloc));
+            loc, MemRefType::get({}, xferOp.getVector().getType()), alloc),
+        ValueRange());
     mapping.map(xferOp.getVector(), load);
     b.clone(*xferOp.getOperation(), mapping);
     b.create<scf::YieldOp>(loc, ValueRange{});
   });
 }
 
 // TODO: Parallelism and threadlocal considerations with a ParallelScope trait.
 static Operation *getAutomaticAllocationScope(Operation *op) {

mlir/lib/TableGen/Operator.cpp

 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
+#include <list>
 
 #define DEBUG_TYPE "mlir-tblgen-operator"
 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 using llvm::DagInit;
 using llvm::DefInit;
auto Operator::operand_end() const -> const_value_iterator {
   return operands.end();
 }
 auto Operator::getOperands() const -> const_value_range {
   return {operand_begin(), operand_end()};
 }
 
 auto Operator::getArg(int index) const -> Argument { return arguments[index]; }
 
-// Mapping from result index to combined argument and result index. Arguments
-// are indexed to match getArg index, while the result indexes are mapped to
-// avoid overlap.
-static int resultIndex(int i) { return -1 - i; }
-
 bool Operator::isVariadic() const {
   return any_of(llvm::concat<const NamedTypeConstraint>(operands, results),
                 [](const NamedTypeConstraint &op) { return op.isVariadic(); });
 }
 
 void Operator::populateTypeInferenceInfo(
     const llvm::StringMap<int> &argumentsAndResultsIndex) {
   // If the type inference op interface is not registered, then do not attempt
if (getTrait("::mlir::OpTrait::SameOperandsAndResultType")) {
     // Check for a non-variable length operand to use as the type anchor.
     auto *operandI = llvm::find_if(arguments, [](const Argument &arg) {
       NamedTypeConstraint *operand = arg.dyn_cast<NamedTypeConstraint *>();
       return operand && !operand->isVariableLength();
     });
     if (operandI == arguments.end())
       return;
 
-    // Map each of the result types to the anchor operation.
+    // All result types are inferred from the operand type.
     int operandIdx = operandI - arguments.begin();
-    resultTypeMapping.resize(getNumResults());
     for (int i = 0; i < getNumResults(); ++i)
-      resultTypeMapping[i].emplace_back(operandIdx);
+      resultTypeMapping.emplace_back(operandIdx, "$_self");
 
     allResultsHaveKnownTypes = true;
     traits.push_back(Trait::create(inferTrait->getDefInit()));
     return;
   }
 
-  // We create equivalence classes of argument/result types where arguments
-  // and results are mapped into the same index space and indices corresponding
-  // to the same type are in the same equivalence class.
-  llvm::EquivalenceClasses<int> ecs;
-  resultTypeMapping.resize(getNumResults());
-  // Captures the argument whose type matches a given result type. Preference
-  // towards capturing operands first before attributes.
-  auto captureMapping = [&](int i) {
-    bool found = false;
-    ecs.insert(resultIndex(i));
-    auto mi = ecs.findLeader(resultIndex(i));
-    for (auto me = ecs.member_end(); mi != me; ++mi) {
-      if (*mi < 0) {
-        auto tc = getResultTypeConstraint(i);
-        if (tc.getBuilderCall()) {
-          resultTypeMapping[i].emplace_back(tc);
-          found = true;
-        }
-        continue;
-      }
+  /// This struct represents a node in this operation's result type inferenece
+  /// graph. Each node has a list of incoming type inference edges `sources`.
+  /// Each edge represents a "source" from which the result type can be
+  /// inferred, either an operand (leaf) or another result (node). When a node
+  /// is known to have a fully-inferred type, `inferred` is set to true.
+  struct ResultTypeInference {
+    /// The list of incoming type inference edges.
+    SmallVector<InferredResultType> sources;
+    /// This flag is set to true when the result type is known to be inferrable.
+    bool inferred = false;
+  };
 
-      resultTypeMapping[i].emplace_back(*mi);
-      found = true;
-    }
-    return found;
-  };
+  // This vector represents the type inference graph, with one node for each
+  // operation result. The nth element is the node for the nth result.
+  SmallVector<ResultTypeInference> inference(getNumResults(), {});
+
+  // For all results whose types are buildable, initialize their type inference
+  // nodes with an edge to themselves. Mark those nodes are fully-inferred.
+  for (auto &[idx, infer] : llvm::enumerate(inference)) {
+    if (getResult(idx).constraint.getBuilderCall()) {
+      infer.sources.emplace_back(InferredResultType::mapResultIndex(idx),
+                                 "$_self");
+      infer.inferred = true;
+    }
 
+  // Use `AllTypesMatch` and `TypesMatchWith` operation traits to build the
+  // result type inference graph.
   for (const Trait &trait : traits) {
     const llvm::Record &def = trait.getDef();
+
     // If the infer type op interface was manually added, then treat it as
     // intention that the op needs special handling.
     // TODO: Reconsider whether to always generate, this is more conservative
     // and keeps existing behavior so starting that way for now.
     if (def.isSubClassOf(
             llvm::formatv("{0}::Trait", inferTypeOpInterface).str()))
       return;
     if (const auto *traitDef = dyn_cast<InterfaceTrait>(&trait))
       if (&traitDef->getDef() == inferTrait)
         return;
 
+    // The `TypesMatchWith` trait represents a 1 -> 1 type inference edge with a
+    // type transformer.
+    if (def.isSubClassOf("TypesMatchWith")) {
+      int target = argumentsAndResultsIndex.lookup(def.getValueAsString("rhs"));
+      // Ignore operand type inference.
+      if (InferredResultType::isArgIndex(target))
+        continue;
+      int resultIndex = InferredResultType::unmapResultIndex(target);
+      ResultTypeInference &infer = inference[resultIndex];
+      // If the type of the result has already been inferred, do nothing.
+      if (infer.inferred)
+        continue;
+      int sourceIndex =
+          argumentsAndResultsIndex.lookup(def.getValueAsString("lhs"));
+      infer.sources.emplace_back(sourceIndex,
+                                 def.getValueAsString("transformer").str());
+      // Locally propagate inferredness.
+      infer.inferred =
+          InferredResultType::isArgIndex(sourceIndex) ||
+          inference[InferredResultType::unmapResultIndex(sourceIndex)].inferred;
+      continue;
+    }
+
     if (!def.isSubClassOf("AllTypesMatch"))
       continue;
 
     auto values = def.getValueAsListOfStrings("values");
-    auto root = argumentsAndResultsIndex.lookup(values.front());
-    for (StringRef str : values)
-      ecs.unionSets(argumentsAndResultsIndex.lookup(str), root);
+    // The `AllTypesMatch` trait represents an N <-> N fanin and fanout. That
+    // is, every result type has an edge from every other type. However, if any
+    // one of the values refers to an operand or a result with a fully-inferred
+    // type, we can infer all other types from that value. Try to find a
+    // fully-inferred type in the list.
+    std::optional<int> fullyInferredIndex;
+    SmallVector<int> resultIndices;
+    for (StringRef name : values) {
+      int index = argumentsAndResultsIndex.lookup(name);
+      if (InferredResultType::isResultIndex(index))
+        resultIndices.push_back(InferredResultType::unmapResultIndex(index));
+      if (InferredResultType::isArgIndex(index) ||
+          inference[InferredResultType::unmapResultIndex(index)].inferred)
+        fullyInferredIndex = index;
+    }
+    if (fullyInferredIndex) {
+      // Make the fully-inferred type the only source for all results that
+      // aren't already inferred -- a 1 -> N fanout.
+      for (int resultIndex : resultIndices) {
+        ResultTypeInference &infer = inference[resultIndex];
+        if (!infer.inferred) {
+          infer.sources.assign(1, {*fullyInferredIndex, "$_self"});
+          infer.inferred = true;
         }
+      }
+    } else {
+      // Add an edge between every result and every other type; N <-> N.
+      for (int resultIndex : resultIndices) {
+        for (int otherResultIndex : resultIndices) {
+          if (resultIndex == otherResultIndex)
+            continue;
+          inference[resultIndex].sources.emplace_back(otherResultIndex,
+                                                      "$_self");
+        }
+      }
+    }
+  }
+
+  // Propagate inferredness until a fixed point.
+  std::list<ResultTypeInference *> worklist;
+  for (ResultTypeInference &infer : inference)
+    if (!infer.inferred)
+      worklist.push_back(&infer);
+  bool changed;
+  do {
+    changed = false;
+    // This is `llvm::make_early_inc_range` but keeps the iterator for erasing.
+    for (auto earlyIncIt = worklist.begin(), cur = earlyIncIt;
+         cur = earlyIncIt++, cur != worklist.end();) {
+      ResultTypeInference &infer = **cur;
+      for (auto &[idx, source] : llvm::enumerate(infer.sources)) {
+        assert(InferredResultType::isResultIndex(source.getIndex()));
+        if (inference[InferredResultType::unmapResultIndex(source.getIndex())]
+                .inferred) {
+          changed = true;
+          infer.inferred = true;
+          // Make this the only source for the result. This breaks any cycles.
+          infer.sources.assign(1, source);
+          worklist.erase(cur);
+          break;
+        }
+      }
+    }
+  } while (changed);
 
-  // Verifies that all output types have a corresponding known input type
-  // and chooses matching operand or attribute (in that order) that
-  // matches it.
-  allResultsHaveKnownTypes =
-      all_of(llvm::seq<int>(0, getNumResults()), captureMapping);
+  allResultsHaveKnownTypes = worklist.empty();
 
   // If the types could be computed, then add type inference trait.
-  if (allResultsHaveKnownTypes)
+  if (allResultsHaveKnownTypes) {
     traits.push_back(Trait::create(inferTrait->getDefInit()));
+    for (const ResultTypeInference &infer : inference)
+      resultTypeMapping.push_back(infer.sources.front());
+  }
 }
 
 void Operator::populateOpStructure() {
   auto &recordKeeper = def.getRecords();
   auto *typeConstraintClass = recordKeeper.getClass("TypeConstraint");
   auto *attrClass = recordKeeper.getClass("Attr");
   auto *derivedAttrClass = recordKeeper.getClass("DerivedAttr");
   auto *opVarClass = recordKeeper.getClass("OpVariable");
if (!resultInit) {
       PrintFatalError(def.getLoc(),
                       Twine("undefined type for result #") + Twine(i));
     }
     auto *resultDef = resultInit->getDef();
     if (resultDef->isSubClassOf(opVarClass))
       resultDef = resultDef->getValueAsDef("constraint");
     results.push_back({name, TypeConstraint(resultDef)});
     if (!name.empty())
-      argumentsAndResultsIndex[name] = resultIndex(i);
+      argumentsAndResultsIndex[name] = InferredResultType::mapResultIndex(i);
 
     // We currently only support VariadicOfVariadic operands.
     if (results.back().constraint.isVariadicOfVariadic()) {
       PrintFatalError(
           def.getLoc(),
           "'VariadicOfVariadic' results are currently not supported");
     }
   }
if (builderList && !builderList->empty()) {
     PrintFatalError(
         def.getLoc(),
         "default builders are skipped and no custom builders provided");
   }
 
   LLVM_DEBUG(print(llvm::dbgs()));
 }
 
-auto Operator::getSameTypeAsResult(int index) const -> ArrayRef<ArgOrType> {
+const InferredResultType &Operator::getInferredResultType(int index) const {
   assert(allResultTypesKnown());
   return resultTypeMapping[index];
 }
 
 ArrayRef<SMLoc> Operator::getLoc() const { return def.getLoc(); }
 
 bool Operator::hasDescription() const {
   return def.getValue("description") != nullptr;

mlir/test/mlir-tblgen/op-result.td

def OpL3 : NS_Op<"op_with_all_types_constraint",
   let results = (outs AnyType:$b);
 }
 
 // CHECK-LABEL: LogicalResult OpL3::inferReturnTypes
 // CHECK-NOT: }
 // CHECK: ::mlir::Type odsInferredType0 = attributes.get("a").cast<::mlir::TypedAttr>().getType();
 // CHECK: inferredReturnTypes[0] = odsInferredType0;
 
+def OpL4 : NS_Op<"two_inference_edges", [
+    TypesMatchWith<"", "a", "b", "infer0($_self)">,
+    TypesMatchWith<"", "b", "c", "infer1($_self)">,
+    TypesMatchWith<"", "input", "a", "fromInput($_self)">]> {
+  let arguments = (ins I32:$input);
+  let results = (outs AnyType:$a, AnyType:$b, AnyType:$c);
+}
+
+// CHECK-LABEL: LogicalResult OpL4::inferReturnTypes
+// CHECK: odsInferredType0 = fromInput(operands[0].getType())
+// CHECK: odsInferredType1 = infer0(odsInferredType0)
+// CHECK: odsInferredType2 = infer1(odsInferredType1)
+// CHECK: inferredReturnTypes[0] = odsInferredType0
+// CHECK: inferredReturnTypes[1] = odsInferredType1
+// CHECK: inferredReturnTypes[2] = odsInferredType2
+
 def OpM : NS_Op<"mix_diff_size_variadic_and_normal_results_op", [AttrSizedResultSegments]> {
   let results = (outs Variadic<AnyTensor>:$output1, AnyTensor:$output2, Optional<AnyTensor>:$output3);
 }
 
 // CHECK-LABEL: OpM::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange output1, ::mlir::Type output2, /*optional*/::mlir::Type output3)
 // CHECK: odsState.addAttribute(getResultSegmentSizesAttrName(odsState.name), odsBuilder.getDenseI32ArrayAttr({static_cast<int32_t>(output1.size()), 1, (output3 ? 1 : 0)}));

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 #include "mlir/TableGen/Class.h"
 #include "mlir/TableGen/CodeGenHelpers.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Interfaces.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/SideEffects.h"
 #include "mlir/TableGen/Trait.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
Method *method = [&]() -> Method * {
     return nullptr;
   }();
   ERROR_IF_PRUNED(method, "inferReturnTypes", op);
   auto &body = method->body();
   body << "  inferredReturnTypes.resize(" << op.getNumResults() << ");\n";
 
   FmtContext fctx;
   fctx.withBuilder("odsBuilder");
+  fctx.addSubst("_ctxt", "context");
   body << "  ::mlir::Builder odsBuilder(context);\n";
 
-  // Preprocess the result types and build all of the types used during
-  // inferrence. This limits the amount of duplicated work when a type is used
-  // to infer multiple others.
-  llvm::DenseMap<Constraint, int> constraintsTypes;
-  llvm::DenseMap<int, int> argumentsTypes;
+  // Process the type inference graph in topological order, starting from types
+  // that are always fully-inferred: operands and results with constructible
+  // types. The type inference graph here will always be a DAG, so this gives
+  // us the correct order for generating the types. -1 is a placeholder to
+  // indicate the type for a result has not been generated.
+  SmallVector<int> constructedIndices(op.getNumResults(), -1);
   int inferredTypeIdx = 0;
+  for (int numResults = op.getNumResults(); inferredTypeIdx != numResults;) {
     for (int i = 0, e = op.getNumResults(); i != e; ++i) {
-    auto type = op.getSameTypeAsResult(i).front();
-
-    // If the type isn't an argument, it refers to a buildable type.
-    if (!type.isArg()) {
-      auto it = constraintsTypes.try_emplace(type.getType(), inferredTypeIdx);
-      if (!it.second)
-        continue;
-
-      // If we haven't seen this constraint, generate a variable for it.
-      body << "  ::mlir::Type odsInferredType" << inferredTypeIdx++ << " = "
-           << tgfmt(*type.getType().getBuilderCall(), &fctx) << ";\n";
-      continue;
-    }
-
-    // Otherwise, this is an argument.
-    int argIndex = type.getArg();
-    auto it = argumentsTypes.try_emplace(argIndex, inferredTypeIdx);
-    if (!it.second)
+      if (constructedIndices[i] >= 0)
         continue;
+      const InferredResultType &infer = op.getInferredResultType(i);
+      std::string typeStr;
       body << "  ::mlir::Type odsInferredType" << inferredTypeIdx++ << " = ";
-
-    // If this is an operand, just index into operand list to access the type.
-    auto arg = op.getArgToOperandOrAttribute(argIndex);
+      if (infer.isArg()) {
+        // If this is an operand, just index into operand list to access the
+        // type.
+        auto arg = op.getArgToOperandOrAttribute(infer.getIndex());
         if (arg.kind() == Operator::OperandOrAttribute::Kind::Operand) {
-      body << "operands[" << arg.operandOrAttributeIndex() << "].getType()";
+          typeStr = ("operands[" + Twine(arg.operandOrAttributeIndex()) +
+                     "].getType()")
+                        .str();
 
           // If this is an attribute, index into the attribute dictionary.
         } else {
           auto *attr =
               op.getArg(arg.operandOrAttributeIndex()).get<NamedAttribute *>();
-      body << "attributes.get(\"" << attr->name
-           << "\").cast<::mlir::TypedAttr>().getType()";
-    }
-    body << ";\n";
+          typeStr = ("attributes.get(\"" + attr->name +
+                     "\").cast<::mlir::TypedAttr>().getType()")
+                        .str();
         }
-
-  // Perform a second pass that handles assigning the inferred types to the
-  // results.
-  for (int i = 0, e = op.getNumResults(); i != e; ++i) {
-    auto types = op.getSameTypeAsResult(i);
-
-    // Append the inferred type.
-    auto type = types.front();
-    body << "  inferredReturnTypes[" << i << "] = odsInferredType"
-         << (type.isArg() ? argumentsTypes[type.getArg()]
-                          : constraintsTypes[type.getType()])
-         << ";\n";
-
-    if (types.size() == 1)
+      } else if (std::optional<StringRef> builder =
+                     op.getResult(infer.getResultIndex())
+                         .constraint.getBuilderCall()) {
+        typeStr = tgfmt(*builder, &fctx).str();
+      } else if (int index = constructedIndices[infer.getResultIndex()];
+                 index >= 0) {
+        typeStr = ("odsInferredType" + Twine(index)).str();
+      } else {
         continue;
-    // TODO: We could verify equality here, but skipping that for verification.
       }
+      body << tgfmt(infer.getTransformer(), &fctx.withSelf(typeStr)) << ";\n";
+      constructedIndices[i] = inferredTypeIdx - 1;
+    }
+  }
+  for (auto [i, index] : llvm::enumerate(constructedIndices))
+    body << "  inferredReturnTypes[" << i << "] = odsInferredType" << index
+         << ";\n";
   body << "  return ::mlir::success();";
 }
 
 void OpEmitter::genParser() {
   if (hasStringAttribute(def, "assemblyFormat"))
     return;
 
   if (!def.getValueAsBit("hasCustomAssemblyFormat"))