[SimpleLoopUnswitch] Fix SCEV invalidation for unswitchTrivialSwitch

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 1bf8ae1 introduced 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.

Differential Revision: https://reviews.llvm.org/D146963

Author: bjope

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

@@ -796,6 +796,13 @@ static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
     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)
@@ -821,10 +828,6 @@ static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   // 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);

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

@@ -0,0 +1,75 @@
+; 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
+}
+