[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

170

Make this an else if?

179

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?

182

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

397

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

176

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

176

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

179

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.

182

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

176

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

176

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

157–158

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)) {
..
}

162

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

static bool findLoopComponents(
   ICmpInst *Compare = dyn_cast<ICmpInst>(BackBranch->getCondition());
   if (!Compare || !IsValidPredicate(Compare->getUnsignedPredicate()) ||
       Compare->hasNUsesOrMore(2)) {
     LLVM_DEBUG(dbgs() << "Could not find valid comparison\n");
     return false;
   }
   IterationInstructions.insert(Compare);
   LLVM_DEBUG(dbgs() << "Found comparison: "; Compare->dump());
 
   // Find increment and limit from the compare
   Increment = nullptr;
+  Value *LatchValue = InductionPHI->getIncomingValueForBlock(Latch);
   if (match(Compare->getOperand(0),
             m_c_Add(m_Specific(InductionPHI), m_ConstantInt<1>()))) {
     Increment = dyn_cast<BinaryOperator>(Compare->getOperand(0));
     Limit = Compare->getOperand(1);
   } else if (Compare->getUnsignedPredicate() == CmpInst::ICMP_NE &&
              match(Compare->getOperand(1),
                    m_c_Add(m_Specific(InductionPHI), m_ConstantInt<1>()))) {
     Increment = dyn_cast<BinaryOperator>(Compare->getOperand(1));
     Limit = Compare->getOperand(0);
+  } else if (match(LatchValue,
+                   m_c_Add(m_Specific(InductionPHI), m_ConstantInt<1>()))) {
+    // The compare may have been altered by another transformation (e.g icmp ult
+    // %inc, limit -> icmp ult %j, limit-1), where limit is a constant. In this
+    // case the increment is obtained from the InductionPHI and the limit is the
+    // RHS of the compare + 1.
+    assert(Compare->getOperand(0) == InductionPHI);
+    Increment = dyn_cast<BinaryOperator>(LatchValue);
+    ConstantInt *RHS = dyn_cast<ConstantInt>(Compare->getOperand(1));
+    ConstantInt *One = ConstantInt::get(RHS->getType(), 1, true);
+    Limit = ConstantInt::get(Compare->getContext(),
+                             RHS->getValue() + One->getValue());
   }
   if (!Increment || Increment->hasNUsesOrMore(3)) {
     LLVM_DEBUG(dbgs() << "Cound not find valid increment\n");
     return false;
   }
   IterationInstructions.insert(Increment);
   LLVM_DEBUG(dbgs() << "Found increment: "; Increment->dump());
   LLVM_DEBUG(dbgs() << "Found limit: "; Limit->dump());
 
   assert(InductionPHI->getNumIncomingValues() == 2);
 
   if (InductionPHI->getIncomingValueForBlock(Latch) != Increment) {
     LLVM_DEBUG(
         dbgs() << "Incoming value from latch is not the increment inst\n");
     return false;
   }
 
   auto *CI = dyn_cast<ConstantInt>(
       InductionPHI->getIncomingValueForBlock(L->getLoopPreheader()));
-  if (!CI || !CI->isZero()) {
+  if (!CI) {
+    LLVM_DEBUG(dbgs() << "PHI value is not zero \n");
+    return false;
+  } else if (!CI->isZero()) {
     LLVM_DEBUG(dbgs() << "PHI value is not zero: "; CI->dump());
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "Successfully found all loop components\n");
   return true;
 }
 
for (User *U : FI.InnerInductionPHI->users()) {
 
     // 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 == InnerLimit) {
+    // 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), where
+    // limit is a constant. Ignore this use as the compare gets removed later
+    // anyway.
+    if (U == FI.InnerBranch->getCondition())
+      continue;
+    else if ((IsAdd || IsAddTrunc) && MatchedItCount == InnerLimit) {
       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
 }
 
+; test_10, test_11 and test_12 are for the case when the inner limit is a
+; constant (e.g. 20), then InstCombine makes the transformation:
+; icmp ult i32 %inc, 20 -> icmp ult i32 %j, 20-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 ], [ %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:
+  %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.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 %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 ]

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 limit is a defined integer (e.g. 20) and the step
+; is 1, InstCombine causes the transformation:
+; icmp ult i32 %inc, 20 -> icmp ult i32 %j, 19.
+; This is an 'unoptimizable' use of the inner induction variable %j but
+; we should still flatten the loop as this compare instruction 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)