[mlir] Fold shape.eq %a, %a to true

[tpopp]

I don't know if this works with whatever Shape Error types exist these days, so please let me know what you think.

[jpienaar]

In D95430#2522385, @tpopp wrote:

I don't know if this works with whatever Shape Error types exist these days, so please let me know what you think.

This goes back to the comment when the op was added: I don't think it as defined makes sense outside of an assuming region. Errors in and bool out I don't know what the consumer could do with the value, it feels like it would be easy to drop error state there (e.g., errors feed into equals, equals get folded and now errors have no consumer anymore and we DCE those).

As defined, this seems consistent.

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

OpFoldResult CstrRequireOp::fold(ArrayRef<Attribute> operands) {
   return operands[0];
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeEqOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeEqOp::fold(ArrayRef<Attribute> operands) {
+  if (lhs() == rhs())
+    return BoolAttr::get(true, getContext());
   auto lhs = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (lhs == nullptr)
     return {};
   auto rhs = operands[1].dyn_cast_or_null<DenseIntElementsAttr>();
   if (rhs == nullptr)
     return {};
   return BoolAttr::get(lhs == rhs, getContext());
 }

mlir/test/Dialect/Shape/canonicalize.mlir

func @shape_eq_fold_0() -> i1 {
   %a = shape.const_shape [1, 2, 3, 4, 5, 6] : !shape.shape
   %b = shape.const_shape [1, 2, 3] : !shape.shape
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }
 
 // -----
 
-// Do not fold `shape_eq` for non-constant shapes.
+// Do not fold `shape_eq` for non-constant different shapes.
 // CHECK-LABEL: @shape_eq_do_not_fold
 // CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1
 func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {
   // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6]
   // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape
   // CHECK: return %[[RESULT]] : i1
   %b = shape.const_shape [4, 5, 6] : !shape.shape
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }
 
+
+// -----
+
+// Fold `shape_eq` for non-constant but same shapes.
+// CHECK-LABEL: @shape_eq_do_fold
+// CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1
+func @shape_eq_do_fold(%a : !shape.shape) -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant true
+  // CHECK: return %[[RESULT]] : i1
+  %result = shape.shape_eq %a, %a : !shape.shape, !shape.shape
+  return %result : i1
+}
+
 // -----
 
 // Fold `mul` for constant sizes.
 // CHECK-LABEL: @fold_mul_size
 func @fold_mul_size() -> !shape.size {
   // CHECK: %[[RESULT:.*]] = shape.const_size 6
   // CHECK: return %[[RESULT]] : !shape.size
   %c2 = shape.const_size 2