[LoopFlatten] Use SCEV and Loop APIs to identify increment and trip count

[RosieSumpter]

Summary

Replace pattern-matching with existing SCEV and Loop APIs as a more
robust way of identifying the loop increment and trip count. Also
rename 'Limit' as 'TripCount' to be consistent with terminology.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

134

I am not sure if an induction phi can have > 2 incoming values. If it can then the assert needs to be a return, and a test case is missing for this.

But if getNumIncomingValues() == 2 is guaranteed by getInductionVariable(), then this is fine. Can you double check by looking at getInductionVariable()?

155–156

Just for clarity, I was wondering if we should rename limit to trip count, because that is what it is at the moment with the restriction that our loops start counting at 0. I think that makes things clearer now that we start using SCEV and request trip counts.

[RosieSumpter]

• Removed assert on the number of incoming phi values
• Rename 'Limit' to 'TripCount' for consistency

[RosieSumpter]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

134

This assert was actually already there, I just moved it up to be where the induction phi is identified. But I agree, since the loop is in simplified form it will only have 2 incoming values so I'll remove this assert.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

74

Perhaps we can document some assumptions and limitations here.

79

This corresponds to a loop induction variable, which we expect to have a start value of 0.

155–156

Perhaps I have a slight preference here for:

f (!L->isCanonical(*SE))
  return false;

so we don't need the Increment = nullptr; here.

155–156

Then, this can just be:

if (Increment->hasNUsesOrMore(3))

155–156

Ah, nice one, isCanonical checks exactly what we were assuming and pattern matching: it checks that the loop starts at 0 and increments by 1. Do we need the isZero() check above?

[RosieSumpter]

• Added comment about induction variable assumptions
• Moved isCanonical check on loop to be higher so we exit earlier if loop is not canonical
• Removed Increment = nullptr
• Removed code which determines trip count in the case where another transformation changes the compare, and return false instead (and added negative test case for this)

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

165

Can this be !=?

165–190

We are not doing this sanity check when the IV is widened.
Can we also check this case and do something like this?

ExtTC = Tripcount->getOperand(0);
if (!isa<ZExt>(ExtTC) || !isa<SExt>(ExtTC) ||
    SE->getSCEV(ExtTC) != SCEVTripCount ) {
  LLVM_DEBUG(dbgs() << "Could not find valid extended trip count\n");
  return false;
}

[RosieSumpter]

Added a check on the trip count when the IV is widened

[RosieSumpter]

Amended the check on the trip count when the IV has been widened

[SjoerdMeijer]

Thanks, nice change that gets rid of quite some pattern matching!

[fhahn]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

109

This is just a limitation to keep the implementation simpler at the moment right? Might be worth clarifying in the comment.

160

Is later code somewhere relying that the increment is an add instruction? What if it is something like an extend?

[RosieSumpter]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

109

Ok sure, there is also a note mentioning this in lines 77/78. I will add some clarification that this is to simplify the implementation before committing.

160

No, we don't rely on it being an Add instruction. Using SCEV, it is now guaranteed that the loop starts at 0 and increments with 1. Most of the pattern matching that we replaced was to ensure that. The increment is also marked as an "IterationInstruction", but that's only for cost-modeling purposes and it doesn't rely on Adds.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

static cl::opt<unsigned> RepeatedInstructionThreshold(
     "loop-flatten-cost-threshold", cl::Hidden, cl::init(2),
     cl::desc("Limit on the cost of instructions that can be repeated due to "
              "loop flattening"));
 
 static cl::opt<bool>
     AssumeNoOverflow("loop-flatten-assume-no-overflow", cl::Hidden,
                      cl::init(false),
                      cl::desc("Assume that the product of the two iteration "
-                              "limits will never overflow"));
+                              "trip counts will never overflow"));
 
 static cl::opt<bool>
     WidenIV("loop-flatten-widen-iv", cl::Hidden,
             cl::init(true),
             cl::desc("Widen the loop induction variables, if possible, so "
                      "overflow checks won't reject flattening"));
 
 struct FlattenInfo {
   Loop *OuterLoop = nullptr;
   Loop *InnerLoop = nullptr;
+  // These PHINodes correspond to loop induction variables, which are expected
+  // to start at zero and increment by one on each loop.
   PHINode *InnerInductionPHI = nullptr;
   PHINode *OuterInductionPHI = nullptr;
-  Value *InnerLimit = nullptr;
-  Value *OuterLimit = nullptr;
+  Value *InnerTripCount = nullptr;
+  Value *OuterTripCount = nullptr;
   BinaryOperator *InnerIncrement = nullptr;
   BinaryOperator *OuterIncrement = nullptr;
   BranchInst *InnerBranch = nullptr;
   BranchInst *OuterBranch = nullptr;
   SmallPtrSet<Value *, 4> LinearIVUses;
   SmallPtrSet<PHINode *, 4> InnerPHIsToTransform;
 
   // Whether this holds the flatten info before or after widening.
   bool Widened = false;
 
   FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
 };
 
