[MLIR] Add OrOp folding rule for constant one operand

[bondhugula]

Summary

Add folding rule for std.or op when an operand has all bits set.

or(x, <all bits set>) -> <all bits set>

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1566–1568

OrOp is a bitwise or, this doesn't look correct for non-i1. Is this intended to capture when all bits are set in one of the values?

[bondhugula]

Fix folder for all ones.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1566–1568

Oops, I missed that; thanks. I actually meant to capture the i1 case but I've now changed this to fold for "all bits set" case.

[bondhugula]

Update comment.

[bondhugula]

Use CHECK-DAG for constants.

[bondhugula]

Update commit summary.

[bondhugula]

Gentle ping to reviewers @rriddle @mehdi_amini @herhut on this trivial revision.

[rriddle]

LG, thanks.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 //===----------------------------------------------------------------------===//
 // OrOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult OrOp::fold(ArrayRef<Attribute> operands) {
   /// or(x, 0) -> x
   if (matchPattern(rhs(), m_Zero()))
     return lhs();
-  /// or(x,x) -> x
+  /// or(x, x) -> x
   if (lhs() == rhs())
     return rhs();
+  /// or(x, <all ones>) -> <all ones>
+  if (auto rhsAttr = operands[1].dyn_cast_or_null<IntegerAttr>())
+    if (rhsAttr.getValue().isAllOnes())
+      return rhsAttr;
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a | b; });
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//

mlir/test/Transforms/canonicalize.mlir

 //CHECK-LABEL: func @or_zero_tensor
 func @or_zero_tensor(%arg0: tensor<4x5xi32>) -> tensor<4x5xi32> {
   // CHECK-NEXT: return %arg0
   %cst = constant dense<0> : tensor<4x5xi32>
   %1 = or %arg0, %cst : tensor<4x5xi32>
   return %1 : tensor<4x5xi32>
 }
 
+// CHECK-LABEL: func @or_all_ones
+func @or_all_ones(%arg0: i1, %arg1: i4) -> (i1, i4) {
+  // CHECK-DAG: %c-1_i4 = constant -1 : i4
+  // CHECK-DAG: %true = constant true
+  %c1_i1 = constant 1 : i1
+  %c15 = constant 15 : i4
+  // CHECK-NEXT: return %true
+  %1 = or %arg0, %c1_i1 : i1
+  %2 = or %arg1, %c15 : i4
+  return %1, %2 : i1, i4
+}
+
 //CHECK-LABEL: func @xor_self
 func @xor_self(%arg0: i32) -> i32 {
   //CHECK-NEXT: %c0_i32 = constant 0
   %1 = xor %arg0, %arg0 : i32
   //CHECK-NEXT: return %c0_i32
   return %1 : i32
 }