[mlir][Shape] Make cstr_eq more like cstr_broadcastable

[bkramer]

Summary

This includes allowing extents and not just shapes.

[jpienaar]

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

805

DeclareOpInterfaceMethods ?

[jpienaar]

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

805

Mmm, this one is weird: is witness considered a buildable type? Then this should be generated as it is a fixed output type.

[bkramer]

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

805

The problem is that tblgen does not generate the methods for ops with variadic operands. I have no idea why, see Operator::populateTypeInferenceInfo.

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

let extraClassDeclaration = [{
     };
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let verifier = [{ return ::verify(*this); }];
 }
 
-def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative]> {
+def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative, InferTypeOpInterface]> {
   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_ShapeType>:$inputs);
+  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
   let results = (outs Shape_WitnessType:$result);
 
-  let assemblyFormat = "$inputs attr-dict";
+  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 = [{

mlir/test/Dialect/Shape/canonicalize.mlir

 
 // -----
 // cstr_eq with non-constant but known equal shapes can be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.cstr_eq %arg0, %arg0, %arg0
+  %0 = shape.cstr_eq %arg0, %arg0, %arg0 : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_eq with equal const_shapes can be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [0, 1] : !shape.shape
   %cs1 = shape.const_shape [0, 1] : !shape.shape
   %cs2 = shape.const_shape [0, 1] : !shape.shape
-  %0 = shape.cstr_eq %cs0, %cs1, %cs2
+  %0 = shape.cstr_eq %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_eq with unequal const_shapes cannot be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_eq
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [0, 1] : !shape.shape
   %cs1 = shape.const_shape [3, 1] : !shape.shape
-  %0 = shape.cstr_eq %cs0, %cs1
+  %0 = shape.cstr_eq %cs0, %cs1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_eq without const_shapes cannot be folded
 // CHECK-LABEL: func @f
 func @f(%arg0: !shape.shape, %arg1: !shape.shape) {
   // CHECK-NEXT: shape.cstr_eq
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.cstr_eq %arg0, %arg1
+  %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_require with constant can be folded
 // CHECK-LABEL: func @cstr_require_fold
 func @cstr_require_fold() {

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
   return %0 : tensor<?xindex>
 }
 
 func @test_constraints() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %true = constant true
   %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
-  %w1 = shape.cstr_eq %0, %1
+  %w1 = shape.cstr_eq %0, %1 : !shape.shape, !shape.shape
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
   %w4 = shape.cstr_require %true, "msg"
   %w_all = shape.assuming_all %w0, %w1, %w2, %w3, %w4
   shape.assuming %w_all -> !shape.shape {
     %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
+func @eq_on_extent_tensors(%lhs : tensor<?xindex>,
+                                      %rhs : tensor<?xindex>) {
+  %w0 = shape.cstr_eq %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
+  return
+}
+
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
   return
 }
 
 func @mul(%size_arg : !shape.size, %index_arg : index) {
   %size_prod = shape.mul %size_arg, %size_arg

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
   // REPLACE-NOT: shape.cstr_eq
   // REPLACE: shape.assuming %[[WITNESS]]
   // CANON-NEXT: test.source
   // CANON-NEXT: return
-  %0 = shape.cstr_eq %arg0, %arg1
+  %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
   %1 = shape.assuming %0 -> index {
     %2 = "test.source"() : () -> (index)
     shape.assuming_yield %2 : index
   }
   return %1 : index
 }
 
 // -----
 // With a non-const value, we cannot fold away the code, but all constraints
 // should be removed still.
 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // CANON-NEXT: test.source
   // CANON-NEXT: return
   %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
-  %1 = shape.cstr_eq %arg0, %arg1
+  %1 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
   %2 = shape.assuming_all %0, %1
   %3 = shape.assuming %0 -> index {
     %4 = "test.source"() : () -> (index)
     shape.assuming_yield %4 : index
   }
   return %3 : index
 }