[MLIR][Shape] Fold `shape.shape_eq`

[frgossen]

Summary

Fold shape.shape_eq.

Depends On D82528

[frgossen]

Remove CHECK-DAGs

[silvas]

LGTM, modulo splitting into two ops and a nit.

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

407

can you use std::equal or check in STLExtras.h if we have a range version? That should make this more compact.

(actually, adding range_equal to STLExtras.h would be cool, then after splitting into two ops this could be just llvm::range_equal(lhs.cast<DenseIntElementsAttr>(), rhs.cast<DenseIntElementsAttr>()))

[frgossen]

Simplify folding

[frgossen]

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

407

Turns out there is actually == defined on these.
Thanks!

[frgossen]

Rebase

[frgossen]

Rebase

[frgossen]

Rebase

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

OpFoldResult EqOp::fold(ArrayRef<Attribute> operands) {
   if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>()) {
     if (auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>()) {
       bool eq = lhs.getValue().eq(rhs.getValue());
       Builder builder(getContext());
       return builder.getI1IntegerAttr(eq);
     }
   }
 
-  // TODO: Fold shape equality.
+  // Fold shape equality.
+  if (auto lhs = operands[0].dyn_cast_or_null<DenseIntElementsAttr>()) {
+    if (auto rhs = operands[1].dyn_cast_or_null<DenseIntElementsAttr>()) {
+      auto iLhs = lhs.begin(), iRhs = rhs.begin();
+      auto endLhs = lhs.end(), endRhs = rhs.end();
+      auto sizeLhs = endLhs - iLhs, sizeRhs = endRhs - iRhs;
+      Builder builder(getContext());
+      if (sizeLhs != sizeRhs)
+        return builder.getI1IntegerAttr(false);
+      while (iLhs < endLhs) {
+        assert(iRhs < endRhs);
+        if (*iLhs++ != *iRhs++)
+          return builder.getI1IntegerAttr(false);
+      }
+      return builder.getI1IntegerAttr(true);
+    }
+  }
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // IndexToSizeOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Shape/canonicalize.mlir

func @eq_size_do_not_fold(%a : !shape.size) -> i1 {
   // CHECK-DAG: %[[B:.*]] = shape.const_size 456
   // CHECK-DAG: %[[RESULT:.*]] = shape.eq %[[A]], %[[B]] : !shape.size
   // CHECK-DAG: return %[[RESULT]] : i1
   %b = shape.const_size 456
   %result = shape.eq %a, %b : !shape.size
   return %result : i1
 }
 
+// -----
+
+// Fold `eq` for equal and constant shapes.
+// CHECK-LABEL: @eq_shape_fold_1
+func @eq_shape_fold_1() -> i1 {
+  // CHECK-DAG: %[[RESULT:.*]] = constant true
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3]
+  %b = shape.const_shape [1, 2, 3]
+  %result = shape.eq %a, %b : !shape.shape
+  return %result : i1
+}
+
+// -----
+
+// Fold `eq` for different but constant shapes of same length.
+// CHECK-LABEL: @eq_shape_fold_0
+func @eq_shape_fold_0() -> i1 {
+  // CHECK-DAG: %[[RESULT:.*]] = constant false
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3]
+  %b = shape.const_shape [4, 5, 6]
+  %result = shape.eq %a, %b : !shape.shape
+  return %result : i1
+}
+
+// -----
+
+// Fold `eq` for different but constant shapes of different length.
+// CHECK-LABEL: @eq_shape_fold_0
+func @eq_shape_fold_0() -> i1 {
+  // CHECK-DAG: %[[RESULT:.*]] = constant false
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3]
+  %b = shape.const_shape [1, 2, 3, 4, 5, 6]
+  %result = shape.eq %a, %b : !shape.shape
+  return %result : i1
+}
+
+// -----
+
+// Do not fold `eq` for non-constant shapes.
+// CHECK-LABEL: @eq_shape_do_not_fold
+// CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1
+func @eq_shape_do_not_fold(%a : !shape.shape) -> i1 {
+  // CHECK-DAG: %[[B:.*]] = shape.const_shape [4, 5, 6]
+  // CHECK-DAG: %[[RESULT:.*]] = shape.eq %[[A]], %[[B]] : !shape.shape
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %b = shape.const_shape [4, 5, 6]
+  %result = shape.eq %a, %b : !shape.shape
+  return %result : i1
+}