[LoopFlatten] Fix IV increment use count

[dmgreen]

Summary

The add from the IV in the inner loop was always checking for 2 uses, the phi and the compare. The compare could be based on the phi though, leaving one valid use of the compare. In the testcase we could be left with the phi and a lcssa phi as the two users, invalidly allowing flattening where we shouldn't.

Fixes 58441

[SjoerdMeijer]

I had a look, but you got there first! Thanks for fixing!

This was reported after enabling LoopFlatten by default. I will wait a couple of days to see if there's no fall out from this change, then I will try flipping the switch again later this week.

[dmgreen]

Thanks - that sounds good to me.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

static bool findLoopComponents(
   LLVM_DEBUG(dbgs() << "Found comparison: "; Compare->dump());
 
   // Find increment and trip count.
   // There are exactly 2 incoming values to the induction phi; one from the
   // pre-header and one from the latch. The incoming latch value is the
   // increment variable.
   Increment =
       cast<BinaryOperator>(InductionPHI->getIncomingValueForBlock(Latch));
-  if (Increment->hasNUsesOrMore(3)) {
+  if ((Compare->getOperand(0) != Increment || !Increment->hasNUses(2)) &&
+      !Increment->hasNUses(1)) {
     LLVM_DEBUG(dbgs() << "Could not find valid increment\n");
     return false;
   }
   // The trip count is the RHS of the compare. If this doesn't match the trip
   // count computed by SCEV then this is because the trip count variable
   // has been widened so the types don't match, or because it is a constant and
   // another transformation has changed the compare (e.g. icmp ult %inc,
   // tripcount -> icmp ult %j, tripcount-1), or both.

llvm/test/Transforms/LoopFlatten/widen-pr58441.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes="loop-flatten" -verify-scev -S < %s | FileCheck %s
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
 ; This shouldn't flatten, as the inner add is used by a lcssa.
 
 define i32 @test() {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i64 6, 6
 ; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER_I:%.*]]
 ; CHECK:       for.cond1.preheader.i:
-; CHECK-NEXT:    [[INDVAR1:%.*]] = phi i64 [ [[INDVAR_NEXT2:%.*]], [[FOR_INC5_I:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[L_011_I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[ADD6_I:%.*]], [[FOR_INC5_I:%.*]] ]
 ; CHECK-NEXT:    br label [[WHILE_COND_I_PREHEADER_I:%.*]]
 ; CHECK:       while.cond.i.preheader.i:
-; CHECK-NEXT:    [[INDVAR:%.*]] = phi i64 [ 0, [[FOR_COND1_PREHEADER_I]] ]
-; CHECK-NEXT:    [[STOREMERGE9_I:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_I]] ]
-; CHECK-NEXT:    [[INDVAR_NEXT:%.*]] = add i64 [[INDVAR]], 1
-; CHECK-NEXT:    [[ADD_I:%.*]] = add nuw nsw i32 [[STOREMERGE9_I]], 1
-; CHECK-NEXT:    [[CMP2_I:%.*]] = icmp ult i64 [[INDVAR]], 5
-; CHECK-NEXT:    br label [[FOR_INC5_I]]
+; CHECK-NEXT:    [[STOREMERGE9_I:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_I]] ], [ [[ADD_I:%.*]], [[WHILE_COND_I_PREHEADER_I]] ]
+; CHECK-NEXT:    [[ADD_I]] = add nuw nsw i32 [[STOREMERGE9_I]], 1
+; CHECK-NEXT:    [[CMP2_I:%.*]] = icmp ult i32 [[STOREMERGE9_I]], 5
+; CHECK-NEXT:    br i1 [[CMP2_I]], label [[WHILE_COND_I_PREHEADER_I]], label [[FOR_INC5_I]]
 ; CHECK:       for.inc5.i:
-; CHECK-NEXT:    [[ADD_I_LCSSA_WIDE:%.*]] = phi i64 [ [[INDVAR_NEXT]], [[WHILE_COND_I_PREHEADER_I]] ]
 ; CHECK-NEXT:    [[ADD_I_LCSSA:%.*]] = phi i32 [ [[ADD_I]], [[WHILE_COND_I_PREHEADER_I]] ]
-; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[ADD_I_LCSSA_WIDE]] to i32
-; CHECK-NEXT:    [[INDVAR_NEXT2]] = add i64 [[INDVAR1]], 1
-; CHECK-NEXT:    [[EXITCOND_NOT_I:%.*]] = icmp eq i64 [[INDVAR_NEXT2]], [[FLATTEN_TRIPCOUNT]]
+; CHECK-NEXT:    [[ADD6_I]] = add nuw nsw i32 [[L_011_I]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT_I:%.*]] = icmp eq i32 [[ADD6_I]], 6
 ; CHECK-NEXT:    br i1 [[EXITCOND_NOT_I]], label [[E_EXIT:%.*]], label [[FOR_COND1_PREHEADER_I]]
 ; CHECK:       e.exit:
