[SimpleLoopUnswitch] Fix SCEV invalidation for unswitchTrivialSwitch

[bjope]

Summary

When doing a trivial unswitch of a switch statement the code need
to "invalidate SCEVs for the outermost loop reached by any of the
exits", as indicated by code comments.

Depending on if we find such an outermost loop or not we can limit
the invalidation to some sub-loops or the full loop-nest. As shown
in the added test case there seem to have been some bugs in the code
that was finding the "outermost loop", so we could end up invalidating
too few loops.

Seems like commit 1bf8ae17f5e2714c8c87978 introduce the bug by
moving the code that invalidates the loops above some of the code
that computed 'OuterL'. This patch fixes that by also moving that
computation of 'OuterL' so that we compute 'OuterL' properly before
we use it for the SCEV invalidation.

[fhahn]

LGTM, thanks for tracking down a test case for this!

llvm/test/Transforms/SimpleLoopUnswitch/update-scev-2.ll

8

I think it is probably better to avoid including a stack trace here, as it may go stable. Just staying it will fail verification would be sufficient.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   Loop *OuterL = &L;
 
   if (DefaultExitBB) {
     // Check the loop containing this exit.
     Loop *ExitL = LI.getLoopFor(DefaultExitBB);
     if (!ExitL || ExitL->contains(OuterL))
       OuterL = ExitL;
   }
+  for (unsigned Index : ExitCaseIndices) {
+    auto CaseI = SI.case_begin() + Index;
+    // Compute the outer loop from this exit.
+    Loop *ExitL = LI.getLoopFor(CaseI->getCaseSuccessor());
+    if (!ExitL || ExitL->contains(OuterL))
+      OuterL = ExitL;
+  }
 
   if (SE) {
     if (OuterL)
       SE->forgetLoop(OuterL);
     else
       SE->forgetTopmostLoop(&L);
   }
 
SmallVector<std::tuple<ConstantInt *, BasicBlock *,
                          SwitchInstProfUpdateWrapper::CaseWeightOpt>,
               4> ExitCases;
   ExitCases.reserve(ExitCaseIndices.size());
   SwitchInstProfUpdateWrapper SIW(SI);
   // We walk the case indices backwards so that we remove the last case first
   // and don't disrupt the earlier indices.
   for (unsigned Index : reverse(ExitCaseIndices)) {
     auto CaseI = SI.case_begin() + Index;
-    // Compute the outer loop from this exit.
-    Loop *ExitL = LI.getLoopFor(CaseI->getCaseSuccessor());
-    if (!ExitL || ExitL->contains(OuterL))
-      OuterL = ExitL;
     // Save the value of this case.
     auto W = SIW.getSuccessorWeight(CaseI->getSuccessorIndex());
     ExitCases.emplace_back(CaseI->getCaseValue(), CaseI->getCaseSuccessor(), W);
     // Delete the unswitched cases.
     SIW.removeCase(CaseI);
   }
 
   // Check if after this all of the remaining cases point at the same

