[LoopFlatten] Fix missed LoopFlatten opportunity

[RosieSumpter]

Summary

When the limit of the inner loop is a known integer, the InstCombine
pass now causes the transformation e.g. imcp ult i32 %inc, tripcount ->
icmp ult %j, tripcount-step (where %j is the inner loop induction
variable and %inc is add %j, step), which is now accounted for when
identifying the trip count of the loop. This is also an acceptable use
of %j (provided the step is 1) so is ignored as long as the compare
that it's used in is also the condition of the inner branch.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

156

Make this an else if?

165

I am wondering what happens if the RHS is not a constant, so instead of this:

%cmp = icmp ult i32 %j, 18

we have something like:

%cmp = icmp ult i32 %j, %N

Do we reject that earlier, and/or do we not have these test cases yet?

168

And I was also wondering if we need to worry about RHS being INT_MAX, in which case the + 1 will overflow.

403

Perhaps here we can have:

// The use is in the compare which is also the condition of the inner
// branch. In this case the compare has been altered by another
// transformation (e.g icmp ult %inc, limit -> icmp ult %j, limit-1).
// Ignore this use as the compare gets removed later anyway.
if (U == FI.InnerBranch->getCondition())
   continue;

To make things a bit simpler.

llvm/test/Transforms/LoopFlatten/loop-flatten-negative.ll

348

You need to CHECK: something here.

379

Same here.

llvm/test/Transforms/LoopFlatten/loop-flatten.ll

589

It's more common to put the CHECK-LABEL just before

; CHECK: entry:

[fhahn]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

162

Those kinds of patterns are very fragile in general. Would it be possible to use SCEV instead, which makes it easy to analyse the induction variables & trip counts in a more robust way?

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

162

You're absolutely right. In fact, there's a quite some pattern matching going on for things like getLoopTest, getLoopIncrement, etc., that you'd probably expect to be present in a loop util helper function or something like that, but they aren't. The pattern matching makes some things easier though (as we are looking for some specific patterns). I feel that moving some things to helpers and using SCEV is major surgery, and if we promise not to add more pattern matching after this, I was hoping to postpone this surgery. I.e., we are on a little mission to get LoopFlatten enabled by default (we have this enabled for years now downstream), and before we do that we wanted to fix 2 minor things:

• this minor extension for a case which you expect to trigger,
• and we need to strengthen an overflow check.

After this, I think a nice follow up project is indeed to see if we can move code to helpers and use SCEV.

[RosieSumpter]

• Changed else to else if
• Amended comments
• Moved 'CHECK-LABEL' statement
• Added test for when incoming phi node value for preheader is a variable (not zero)
• Fixed debug output for when incoming phi node value can't be cast to ConstantInt (gives null which causes a crash when dump() is used)

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

165

It shouldn't get here unless it is constant since this only catches the case where the compare which is the condition of the back edge has been changed from e.g.

%cmp = icmp ult i32 %inc, 19

to

%cmp = icmp ult i32 %j, 18

If the RHS is unknown e.g.

%cmp = icmp ult i32 %inc, %N

it won't have been changed and so will be caught by one of the previous pattern matches.

I've altered the comments so that it's clearer that this match is only for the case when limit is a constant, and also added an assert to ensure that the LHS of the compare is the induction phi.

168

So here the limit is RHS+1 because the another transformation changed the compare from

%cmp = icmp ult i32 %inc, limit

to

%cmp = icmp ult i32 %j, limit-1

(where limit is a known constant) so the only way RHS=limit-1 would be INT_MAX is if limit had already overflowed, so I think we don't need to check for it here?

llvm/test/Transforms/LoopFlatten/loop-flatten-negative.ll

348

The only CHECK in this file is at the beginning:

CHECK-NOT: Checks all passed, doing the transformation

which checks that all these tests fail doesn't it? Or should I being doing more specific checks here?

[fhahn]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

162

I was hoping to postpone this surgery. I.e., we are on a little mission to get LoopFlatten enabled by default

I understand the motivation, but I am not sure it is the best way forward. If SCEV is the the right tool to use, I think it would be great to make the switch before enabling it by default; I am not too familiar with the details, but it looks like the pattern matching of induction variables and compares to get the loop trip counts seems like something that SCEV already provides.

