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

[frgossen]

Summary

Fold shape.size_eq.

Depends On D82628

[frgossen]

Remove CHECK-DAGs

[silvas]

LGTM modulo splitting into two separate ops.

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

541

This will be much cleaner with two separate ops.

[rriddle]

mlir/include/mlir/IR/Builders.h

112 ↗

(On Diff #273386)

I would rather use bool here. This method also isn't strictly necessary given that it does exactly the same thing as getBoolAttr.

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

544

This looks much longer than just doing:

BoolAttr::get(eq, getContext());

[frgossen]

Address comments

[frgossen]

Rebase

[rriddle]

mlir/test/Dialect/Shape/canonicalize.mlir

755

Can you simplify this to just checking something like:

// CHECK: %[[RESULT:.*]] = shape.size_eq
// CHECK: return %[[RESULT]]

?

[frgossen]

Address comment

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

let description = [{
     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)";
+  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

 
 //===----------------------------------------------------------------------===//
 // 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());
 }
 
 //===----------------------------------------------------------------------===//
+// SizeEqOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult SizeEqOp::fold(ArrayRef<Attribute> operands) {
+  auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>();
+  if (lhs == nullptr)
+    return {};
+  auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>();
+  if (rhs == nullptr)
+    return {};
+  bool equal = lhs.getValue().eq(rhs.getValue());
+  return BoolAttr::get(equal, 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;

mlir/test/Dialect/Shape/canonicalize.mlir

func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor<index>) {
   %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
   %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
+// Fold `size_eq` for equal and constant sizes.
+// CHECK-LABEL: @size_eq_fold_1
+func @size_eq_fold_1() -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant true
+  // CHECK: return %[[RESULT]] : i1
+  %a = shape.const_size 123
+  %b = shape.const_size 123
+  %result = shape.size_eq %a, %b : !shape.size, !shape.size
+  return %result : i1
+}
+
+// -----
+
+// Fold `size_eq` for different but constant sizes.
+// CHECK-LABEL: @size_eq_fold_0
+func @size_eq_fold_0() -> i1 {
+  // CHECK: %[[RESULT:.*]] = constant false
+  // CHECK: return %[[RESULT]] : i1
+  %a = constant 123 : index
+  %b = shape.const_size 456
+  %result = shape.size_eq %a, %b : index, !shape.size
+  return %result : i1
+}
+
+// -----
+
+// Do not fold `size_eq` for non-constant sizes.
+// CHECK-LABEL: @size_eq_do_not_fold
+func @size_eq_do_not_fold(%a : !shape.size) -> i1 {
+  // CHECK: %[[RESULT:.*]] = shape.size_eq
+  // CHECK: return %[[RESULT]] : i1
+  %b = shape.const_size 456
+  %result = shape.size_eq %a, %b : !shape.size, !shape.size
+  return %result : i1
+}
+
+// -----
+
 // 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] : !shape.shape
   %b = shape.const_shape [1, 2, 3] : tensor<?xindex>
   %result = shape.shape_eq %a, %b : !shape.shape, tensor<?xindex>