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

Differential Revision: https://reviews.llvm.org/D95430

Author: tpopp

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

@@ -572,6 +572,8 @@ OpFoldResult CstrRequireOp::fold(ArrayRef<Attribute> operands) {
 //===----------------------------------------------------------------------===//
 
 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 {};

mlir/test/Dialect/Shape/canonicalize.mlir

@@ -787,7 +787,7 @@ func @shape_eq_fold_0() -> 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 {
@@ -799,6 +799,19 @@ func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {
   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.