[MLIR] Generalize select to arithmetic canonicalization

[wsmoses]

Summary

Given a select whose result is an i1, we can eliminate the conditional in the select completely by adding a few arithmetic operations.

[bondhugula]

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

833–840

Always avoid nesting multiple IR-creating calls. It can lead to an unpredictable "order of evaluation" dependent output that is sensitive to compilers, their versions, and options.

mlir/test/Dialect/Standard/canonicalize.mlir

67–69

An unstable/unknown order isn't desirable.

[wsmoses]

Address comments

[kiranchandramohan]

Given a select whose result is an i1, we can eliminate the conditional branch completely by adding a few arithmetic operations.

Nit: Could you clarify which conditional branch you are referring to here? Is it the conditional branch in machine instruction for those architectures without a select instruction?

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

816

Nit: This comment probably needs an update.

[wsmoses]

Address nits

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

static LogicalResult verify(ReturnOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
 // Transforms a select to a not, where relevant.
 //
-//  select %arg, %false, %true
+//  select %arg, %x, %y : i1
 //
 //  becomes
 //
-//  xor %arg, %true
-struct SelectToNot : public OpRewritePattern<SelectOp> {
+//  and(%arg, %x) or and(!%arg, %y)
+struct SelectI1Simplify : 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<arith::XOrIOp>(op, op.getCondition(),
+    Value falseConstant =
+        rewriter.create<arith::ConstantIntOp>(op.getLoc(), true, 1);
+    Value notCondition = rewriter.create<arith::XOrIOp>(
+        op.getLoc(), op.getCondition(), falseConstant);
+
+    Value trueVal = rewriter.create<arith::AndIOp>(
+        op.getLoc(), op.getCondition(), op.getTrueValue());
+    Value falseVal = rewriter.create<arith::AndIOp>(op.getLoc(), notCondition,
                                                     op.getFalseValue());
+    rewriter.replaceOpWithNewOp<arith::OrIOp>(op, trueVal, falseVal);
     return success();
   }
 };
 
 void SelectOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
-  results.insert<SelectToNot>(context);
+  results.insert<SelectI1Simplify>(context);
 }
 
 OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
   auto trueVal = getTrueValue();
   auto falseVal = getFalseValue();
   if (trueVal == falseVal)
     return trueVal;
 

mlir/test/Dialect/Standard/canonicalize.mlir

 ^exit:
   return
 }
 
 // -----
 
 // CHECK-LABEL: @selToNot
 //       CHECK:       %[[trueval:.+]] = arith.constant true
-//       CHECK:       %{{.+}} = arith.xori %arg0, %[[trueval]] : i1
+//       CHECK:       %[[res:.+]] = arith.xori %arg0, %[[trueval]] : i1
+//       CHECK:   return %[[res]]
 func @selToNot(%arg0: i1) -> i1 {
   %true = arith.constant true
   %false = arith.constant false
   %res = select %arg0, %false, %true : i1
   return %res : i1
 }
+
+// CHECK-LABEL: @selToArith
+//       CHECK-NEXT:       %[[trueval:.+]] = arith.constant true
+//       CHECK-NEXT:       %[[notcmp:.+]] = arith.xori %arg0, %[[trueval]] : i1
+//       CHECK-NEXT:       %[[condtrue:.+]] = arith.andi %arg0, %arg1 : i1
+//       CHECK-NEXT:       %[[condfalse:.+]] = arith.andi %[[notcmp]], %arg2 : i1
+//       CHECK-NEXT:       %[[res:.+]] = arith.ori %[[condtrue]], %[[condfalse]] : i1
+//       CHECK:   return %[[res]]
+func @selToArith(%arg0: i1, %arg1 : i1, %arg2 : i1) -> i1 {
+  %res = select %arg0, %arg1, %arg2 : i1
+  return %res : i1
+}