[MLIR] Canonicalize/fold select %x, 1, 0 to extui

[wsmoses]

Summary

Two canonicalizations for select %x, 1, 0

If the return type is i1, return simply the condition %x, otherwise extui %x to the return type.

[mehdi_amini]

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

990

Can we also support select %arg, %c0, %c1?
Maybe extui (xor %arg, true) ?

990

That adds an operation in the chain, but it also turns the ternary in a fairly simple binary sequence.

[wsmoses]

Add not case

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

if (!op.getType().isInteger(1))
       return failure();
 
     rewriter.replaceOpWithNewOp<arith::XOrIOp>(op, op.getCondition(),
                                                op.getFalseValue());
     return success();
   }
 };
 
+//  select %arg, %c1, %c0 => extui %arg
+struct SelectToExtUI : public OpRewritePattern<SelectOp> {
+  using OpRewritePattern<SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SelectOp op,
+                                PatternRewriter &rewriter) const override {
+    // Cannot extui i1 to i1, or i1 to f32
+    if (!op.getType().isa<IntegerType>() || op.getType().isInteger(1))
+      return failure();
+
+    // select %x, c1, %c0 => extui %arg
+    if (matchPattern(op.getTrueValue(), m_One()))
+      if (matchPattern(op.getFalseValue(), m_Zero())) {
+        rewriter.replaceOpWithNewOp<arith::ExtUIOp>(op, op.getType(),
+                                                    op.getCondition());
+        return success();
+      }
+
+    // select %x, c0, %c1 => extui (xor %arg, true)
+    if (matchPattern(op.getTrueValue(), m_Zero()))
+      if (matchPattern(op.getFalseValue(), m_One())) {
+        rewriter.replaceOpWithNewOp<arith::ExtUIOp>(
+            op, op.getType(),
+            rewriter.create<arith::XOrIOp>(
+                op.getLoc(), op.getCondition(),
+                rewriter.create<arith::ConstantIntOp>(
+                    op.getLoc(), 1, op.getCondition().getType())));
+        return success();
+      }
+
+    return failure();
+  }
+};
+
 void SelectOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
-  results.insert<SelectToNot>(context);
+  results.insert<SelectToNot, SelectToExtUI>(context);
 }
 
 OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
   auto trueVal = getTrueValue();
   auto falseVal = getFalseValue();
   if (trueVal == falseVal)
     return trueVal;
 
   auto condition = getCondition();
 
   // select true, %0, %1 => %0
   if (matchPattern(condition, m_One()))
     return trueVal;
 
   // select false, %0, %1 => %1
   if (matchPattern(condition, m_Zero()))
     return falseVal;
 
+  // select %x, true, false => %x
+  if (getType().isInteger(1))
+    if (matchPattern(getTrueValue(), m_One()))
+      if (matchPattern(getFalseValue(), m_Zero()))
+        return condition;
+
   if (auto cmp = dyn_cast_or_null<arith::CmpIOp>(condition.getDefiningOp())) {
     auto pred = cmp.getPredicate();
     if (pred == arith::CmpIPredicate::eq || pred == arith::CmpIPredicate::ne) {
       auto cmpLhs = cmp.getLhs();
       auto cmpRhs = cmp.getRhs();
 
       // %0 = arith.cmpi eq, %arg0, %arg1
       // %1 = select %0, %arg0, %arg1 => %arg1

mlir/test/Dialect/Standard/canonicalize.mlir

 func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
   %0 = arith.cmpi ne, %arg0, %arg1 : i64
   %1 = select %0, %arg0, %arg1 : i64
   return %1 : i64
 }
 
 // -----
 
+// CHECK-LABEL: @select_extui
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 : i1 to i64
+//       CHECK:   return %[[res]]
+func @select_extui(%arg0: i1) -> i64 {
+  %c0_i64 = arith.constant 0 : i64
+  %c1_i64 = arith.constant 1 : i64
+  %res = select %arg0, %c1_i64, %c0_i64 : i64
+  return %res : i64
+}
+
+// CHECK-LABEL: @select_extui2
+// CHECK-DAG:  %true = arith.constant true
+// CHECK-DAG:  %[[xor:.+]] = arith.xori %arg0, %true : i1
+// CHECK-DAG:  %[[res:.+]] = arith.extui %[[xor]] : i1 to i64
+//       CHECK:   return %[[res]]
+func @select_extui2(%arg0: i1) -> i64 {
+  %c0_i64 = arith.constant 0 : i64
+  %c1_i64 = arith.constant 1 : i64
+  %res = select %arg0, %c0_i64, %c1_i64 : i64
+  return %res : i64
+}
+
+// -----
+
+// CHECK-LABEL: @select_extui_i1
+//  CHECK-NEXT:   return %arg0
+func @select_extui_i1(%arg0: i1) -> i1 {
+  %c0_i1 = arith.constant false
+  %c1_i1 = arith.constant true
+  %res = select %arg0, %c1_i1, %c0_i1 : i1
+  return %res : i1
+}
+
+// -----
+
 // CHECK-LABEL: @branchCondProp
 //       CHECK:       %[[trueval:.+]] = arith.constant true
 //       CHECK:       %[[falseval:.+]] = arith.constant false
 //       CHECK:       "test.consumer1"(%[[trueval]]) : (i1) -> ()
 //       CHECK:       "test.consumer2"(%[[falseval]]) : (i1) -> ()
 func @branchCondProp(%arg0: i1) {
   cond_br %arg0, ^trueB, ^falseB