[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

1075

nit: Remove the mlir:: here and below.

1080

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?

1082

1083–1084

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

1084

Isn't there a rewriter.getBoolAttr?

1094

1095

Same comment here.

1101

?

mlir/test/Dialect/Standard/canonicalize.mlir

408–409

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

1103–1115

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;
+    Type ty = rewriter.getI1Type();
+
+    // These variables serve to prevent creating duplicate constants
+    // and hold constant true or false values.
+    Value constantTrue = nullptr;
+    Value constantFalse = nullptr;
+
+    // TODO These checks can be expanded to encompas any use with only
+    // either the true of false edge as a predecessor. For now, we fall
+    // back to checking the single predecessor is given by the true/fasle
+    // destination, thereby ensuring that only that edge can reach the
+    // op.
+    if (condbr.getTrueDest()->getSinglePredecessor()) {
+      for (OpOperand &use :
+           llvm::make_early_inc_range(condbr.condition().getUses())) {
+        if (use.getOwner()->getBlock() == condbr.getTrueDest()) {
+          replaced = true;
+
+          if (!constantTrue)
+            constantTrue = rewriter.create<mlir::ConstantOp>(
+                condbr.getLoc(), ty, rewriter.getBoolAttr(true));
+
+          rewriter.updateRootInPlace(use.getOwner(),
+                                     [&] { use.set(constantTrue); });
+        }
+      }
+    }
+    if (condbr.getFalseDest()->getSinglePredecessor()) {
+      for (OpOperand &use :
+           llvm::make_early_inc_range(condbr.condition().getUses())) {
+        if (use.getOwner()->getBlock() == condbr.getFalseDest()) {
+          replaced = true;
+
+          if (!constantFalse)
+            constantFalse = rewriter.create<mlir::ConstantOp>(
+                condbr.getLoc(), ty, rewriter.getBoolAttr(false));
+
+          rewriter.updateRootInPlace(use.getOwner(),
+                                     [&] { use.set(constantFalse); });
+        }
+      }
+    }
+    return success(replaced);
+  }
+};
 } // 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();
 }

mlir/test/Dialect/Standard/canonicalize.mlir

func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
   return %1 : i64
 }
 
 // -----
 
 // CHECK-LABEL: @indexCastOfSignExtend
 //       CHECK:   %[[res:.+]] = index_cast %arg0 : i8 to index
 //       CHECK:   return %[[res]]
 func @indexCastOfSignExtend(%arg0: i8) -> index {
   %ext = sexti %arg0 : i8 to i16
   %idx = index_cast %ext : i16 to index
   return %idx : index
 }
 
 // CHECK-LABEL: @signExtendConstant
 //       CHECK:   %[[cres:.+]] = constant -2 : i16
 //       CHECK:   return %[[cres]]
 func @signExtendConstant() -> i16 {
   %c-2 = constant -2 : i8
   %ext = sexti %c-2 : i8 to i16
   return %ext : i16
 }
 
 // CHECK-LABEL: @truncConstant
 //       CHECK:   %[[cres:.+]] = constant -2 : i16
 //       CHECK:   return %[[cres]]
 func @truncConstant(%arg0: i8) -> i16 {
   %c-2 = constant -2 : i32
   %tr = trunci %c-2 : i32 to i16
   return %tr : i16
 }
 
-// -----
-
 // CHECK-LABEL: @tripleAddAdd
 //       CHECK:   %[[cres:.+]] = constant 59 : index 
 //       CHECK:   %[[add:.+]] = addi %arg0, %[[cres]] : index 
 //       CHECK:   return %[[add]]
 func @tripleAddAdd(%arg0: index) -> index {
   %c17 = constant 17 : index
   %c42 = constant 42 : index
   %add1 = addi %c17, %arg0 : index

 //       CHECK:   %[[cres:.+]] = cmpi ult, %arg0, %arg1 : i8
 //       CHECK:   return %[[cres]]
 func @notCmpUGE(%arg0: i8, %arg1: i8) -> i1 {
   %true = constant true
   %cmp = cmpi "uge", %arg0, %arg1 : i8
   %ncmp = xor %cmp, %true : i1
   return %ncmp : i1
 }
+
+// -----
+
+// CHECK-LABEL: @branchCondProp
+//       CHECK:       %[[trueval:.+]] = constant true
+//       CHECK:       %[[falseval:.+]] = constant false
+//       CHECK:       "test.consumer1"(%[[trueval]]) : (i1) -> ()
+//       CHECK:       "test.consumer2"(%[[falseval]]) : (i1) -> ()
+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
+}