[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:

I am not sure then how to fix this?

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