-; CHECK-NEXT:    [[ADD_I_LCSSA_LCSSA:%.*]] = phi i32 [ [[TMP0]], [[FOR_INC5_I]] ]
+; CHECK-NEXT:    [[ADD_I_LCSSA_LCSSA:%.*]] = phi i32 [ [[ADD_I_LCSSA]], [[FOR_INC5_I]] ]
 ; CHECK-NEXT:    ret i32 [[ADD_I_LCSSA_LCSSA]]
 ;
 entry:
   br label %for.cond1.preheader.i
 
 for.cond1.preheader.i:                            ; preds = %for.inc5.i, %entry
   %l.011.i = phi i32 [ 0, %entry ], [ %add6.i, %for.inc5.i ]
   br label %while.cond.i.preheader.i
e.exit:                                           ; preds = %for.inc5.i
   %add.i.lcssa.lcssa = phi i32 [ %add.i.lcssa, %for.inc5.i ]
   ret i32 %add.i.lcssa.lcssa
 }
 
 
 define i32 @test64() {
 ; CHECK-LABEL: @test64(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i64 6, 6
 ; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER_I:%.*]]
 ; CHECK:       for.cond1.preheader.i:
 ; CHECK-NEXT:    [[L_011_I:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[ADD6_I:%.*]], [[FOR_INC5_I:%.*]] ]
 ; CHECK-NEXT:    br label [[WHILE_COND_I_PREHEADER_I:%.*]]
 ; CHECK:       while.cond.i.preheader.i:
-; CHECK-NEXT:    [[STOREMERGE9_I:%.*]] = phi i64 [ 0, [[FOR_COND1_PREHEADER_I]] ]
-; CHECK-NEXT:    [[ADD_I:%.*]] = add nuw nsw i64 [[STOREMERGE9_I]], 1
+; CHECK-NEXT:    [[STOREMERGE9_I:%.*]] = phi i64 [ 0, [[FOR_COND1_PREHEADER_I]] ], [ [[ADD_I:%.*]], [[WHILE_COND_I_PREHEADER_I]] ]
+; CHECK-NEXT:    [[ADD_I]] = add nuw nsw i64 [[STOREMERGE9_I]], 1
 ; CHECK-NEXT:    [[CMP2_I:%.*]] = icmp ult i64 [[STOREMERGE9_I]], 5
-; CHECK-NEXT:    br label [[FOR_INC5_I]]
+; CHECK-NEXT:    br i1 [[CMP2_I]], label [[WHILE_COND_I_PREHEADER_I]], label [[FOR_INC5_I]]
 ; CHECK:       for.inc5.i:
 ; CHECK-NEXT:    [[ADD_I_LCSSA_WIDEN:%.*]] = phi i64 [ [[ADD_I]], [[WHILE_COND_I_PREHEADER_I]] ]
 ; CHECK-NEXT:    [[ADD_I_LCSSA:%.*]] = trunc i64 [[ADD_I_LCSSA_WIDEN]] to i32
 ; CHECK-NEXT:    [[ADD6_I]] = add nuw nsw i64 [[L_011_I]], 1
-; CHECK-NEXT:    [[EXITCOND_NOT_I:%.*]] = icmp eq i64 [[ADD6_I]], [[FLATTEN_TRIPCOUNT]]
+; CHECK-NEXT:    [[EXITCOND_NOT_I:%.*]] = icmp eq i64 [[ADD6_I]], 6
 ; CHECK-NEXT:    br i1 [[EXITCOND_NOT_I]], label [[E_EXIT:%.*]], label [[FOR_COND1_PREHEADER_I]]
 ; CHECK:       e.exit:
 ; CHECK-NEXT:    [[ADD_I_LCSSA_LCSSA:%.*]] = phi i32 [ [[ADD_I_LCSSA]], [[FOR_INC5_I]] ]
 ; CHECK-NEXT:    ret i32 [[ADD_I_LCSSA_LCSSA]]
 ;
 entry:
   br label %for.cond1.preheader.i