llvm/test/Transforms/SimpleLoopUnswitch/update-scev-2.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt < %s -passes='print<scalar-evolution>,simple-loop-unswitch<no-nontrivial>,print<scalar-evolution>' -disable-output 2>&1 | FileCheck -check-prefix CHECK-SCEV %s
+; RUN: opt < %s -passes='function(loop-deletion,simple-loop-unswitch<no-nontrivial>),verify<scalar-evolution>' -S | FileCheck -check-prefix CHECK-IR %s
+
+; These tests used to hit an assertion failure in llvm::ScalarEvolution::BackedgeTakenInfo::getExact:
+;    "We should only have known counts for exiting blocks that dominate latch!"' failed.
+;
+; Verify that we no longer hit that assert, while we expect to have forgotten SCEV for the outer loop.
+; Also added checks to show and verify the IR transformation.
+
+
+; Before the unswitch:
+;
+; CHECK-SCEV: Classifying expressions for: @f4
+; CHECK-SCEV-NEXT:   %j.0 = phi i16 [ 0, %entry ], [ %0, %sw.bb2 ]
+; CHECK-SCEV-NEXT:   -->  {0,+,1}<nuw><nsw><%lbl1> U: [0,3) S: [0,3)               Exits: 2                LoopDispositions: { %lbl1: Computable, %lbl2: Invariant }
+; CHECK-SCEV-NEXT:   %0 = add i16 %j.0, 1
+; CHECK-SCEV-NEXT:   -->  {1,+,1}<nuw><nsw><%lbl1> U: [1,4) S: [1,4)               Exits: 3                LoopDispositions: { %lbl1: Computable, %lbl2: Invariant }
+; CHECK-SCEV-DAG: Loop %lbl2: Unpredictable backedge-taken count.
+; CHECK-SCEV-DAG: Loop %lbl1: backedge-taken count is 2
+;
+; After the unswitch:
+;
+; CHECK-SCEV: Classifying expressions for: @f4
+; CHECK-SCEV-NEXT:  %j.0 = phi i16 [ 0, %entry ], [ %0, %sw.bb2 ]
+; CHECK-SCEV-NEXT:  -->  {0,+,1}<nuw><nsw><%lbl1> U: [0,-32768) S: [0,-32768)             Exits: <<Unknown>>              LoopDispositions: { %lbl1: Computable }
+; CHECK-SCEV-NEXT:  %0 = add i16 %j.0, 1
+; CHECK-SCEV-NEXT:  -->  {1,+,1}<nuw><nsw><%lbl1> U: [1,-32768) S: [1,-32768)             Exits: <<Unknown>>              LoopDispositions: { %lbl1: Computable }
+; CHECK-SCEV-DAG: Loop %lbl1: Unpredictable backedge-taken count.
+; CHECK-SCEV-DAG: Loop %lbl2: <multiple exits> Unpredictable backedge-taken count.
+
+
+define i16 @f4() {
+; CHECK-IR-LABEL: define i16 @f4() {
+; CHECK-IR-NEXT:  entry:
+; CHECK-IR-NEXT:    br label [[LBL1:%.*]]
+; CHECK-IR:       lbl1:
+; CHECK-IR-NEXT:    [[J_0:%.*]] = phi i16 [ 0, [[ENTRY:%.*]] ], [ [[TMP0:%.*]], [[SW_BB2:%.*]] ]
+; CHECK-IR-NEXT:    switch i16 [[J_0]], label [[LBL1_SPLIT:%.*]] [
+; CHECK-IR-NEXT:    i16 0, label [[SW_BB2]]
+; CHECK-IR-NEXT:    i16 1, label [[SW_BB2]]
+; CHECK-IR-NEXT:    i16 2, label [[LBL3:%.*]]
+; CHECK-IR-NEXT:    ]
+; CHECK-IR:       lbl1.split:
+; CHECK-IR-NEXT:    br label [[LBL2:%.*]]
+; CHECK-IR:       lbl2:
+; CHECK-IR-NEXT:    br label [[LBL2]]
+; CHECK-IR:       sw.bb2:
+; CHECK-IR-NEXT:    [[TMP0]] = add i16 [[J_0]], 1
+; CHECK-IR-NEXT:    br label [[LBL1]]
+; CHECK-IR:       lbl3:
+; CHECK-IR-NEXT:    ret i16 0
+;
+entry:
+  br label %lbl1
+
+lbl1:                                             ; preds = %sw.bb2, %entry
+  %j.0 = phi i16 [ 0, %entry ], [ %0, %sw.bb2 ]
+  br label %lbl2
+
+lbl2:                                             ; preds = %lbl2, %lbl1
+  switch i16 %j.0, label %lbl2 [
+  i16 0, label %sw.bb2
+  i16 1, label %sw.bb2
+  i16 2, label %lbl3
+  ]
+
+sw.bb2:                                           ; preds = %lbl2, %lbl2
+  %0 = add i16 %j.0, 1
+  br label %lbl1
+
+lbl3:                                             ; preds = %lbl2
+  ret i16 0
+}
+