[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

390

This will be much cleaner with two separate ops.

[rriddle]

mlir/include/mlir/IR/Builders.h

112

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

393

This looks much longer than just doing:

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

[frgossen]

Address comments

[frgossen]

Rebase

[rriddle]

mlir/test/Dialect/Shape/canonicalize.mlir

476

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 = [{
     they are equal.
     Shape operands are equal if their rank and all extents are equal.
   }];
 
   let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
   let results = (outs I1:$result);
 
   let assemblyFormat = "$lhs `,` $rhs `:` type($lhs) 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/include/mlir/IR/Builders.h

public:
   // ranked tensor of them. Returns null attribute otherwise.
   Attribute getZeroAttr(Type type);
 
   // Convenience methods for fixed types.
   FloatAttr getF16FloatAttr(float value);
   FloatAttr getF32FloatAttr(float value);
   FloatAttr getF64FloatAttr(double value);
 
+  IntegerAttr getI1IntegerAttr(int8_t value);
   IntegerAttr getI8IntegerAttr(int8_t value);
   IntegerAttr getI16IntegerAttr(int16_t value);
   IntegerAttr getI32IntegerAttr(int32_t value);
   IntegerAttr getI64IntegerAttr(int64_t value);
   IntegerAttr getIndexAttr(int64_t value);
 
   /// Signed and unsigned integer attribute getters.
   IntegerAttr getSI32IntegerAttr(int32_t value);

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

 
 //===----------------------------------------------------------------------===//
 // ConstWitnessOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
 
 //===----------------------------------------------------------------------===//
+// EqOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult EqOp::fold(ArrayRef<Attribute> operands) {
+
+  // Fold size equality.
+  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.
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // 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/lib/IR/Builders.cpp

 IntegerAttr Builder::getI16IntegerAttr(int16_t value) {
   return IntegerAttr::get(getIntegerType(16), APInt(16, value));
 }
 
 IntegerAttr Builder::getI8IntegerAttr(int8_t value) {
   return IntegerAttr::get(getIntegerType(8), APInt(8, value));
 }
 
+IntegerAttr Builder::getI1IntegerAttr(int8_t value) {
+  return IntegerAttr::get(getI1Type(), value);
+}
+
 IntegerAttr Builder::getIntegerAttr(Type type, int64_t value) {
   if (type.isIndex())
     return IntegerAttr::get(type, APInt(64, value));
   return IntegerAttr::get(
       type, APInt(type.getIntOrFloatBitWidth(), value, type.isSignedInteger()));
 }
 
 IntegerAttr Builder::getIntegerAttr(Type type, const APInt &value) {

mlir/test/Dialect/Shape/canonicalize.mlir

 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_broadcastable %arg0, %arg0
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
+
+// -----
+
+// Fold `eq` for equal and constant sizes.
+// CHECK-LABEL: @eq_size_fold_1
+func @eq_size_fold_1() -> i1 {
+  // CHECK-DAG: %[[RESULT:.*]] = constant true
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %a = shape.const_size 123
+  %b = shape.const_size 123
+  %result = shape.eq %a, %b : !shape.size
+  return %result : i1
+}
+
+// -----
+
+// Fold `eq` for different but constant sizes.
+// CHECK-LABEL: @eq_size_fold_0
+func @eq_size_fold_0() -> i1 {
+  // CHECK-DAG: %[[RESULT:.*]] = constant false
+  // CHECK-DAG: return %[[RESULT]] : i1
+  %a = shape.const_size 123
+  %b = shape.const_size 456
+  %result = shape.eq %a, %b : !shape.size
+  return %result : i1
+}
+
+// -----
+
+// Do not fold `eq` for non-constant sizes.
+// CHECK-LABEL: @eq_size_do_not_fold
+// CHECK-SAME: (%[[A:.*]]: !shape.size) -> i1
+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
+}
+