From 46aac949bcfef83d03bffd849955c84c6a8e7c31 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Thu, 5 Oct 2023 11:19:26 +0200
Subject: [PATCH] [GVN] Remove users from ICF when RAUWing loads

When performing store to load forwarding, replacing users of the
load may turn an indirect call into one with a known callee, in
which case it might become willreturn, invalidating cached ICF
information. Avoid this by removing users.

This is a bit more aggressive than strictly necessary (e.g. this
shouldn't be necessary when doing load-load CSE), but better safe
than sorry.

Fixes https://github.com/llvm/llvm-project/issues/48805.
---
 llvm/lib/Transforms/Scalar/GVN.cpp  |  3 ++
 llvm/test/Transforms/GVN/pr48805.ll | 49 +++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 llvm/test/Transforms/GVN/pr48805.ll

diff --git a/llvm/lib/Transforms/Scalar/GVN.cpp b/llvm/lib/Transforms/Scalar/GVN.cpp
index bc54846ccf0a..6a97137be3f5 100644
--- a/llvm/lib/Transforms/Scalar/GVN.cpp
+++ b/llvm/lib/Transforms/Scalar/GVN.cpp
@@ -1472,6 +1472,7 @@ void GVNPass::eliminatePartiallyRedundantLoad(
   // Perform PHI construction.
   Value *V = ConstructSSAForLoadSet(Load, ValuesPerBlock, *this);
   // ConstructSSAForLoadSet is responsible for combining metadata.
+  ICF->replaceUsersOf(Load);
   Load->replaceAllUsesWith(V);
   if (isa<PHINode>(V))
     V->takeName(Load);
@@ -1891,6 +1892,7 @@ bool GVNPass::processNonLocalLoad(LoadInst *Load) {
     // Perform PHI construction.
     Value *V = ConstructSSAForLoadSet(Load, ValuesPerBlock, *this);
     // ConstructSSAForLoadSet is responsible for combining metadata.
+    ICF->removeUsersOf(Load);
     Load->replaceAllUsesWith(V);
 
     if (isa<PHINode>(V))
@@ -2165,6 +2167,7 @@ bool GVNPass::processLoad(LoadInst *L) {
   Value *AvailableValue = AV->MaterializeAdjustedValue(L, L, *this);
 
   // MaterializeAdjustedValue is responsible for combining metadata.
+  ICF->removeUsersOf(L);
   L->replaceAllUsesWith(AvailableValue);
   markInstructionForDeletion(L);
   if (MSSAU)
diff --git a/llvm/test/Transforms/GVN/pr48805.ll b/llvm/test/Transforms/GVN/pr48805.ll
new file mode 100644
index 000000000000..e3fdc280255f
--- /dev/null
+++ b/llvm/test/Transforms/GVN/pr48805.ll
@@ -0,0 +1,49 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -S -passes=gvn < %s | FileCheck %s
+
+declare void @willreturn() nounwind willreturn
+
+declare void @capture(ptr)
+
+; Make sure ICF is invalidated when the callee becomes known.
+
+define i64 @test(ptr %p) {
+; CHECK-LABEL: define i64 @test(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca [2 x ptr], align 8
+; CHECK-NEXT:    [[A2:%.*]] = getelementptr ptr, ptr [[A]], i64 1
+; CHECK-NEXT:    call void @capture(ptr [[A]])
+; CHECK-NEXT:    br i1 false, label [[IF:%.*]], label [[ENTRY_EXIT_CRIT_EDGE:%.*]]
+; CHECK:       entry.exit_crit_edge:
+; CHECK-NEXT:    [[RES_PRE:%.*]] = load i64, ptr [[A2]], align 8
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    [[P1:%.*]] = load ptr, ptr [[A2]], align 8
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RES:%.*]] = phi i64 [ [[RES_PRE]], [[ENTRY_EXIT_CRIT_EDGE]] ], [ undef, [[IF]] ]
+; CHECK-NEXT:    store ptr @willreturn, ptr [[P]], align 8
+; CHECK-NEXT:    tail call void @willreturn()
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+entry:
+  %a = alloca [2 x ptr], align 8
+  %a2 = getelementptr ptr, ptr %a, i64 1
+  call void @capture(ptr %a)
+  br i1 false, label %if, label %exit
+
+if:
+  %p0 = load ptr, ptr %a, align 8
+  %p1 = load ptr, ptr %a2, align 8
+  br label %exit
+
+exit:
+  store ptr @willreturn, ptr %p
+  %p2 = load ptr, ptr %a, align 8
+  %pgocount.i = load i64, ptr %p2, align 8
+  %fn = load ptr, ptr %p
+  tail call void %fn()
+  %res = load i64, ptr %a2, align 8
+  ret i64 %res
+}