[mlir][shape] Update meet to handle all size & shape types

[jpienaar]

Summary

Also tighten up return type inference & compatibility functions.

[yaochengji]

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

1318

isExtentTensorType could be reused here.

1347

SizeType, IndexType, ShapeType have already been checked, so better to use isExtentTensorType instead of isShapeType.

[jpienaar]

Use isExtentTensorType helper.

[yaochengji]

LGTM to me after the CI passes.

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

def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
 def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
                                                Shape_ExtentTensorType],
                                               "shape or extent tensor">;
 
 def Shape_SizeOrIndexType : AnyTypeOf<[Shape_SizeType, Index], "size or index">;
 
+// Any type representing a shape or size/dim.
+def Shape_AnyShapeOrSizeType : AnyTypeOf<
+    [Shape_SizeOrIndexType, Shape_ShapeOrExtentTensorType],
+  "any shape or size">;
+
 def Shape_WitnessType : Shape_Type<"Witness", "witness"> {
   let description = [{
     A witness is a structural device in the compiler to maintain ordering of
     code relying on information obtained from passing assertions. Witnesses do
     not represent any physical data.
 
     "cstr_" operations will return witnesses and be lowered into assertion logic
     when not resolvable at compile time.

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

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_MeetOp : Shape_Op<"meet",
     [Commutative, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
-  let summary = "Returns the least general shape.shape of its operands";
+  let summary = "Returns the least general shape or size of its operands";
   let description = [{
-    An operation that computes the least general shape of input operands.
+    An operation that computes the least general shape or dim of input operands.
     This effectively asserts that corresponding static dimensions are equal.
-    The behavior is to match each element of the `shape.shape` and propagate the
+    The behavior is to match each element of the shape/size and propagate the
     most restrictive information, returning an invalid shape if there are
     contradictory requirements. E.g., using pseudo code
 
     ```
     shape.meet([*], [*]) -> [*]
     shape.meet([*], [1, ?]) -> [1, ?]
     shape.meet([1, 2], [1, ?]) -> [1, 2]
     shape.meet([*], [1, 2]) -> [1, 2]
     shape.meet([], []) -> []
     shape.meet([], [*]) -> []
     shape.meet([], [?, ?]) -> [invalid]
     shape.meet([1, ?], [2, ?, ?]) -> [invalid]
     ```
 
     `shape.meet` also allows specifying an optional error string, that may be
     used to return an error to the user upon mismatch of dimensions.
 
     ```mlir
     %c = shape.meet %a, %b, error="<reason>" : !shape.shape, !shape.shape -> !shape.shape
     ```
   }];
 
-  let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
-                   OptionalAttr<StrAttr>:$error);
-  let results = (outs Shape_ShapeOrSizeType:$result);
+  let arguments = (ins
+    Shape_AnyShapeOrSizeType:$arg0,
+    Shape_AnyShapeOrSizeType:$arg1,
+  let results = (outs Shape_AnyShapeOrSizeType:$result);
 
   let assemblyFormat = [{
     $arg0 `,` $arg1 (`,` `error` `=` $error^)? attr-dict `:`
       type($arg0) `,` type($arg1) `->` type($result)
   }];
 
   let extraClassDeclaration = [{
     // Returns when two result types are compatible for this op; method used by

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

 //===----------------------------------------------------------------------===//
 // MeetOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::shape::MeetOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
-  inferredReturnTypes.assign({operands[0].getType()});
+  if (operands.empty())
+    return failure();
+
+  auto isShapeType = [](Type arg) {
+    if (arg.isa<ShapeType>())
+      return true;
+    return isExtentTensorType(arg);
+  };
+
+  ValueRange::type_range types = operands.getTypes();
+  Type acc = types.front();
+  for (auto t : drop_begin(types)) {
+    Type l = acc, r = t;
+    if (!l.isa<ShapeType, SizeType>())
+      std::swap(l, r);
+
+    // Handle sizes, propagate error type if present.
+    if (l.isa<SizeType>()) {
+      if (r.isa<SizeType, IndexType>())
+        acc = l;
+      else
+        return emitOptionalError(location, "requires all sizes or shapes");
+    } else if (l.isa<IndexType>()) {
+      if (r.isa<IndexType>())
+        acc = r;
+      else
+        return emitOptionalError(location, "requires all sizes or shapes");
+    } else if (l.isa<ShapeType>()) {
+      // Handle shapes, propagate error type if present.
+      if (isShapeType(r))
+        acc = l;
+      else
+        return emitOptionalError(location, "requires all sizes or shapes");
+    } else if (isExtentTensorType(l)) {
+      auto rank1 = l.cast<RankedTensorType>().getShape()[0];
+      auto rank2 = r.cast<RankedTensorType>().getShape()[0];
+      if (ShapedType::isDynamic(rank1))
+        acc = l;
+      else if (ShapedType::isDynamic(rank2))
+        acc = r;
+      else if (rank1 != rank2)
+        return emitOptionalError(location, "unequal shape cardinality");
+      else
+        acc = l;
+    }
+  }
+  inferredReturnTypes.assign({acc});
   return success();
 }
 
 bool mlir::shape::MeetOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   if (l.size() != 1 || r.size() != 1)
     return false;
   if (l == r)
     return true;
 
   Type lhs = l.front();
   Type rhs = r.front();
 
-  if (lhs != rhs)
-    return false;
+  if (!lhs.isa<ShapeType, SizeType>())
+    std::swap(lhs, rhs);
 
-  if (lhs.isa<SizeType>() || lhs.isa<ShapeType>())
-    return true;
+  if (lhs.isa<SizeType>())
+    return rhs.isa<SizeType, IndexType>();
+  if (lhs.isa<ShapeType>())
+    return rhs.isa<ShapeType, TensorType>();
 
   if (succeeded(verifyCompatibleShapes({lhs, rhs})))
     return true;
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp

mlir/test/Dialect/Shape/invalid.mlir

 
 // Test that type inference flags the wrong return type.
 
 func.func @const_shape() {
   // expected-error@+1 {{'tensor<3xindex>' are incompatible with return type(s) of operation 'tensor<2xindex>'}}
   %0 = shape.const_shape [4, 5, 6] : tensor<2xindex>
   return
 }
+
+// -----
+
+func.func @invalid_meet(%arg0 : !shape.shape, %arg1 : index) -> index {
+  // expected-error@+1 {{requires all sizes or shapes}}
+  %result = shape.meet %arg0, %arg1 : !shape.shape, index -> index
+  return %result : index
+}
+
+// -----
+
+func.func @invalid_meet(%arg0 : tensor<2xindex>, %arg1 : tensor<3xindex>) -> tensor<?xindex> {
+  // expected-error@+1 {{unequal shape cardinality}}
+  %result = shape.meet %arg0, %arg1 : tensor<2xindex>, tensor<3xindex> -> tensor<?xindex>
+  return %result : tensor<?xindex>
+}
+

mlir/test/Dialect/Shape/ops.mlir

 
 func.func @size_lower_bounded_by_constant(%a: !shape.size) -> !shape.size {
   %0 = shape.const_size 9
   %1 = shape.min %a, %0 : !shape.size, !shape.size -> !shape.size
   %2 = shape.meet %0, %1, error="lower bound element-wise exceeded" :
     !shape.size, !shape.size -> !shape.size
   return %2 : !shape.size
 }
+
+func.func @meet_index(%arg0 : index, %arg1 : index) -> index {
+  %result = shape.meet %arg0, %arg1 : index, index -> index
+  return %result : index
+}
+