[Mlir-commits] [mlir] 1e39d32 - [MLIR] Add OrOp folding rule for constant one operand

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-10-07T08:05:39+05:30
New Revision: 1e39d32c5a119c8f49eb56e2cdd9707022fa06ab

URL: https://github.com/llvm/llvm-project/commit/1e39d32c5a119c8f49eb56e2cdd9707022fa06ab
DIFF: https://github.com/llvm/llvm-project/commit/1e39d32c5a119c8f49eb56e2cdd9707022fa06ab.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp
    mlir/test/Transforms/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 6f95bf03e91d..64a248717aa3 100644
--- 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 a510be43669a..aba52519b95e 100644
--- 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