[SCEV] Compute trip multiple for multiple exit loops

[reames]

Summary

This patch implements getSmallConstantTripMultiple(L) correctly for multiple exit loops. The previous implementation was both imprecise, and violated the specified behavior of the method. This was fine in practice, because it turns out the function was both dead in real code, and not tested for the multiple exit case.

[nikic]

LGTM

llvm/lib/Analysis/ScalarEvolution.cpp

6985

Could be written as return Res.getValueOr(1). Feel free to keep it keep as-is though.

llvm/include/llvm/Analysis/ScalarEvolution.h

public:
   /// always a multiple of the returned value. Returns 1 if the trip count is
   /// unknown or not guaranteed to be the multiple of a constant., Will also
   /// return 1 if the trip count is very large (>= 2^32).
   /// Note that the argument is an exit count for loop L, NOT a trip count.
   unsigned getSmallConstantTripMultiple(const Loop *L,
                                         const SCEV *ExitCount);
 
   /// Returns the largest constant divisor of the trip count of the
-  /// loop if it is a single-exit loop and we can compute a small maximum for
-  /// that loop.
-  ///
-  /// Implemented in terms of the \c getSmallConstantTripMultiple overload with
-  /// the single exiting block passed to it. See that routine for details.
+  /// loop.  Will return 1 if no trip count could be computed, or if a
+  /// divisor could not be found.
   unsigned getSmallConstantTripMultiple(const Loop *L);
 
   /// Returns the largest constant divisor of the trip count of this loop as a
   /// normal unsigned value, if possible. This means that the actual trip
   /// count is always a multiple of the returned value (don't forget the trip
   /// count could very well be zero as well!). As explained in the comments
   /// for getSmallConstantTripCount, this assumes that control exits the loop
   /// via ExitingBlock.

llvm/lib/Analysis/ScalarEvolution.cpp

 
 unsigned ScalarEvolution::getSmallConstantMaxTripCount(const Loop *L) {
   const auto *MaxExitCount =
       dyn_cast<SCEVConstant>(getConstantMaxBackedgeTakenCount(L));
   return getConstantTripCount(MaxExitCount);
 }
 
 unsigned ScalarEvolution::getSmallConstantTripMultiple(const Loop *L) {
-  if (BasicBlock *ExitingBB = L->getExitingBlock())
-    return getSmallConstantTripMultiple(L, ExitingBB);
+  SmallVector<BasicBlock *, 8> ExitingBlocks;
+  L->getExitingBlocks(ExitingBlocks);
 
-  // No trip multiple information for multiple exits.
-  return 0;
+  Optional<unsigned> Res = None;
+  for (auto *ExitingBB : ExitingBlocks) {
+    unsigned Multiple = getSmallConstantTripMultiple(L, ExitingBB);
+    if (!Res)
+      Res = Multiple;
+    Res = (unsigned)GreatestCommonDivisor64(*Res, Multiple);
+  }
+  return Res.getValueOr(1);
 }
 
 unsigned ScalarEvolution::getSmallConstantTripMultiple(const Loop *L,
                                                        const SCEV *ExitCount) {
   if (ExitCount == getCouldNotCompute())
     return 1;
 
   // Get the trip count
   const SCEV *TCExpr = getTripCountFromExitCount(ExitCount);
 

llvm/test/Analysis/ScalarEvolution/tripmultiple_calculation.ll

l3:
   %inc = add i32 %x.0, 1
   %exitcond = icmp eq i32 %inc, 3
   br i1 %exitcond, label %exit, label %l3
 
 exit:
   ret i32 0
 }
 
+; If there are multiple exits, the result is the GCD of the multiples
+; of each individual exit (since we don't know which is taken).
+
+; CHECK: Loop %l4: Trip multiple is 50
+
+define i32 @foo5(i32 %n) {
+entry:
+  br label %l4
+
+l4:
+  %x.0 = phi i32 [ 0, %entry ], [ %inc, %l4-latch ]
+  call void @f()
+  %inc = add i32 %x.0, 1
+  %earlycond = icmp eq i32 %inc, 150
+  br i1 %earlycond, label %exit, label %l4-latch
+
+l4-latch:
+  %exitcond = icmp eq i32 %inc, 200
+  br i1 %exitcond, label %exit, label %l4
+
+exit:
+  ret i32 0
+}
+