[IndVarSimplify] Notify top most loop to drop cached exit counts

[guopeilin]

Summary

Some nested loops may share the same ExitingBB, so after we finishing FoldExit,
we need to notify OuterLoop and SCEV to drop any stored trip count.

[guopeilin]

https://bugs.llvm.org/show_bug.cgi?id=48158

[mkazantsev]

The strategy we use across the optimizer is to drop cached values in the transform that invalidated them, not before their users (as the users are multiple). Do we know who broke it in this case?

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1447

This whole logic is completely expensive as it visits the exit blocks multiple times. Why not just call forgetTopmostLoop(L)?

[mkazantsev]

I'm not really getting how wrong cached exit count may possibly lead to hang. Could you please explain this problem more?

[guopeilin]

In D91325#2392962, @mkazantsev wrote:

I'm not really getting how wrong cached exit count may possibly lead to hang. Could you please explain this problem more?

From the cfg we can see that there exist a nested loop, and basic block if.then is a shared exitingBB between inner and outer loop.
When we start running IndVarSimplify pass on the inner loop, we can get block if.then's ExitCount by function ScalarEvolution::getExitCount(), also we can get inner loop's MaxExitCount by function getMaxBackedgeTakenCount.
If a exitingBB's exit count less than loop's max backedge taken count, we can say that this loop never exit from this exitingBB. So we can FoldExit, that is we replace the condition of exitingBB's BranchInst with a ConstantInt. So after we done IndVarSimplify pass on the inner loop, the if.the will be simplified into a uncondtional block.
However, when we begin running IndVarSimplify on outer loop, if.then also is outer loop's exitingBB, and we don't know it has been simplified. So we will get the cached exit count (which is totally wrong), this count is used for compute MaxBackedgeTakenCount of outer loop.
Unfortunately this MaxBackedgeTakenCount is less than block for.inc15's exit count. So IndVarSimplify think outer loop never exit in for.inc15, we can then fold this exitingBB, finally make outer loop endless.

Before IndVarSimplify

After IndVarSimplify

[mkazantsev]

Thanks for explaining, I didn't realize we end up having wrong exact exit count.

I agree with the motivation, but I think that the fix should be unconditional drop of caches in the end of optimizeLoopExits if Changed == true.

[guopeilin]

use forgetTopmostLoop() directly instead of visiting exitingBB among nested loops

[guopeilin]

Yse, So I use forgetTopmostLoop() after optimizeLoopExits directly, I think it's more concise.

In D91325#2396652, @mkazantsev wrote:

Thanks for explaining, I didn't realize we end up having wrong exact exit count.

I agree with the motivation, but I think that the fix should be unconditional drop of caches in the end of optimizeLoopExits if Changed == true.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1447

The reason why I chose to visit exit blocks is that I think it would be expensive if we drop out all cached value for those nested loops that do not share exiting BB.

[mkazantsev]

LGTM, thanks.

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1781–1783

As a follow-up, could you please also replace this to forgetTopmostLoop? Might be causing same type of problem I guess.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

