[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

411

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

411

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

[frgossen]

Rebase

[frgossen]

Rebase

[frgossen]

Rebase

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

let description = [{
     also unknown.
     Error values propagate in the same way.
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs I1:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
+  let hasFolder = 1;
 }
 
 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/lib/Dialect/Shape/IR/Shape.cpp

 
 //===----------------------------------------------------------------------===//
 // ConstWitnessOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
 
 //===----------------------------------------------------------------------===//
+// ShapeEqOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ShapeEqOp::fold(ArrayRef<Attribute> operands) {
+  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());
+}
+
+//===----------------------------------------------------------------------===//
 // IndexToSizeOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // FromExtentsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   if (llvm::any_of(operands, [](Attribute a) { return !a; }))

mlir/test/Dialect/Shape/canonicalize.mlir

 func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> !shape.size {
 // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32>
 // CHECK-DAG: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
 // CHECK-DAG: return %[[SIZE]] : !shape.size
 %shape = shape.shape_of %arg : tensor<*xf32>
 %rank = shape.rank %shape
 return %rank : !shape.size
 }
+
+// -----
+
+// Fold `shape_eq` for equal and constant shapes.
+// CHECK-LABEL: @shape_eq_fold_1
+func @shape_eq_fold_1() -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant true
+  // CHECK: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3]
+  %b = shape.const_shape [1, 2, 3]
+  %result = shape.shape_eq %a, %b
+  return %result : i1
+}
+
+// -----
+
+// Fold `shape_eq` for different but constant shapes of same length.
+// CHECK-LABEL: @shape_eq_fold_0
+func @shape_eq_fold_0() -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant false
+  // CHECK: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3]
+  %b = shape.const_shape [4, 5, 6]
+  %result = shape.shape_eq %a, %b
+  return %result : i1
+}
+
+// -----
+
+// Fold `shape_eq` for different but constant shapes of different length.
+// CHECK-LABEL: @shape_eq_fold_0
+func @shape_eq_fold_0() -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant false
+  // CHECK: return %[[RESULT]] : i1
+  %a = shape.const_shape [1, 2, 3, 4, 5, 6]
+  %b = shape.const_shape [1, 2, 3]
+  %result = shape.shape_eq %a, %b
+  return %result : i1
+}
+
+// -----
+
+// Do not fold `shape_eq` for non-constant 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]]
+  // CHECK: return %[[RESULT]] : i1
+  %b = shape.const_shape [4, 5, 6]
+  %result = shape.shape_eq %a, %b
+  return %result : i1
+}