[MLIR] Add OrOp folding rule for constant one operand

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

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

Differential Revision: https://reviews.llvm.org/D111206

diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 6f95bf0..64a2487 100644 (file)
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -1560,9 +1560,13 @@ 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; });

diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index a510be4..aba5251 100644 (file)
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -287,6 +287,18 @@ func @or_zero_tensor(%arg0: tensor<4x5xi32>) -> 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