After it is enabled by default, the motivation for a substantial refactoring might be smaller than before enabling it. When enabling it by default, it is also beneficial for the pass to be as general as possible, so fundamental problems can be flushed out early and the compile-time/code-size impact is also representative on a large set of inputs.

Also, not using SCEV makes it harder to reason about the pass, because we can't rely on trusting SCEV to correctly handle all the difficult cases with respect to overflows and such when it comes to trip counts and inductions. We instead need to verify the manual patterns and their interaction with the dependent code. For example, it seems like using SCEV would allow us to be more confident when reasoning about some of the issues in the comments below.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

162

Alright, fair enough, let's then first see what using SCEV exactly means and would involve.

I am not too familiar with the details, but it looks like the pattern matching of induction variables and compares to get the loop trip counts seems like something that SCEV already provides.

Yeah, the trip count is the easy part, but it's also the other loop components like the increment, test, etc., which are not provided by existing APIs. Thus, with major surgery, I probably also meant getting rid of findLoopComponents and adding something similar to Loop or LoopUtils, possibly using SCEV more.

This little pattern match extension here might still be innocent and the way to go, but like I said, let's review SCEV usage first then.

[RosieSumpter]

Other patches have refactored parts of LoopFlatten's code to use SCEV and Loop instead of pattern matching (see https://reviews.llvm.org/D106580 and https://reviews.llvm.org/D106256) so this diff has been updated to be on top of these changes.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

177–178

Nit: every time I need to read this if a couple of times, but I think it makes sense. Just for readability, to reduce indentation inside the if a little bit, I was wondering if this would help:

if (SE->getSCEV(TripCount) != SCEVTripCount) && !IsWidened) {
..
} else if (SE->getSCEV(TripCount) != SCEVTripCount)) {
..
}

206

Nit: what was this true argument again? I would guess thought that we don't need...

llvm/test/Transforms/LoopFlatten/loop-flatten-negative.ll

348

Ah, sorry, I had missed that. It's indeed fine like this.

[RosieSumpter]

• Replaced checks on RHS of compare with asserts
• Slight rewriting to improve readability

[SjoerdMeijer]

Thanks, much easier to read now, LGTM!

[RosieSumpter]

• Added check and test for trip count

[SjoerdMeijer]

I've ran downstream testing for this patch, and that came back okay. LGTM.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