for (unsigned i = 1; i < ExitingBlocks.size(); i++) {
     assert(DT->dominates(ExitingBlocks[i-1], ExitingBlocks[i]));
   }
 #endif
 
   bool Changed = false;
   bool SkipLastIter = false;
   SmallSet<const SCEV*, 8> DominatingExitCounts;
   for (BasicBlock *ExitingBB : ExitingBlocks) {
     const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
     if (isa<SCEVCouldNotCompute>(ExitCount)) {
       // Okay, we do not know the exit count here. Can we at least prove that it
       // will remain the same within iteration space?
       auto *BI = cast<BranchInst>(ExitingBB->getTerminator());
       auto OptimizeCond = [&](bool Inverted, bool SkipLastIter) {
         return optimizeLoopExitWithUnknownExitCount(
             L, BI, ExitingBB, MaxExitCount, Inverted, SkipLastIter, SE,
             Rewriter, DeadInsts);
#endif
 
   // Eliminate redundant IV cycles.
   NumElimIV += Rewriter.replaceCongruentIVs(L, DT, DeadInsts);
 
   // Try to eliminate loop exits based on analyzeable exit counts
   if (optimizeLoopExits(L, Rewriter))  {
     Changed = true;
     // Given we've changed exit counts, notify SCEV
-    SE->forgetLoop(L);
+    // Some nested loops may share same folded exit basic block,
+    // thus we need to notify top most loop.
+    SE->forgetTopmostLoop(L);
   }
 
   // Try to form loop invariant tests for loop exits by changing how many
   // iterations of the loop run when that is unobservable.
   if (predicateLoopExits(L, Rewriter)) {
     Changed = true;
     // Given we've changed exit counts, notify SCEV
     SE->forgetLoop(L);

llvm/test/Transforms/IndVarSimplify/shared-exit-between-nested-loop.ll

+; RUN: opt -indvars -S < %s | FileCheck %s
+
+@__const.e.f = private unnamed_addr constant <{ <{ [2 x i32], [2 x i32], [8 x [2 x i32]] }>, [10 x [2 x i32]] }> <{ <{ [2 x i32], [2 x i32], [8 x [2 x i32]] }> <{ [2 x i32] [i32 4, i32 8], [2 x i32] [i32 3, i32 10], [8 x [2 x i32]] zeroinitializer }>, [10 x [2 x i32]] [[2 x i32] zeroinitializer, [2 x i32] zeroinitializer, [2 x i32] [i32 1, i32 5], [2 x i32] [i32 2080555007, i32 0], [2 x i32] zeroinitializer, [2 x i32] zeroinitializer, [2 x i32] zeroinitializer, [2 x i32] zeroinitializer, [2 x i32] zeroinitializer, [2 x i32] zeroinitializer] }>, align 4
+
+define dso_local i8 @main() local_unnamed_addr {
+entry:
+  br label %for.cond1.preheader
+
+for.cond1.preheader:                              ; preds = %for.inc15, %entry
+  %storemerge32 = phi i32 [ 0, %entry ], [ %inc16, %for.inc15 ]
+  br label %for.cond4
+
+for.cond4:                                        ; preds = %for.inc, %for.cond1.preheader
+  %storemerge2531 = phi i32 [ 0, %for.cond1.preheader ], [ %inc, %for.inc ]
+  %tobool = icmp eq i32 1, 0
+  br i1 %tobool, label %cleanup, label %if.then
+
+if.then:                                          ; preds = %for.cond4
+; CHECK:        if.then:                                          ; preds = %for.cond4
+; CHECK-NEXT:     br i1 false, label %cleanup, label %for.inc
+  %add = add nuw nsw i32 %storemerge2531, 2
+  %idxprom8 = zext i32 %add to i64
+  %arrayidx10 = getelementptr inbounds <{ <{ [2 x i32], [2 x i32], [8 x [2 x i32]] }>, [10 x [2 x i32]] }>, <{ <{ [2 x i32], [2 x i32], [8 x [2 x i32]] }>, [10 x [2 x i32]] }>* @__const.e.f, i64 0, i32 1, i64 %idxprom8, i64 0
+  %0 = load i32, i32* %arrayidx10, align 4
+  %tobool11 = icmp eq i32 %0, 0
+  br i1 %tobool11, label %cleanup, label %for.inc
+
+for.inc:                                          ; preds = %if.then
+  %inc = add nuw nsw i32 %storemerge2531, 1
+  %cmp2 = icmp ult i32 %inc, 2
+  br i1 %cmp2, label %for.cond4, label %for.inc15
+
+cleanup:                                          ; preds = %if.then, %for.cond4
+  br label %return
+
+for.inc15:                                        ; preds = %for.inc
+; CHECK:        for.inc15:
+; CHECK-NEXT:     %inc16 = add nuw nsw i32 %storemerge32, 1
+; CHECK-NEXT:     %cmp = icmp eq i32 %inc16, 4
+; CHECK-NEXT:     br i1 %cmp, label %for.cond18thread-pre-split, label %for.cond1.preheader
+  %inc16 = add nuw nsw i32 %storemerge32, 1
+  %cmp = icmp eq i32 %inc16, 4
+  br i1 %cmp, label %for.cond18thread-pre-split, label %for.cond1.preheader
+
+for.cond18thread-pre-split:                       ; preds = %for.inc15
+  br i1 1, label %return.loopexit, label %for.inc21.lr.ph
+
+for.inc21.lr.ph:                                  ; preds = %for.cond18thread-pre-split
+  br label %for.inc21
+
+for.inc21:                                        ; preds = %for.inc21, %for.inc21.lr.ph
+  %1 = phi i32 [ undef, %for.inc21.lr.ph ], [ %inc22, %for.inc21 ]
+  %inc22 = add nsw i32 %1, 1
+  br i1 true, label %for.cond18.return.loopexit_crit_edge, label %for.inc21
+
+for.cond18.return.loopexit_crit_edge:             ; preds = %for.inc21
+  br label %return.loopexit
+
+return.loopexit:                                  ; preds = %for.cond18.return.loopexit_crit_edge, %for.cond18thread-pre-split
+  br label %return
+
+return:                                           ; preds = %return.loopexit, %cleanup
+  ret i8 undef
+}