[JumpThreading] Invalidate LVI after `combineMetadataForCSE`.

(cherry picked from commit 7ded71b1e43fff0be3acb74038bfea87f38d5cfa)

4 files changed, 25 insertions, 5 deletions

diff --git a/llvm/include/llvm/Analysis/LazyValueInfo.h b/llvm/include/llvm/Analysis/LazyValueInfo.h
index b109b7f7e65a..7b2bfdac75a8 100644
--- a/llvm/include/llvm/Analysis/LazyValueInfo.h
+++ b/llvm/include/llvm/Analysis/LazyValueInfo.h
@@ -115,6 +115,9 @@ public:
   /// PredBB to OldSucc to be from PredBB to NewSucc instead.
   void threadEdge(BasicBlock *PredBB, BasicBlock *OldSucc, BasicBlock *NewSucc);
 
+  /// Remove information related to this value from the cache.
+  void forgetValue(Value *V);
+
   /// Inform the analysis cache that we have erased a block.
   void eraseBlock(BasicBlock *BB);
 
diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index 33651783cb17..2ba6036056d9 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -465,6 +465,10 @@ public:
     F.print(OS, &Writer);
   }
 
+  /// This is part of the update interface to remove information related to this
+  /// value from the cache.
+  void forgetValue(Value *V) { TheCache.eraseValue(V); }
+
   /// This is part of the update interface to inform the cache
   /// that a block has been deleted.
   void eraseBlock(BasicBlock *BB) {
@@ -1969,6 +1973,11 @@ void LazyValueInfo::threadEdge(BasicBlock *PredBB, BasicBlock *OldSucc,
   }
 }
 
+void LazyValueInfo::forgetValue(Value *V) {
+  if (PImpl)
+    getImpl(PImpl, AC, nullptr).forgetValue(V);
+}
+
 void LazyValueInfo::eraseBlock(BasicBlock *BB) {
   if (PImpl) {
     getImpl(PImpl, AC, BB->getModule()).eraseBlock(BB);
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 24390f1b54f6..5b8f1b00dc03 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -1269,6 +1269,7 @@ bool JumpThreadingPass::simplifyPartiallyRedundantLoad(LoadInst *LoadI) {
     if (IsLoadCSE) {
       LoadInst *NLoadI = cast<LoadInst>(AvailableVal);
       combineMetadataForCSE(NLoadI, LoadI, false);
+      LVI->forgetValue(NLoadI);
     };
 
     // If the returned value is the load itself, replace with poison. This can
@@ -1461,6 +1462,7 @@ bool JumpThreadingPass::simplifyPartiallyRedundantLoad(LoadInst *LoadI) {
 
   for (LoadInst *PredLoadI : CSELoads) {
     combineMetadataForCSE(PredLoadI, LoadI, true);
+    LVI->forgetValue(PredLoadI);
   }
 
   LoadI->replaceAllUsesWith(PN);
diff --git a/llvm/test/Transforms/JumpThreading/invalidate-lvi.ll b/llvm/test/Transforms/JumpThreading/invalidate-lvi.ll
index 9c5cbfac62d9..27191d6f54c2 100644
--- a/llvm/test/Transforms/JumpThreading/invalidate-lvi.ll
+++ b/llvm/test/Transforms/JumpThreading/invalidate-lvi.ll
@@ -12,15 +12,21 @@ define void @foo(i1 %0) {
 ; CHECK-NEXT:    [[V:%.*]] = alloca i64, align 8
 ; CHECK-NEXT:    call void @set_value(ptr [[V]])
 ; CHECK-NEXT:    [[L1:%.*]] = load i64, ptr [[V]], align 8
-; CHECK-NEXT:    br i1 [[TMP0]], label [[BB0:%.*]], label [[BB4:%.*]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[BB0:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb0:
 ; CHECK-NEXT:    [[C1:%.*]] = icmp eq i64 [[L1]], 0
-; CHECK-NEXT:    br i1 [[C1]], label [[BB1:%.*]], label [[BB4]]
-; CHECK:       bb1:
+; CHECK-NEXT:    br i1 [[C1]], label [[BB2_THREAD:%.*]], label [[BB2]]
+; CHECK:       bb2.thread:
 ; CHECK-NEXT:    store i64 0, ptr [[V]], align 8
-; CHECK-NEXT:    br label [[BB4]]
+; CHECK-NEXT:    br label [[BB4:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[L2:%.*]] = phi i64 [ [[L1]], [[BB0]] ], [ [[L1]], [[START:%.*]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[L2]], 2
+; CHECK-NEXT:    br i1 [[TMP1]], label [[BB3:%.*]], label [[BB4]]
+; CHECK:       bb3:
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    ret void
 ; CHECK:       bb4:
-; CHECK-NEXT:    [[L2:%.*]] = phi i64 [ 0, [[BB1]] ], [ [[L1]], [[BB0]] ], [ [[L1]], [[START:%.*]] ]
 ; CHECK-NEXT:    ret void
 ;
 start: