[llvm] [LoopVectorize] Don't speculate an early-exit trip count that may cause UB (PR #219893)

[Madhur Amilkanthwar via llvm-commits]

================
@@ -10,16 +11,55 @@ declare i32 @llvm.cttz.i32(i32, i1 immarg)
 ; well defined on that path must not be computed unconditionally in the
 ; preheader.
 
-; FIXME: %ct is poison when %x is 0, so %d may be poison. The udiv is currently
-; speculated into the preheader, where it can divide by poison, which is UB. The
-; original loop returns 0 without evaluating %d when %skip is true, so this is a
-; miscompile and the loop must not be vectorized.
+; %ct is poison when %x is 0, so %d may be poison and speculating (63 /u %d)
+; would divide by poison, which is UB. The original loop returns 0 without ever
+; evaluating %d when %skip is true, so do not vectorize.
 define i32 @udiv_by_poison_step(i32 %x, i1 %skip) {
 ; CHECK-LABEL: define i32 @udiv_by_poison_step(
 ; CHECK-SAME: i32 [[X:%.*]], i1 [[SKIP:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    [[CT:%.*]] = call i32 @llvm.cttz.i32(i32 [[X]], i1 true)
 ; CHECK-NEXT:    [[D:%.*]] = add nuw nsw i32 [[CT]], 1
+; CHECK-NEXT:    br label %[[LOOP1:.*]]
+; CHECK:       [[LOOP1]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[INC:%.*]], %[[LATCH:.*]] ]
+; CHECK-NEXT:    br i1 [[SKIP]], label %[[EXIT:.*]], label %[[LATCH]]
+; CHECK:       [[LATCH]]:
+; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I]], [[D]]
+; CHECK-NEXT:    [[DONE:%.*]] = icmp uge i32 [[INC]], 64
+; CHECK-NEXT:    br i1 [[DONE]], label %[[EXIT]], label %[[LOOP1]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[R:%.*]] = phi i32 [ 0, %[[LOOP1]] ], [ [[INC]], %[[LATCH]] ]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+entry:
+  %ct = call i32 @llvm.cttz.i32(i32 %x, i1 true)
----------------
madhur13490 wrote:

Agreed, it's not really about immediate UB — any poison symbolic max BTC in the preheader is the problem, since it feeds the min-iters branch. Today's `isGuaranteedNotToCauseUB` only catches a udiv with a bad divisor, so a poison BTC without a udiv still vectorizes.

I tried switching to isGuaranteedNotToBePoison. It fixes that case, but regresses 7 early-exit tests that stop vectorizing: AArch64/simple_early_exit.ll, early_exit_legality.ll, early_exit_store_legality.ll, fold-epilogue-tail.ll, single-early-exit-deref-assumptions.ll, single_early_exit.ll, vect.stats.ll. It's too blunt — it bails whenever the BTC has a value it can't prove non-poison (a loop bound that's an arg or a load), which is most variable-trip-count loops.

Following the existing TODO, I think the fix is to freeze the expanded BTC instead of bailing, so poison can't reach the branch but these loops still vectorize. Sound good, or would you prefer a narrower predicate? I can prototype the freeze.



https://github.com/llvm/llvm-project/pull/219893