[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

diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 48a539a..ebb7c37 100644 (file)
--- 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 837986f..f6abfc4 100644 (file)
--- 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 f3b9bdf..d2ef830 100644 (file)
--- 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
+}