Have isKnownNotFullPoison be smarter around control flow

[sanjoy]

Summary

(... while still not using a PostDomTree)

The way we use isKnownNotFullPoison from SCEV today, the new CFG walking
logic will not trigger for any realistic cases -- it will kick in only
for situations where we could have merged the contiguous basic blocks
anyway[0], since the poison generating instruction dominates all of its
non-PHI uses (which are the only uses we consider right now).

However, having this change in place will allow a later bugfix to break
fewer llvm-lit tests.

[0]: i.e. cases where block A branches to block B and B is A's only
successor and A is B's only predecessor.

[broune]

lib/Analysis/ValueTracking.cpp

3413

As far as I can tell, this code for Br is not in the official repository, yet it is not listed as a changed line in the review.

https://github.com/llvm-mirror/llvm/blob/615a6cf40aa9c7b280fd665e4ad66d4e0d5395fe/lib/Analysis/ValueTracking.cpp#L3522

Maybe I'm just confused about how diffs are shown or this was added very recently?

3442

Would it make sense to have a limit on how many instructions/basic blocks this will consider?

3462

Could there be braces for this multi-line if?

[sanjoy]

lib/Analysis/ValueTracking.cpp

3413

This is my fault. I had this revision based off of D19211 but didn't note it in the revision.

[sanjoy]

lib/Analysis/ValueTracking.cpp

3442

Added a limit of MaxDepth blocks (for the lack of a better limit).

lib/Analysis/ValueTracking.cpp