-// Finds the induction variable, increment and limit for a simple loop that we
-// can flatten.
+// Finds the induction variable, increment and trip count for a simple loop that
+// we can flatten.
 static bool findLoopComponents(
     Loop *L, SmallPtrSetImpl<Instruction *> &IterationInstructions,
-    PHINode *&InductionPHI, Value *&Limit, BinaryOperator *&Increment,
-    BranchInst *&BackBranch, ScalarEvolution *SE) {
+    PHINode *&InductionPHI, Value *&TripCount, BinaryOperator *&Increment,
+    BranchInst *&BackBranch, ScalarEvolution *SE, bool IsWidened) {
   LLVM_DEBUG(dbgs() << "Finding components of loop: " << L->getName() << "\n");
 
   if (!L->isLoopSimplifyForm()) {
     LLVM_DEBUG(dbgs() << "Loop is not in normal form\n");
     return false;
   }
 
+  // Currently, to simplify the implementation, the Loop induction variable must
+  // start at zero and increment with a step size of one.
+  if (!L->isCanonical(*SE)) {
+    LLVM_DEBUG(dbgs() << "Loop is not canonical\n");
+    return false;
+  }
+
   // There must be exactly one exiting block, and it must be the same at the
   // latch.
   BasicBlock *Latch = L->getLoopLatch();
   if (L->getExitingBlock() != Latch) {
     LLVM_DEBUG(dbgs() << "Exiting and latch block are different\n");
     return false;
   }
 
   // Find the induction PHI. If there is no induction PHI, we can't do the
   // transformation. TODO: could other variables trigger this? Do we have to
   // search for the best one?
   InductionPHI = L->getInductionVariable(*SE);
   if (!InductionPHI) {
     LLVM_DEBUG(dbgs() << "Could not find induction PHI\n");
     return false;
   }
   LLVM_DEBUG(dbgs() << "Found induction PHI: "; InductionPHI->dump());
 
   bool ContinueOnTrue = L->contains(Latch->getTerminator()->getSuccessor(0));
   auto IsValidPredicate = [&](ICmpInst::Predicate Pred) {
     if (ContinueOnTrue)
       return Pred == CmpInst::ICMP_NE || Pred == CmpInst::ICMP_ULT;
     else
       return Pred == CmpInst::ICMP_EQ;
   };
 
   // Find Compare and make sure it is valid. getLatchCmpInst checks that the
   // back branch of the latch is conditional.
   ICmpInst *Compare = L->getLatchCmpInst();
   if (!Compare || !IsValidPredicate(Compare->getUnsignedPredicate()) ||
       Compare->hasNUsesOrMore(2)) {
     LLVM_DEBUG(dbgs() << "Could not find valid comparison\n");
     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 limit from the compare
-  Increment = nullptr;
-  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);
-  }
-  if (!Increment || Increment->hasNUsesOrMore(3)) {
-    LLVM_DEBUG(dbgs() << "Cound not find valid increment\n");
+  // 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).
+  TripCount = Compare->getOperand(1);
+  const SCEV *SCEVTripCount =
+      SE->getTripCountFromExitCount(SE->getBackedgeTakenCount(L));
+  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 (!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 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()) {
-    LLVM_DEBUG(dbgs() << "PHI value is not zero: "; CI->dump());
-    return false;
-  }
+  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 (auto &I : *B) {
       // a fall-through, so adds no cost.
       BranchInst *Br = dyn_cast<BranchInst>(&I);
       if (Br && Br->isUnconditional() &&
           Br->getSuccessor(0) == FI.InnerLoop->getHeader())
         continue;
       // Multiplies of the outer iteration variable and inner iteration
       // count will be optimised out.
       if (match(&I, m_c_Mul(m_Specific(FI.OuterInductionPHI),
-                            m_Specific(FI.InnerLimit))))
+                            m_Specific(FI.InnerTripCount))))
         continue;
       InstructionCost Cost =
           TTI->getUserCost(&I, TargetTransformInfo::TCK_SizeAndLatency);
       LLVM_DEBUG(dbgs() << "Cost " << Cost << ": "; I.dump());
       RepeatedInstrCost += Cost;
     }
   }
 
   LLVM_DEBUG(dbgs() << "Cost of instructions that will be repeated: "
                     << RepeatedInstrCost << "\n");
   // Bail out if flattening the loops would cause instructions in the outer
   // loop but not in the inner loop to be executed extra times.
   if (RepeatedInstrCost > RepeatedInstructionThreshold) {
     LLVM_DEBUG(dbgs() << "checkOuterLoopInsts: not profitable, bailing.\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "checkOuterLoopInsts: OK\n");
   return true;
 }
 
 static bool checkIVUsers(FlattenInfo &FI) {
   // We require all uses of both induction variables to match this pattern:
   //
-  //   (OuterPHI * InnerLimit) + InnerPHI
+  //   (OuterPHI * InnerTripCount) + InnerPHI
   //
   // Any uses of the induction variables not matching that pattern would
   // require a div/mod to reconstruct in the flattened loop, so the
   // transformation wouldn't be profitable.
 
-  Value *InnerLimit = FI.InnerLimit;
+  Value *InnerTripCount = FI.InnerTripCount;
   if (FI.Widened &&
-      (isa<SExtInst>(InnerLimit) || isa<ZExtInst>(InnerLimit)))
-    InnerLimit = cast<Instruction>(InnerLimit)->getOperand(0);
+      (isa<SExtInst>(InnerTripCount) || isa<ZExtInst>(InnerTripCount)))
+    InnerTripCount = cast<Instruction>(InnerTripCount)->getOperand(0);
 
   // Check that all uses of the inner loop's induction variable match the
   // expected pattern, recording the uses of the outer IV.
   SmallPtrSet<Value *, 4> ValidOuterPHIUses;
   for (User *U : FI.InnerInductionPHI->users()) {
     if (U == FI.InnerIncrement)
       continue;
 
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) {
+    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_DEBUG(dbgs() << "checkIVUsers: OK\n";
              for (Value *V : FI.LinearIVUses) {
                dbgs() << "  ";
                V->dump();
              });
   return true;
 }
 
 // Return an OverflowResult dependant on if overflow of the multiplication of
-// InnerLimit and OuterLimit can be assumed not to happen.
+// InnerTripCount and OuterTripCount can be assumed not to happen.
 static OverflowResult checkOverflow(FlattenInfo &FI, DominatorTree *DT,
                                     AssumptionCache *AC) {
   Function *F = FI.OuterLoop->getHeader()->getParent();
   const DataLayout &DL = F->getParent()->getDataLayout();
 
   // For debugging/testing.
   if (AssumeNoOverflow)
     return OverflowResult::NeverOverflows;
 
   // Check if the multiply could not overflow due to known ranges of the
   // input values.
   OverflowResult OR = computeOverflowForUnsignedMul(
-      FI.InnerLimit, FI.OuterLimit, DL, AC,
+      FI.InnerTripCount, FI.OuterTripCount, DL, AC,
       FI.OuterLoop->getLoopPreheader()->getTerminator(), DT);
   if (OR != OverflowResult::MayOverflow)
     return OR;
 
   for (Value *V : FI.LinearIVUses) {
     for (Value *U : V->users()) {
       if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
         // The IV is used as the operand of a GEP, and the IV is at least as
static OverflowResult checkOverflow(FlattenInfo &FI, DominatorTree *DT,
 
   return OverflowResult::MayOverflow;
 }
 
 static bool CanFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                                ScalarEvolution *SE, AssumptionCache *AC,
                                const TargetTransformInfo *TTI) {
   SmallPtrSet<Instruction *, 8> IterationInstructions;
-  if (!findLoopComponents(FI.InnerLoop, IterationInstructions, FI.InnerInductionPHI,
-                          FI.InnerLimit, FI.InnerIncrement, FI.InnerBranch, SE))
+  if (!findLoopComponents(FI.InnerLoop, IterationInstructions,
+                          FI.InnerInductionPHI, FI.InnerTripCount,
+                          FI.InnerIncrement, FI.InnerBranch, SE, FI.Widened))
     return false;
-  if (!findLoopComponents(FI.OuterLoop, IterationInstructions, FI.OuterInductionPHI,
-                          FI.OuterLimit, FI.OuterIncrement, FI.OuterBranch, SE))
+  if (!findLoopComponents(FI.OuterLoop, IterationInstructions,
+                          FI.OuterInductionPHI, FI.OuterTripCount,
+                          FI.OuterIncrement, FI.OuterBranch, SE, FI.Widened))
     return false;
 
-  // Both of the loop limit values must be invariant in the outer loop
+  // Both of the loop trip count values must be invariant in the outer loop
   // (non-instructions are all inherently invariant).
-  if (!FI.OuterLoop->isLoopInvariant(FI.InnerLimit)) {
-    LLVM_DEBUG(dbgs() << "inner loop limit not invariant\n");
+  if (!FI.OuterLoop->isLoopInvariant(FI.InnerTripCount)) {
+    LLVM_DEBUG(dbgs() << "inner loop trip count not invariant\n");
     return false;
   }
-  if (!FI.OuterLoop->isLoopInvariant(FI.OuterLimit)) {
-    LLVM_DEBUG(dbgs() << "outer loop limit not invariant\n");
+  if (!FI.OuterLoop->isLoopInvariant(FI.OuterTripCount)) {
+    LLVM_DEBUG(dbgs() << "outer loop trip count not invariant\n");
     return false;
   }
 
   if (!checkPHIs(FI, TTI))
     return false;
 
   // FIXME: it should be possible to handle different types correctly.
   if (FI.InnerInductionPHI->getType() != FI.OuterInductionPHI->getType())
LLVM_DEBUG(dbgs() << "Checks all passed, doing the transformation\n");
     using namespace ore;
     OptimizationRemark Remark(DEBUG_TYPE, "Flattened", FI.InnerLoop->getStartLoc(),
                               FI.InnerLoop->getHeader());
     OptimizationRemarkEmitter ORE(F);
     Remark << "Flattened into outer loop";
     ORE.emit(Remark);
   }
 
-  Value *NewTripCount =
-      BinaryOperator::CreateMul(FI.InnerLimit, FI.OuterLimit, "flatten.tripcount",
-                                FI.OuterLoop->getLoopPreheader()->getTerminator());
+  Value *NewTripCount = BinaryOperator::CreateMul(
+      FI.InnerTripCount, FI.OuterTripCount, "flatten.tripcount",
   LLVM_DEBUG(dbgs() << "Created new trip count in preheader: ";
              NewTripCount->dump());
 
   // Fix up PHI nodes that take values from the inner loop back-edge, which
   // we are about to remove.
   FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
 
   // The old Phi will be optimised away later, but for now we can't leave
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   auto &DL = M->getDataLayout();
   auto *InnerType = FI.InnerInductionPHI->getType();
   auto *OuterType = FI.OuterInductionPHI->getType();
   unsigned MaxLegalSize = DL.getLargestLegalIntTypeSizeInBits();
   auto *MaxLegalType = DL.getLargestLegalIntType(M->getContext());
 
   // If both induction types are less than the maximum legal integer width,
   // promote both to the widest type available so we know calculating
-  // (OuterLimit * InnerLimit) as the new trip count is safe.
+  // (OuterTripCount * InnerTripCount) as the new trip count is safe.
   if (InnerType != OuterType ||
       InnerType->getScalarSizeInBits() >= MaxLegalSize ||
       MaxLegalType->getScalarSizeInBits() < InnerType->getScalarSizeInBits() * 2) {
     LLVM_DEBUG(dbgs() << "Can't widen the IV\n");
     return false;
   }
 
   SCEVExpander Rewriter(*SE, DL, "loopflatten");

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) {
+entry:
+  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
+
+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
+
+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
+
+for.cond.cleanup:
+  %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 ]