Teach ValueTracking that compare-and-branch on poison is UB

[sanjoy]

Summary

This lets us optimize the interesting case of

do {
  I = PHI(I_Begin, I_Next);
  I_Next = I +nsw 1
  Use(Sext(I_Next))
} while (I_Next < N);

Note: we "already get" the above case, but the way we do so is
incorrect. Making our poison value analysis smarter about this will let
us get rid of the buggy behavior (in a later patch) without regressing
too much performance across the board.

[majnemer]

lib/Analysis/ValueTracking.cpp

3413–3416

Our langref states:

An instruction control-depends on a terminator instruction if the terminator instruction has multiple successors and the instruction is always executed when control transfers to one of the successors, and may not be executed when control is transferred to another.

The code, as written, sounds like it will do the wrong thing if the branch instruction branches to the same basic block regardless of the condition operand. Maybe this is just a bug in the langref? Equally possible is that I misunderstand the behavior that this function is trying to suss out.

[sanjoy]

lib/Analysis/ValueTracking.cpp

3413–3416

You're right; as written this code is not correct and it can't easily be fixed either. Going by the langref, this code fragment (the motivation for this change)

loop:
  %iv = phi(0, %iv.inc)
  %iv.inc = add nsw %iv, 1
  %cmp = %iv.inc s< %L
  br i1 %cmp, label %leave, label %loop
leave:
  ...

is well defined, even when %iv.inc is poison.

/sadtrombone

[sanjoy]

As @majnemer pointed out, this change is wrong.

lib/Analysis/ValueTracking.cpp

case Instruction::Mul: {
             // the other operand that is poison.
             return !CI->isZero();
           }
         }
       }
       return false;
     }
 
+    case Instruction::ICmp:
+      // Comparing poison with any value yields poison.  This is why, for
+      // instance, x s< (x +nsw 1) can be folded to true.
+      return true;
+
     case Instruction::GetElementPtr:
       // A GEP implicitly represents a sequence of additions, subtractions,
       // truncations, sign extensions and multiplications. The multiplications
       // are by the non-zero sizes of some set of types, so we do not have to be
       // concerned with multiplication by zero. If the GEP is in-bounds, then
       // these operations are implicitly no-signed-wrap so poison is propagated
       // by the arguments above for Add, Sub, Trunc, SExt and Mul.
       return cast<GEPOperator>(I)->isInBounds();
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

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

; CHECK: --> {%offset,+,1}<nuw>
   %f = load float, float* %ptr, align 4
   %exitcond = icmp eq i32 %nexti, %numIterations
   br i1 %exitcond, label %exit, label %loop
 
 exit:
   ret void
 }
 
+; Example where an add should get the nuw flag: the UB-on-poison in
+; this example is not due to a memory operation, but due to a
+; compare-and-branch.
+define void @test-add-nuw-from-icmp-and-branch(
+    float* %input, i32 %offset, i32 %numIterations) {
+; CHECK-LABEL: @test-add-nuw-from-icmp-and-branch
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %loop ], [ 0, %entry ]
+
+; CHECK: %index32 =
+; CHECK: --> {%offset,+,1}<nuw>
+  %index32 = add nuw i32 %i, %offset
+
+  %nexti = add nuw i32 %i, 1
+  %exitcond = icmp eq i32 %index32, %numIterations
+  br i1 %exitcond, label %exit, label %loop
+
+exit:
+  ret void
+}
+
+; Like @test-add-nuw-from-icmp-and-branch, but we're testing just the
+; icmp logic here (this pattern is unlikely to appear in real code).
+define void @test-add-nuw-from-icmp(float* %input, i32 %offset,
+                                    i32 %numIterations) {
+; CHECK-LABEL: @test-add-nuw-from-icmp
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %loop ], [ 0, %entry ]
+
+; CHECK: %index32 =
+; CHECK: --> {%offset,+,1}<nuw>
+  %index32 = add nuw i32 %i, %offset
+  %cmp = icmp sgt i32 %index32, 0
+  %cmp.idx = sext i1 %cmp to i32
+
+  %ptr = getelementptr inbounds float, float* %input, i32 %cmp.idx
+  %nexti = add nuw i32 %i, 1
+  %f = load float, float* %ptr, align 4
+  %exitcond = icmp eq i32 %nexti, %numIterations
+  br i1 %exitcond, label %exit, label %loop
+
+exit:
+  ret void
+}
+
 ; With no load to trigger UB from poison, we cannot infer nsw.
 define void @test-add-no-load(float* %input, i32 %offset, i32 %numIterations) {
 ; CHECK-LABEL: @test-add-no-load
 entry:
   br label %loop
 loop:
   %i = phi i32 [ %nexti, %loop ], [ 0, %entry ]