[MLIR][Shape] Add `shape.size_eq` to the shape dialect

[frgossen]

Summary

The new operation shape.size_eq compares equality of the shape dialect's size
values !shape.size and also supports index types.

[silvas]

LGTM modulo incorporating our decision on the error case discussion from shape_eq.

[frgossen]

Address error handling

[frgossen]

@jpienaar , can I land this?

[jpienaar]

LG, modulo I think the error discussion on the shape_eq one wasn't updated there or used here. But we can do those as uniform follow up too.

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

156

I don't understand the equality part of errors. The behavior elsewhere is that an error is propagated, here it sounds like you can compare if the size (that was supposed to be there) is equal, and returns true if there are no sizes. I don't think we've defined equality of errors either (e.g., if I have two errors due to being non-broadcastable are they equal? equal only when the operands are the same too?). But it is above already. And we can cycle back on both.

I agree with @herhut on when shape_eq was added about UB (in https://reviews.llvm.org/D82528) or my follow up question about usage only where non-error.

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

def Shape_ShapeEqOp : Shape_Op<"shape_eq", [Commutative, NoSideEffect]> {
   let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
                        Shape_ShapeOrExtentTensorType:$rhs);
   let results = (outs I1:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
   let hasFolder = 1;
 }
 
+def Shape_SizeEqOp : Shape_Op<"size_eq", [Commutative, NoSideEffect]> {
+  let summary = "Returns whether the input sizes are equal";
+  let description = [{
+    Takes two size or index operands and determines whether they are equal.
+    Index values are regarded as equivalent to non-error size values. Error
+    values can be tested for equality like any other size value, meaning that
+    error value is equal to itself.
+  }];
+
+  let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
+}
+
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
   let summary = "Creates a shape from extents";
   let description = [{
     Creates a shape from multiple SSA values representing the extents of
     the shape.
 
     ```mlir
     // Rank 2 shape.

mlir/test/Dialect/Shape/ops.mlir

 
 func @any_on_extent_tensors(%a : tensor<?xindex>,
                             %b : tensor<?xindex>,
                             %c : tensor<?xindex>) -> tensor<?xindex> {
   %result = shape.any %a, %b, %c
       : tensor<?xindex>, tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
   return %result : tensor<?xindex>
 }
+
+func @size_eq_on_sizes(%a : !shape.size, %b : !shape.size) -> i1 {
+  %result = shape.size_eq %a, %b : !shape.size, !shape.size
+  return %result : i1
+}
+
+func @size_eq_on_indices(%a : index, %b : index) -> i1 {
+  %result = shape.size_eq %a, %b : index, index
+  return %result : i1
+}
+
+func @size_eq(%a : !shape.size, %b : index) -> i1 {
+  %result = shape.size_eq %a, %b : !shape.size, index
+  return %result : i1
+}