[mlir][ODS] Infer return types if the operands are variadic but the results are not

[bkramer]

Summary

Clean up code that worked around this limitation.

[jpienaar]

Nice, thanks!

mlir/include/mlir/TableGen/Operator.h

106

hasVariadicResults (to match naming above)

[bkramer]

mlir/include/mlir/TableGen/Operator.h

106

That's actually not the same. This is also checking for Optional<> in the results, while isVariadic only looks for Variadic<>.

[bkramer]

Use existing method to check for variadic results and clarify what variable length means.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_DivOp : Shape_Op<"div", [NoSideEffect,
 
   let extraClassDeclaration = [{
     // Returns when two result types are compatible for this op; method used by
     // InferTypeOpInterface
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
-def Shape_ShapeEqOp : Shape_Op<"shape_eq",
-    [NoSideEffect, Commutative, InferTypeOpInterface]> {
+def Shape_ShapeEqOp : Shape_Op<"shape_eq", [NoSideEffect, Commutative]> {
   let summary = "Returns whether the input shapes or extent tensors are equal";
   let description = [{
     Takes one or more shape or extent tensor operands and determines whether
     they are equal. When extent tensors are compared to shapes they are regarded
     as their equivalent non-error shapes. Error shapes can be tested for
     equality like any other shape value, meaning that the error value is equal
     to itself.
   }];
 
   let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
   let results = (outs I1:$result);
 
   // Convenience builder alias for the binary version.
   let builders = [
   OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
     [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
   ];
-  let extraClassDeclaration = [{
-    // TODO: This should really be automatic. Figure out how to not need this defined.
-    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
-    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
-    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
-    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
-      inferredReturnTypes.push_back(::mlir::IntegerType::get(context,
-                                                             /*width=*/1));
-      return success();
-    };
-  }];
 
   let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
   let hasFolder = 1;
 }
 
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
   let summary = "Creates a shape from extents";
   let description = [{
def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   }];
 
   let arguments = (ins 1DTensorOf<[Index]>:$input);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$input attr-dict `:` type($input)";
 }
 
-def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable",
-                                       [Commutative, InferTypeOpInterface]> {
+def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable", [Commutative]> {
   let summary = "Determines if 2+ shapes can be successfully broadcasted";
   let description = [{
     Given multiple input shapes or extent tensors, return a predicate specifying
     if they are broadcastable. This broadcastable follows the same logic as what
     shape.broadcast documents.
 
     Concretely, shape.is_broadcastable returning true implies that
     shape.broadcast will not give an error, and shape.cstr_broadcastable will
def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable", [Commutative]> {
 
   let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
   let results = (outs I1:$result);
 
   let builders = [
   OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
     [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
   ];
-  let extraClassDeclaration = [{
-    // TODO: This should really be automatic. Figure out how to not need this defined.
-    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
-    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
-    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
-    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
-      inferredReturnTypes.push_back(::mlir::IntegerType::get(context,
-                                                             /*width=*/1));
-      return success();
-    };
-  }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 
   let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
 }
 
 def Shape_RankOp : Shape_Op<"rank",
let description = [{
     %wf = shape.assuming_all %w0, %w1 // Failure
     %wt = shape.assuming_all %w0, %w2 // Passing
     ```
   }];
 
   let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
 
-  // Only needed while tablegen is unable to generate this for ops with variadic
-  // arguments.
-  let builders = [
-    OpBuilder<(ins "ValueRange":$inputs)>,
-  ];
-
   let assemblyFormat = "$inputs attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Shape_AssumingOp : Shape_Op<"assuming", [
def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
 
   let builders = [
     OpBuilder<(ins), [{ /* nothing to do */ }]>,
   ];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
-def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable",
-                                         [Commutative, InferTypeOpInterface]> {
+def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
   let summary = "Determines if 2+ shapes can be successfully broadcasted";
   let description = [{
     Given input shapes or extent tensors, return a witness specifying if they
     are broadcastable. This broadcastable follows the same logic as what
     shape.broadcast documents.
 
     "cstr" operations represent runtime assertions.
 
def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
 
   let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
 
   let builders = [
   OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
     [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
   ];
 
-  let extraClassDeclaration = [{
-    // TODO: This should really be automatic. Figure out how to not need this defined.
-    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
-    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
-    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
-    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
-      inferredReturnTypes.push_back(::mlir::shape::WitnessType::get(context));
-      return success();
-    };
-  }];
-
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
-def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative, InferTypeOpInterface]> {
+def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative]> {
   let summary = "Determines if all input shapes are equal";
   let description = [{
     Given 1 or more input shapes, determine if all shapes are the exact same.
 
     "cstr" operations represent runtime assertions.
 
     Example:
     ```mlir
     %w0 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
     %w1 = shape.cstr_eq [2,2], [1,2] // Failure
     ```
   }];
   let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
 
-  let extraClassDeclaration = [{
-    // TODO: This should really be automatic. Figure out how to not need this defined.
-    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
-    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
-    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
-    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
-      inferredReturnTypes.push_back(::mlir::shape::WitnessType::get(context));
-      return success();
-    };
-  }];
-
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def Shape_ConstWitnessOp : Shape_Op<"const_witness", [ConstantLike, NoSideEffect]> {
   let summary = "An operation that returns a statically known witness value";
   let description = [{
   This operation represents a statically known witness result. This can be

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

let description = [{
     Example:
 
     ```mlir
     %v = vector.insert_map %ev %v[%id] : vector<1xf32> into vector<32xf32>
     %v1 = vector.insert_map %ev1, %v1[%arg0, %arg1] : vector<2x4x1xf32>
       into vector<64x4x32xf32>
     ```
   }];
-  let builders = [OpBuilder<(ins "Value":$vector, "Value":$dest,
-                                "ValueRange":$ids)>];
   let extraClassDeclaration = [{
     VectorType getSourceVectorType() {
       return vector().getType().cast<VectorType>();
     }
     VectorType getResultType() {
       return getResult().getType().cast<VectorType>();
     }
     // Return a map indicating the dimension mapping to the given ids.

mlir/include/mlir/TableGen/Operator.h

public:
   using value_iterator = NamedTypeConstraint *;
   using const_value_iterator = const NamedTypeConstraint *;
   using value_range = llvm::iterator_range<value_iterator>;
   using const_value_range = llvm::iterator_range<const_value_iterator>;
 
   // Returns true if this op has variable length operands or results.
   bool isVariadic() const;
 
+  // Returns true if this op has variable length results.
+  bool hasVariableLengthResults() const;
+
   // Returns true if default builders should not be generated.
   bool skipDefaultBuilders() const;
 
   // Op result iterators.
   const_value_iterator result_begin() const;
   const_value_iterator result_end() const;
   const_value_range getResults() const;
 

mlir/lib/Dialect/Shape/IR/Shape.cpp

 LogicalResult AssumingAllOp::verify() {
   // Ensure that AssumingAllOp contains at least one operand
   if (getNumOperands() == 0)
     return emitOpError("no operands specified");
 
   return success();
 }
 
-void AssumingAllOp::build(OpBuilder &b, OperationState &state,
-                          ValueRange inputs) {
-  build(b, state, b.getType<WitnessType>(), inputs);
-}
-
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
   if (getShapes().size() == 1) {
     // Otherwise, we need a cast which would be a canonicalization, not folding.
     if (getShapes().front().getType() != getType())

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

if (position().empty())
     return source();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // InsertMapOp
 //===----------------------------------------------------------------------===//
 
-void InsertMapOp::build(OpBuilder &builder, OperationState &result,
-                        Value vector, Value dest, ValueRange ids) {
-  InsertMapOp::build(builder, result, dest.getType(), vector, dest, ids);
-}
-
 LogicalResult InsertMapOp::verify() {
   if (getSourceVectorType().getRank() != getResultType().getRank())
     return emitOpError("expected source and destination vectors of same rank");
   unsigned numId = 0;
   for (unsigned i = 0, e = getResultType().getRank(); i < e; i++) {
     if (getResultType().getDimSize(i) % getSourceVectorType().getDimSize(i) !=
         0)
       return emitOpError(

mlir/lib/TableGen/Operator.cpp

 // 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(); });
 }
 
+bool Operator::hasVariableLengthResults() const {
+  return any_of(results, [](const NamedTypeConstraint &op) {
+    return op.isVariableLength();
+  });
+}
+
 void Operator::populateTypeInferenceInfo(
     const llvm::StringMap<int> &argumentsAndResultsIndex) {
   // If the type inference op interface is not registered, then do not attempt
   // to determine if the result types an be inferred.
   auto &recordKeeper = def.getRecords();
   auto *inferTrait = recordKeeper.getDef(inferTypeOpInterface);
   allResultsHaveKnownTypes = false;
   if (!inferTrait)
     return;
 
   // If there are no results, the skip this else the build method generated
   // overlaps with another autogenerated builder.
   if (getNumResults() == 0)
     return;
 
-  // Skip for ops with variadic operands/results.
-  // TODO: This can be relaxed.
-  if (isVariadic())
+  // Skip for ops with variadic results.
+  if (hasVariableLengthResults())
     return;
 
   // Skip cases currently being custom generated.
   // TODO: Remove special cases.
   if (getTrait("::mlir::OpTrait::SameOperandsAndResultType"))
     return;
 
   // We create equivalence classes of argument/result types where arguments