[ValueTracking] Let isGuaranteedNotToBeUndefOrPoison look into branch conditions of dominating blocks' terminators

Summary:
```
  br i1 c, BB1, BB2:
BB1:
  use1(c)
BB2:
  use2(c)
```

In BB1 and BB2, c is never undef or poison because otherwise the branch would have triggered UB.

Checked with Alive2

Reviewers: xbolva00, spatel, lebedev.ri, reames, jdoerfert, nlopes, sanjoy

Reviewed By: reames

Subscribers: jdoerfert, hiraditya, llvm-commits

Tags: #llvm

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

Author: aqjune

llvm/include/llvm/Analysis/ValueTracking.h

@@ -569,7 +569,13 @@ class Value;
 
   /// Return true if this function can prove that V is never undef value
   /// or poison value.
-  bool isGuaranteedNotToBeUndefOrPoison(const Value *V);
+  //
+  /// If CtxI and DT are specified this method performs flow-sensitive analysis
+  /// and returns true if it is guaranteed to be never undef or poison
+  /// immediately before the CtxI.
+  bool isGuaranteedNotToBeUndefOrPoison(const Value *V,
+                                        const Instruction *CtxI = nullptr,
+                                        const DominatorTree *DT = nullptr);
 
   /// Specific patterns of select instructions we can match.
   enum SelectPatternFlavor {

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -5361,16 +5361,17 @@ Value *llvm::SimplifyCall(CallBase *Call, const SimplifyQuery &Q) {
 }
 
 /// Given operands for a Freeze, see if we can fold the result.
-static Value *SimplifyFreezeInst(Value *Op0) {
+static Value *SimplifyFreezeInst(Value *Op0, const SimplifyQuery &Q) {
   // Use a utility function defined in ValueTracking.
-  if (llvm::isGuaranteedNotToBeUndefOrPoison(Op0))
+  if (llvm::isGuaranteedNotToBeUndefOrPoison(Op0, Q.CxtI, Q.DT))
     return Op0;
+
   // We have room for improvement.
   return nullptr;
 }
 
 Value *llvm::SimplifyFreezeInst(Value *Op0, const SimplifyQuery &Q) {
-  return ::SimplifyFreezeInst(Op0);
+  return ::SimplifyFreezeInst(Op0, Q);
 }
 
 /// See if we can compute a simplified version of this instruction.

llvm/lib/Analysis/ValueTracking.cpp

@@ -4525,7 +4525,9 @@ bool llvm::isOverflowIntrinsicNoWrap(const WithOverflowInst *WO,
   return llvm::any_of(GuardingBranches, AllUsesGuardedByBranch);
 }
 
-bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V) {
+bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V,
+                                            const Instruction *CtxI,
+                                            const DominatorTree *DT) {
   // If the value is a freeze instruction, then it can never
   // be undef or poison.
   if (isa<FreezeInst>(V))
@@ -4558,6 +4560,29 @@ bool llvm::isGuaranteedNotToBeUndefOrPoison(const Value *V) {
       return true;
   }
 
+  if (!CtxI || !DT)
+    return false;
+
+  // If V is used as a branch condition before reaching CtxI, V cannot be
+  // undef or poison.
+  //   br V, BB1, BB2
+  // BB1:
+  //   CtxI ; V cannot be undef or poison here
+  auto Dominator = DT->getNode(CtxI->getParent())->getIDom();
+  while (Dominator) {
+    auto *TI = Dominator->getBlock()->getTerminator();
+
+    if (auto BI = dyn_cast<BranchInst>(TI)) {
+      if (BI->isConditional() && BI->getCondition() == V)
+        return true;
+    } else if (auto SI = dyn_cast<SwitchInst>(TI)) {
+      if (SI->getCondition() == V)
+        return true;
+    }
+
+    Dominator = Dominator->getIDom();
+  }
+
   return false;
 }
 

llvm/test/Transforms/InstSimplify/freeze.ll

@@ -18,3 +18,66 @@ define i32 @make_const() {
   %x = freeze i32 10
   ret i32 %x
 }
+
+define i1 @brcond(i1 %c, i1 %c2) {
+; CHECK-LABEL: @brcond(
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    br i1 [[C2:%.*]], label [[A2:%.*]], label [[B]]
+; CHECK:       A2:
+; CHECK-NEXT:    ret i1 [[C]]
+; CHECK:       B:
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  br i1 %c, label %A, label %B
+A:
+  br i1 %c2, label %A2, label %B
+A2:
+  %f1 = freeze i1 %c
+  ret i1 %f1
+B:
+  %f2 = freeze i1 %c
+  ret i1 %f2
+}
+
+define i32 @brcond_switch(i32 %x) {
+; CHECK-LABEL: @brcond_switch(
+; CHECK-NEXT:    switch i32 [[X:%.*]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    i32 0, label [[A:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       A:
+; CHECK-NEXT:    ret i32 [[X]]
+; CHECK:       EXIT:
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  switch i32 %x, label %EXIT [ i32 0, label %A ]
+A:
+  %fr1 = freeze i32 %x
+  ret i32 %fr1
+EXIT:
+  %fr2 = freeze i32 %x
+  ret i32 %fr2
+}
+
+define i1 @brcond_noopt(i1 %c, i1 %c2) {
+; CHECK-LABEL: @brcond_noopt(
+; CHECK-NEXT:    [[F:%.*]] = freeze i1 [[C:%.*]]
+; CHECK-NEXT:    call void @f1(i1 [[F]])
+; CHECK-NEXT:    call void @f2()
+; CHECK-NEXT:    br i1 [[C]], label [[A:%.*]], label [[B:%.*]]
+; CHECK:       A:
+; CHECK-NEXT:    ret i1 false
+; CHECK:       B:
+; CHECK-NEXT:    ret i1 true
+;
+  %f = freeze i1 %c
+  call void @f1(i1 %f) ; cannot optimize i1 %f to %c
+  call void @f2()      ; .. because if f2() exits, `br %c` cannot be reached
+  br i1 %c, label %A, label %B
+A:
+  ret i1 0
+B:
+  ret i1 1
+}
+declare void @f1(i1)
+declare void @f2()