[MLIR] Conditional Branch Argument Propagation

[wsmoses]

Summary

In an operation in the true/false dest of a branch,
one can assume that the operation itself was true/false if
only that edge can reach the operation.

[rriddle]

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

961

nit: Remove the mlir:: here and below.

966

Ideally this would be based on a DominanceInfo that gets passed in as this only works for local cases. Can you add a TODO/FIXME on this pattern for upgrading when we have access to a DominanceInfo here?

968

969–970

Do you need to create a new ConstantOp for every use? Can you just check to see if one has already been created?

970

Isn't there a rewriter.getBoolAttr?

980

981

Same comment here.

987

?

mlir/test/Dialect/Standard/canonicalize.mlir

315–316

Can you clean this up to check only what is necessary? We shouldn't be checking comments, blocks, etc. We should only need to check that constants get generated and those constants feed into the consumers. Any SSA values should also be using regex captures, and not the literal value from the printer.

[wsmoses]

Remove extra mlir namespace

[wsmoses]

Simplify test and address remaining comments

[wsmoses]

Add comment stating this can be expanded later [as a todo]

[ftynse]

LGTM with River's comments addressed.

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

989–1001

It would be nice to have this code factored out into a function, it's almost the same as above and the overall pattern looks useful.

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

if (currentBlock == predBranch.trueDest())
       rewriter.replaceOpWithNewOp<BranchOp>(condbr, condbr.trueDest(),
                                             condbr.trueDestOperands());
     else
       rewriter.replaceOpWithNewOp<BranchOp>(condbr, condbr.falseDest(),
                                             condbr.falseDestOperands());
     return success();
   }
 };
+
+///   cond_br %arg0, ^trueB, ^falseB
+///
+/// ^trueB:
+///   "test.consumer1"(%arg0) : (i1) -> ()
+///    ...
+///
+/// ^falseB:
+///   "test.consumer2"(%arg0) : (i1) -> ()
+///   ...
+///
+/// ->
+///
+///   cond_br %arg0, ^trueB, ^falseB
+/// ^trueB:
+///   "test.consumer1"(%true) : (i1) -> ()
+///   ...
+///
+/// ^falseB:
+///   "test.consumer2"(%false) : (i1) -> ()
+///   ...
+struct CondBranchTruthPropagation : public OpRewritePattern<CondBranchOp> {
+  using OpRewritePattern<CondBranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CondBranchOp condbr,
+                                PatternRewriter &rewriter) const override {
+    // Check that we have a single distinct predecessor.
+    bool replaced = false;
+    mlir::Type ty = rewriter.getI1Type();
+
+    if (condbr.getTrueDest()->getSinglePredecessor()) {
+      for (OpOperand &use :
+           llvm::make_early_inc_range(condbr.condition().getUses())) {
+        if (use.getOwner()->getBlock() == condbr.getTrueDest()) {
+          replaced = true;
+          rewriter.updateRootInPlace(use.getOwner(), [&]() {
+            use.set(rewriter.create<mlir::ConstantOp>(
+                condbr.getLoc(), ty, rewriter.getIntegerAttr(ty, 1)));
+          });
+        }
+      }
+    }
+    if (condbr.getFalseDest()->getSinglePredecessor()) {
+      for (OpOperand &use :
+           llvm::make_early_inc_range(condbr.condition().getUses())) {
+        if (use.getOwner()->getBlock() == condbr.getFalseDest()) {
+          replaced = true;
+          rewriter.updateRootInPlace(use.getOwner(), [&]() {
+            use.set(rewriter.create<mlir::ConstantOp>(
+                condbr.getLoc(), ty, rewriter.getIntegerAttr(ty, 0)));
+          });
+        }
+      }
+    }
+    return replaced ? success() : failure();
+  }
+};
 } // end anonymous namespace
 
 void CondBranchOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
               SimplifyCondBranchIdenticalSuccessors,
-              SimplifyCondBranchFromCondBranchOnSameCondition>(context);
+              SimplifyCondBranchFromCondBranchOnSameCondition,
+              CondBranchTruthPropagation>(context);
 }
 
 Optional<MutableOperandRange>
 CondBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return index == trueIndex ? trueDestOperandsMutable()
                             : falseDestOperandsMutable();
 }
 
 Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   if (IntegerAttr condAttr = operands.front().dyn_cast_or_null<IntegerAttr>())
     return condAttr.getValue().isOneValue() ? trueDest() : falseDest();

mlir/test/Dialect/Standard/canonicalize.mlir

 
 // CHECK-LABEL: @select_cmp_ne_select
 //       CHECK:   return %arg0
 func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
   %0 = cmpi ne, %arg0, %arg1 : i64
   %1 = select %0, %arg0, %arg1 : i64
   return %1 : i64
 }
+
+// -----
+
+// CHECK-LABEL: @branchCondProp
+//       CHECK:       %true = constant true
+//       CHECK:       %false = constant false
+//       CHECK:       cond_br %arg0, ^bb1, ^bb2
+//       CHECK:     ^bb1:  // pred: ^bb0
+//       CHECK:       "test.consumer1"(%true) : (i1) -> ()
+//       CHECK:       br ^bb3
+//       CHECK:     ^bb2:  // pred: ^bb0
+//       CHECK:       "test.consumer2"(%false) : (i1) -> ()
+//       CHECK:       br ^bb3
+//       CHECK:     ^bb3:  // 2 preds: ^bb1, ^bb2
+//       CHECK:       return
+func @branchCondProp(%arg0: i1) {
+  cond_br %arg0, ^trueB, ^falseB
+
+^trueB:
+  "test.consumer1"(%arg0) : (i1) -> ()
+  br ^exit  
+
+^falseB:
+  "test.consumer2"(%arg0) : (i1) -> ()
+  br ^exit
+
+^exit:
+  return
+}