[SimpleLoopUnswitch] Forget loop if its one of user was replaced

[StephenFan]

Summary

Replacing user in loop may change the scalar evolution in this loop. So
we need to forget this loop.

Fixes: https://github.com/llvm/llvm-project/issues/58868

[fhahn]

I think this should already be fixed slightly differently with 7854a1abfd9a

[StephenFan]

In D137910#3923641, @fhahn wrote:

I think this should already be fixed slightly differently with 7854a1abfd9a

Thanks!

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 /// update the `AssumptionCache` provided in `AC`.
 static BasicBlock *buildClonedLoopBlocks(
     Loop &L, BasicBlock *LoopPH, BasicBlock *SplitBB,
     ArrayRef<BasicBlock *> ExitBlocks, BasicBlock *ParentBB,
     BasicBlock *UnswitchedSuccBB, BasicBlock *ContinueSuccBB,
     const SmallDenseMap<BasicBlock *, BasicBlock *, 16> &DominatingSucc,
     ValueToValueMapTy &VMap,
     SmallVectorImpl<DominatorTree::UpdateType> &DTUpdates, AssumptionCache &AC,
-    DominatorTree &DT, LoopInfo &LI, MemorySSAUpdater *MSSAU) {
+    DominatorTree &DT, LoopInfo &LI, MemorySSAUpdater *MSSAU, ScalarEvolution *SE) {
   SmallVector<BasicBlock *, 4> NewBlocks;
   NewBlocks.reserve(L.getNumBlocks() + ExitBlocks.size());
 
   // We will need to clone a bunch of blocks, wrap up the clone operation in
   // a helper.
   auto CloneBlock = [&](BasicBlock *OldBB) {
     // Clone the basic block and insert it before the new preheader.
     BasicBlock *NewBB = CloneBasicBlock(OldBB, VMap, ".us", OldBB->getParent());
for (auto ZippedInsts : llvm::zip_first(
           (isa<PHINode>(I) || isa<LandingPadInst>(I) || isa<CatchPadInst>(I)) &&
           "Bad instruction in exit block!");
       // We should have a value map between the instruction and its clone.
       assert(VMap.lookup(&I) == &ClonedI && "Mismatch in the value map!");
 
       auto *MergePN =
           PHINode::Create(I.getType(), /*NumReservedValues*/ 2, ".us-phi",
                           &*MergeBB->getFirstInsertionPt());
+
+      // Replacing user in loop may change the scalar evolution in this loop.
+      for (auto *User : I.users())
+        if (auto *Inst = dyn_cast<Instruction>(User))
+          if (Loop *L = LI.getLoopFor(Inst->getParent()))
+            SE->forgetLoop(L);
+
       I.replaceAllUsesWith(MergePN);
       MergePN->addIncoming(&I, ExitBB);
       MergePN->addIncoming(&ClonedI, ClonedExitBB);
     }
   }
 
   // Rewrite the instructions in the cloned blocks to refer to the instructions
   // in the cloned blocks. We have to do this as a second pass so that we have
static void unswitchNontrivialInvariants(
   // Clone the loop for each unswitched successor.
   SmallVector<std::unique_ptr<ValueToValueMapTy>, 4> VMaps;
   VMaps.reserve(UnswitchedSuccBBs.size());
   SmallDenseMap<BasicBlock *, BasicBlock *, 4> ClonedPHs;
   for (auto *SuccBB : UnswitchedSuccBBs) {
     VMaps.emplace_back(new ValueToValueMapTy());
     ClonedPHs[SuccBB] = buildClonedLoopBlocks(
         L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB, RetainedSuccBB,
-        DominatingSucc, *VMaps.back(), DTUpdates, AC, DT, LI, MSSAU);
+        DominatingSucc, *VMaps.back(), DTUpdates, AC, DT, LI, MSSAU, SE);
   }
 
   // Drop metadata if we may break its semantics by moving this instr into the
   // split block.
   if (TI.getMetadata(LLVMContext::MD_make_implicit)) {
     if (DropNonTrivialImplicitNullChecks)
       // Do not spend time trying to understand if we can keep it, just drop it
       // to save compile time.

llvm/test/Transforms/SimpleLoopUnswitch/pr58868.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='loop-mssa(loop-deletion,simple-loop-unswitch<nontrivial;trivial>),loop(indvars)' < %s -S | FileCheck %s
+
+define void @f(i1 %0) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[TMP0:%.*]], label [[ENTRY_SPLIT_US:%.*]], label [[ENTRY_SPLIT:%.*]]
+; CHECK:       entry.split.us:
+; CHECK-NEXT:    br label [[LBL_US:%.*]]
+; CHECK:       lbl.us:
+; CHECK-NEXT:    br label [[CONT_US:%.*]]
+; CHECK:       cont.us:
+; CHECK-NEXT:    [[F_B_PROMOTED_US:%.*]] = load i16, i16* null, align 1
+; CHECK-NEXT:    br i1 true, label [[LBL_LOOPEXIT_US:%.*]], label [[CONT_SPLIT_SPLIT_US:%.*]]
+; CHECK:       lbl.loopexit.us:
+; CHECK-NEXT:    br label [[LBL_US]]
+; CHECK:       cont.split.split.us:
+; CHECK-NEXT:    [[F_B_PROMOTED_LCSSA_US:%.*]] = phi i16 [ [[F_B_PROMOTED_US]], [[CONT_US]] ]
+; CHECK-NEXT:    br label [[CONT_SPLIT:%.*]]
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label [[LBL:%.*]]
+; CHECK:       lbl.loopexit:
+; CHECK-NEXT:    unreachable
+; CHECK:       lbl:
+; CHECK-NEXT:    br label [[HANDLER_TYPE_MISMATCH:%.*]]
+; CHECK:       handler.type_mismatch:
+; CHECK-NEXT:    br label [[CONT:%.*]]
+; CHECK:       cont:
+; CHECK-NEXT:    [[F_B_PROMOTED:%.*]] = load i16, i16* null, align 1
+; CHECK-NEXT:    br i1 false, label [[LBL_LOOPEXIT:%.*]], label [[CONT_SPLIT_SPLIT:%.*]]
+; CHECK:       cont.split.split:
+; CHECK-NEXT:    [[F_B_PROMOTED_LCSSA:%.*]] = phi i16 [ [[F_B_PROMOTED]], [[CONT]] ]
+; CHECK-NEXT:    br label [[CONT_SPLIT]]
+; CHECK:       cont.split:
+; CHECK-NEXT:    [[DOTUS_PHI:%.*]] = phi i16 [ [[F_B_PROMOTED_LCSSA]], [[CONT_SPLIT_SPLIT]] ], [ [[F_B_PROMOTED_LCSSA_US]], [[CONT_SPLIT_SPLIT_US]] ]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INC4:%.*]] = phi i16 [ [[DOTUS_PHI]], [[CONT_SPLIT]] ], [ [[INC:%.*]], [[FOR_INC:%.*]] ]
+; CHECK-NEXT:    br label [[FOR_INC]]
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[INC]] = add nsw i16 [[INC4]], 1
+; CHECK-NEXT:    store i16 0, i16* null, align 1
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ne i16 [[INC4]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[FOR_BODY]], label [[CRIT_EDGE3:%.*]]
+; CHECK:       crit_edge3:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %lbl
+
+lbl:                                              ; preds = %for.body, %entry
+  br i1 %0, label %cont, label %handler.type_mismatch
+
+handler.type_mismatch:                            ; preds = %lbl
+  br label %cont
+
+cont:                                             ; preds = %handler.type_mismatch, %lbl
+  %f.b.promoted = load i16, i16* null, align 1
+  br label %for.body
+
+for.body:                                         ; preds = %for.inc, %cont
+  %inc4 = phi i16 [ %f.b.promoted, %cont ], [ %inc, %for.inc ]
+  br i1 %0, label %lbl, label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i16 %inc4, 1
+  store i16 0, i16* null, align 1
+  %tobool = icmp ne i16 %inc4, 0
+  br i1 %tobool, label %for.body, label %crit_edge3
+
+crit_edge3:                                       ; preds = %for.inc
+  ret void
+}