case Instruction::AtomicRMW:
       return cast<AtomicRMWInst>(I)->getPointerOperand();
 
     case Instruction::UDiv:
     case Instruction::SDiv:
     case Instruction::URem:
     case Instruction::SRem:
       return I->getOperand(1);
 
     case Instruction::Br:
       if (cast<BranchInst>(I)->isConditional())
         return cast<BranchInst>(I)->getCondition();
       return nullptr;
 
     default:
       return nullptr;
   }
 }
 
 bool llvm::isKnownNotFullPoison(const Instruction *PoisonI) {
   // We currently only look for uses of poison values within the same basic
   // block, as that makes it easier to guarantee that the uses will be
   // executed given that PoisonI is executed.
   //
   // FIXME: Expand this to consider uses beyond the same basic block. To do
   // this, look out for the distinction between post-dominance and strong
   // post-dominance.
   const BasicBlock *BB = PoisonI->getParent();
 
   // Set of instructions that we have proved will yield poison if PoisonI
   // does.
   SmallSet<const Value *, 16> YieldsPoison;
+  SmallSet<const BasicBlock *, 4> Visited;
   YieldsPoison.insert(PoisonI);
+  Visited.insert(PoisonI->getParent());
 
-  for (BasicBlock::const_iterator I = PoisonI->getIterator(), E = BB->end();
-       I != E; ++I) {
-    if (&*I != PoisonI) {
-      const Value *NotPoison = getGuaranteedNonFullPoisonOp(&*I);
-      if (NotPoison != nullptr && YieldsPoison.count(NotPoison)) return true;
-      if (!isGuaranteedToTransferExecutionToSuccessor(&*I))
+  BasicBlock::const_iterator Begin = PoisonI->getIterator(), End = BB->end();
+
+  while (true) {
+    for (auto &I : make_range(Begin, End)) {
+      if (&I != PoisonI) {
+        const Value *NotPoison = getGuaranteedNonFullPoisonOp(&I);
+        if (NotPoison != nullptr && YieldsPoison.count(NotPoison))
+          return true;
+        if (!isGuaranteedToTransferExecutionToSuccessor(&I))
           return false;
       }
 
       // Mark poison that propagates from I through uses of I.
-    if (YieldsPoison.count(&*I)) {
-      for (const User *User : I->users()) {
+      if (YieldsPoison.count(&I)) {
+        for (const User *User : I.users()) {
           const Instruction *UserI = cast<Instruction>(User);
-        if (UserI->getParent() == BB && propagatesFullPoison(UserI))
+          if (propagatesFullPoison(UserI))
             YieldsPoison.insert(User);
         }
       }
     }
+
+    if (auto *NextBB = BB->getSingleSuccessor())
+      if (Visited.insert(NextBB).second) {
+        BB = NextBB;
+        Begin = BB->getFirstNonPHI()->getIterator();
+        End = BB->end();
+        continue;
+      }
+
+    break;
+  };
   return false;
 }
 
 static bool isKnownNonNaN(Value *V, FastMathFlags FMF) {
   if (FMF.noNaNs())
     return true;
 
   if (auto *C = dyn_cast<ConstantFP>(V))

test/Analysis/ScalarEvolution/flags-from-poison.ll

 define void @test-add-not-header2(float* %input, i32 %offset, i32 %numIterations) {
 ; CHECK-LABEL: @test-add-not-header2
 entry:
   br label %loop
 loop:
   %i = phi i32 [ %nexti, %loop2 ], [ 0, %entry ]
 
 ; CHECK: %index32 =
+; CHECK: --> {%offset,+,1}<nsw>
+  %index32 = add nsw i32 %i, %offset
+
+  %ptr = getelementptr inbounds float, float* %input, i32 %index32
+  %nexti = add nsw i32 %i, 1
+  br label %loop2
+loop2:
+  %f = load float, float* %ptr, align 4
+  %exitcond = icmp eq i32 %nexti, %numIterations
+  br i1 %exitcond, label %exit, label %loop
+exit:
+  ret void
+}
+
+; Similar to test-add-not-header, but in this case the load
+; instruction may not be executed.
+define void @test-add-not-header3(float* %input, i32 %offset, i32 %numIterations,
+                                 i1* %cond_buf) {
+; CHECK-LABEL: @test-add-not-header3
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %loop2 ], [ 0, %entry ]
+
+; CHECK: %index32 =
 ; CHECK: --> {%offset,+,1}<nw>
   %index32 = add nsw i32 %i, %offset
 
   %ptr = getelementptr inbounds float, float* %input, i32 %index32
   %nexti = add nsw i32 %i, 1
+  %cond = load volatile i1, i1* %cond_buf
+  br i1 %cond, label %loop2, label %exit
+loop2:
+  %f = load float, float* %ptr, align 4
+  %exitcond = icmp eq i32 %nexti, %numIterations
+  br i1 %exitcond, label %exit, label %loop
+exit:
+  ret void
+}
+
+; Same thing as test-add-not-header2, except we have a few extra
+; blocks.
+define void @test-add-not-header4(float* %input, i32 %offset, i32 %numIterations) {
+; CHECK-LABEL: @test-add-not-header4
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %loop2 ], [ 0, %entry ]
+
+; CHECK: %index32 =
+; CHECK: --> {%offset,+,1}<nsw>
+  %index32 = add nsw i32 %i, %offset
+
+  %ptr = getelementptr inbounds float, float* %input, i32 %index32
+  %nexti = add nsw i32 %i, 1
+  br label %loop3
+loop3:
+  br label %loop4
+loop4:
   br label %loop2
 loop2:
   %f = load float, float* %ptr, align 4
   %exitcond = icmp eq i32 %nexti, %numIterations
   br i1 %exitcond, label %exit, label %loop
 exit:
   ret void
 }
 
+; Demonstrate why we need a Visited set in llvm::isKnownNotFullPoison.
+define void @test-add-not-header5(float* %input, i32 %offset) {
+; CHECK-LABEL: @test-add-not-header5
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %loop ], [ 0, %entry ]
+
+; CHECK: %index32 =
+; CHECK: --> {%offset,+,1}<nw>
+  %index32 = add nsw i32 %i, %offset
+
+  %ptr = getelementptr inbounds float, float* %input, i32 %index32
+  %nexti = add nsw i32 %i, 1
+  br label %loop
+
+exit:
+  ret void
+}
+
 ; The call instruction makes it not guaranteed that the add will be
 ; executed, since it could run forever or throw an exception, so we
 ; cannot assume that the UB is realized.
 define void @test-add-call(float* %input, i32 %offset, i32 %numIterations) {
 ; CHECK-LABEL: @test-add-call
 entry:
   br label %loop
 loop: