[Mlir-commits] [mlir] 35c0ab7 - [MLIR] Simplify select to a not

[William S. Moses]

Author: William S. Moses
Date: 2021-06-28T11:00:02-04:00
New Revision: 35c0ab72fc20fcd47adda07f738338733d2c49e3

URL: https://github.com/llvm/llvm-project/commit/35c0ab72fc20fcd47adda07f738338733d2c49e3
DIFF: https://github.com/llvm/llvm-project/commit/35c0ab72fc20fcd47adda07f738338733d2c49e3.diff

LOG: [MLIR] Simplify select to a not

Given a select that returns the logical negation of the condition, replace it with a not of the condition.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp
    mlir/test/Dialect/Standard/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 48a539a387bd..ebb7c37703c1 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -1454,6 +1454,7 @@ def SelectOp : Std_Op<"select", [NoSideEffect,
       Value getFalseValue() { return false_value(); }
   }];
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 

diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 837986fc0353..f6abfc4060d0 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -1496,6 +1496,37 @@ static LogicalResult verify(ReturnOp op) {
 // SelectOp
 //===----------------------------------------------------------------------===//
 
+// Transforms a select to a not, where relevant.
+//
+//  select %arg, %false, %true
+//
+//  becomes
+//
+//  xor %arg, %true
+struct SelectToNot : public OpRewritePattern<SelectOp> {
+  using OpRewritePattern<SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SelectOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!matchPattern(op.getTrueValue(), m_Zero()))
+      return failure();
+
+    if (!matchPattern(op.getFalseValue(), m_One()))
+      return failure();
+
+    if (!op.getType().isInteger(1))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<XOrOp>(op, op.condition(), op.getFalseValue());
+    return success();
+  }
+};
+
+void SelectOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                           MLIRContext *context) {
+  results.insert<SelectToNot>(context);
+}
+
 OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
   auto trueVal = getTrueValue();
   auto falseVal = getFalseValue();

diff  --git a/mlir/test/Dialect/Standard/canonicalize.mlir b/mlir/test/Dialect/Standard/canonicalize.mlir
index f3b9bdf9c54f..d2ef830537f9 100644
--- a/mlir/test/Dialect/Standard/canonicalize.mlir
+++ b/mlir/test/Dialect/Standard/canonicalize.mlir
@@ -319,3 +319,15 @@ func @branchCondProp(%arg0: i1) {
 ^exit:
   return
 }
+
+// -----
+
+// CHECK-LABEL: @selToNot
+//       CHECK:       %[[trueval:.+]] = constant true
+//       CHECK:       %{{.+}} = xor %arg0, %[[trueval]] : i1
+func @selToNot(%arg0: i1) -> i1 {
+  %true = constant true
+  %false = constant false
+  %res = select %arg0, %false, %true : i1
+  return %res : i1
+}