Broaden the definition of a "widenable branch"

As a reminder, a "widenable branch" is the pattern "br i1 (and i1 X, WC()), label %taken, label %untaken" where "WC" is the widenable condition intrinsics. The semantics of such a branch (derived from the semantics of WC) is that a new condition can be added into the condition arbitrarily without violating legality.

Broaden the definition in two ways:
    Allow swapped operands to the br (and X, WC()) form
    Allow widenable branch w/trivial condition (i.e. true) which takes form of br i1 WC()

The former is just general robustness (e.g. for X = non-instruction this is what instcombine produces). The later is specifically important as partial unswitching of a widenable range check produces exactly this form above the loop.

Differential Revision: https://reviews.llvm.org/D70502

1 files changed, 39 insertions, 11 deletions

diff --git a/llvm/lib/Transforms/Utils/GuardUtils.cpp b/llvm/lib/Transforms/Utils/GuardUtils.cpp
index b8c4c76..241bfbf 100644
--- a/llvm/lib/Transforms/Utils/GuardUtils.cpp
+++ b/llvm/lib/Transforms/Utils/GuardUtils.cpp
@@ -15,10 +15,12 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 static cl::opt<uint32_t> PredicatePassBranchWeight(
     "guards-predicate-pass-branch-weight", cl::Hidden, cl::init(1 << 20),
@@ -80,23 +82,49 @@ void llvm::makeGuardControlFlowExplicit(Function *DeoptIntrinsic,
 void llvm::widenWidenableBranch(BranchInst *WidenableBR, Value *NewCond) {
   assert(isWidenableBranch(WidenableBR) && "precondition");
 
-  Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
-  // Condition is only guaranteed to dominate branch
-  WCAnd->moveBefore(WidenableBR);
-  Value *OldCond = WCAnd->getOperand(0);
-  IRBuilder<> B(WCAnd);
-  WCAnd->setOperand(0, B.CreateAnd(NewCond, OldCond));
+  // The tempting trivially option is to produce something like this:
+  // br (and oldcond, newcond) where oldcond is assumed to contain a widenable
+  // condition, but that doesn't match the pattern parseWidenableBranch expects
+  // so we have to be more sophisticated.
 
+  if (match(WidenableBR->getCondition(),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>())) {
+    IRBuilder<> B(WidenableBR);
+    WidenableBR->setCondition(B.CreateAnd(NewCond,
+                                          WidenableBR->getCondition()));
+  } else {
+    Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
+    // Condition is only guaranteed to dominate branch
+    WCAnd->moveBefore(WidenableBR);
+    IRBuilder<> B(WCAnd);
+    const bool Op0IsWC =
+      match(WCAnd->getOperand(0),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>());
+    const unsigned CondOpIdx = Op0IsWC ? 1 : 0;
+    Value *OldCond = WCAnd->getOperand(CondOpIdx);
+    NewCond = B.CreateAnd(NewCond, OldCond);
+    WCAnd->setOperand(CondOpIdx, NewCond);
+  }
   assert(isWidenableBranch(WidenableBR) && "preserve widenabiliy");
 }
 
 void llvm::setWidenableBranchCond(BranchInst *WidenableBR, Value *NewCond) {
   assert(isWidenableBranch(WidenableBR) && "precondition");
 
-  Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
-  // Condition is only guaranteed to dominate branch
-  WCAnd->moveBefore(WidenableBR);
-  WCAnd->setOperand(0, NewCond);
-
+  if (match(WidenableBR->getCondition(),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>())) {
+    IRBuilder<> B(WidenableBR);
+    WidenableBR->setCondition(B.CreateAnd(NewCond,
+                                          WidenableBR->getCondition()));
+  } else {
+    Instruction *WCAnd = cast<Instruction>(WidenableBR->getCondition());
+    // Condition is only guaranteed to dominate branch
+    WCAnd->moveBefore(WidenableBR);
+    const bool Op0IsWC =
+      match(WCAnd->getOperand(0),
+            m_Intrinsic<Intrinsic::experimental_widenable_condition>());
+    const unsigned CondOpIdx = Op0IsWC ? 1 : 0;
+    WCAnd->setOperand(CondOpIdx, NewCond);
+  }
   assert(isWidenableBranch(WidenableBR) && "preserve widenabiliy");
 }