if (!Compare || !IsValidPredicate(Compare->getUnsignedPredicate()) ||
     return false;
   }
   BackBranch = cast<BranchInst>(Latch->getTerminator());
   IterationInstructions.insert(BackBranch);
   LLVM_DEBUG(dbgs() << "Found back branch: "; BackBranch->dump());
   IterationInstructions.insert(Compare);
   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 =
       dyn_cast<BinaryOperator>(InductionPHI->getIncomingValueForBlock(Latch));
   if (Increment->hasNUsesOrMore(3)) {
     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 either because the trip count variable
   // has been widened (then leave the trip count as it is), or because it is a
   // constant and another transformation has changed the compare, e.g.
-  // icmp ult %inc, tripcount -> icmp ult %j, tripcount-1, then we don't flatten
-  // the loop (yet).
+  // icmp ult %inc, tripcount -> icmp ult %j, tripcount-1.
   TripCount = Compare->getOperand(1);
   const SCEV *BackedgeTakenCount = SE->getBackedgeTakenCount(L);
   if (isa<SCEVCouldNotCompute>(BackedgeTakenCount)) {
     LLVM_DEBUG(dbgs() << "Backedge-taken count is not predictable\n");
     return false;
   }
   const SCEV *SCEVTripCount = SE->getTripCountFromExitCount(BackedgeTakenCount);
-  if (SE->getSCEV(TripCount) != SCEVTripCount) {
-    if (!IsWidened) {
-      LLVM_DEBUG(dbgs() << "Could not find valid trip count\n");
-      return false;
-    }
-    auto TripCountInst = dyn_cast<Instruction>(TripCount);
+  if (SE->getSCEV(TripCount) != SCEVTripCount && !IsWidened) {
+    ConstantInt *RHS = dyn_cast<ConstantInt>(TripCount);
+    // If the IV hasn't been widened, the only way the RHS of the Compare can be
+    // different from the SCEV trip count is if it is a constant which has been
+    // changed by another transformation.
+    assert(RHS && "Expected RHS of compare to be constant");
+    // The L->isCanonical check above ensures we only get here if the loop
+    // increments by 1 on each iteration, so the RHS of the Compare is
+    // tripcount-1 (i.e equivalent to the backedge taken count).
+    assert(SE->getSCEV(RHS) == BackedgeTakenCount &&
+           "Expected RHS of compare to be equal to the backedge taken count");
+    ConstantInt *One = ConstantInt::get(RHS->getType(), 1);
+    TripCount = ConstantInt::get(TripCount->getContext(),
+                                 RHS->getValue() + One->getValue());
+  } else if (SE->getSCEV(TripCount) != SCEVTripCount) {
+    auto *TripCountInst = dyn_cast<Instruction>(TripCount);
     if (!TripCountInst) {
       LLVM_DEBUG(dbgs() << "Could not find valid extended trip count\n");
       return false;
     }
     if ((!isa<ZExtInst>(TripCountInst) && !isa<SExtInst>(TripCountInst)) ||
         SE->getSCEV(TripCountInst->getOperand(0)) != SCEVTripCount) {
       LLVM_DEBUG(dbgs() << "Could not find valid extended trip count\n");
       return false;
     }
   }
   IterationInstructions.insert(Increment);
   LLVM_DEBUG(dbgs() << "Found increment: "; Increment->dump());
   LLVM_DEBUG(dbgs() << "Found trip count: "; TripCount->dump());
 
   LLVM_DEBUG(dbgs() << "Successfully found all loop components\n");
   return true;
 }
 
 static bool checkPHIs(FlattenInfo &FI, const TargetTransformInfo *TTI) {
   // All PHIs in the inner and outer headers must either be:
   // - The induction PHI, which we are going to rewrite as one induction in
for (User *U : FI.InnerInductionPHI->users()) {
     // After widening the IVs, a trunc instruction might have been introduced, so
     // look through truncs.
     if (isa<TruncInst>(U)) {
       if (!U->hasOneUse())
         return false;
       U = *U->user_begin();
     }
 
+    // If the use is in the compare (which is also the condition of the inner
+    // branch) then the compare has been altered by another transformation e.g
+    // icmp ult %inc, tripcount -> icmp ult %j, tripcount-1, where tripcount is
+    // a constant. Ignore this use as the compare gets removed later anyway.
+    if (U == FI.InnerBranch->getCondition())
+      continue;
+
     LLVM_DEBUG(dbgs() << "Found use of inner induction variable: "; U->dump());
 
     Value *MatchedMul;
     Value *MatchedItCount;
     bool IsAdd = match(U, m_c_Add(m_Specific(FI.InnerInductionPHI),
                                   m_Value(MatchedMul))) &&
                  match(MatchedMul, m_c_Mul(m_Specific(FI.OuterInductionPHI),
                                            m_Value(MatchedItCount)));
 
     // Matches the same pattern as above, except it also looks for truncs
     // on the phi, which can be the result of widening the induction variables.
     bool IsAddTrunc = match(U, m_c_Add(m_Trunc(m_Specific(FI.InnerInductionPHI)),
                                        m_Value(MatchedMul))) &&
                       match(MatchedMul,
                             m_c_Mul(m_Trunc(m_Specific(FI.OuterInductionPHI)),
                             m_Value(MatchedItCount)));
 
     if ((IsAdd || IsAddTrunc) && MatchedItCount == InnerTripCount) {
       LLVM_DEBUG(dbgs() << "Use is optimisable\n");
       ValidOuterPHIUses.insert(MatchedMul);
       FI.LinearIVUses.insert(U);
     } else {
       LLVM_DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
       return false;
     }

llvm/test/Transforms/LoopFlatten/loop-flatten-negative.ll

for.inc6:                                         ; preds = %for.body3
   %inc7 = add nuw nsw i32 %i.018, 1
   %exitcond19 = icmp ne i32 %inc7, 10
   br i1 %exitcond19, label %for.body, label %for.end8
 
 for.end8:                                         ; preds = %for.inc6
   ret i32 10
 }
 
-; When the loop trip count is a constant (e.g. 20) and the step size is
-; 1, InstCombine causes the transformation icmp ult i32 %inc, 20 ->
-; icmp ult i32 %j, 19. In this case a valid trip count is not found so
-; the loop is not flattened. 
-define i32 @test9(i32* nocapture %A) {
+; test_10, test_11 and test_12 are for the case when the
+; inner trip count is a constant, then the InstCombine pass
+; makes the transformation icmp ult i32 %inc, tripcount ->
+; icmp ult i32 %j, tripcount-step.
+
+; test_10: The step is not 1.
+define i32 @test_10(i32* nocapture %A) {
 entry:
   br label %for.cond1.preheader
 
 for.cond1.preheader:
-  %i.017 = phi i32 [ 0, %entry ], [ %inc6, %for.cond.cleanup3 ]
+  %i.017 = phi i32 [ 0, %entry ], [ %inc, %for.cond.cleanup3 ]
   %mul = mul i32 %i.017, 20
   br label %for.body4
 
+for.body4:
+  %j.016 = phi i32 [ 0, %for.cond1.preheader ], [ %add5, %for.body4 ]
+  %add = add i32 %j.016, %mul
+  %arrayidx = getelementptr inbounds i32, i32* %A, i32 %add
+  store i32 30, i32* %arrayidx, align 4
+  %add5 = add nuw nsw i32 %j.016, 2
+  %cmp2 = icmp ult i32 %j.016, 18
+  br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
+
 for.cond.cleanup3:
-  %inc6 = add i32 %i.017, 1
-  %cmp = icmp ult i32 %inc6, 11
+  %inc = add i32 %i.017, 1
+  %cmp = icmp ult i32 %inc, 11
   br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup
 
+for.cond.cleanup:
+  %0 = load i32, i32* %A, align 4
+  ret i32 %0
+}
+
+; test_11: The inner inducation variable is used in a compare which
+; isn't the condition of the inner branch.
+define i32 @test_11(i32* nocapture %A) {
+entry:
+  br label %for.cond1.preheader
+
+for.cond1.preheader:
+  %i.020 = phi i32 [ 0, %entry ], [ %inc7, %for.cond.cleanup3 ]
+  %mul = mul i32 %i.020, 20
+  br label %for.body4
+
 for.body4:
-  %j.016 = phi i32 [ 0, %for.cond1.preheader ], [ %inc, %for.body4 ]
-  %add = add i32 %j.016, %mul
+  %j.019 = phi i32 [ 0, %for.cond1.preheader ], [ %inc, %for.body4 ]
+  %cmp5 = icmp ult i32 %j.019, 5
+  %cond = select i1 %cmp5, i32 30, i32 15
+  %add = add i32 %j.019, %mul
   %arrayidx = getelementptr inbounds i32, i32* %A, i32 %add
-  store i32 30, i32* %arrayidx, align 4
-  %inc = add nuw nsw i32 %j.016, 1
-  %cmp2 = icmp ult i32 %j.016, 19
+  store i32 %cond, i32* %arrayidx, align 4
+  %inc = add nuw nsw i32 %j.019, 1
+  %cmp2 = icmp ult i32 %j.019, 19
   br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
 
+for.cond.cleanup3:
+  %inc7 = add i32 %i.020, 1
+  %cmp = icmp ult i32 %inc7, 11
+  br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup
+
 for.cond.cleanup:
   %0 = load i32, i32* %A, align 4
   ret i32 %0
 }
 
+; test_12: Incoming phi node value for preheader is a variable
+define i32 @test_12(i32* %A) {
+entry:
+  br label %while.cond1.preheader
+
+while.cond1.preheader:
+  %j.017 = phi i32 [ 0, %entry ], [ %j.1, %while.end ]
+  %i.016 = phi i32 [ 0, %entry ], [ %inc4, %while.end ]
+  %mul = mul i32 %i.016, 20
+  %cmp214 = icmp ult i32 %j.017, 20
+  br i1 %cmp214, label %while.body3.preheader, label %while.end
+
+while.body3.preheader:
+  br label %while.body3
+
+while.body3:
+  %j.115 = phi i32 [ %inc, %while.body3 ], [ %j.017, %while.body3.preheader ]
+  %add = add i32 %j.115, %mul
+  %arrayidx = getelementptr inbounds i32, i32* %A, i32 %add
+  store i32 30, i32* %arrayidx, align 4
+  %inc = add nuw nsw i32 %j.115, 1
+  %cmp2 = icmp ult i32 %j.115, 19
+  br i1 %cmp2, label %while.body3, label %while.end.loopexit
+
+while.end.loopexit:
+  %inc.lcssa = phi i32 [ %inc, %while.body3 ]
+  br label %while.end
+  
+while.end:
+  %j.1 = phi i32 [ %j.017, %while.cond1.preheader], [ %inc.lcssa, %while.end.loopexit ]
+  %inc4 = add i32 %i.016, 1
+  %cmp = icmp ult i32 %inc4, 11
+  br i1 %cmp, label %while.cond1.preheader, label %while.end5
+
+while.end5:
+  %0 = load i32, i32* %A, align 4
+  ret i32 %0
+}
+
 ; Outer loop conditional phi
 define i32 @e() {
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.end16
   %f.033 = phi i32 [ 0, %entry ], [ %inc18, %for.end16 ]
   %g.032 = phi i32 [ undef, %entry ], [ %g.3.lcssa, %for.end16 ]

llvm/test/Transforms/LoopFlatten/loop-flatten.ll

 ; CHECK:   %inc7 = add nuw nsw i32 %i.018, 1
 ; CHECK:   %exitcond19 = icmp eq i32 %inc7, %flatten.tripcount
 ; CHECK:   br i1 %exitcond19, label %for.end8, label %for.body
 
 for.end8:                                         ; preds = %for.inc6
   ret i32 10
 }
 
+; When the inner loop trip count is a constant and the step
+; is 1, the InstCombine pass causes the transformation e.g.
+; icmp ult i32 %inc, 20 -> icmp ult i32 %j, 19. This doesn't
+; match the pattern (OuterPHI * InnerTripCount) + InnerPHI but
+; we should still flatten the loop as the compare is removed
+; later anyway.
+define i32 @test9(i32* nocapture %A) {
+entry:
+  br label %for.cond1.preheader
+; CHECK-LABEL: test9
+; CHECK: entry:
+; CHECK: %flatten.tripcount = mul i32 20, 11
+; CHECK: br label %for.cond1.preheader
+
+for.cond1.preheader:
+  %i.017 = phi i32 [ 0, %entry ], [ %inc6, %for.cond.cleanup3 ]
+  %mul = mul i32 %i.017, 20
+  br label %for.body4
+; CHECK: for.cond1.preheader:
+; CHECK:   %i.017 = phi i32 [ 0, %entry ], [ %inc6, %for.cond.cleanup3 ]
+; CHECK:   %mul = mul i32 %i.017, 20
+; CHECK:   br label %for.body4
+
+for.cond.cleanup3:
+  %inc6 = add i32 %i.017, 1
+  %cmp = icmp ult i32 %inc6, 11
+  br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup
+; CHECK: for.cond.cleanup3:
+; CHECK:   %inc6 = add i32 %i.017, 1
+; CHECK:   %cmp = icmp ult i32 %inc6, %flatten.tripcount
+; CHECK:   br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup
+
+for.body4:
+  %j.016 = phi i32 [ 0, %for.cond1.preheader ], [ %inc, %for.body4 ]
+  %add = add i32 %j.016, %mul
+  %arrayidx = getelementptr inbounds i32, i32* %A, i32 %add
+  store i32 30, i32* %arrayidx, align 4
+  %inc = add nuw nsw i32 %j.016, 1
+  %cmp2 = icmp ult i32 %j.016, 19
+  br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
+; CHECK: for.body4
+; CHECK:   %j.016 = phi i32 [ 0, %for.cond1.preheader ]
+; CHECK:   %add = add i32 %j.016, %mul
+; CHECK:   %arrayidx = getelementptr inbounds i32, i32* %A, i32 %i.017
+; CHECK:   store i32 30, i32* %arrayidx, align 4
+; CHECK:   %inc = add nuw nsw i32 %j.016, 1
+; CHECK:   %cmp2 = icmp ult i32 %j.016, 19
+; CHECK:   br label %for.cond.cleanup3
+
+for.cond.cleanup:
+  %0 = load i32, i32* %A, align 4
+  ret i32 %0
+}
 
 declare i32 @func(i32)