[llvm] [Attributor] Prevent infinite loop in AAGlobalValueInfoFloating (PR #94941)

[Ethan Luis McDonough via llvm-commits]

https://github.com/EthanLuisMcDonough updated https://github.com/llvm/llvm-project/pull/94941

>From c8655724bdf26aaac8a39d5fcf05bc1cfbab64a3 Mon Sep 17 00:00:00 2001
From: Ethan Luis McDonough <ethanluismcdonough at gmail.com>
Date: Mon, 10 Jun 2024 03:38:46 -0500
Subject: [PATCH 1/3] Fix infinite loop in attributor for recursive globals

---
 .../Transforms/IPO/AttributorAttributes.cpp   |  2 +-
 .../Attributor/recursive_globals.ll           | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/Attributor/recursive_globals.ll

diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 549d03645f939..8c268e3faf177 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -12031,7 +12031,7 @@ struct AAGlobalValueInfoFloating : public AAGlobalValueInfo {
                 SmallVectorImpl<const Value *> &Worklist) {
     Instruction *UInst = dyn_cast<Instruction>(U.getUser());
     if (!UInst) {
-      Follow = true;
+      Follow = !Uses.contains(&U);
       return true;
     }
 
diff --git a/llvm/test/Transforms/Attributor/recursive_globals.ll b/llvm/test/Transforms/Attributor/recursive_globals.ll
new file mode 100644
index 0000000000000..1af6383b063f4
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/recursive_globals.ll
@@ -0,0 +1,24 @@
+; RUN: opt < %s -passes=attributor
+
+ at glob1 = global { ptr, ptr } { ptr @glob1, ptr @fnc1 }
+ at glob2 = global { ptr, ptr } { ptr @glob3, ptr @fnc2 }
+ at glob3 = global { ptr, ptr } { ptr @glob2, ptr @fnc2 }
+
+define internal void @fnc1() {
+  ret void
+}
+
+define internal void @fnc2() {
+  ret void
+}
+
+define internal void @indr_caller(ptr %0) {
+  call void %0()
+  ret void
+}
+
+define void @main() {
+  call void @indr_caller(ptr @fnc1)
+  call void @indr_caller(ptr @fnc2)
+  ret void
+}

>From b6ef071dd4af3c19e881f9bb396e467f98da5fa2 Mon Sep 17 00:00:00 2001
From: Ethan Luis McDonough <ethanluismcdonough at gmail.com>
Date: Mon, 10 Jun 2024 15:01:19 -0500
Subject: [PATCH 2/3] Address requested changes

---
 llvm/lib/Transforms/IPO/Attributor.cpp           | 6 +++++-
 llvm/lib/Transforms/IPO/AttributorAttributes.cpp | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index b6866580ccd3f..39e69219e1c7b 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -1749,6 +1749,10 @@ bool Attributor::checkForAllCallees(
   return Pred(Callees.getArrayRef());
 }
 
+bool noRevisitUser(const User *user) {
+  return isa<PHINode>(user) || !isa<Instruction>(user);
+}
+
 bool Attributor::checkForAllUses(
     function_ref<bool(const Use &, bool &)> Pred,
     const AbstractAttribute &QueryingAA, const Value &V,
@@ -1796,7 +1800,7 @@ bool Attributor::checkForAllUses(
 
   while (!Worklist.empty()) {
     const Use *U = Worklist.pop_back_val();
-    if (isa<PHINode>(U->getUser()) && !Visited.insert(U).second)
+    if (noRevisitUser(U->getUser()) && !Visited.insert(U).second)
       continue;
     DEBUG_WITH_TYPE(VERBOSE_DEBUG_TYPE, {
       if (auto *Fn = dyn_cast<Function>(U->getUser()))
diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 8c268e3faf177..549d03645f939 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -12031,7 +12031,7 @@ struct AAGlobalValueInfoFloating : public AAGlobalValueInfo {
                 SmallVectorImpl<const Value *> &Worklist) {
     Instruction *UInst = dyn_cast<Instruction>(U.getUser());
     if (!UInst) {
-      Follow = !Uses.contains(&U);
+      Follow = true;
       return true;
     }
 

>From 530ec6d11462f7801c9cfc819b0deda876ca3d2d Mon Sep 17 00:00:00 2001
From: Ethan Luis McDonough <ethanluismcdonough at gmail.com>
Date: Mon, 10 Jun 2024 15:18:25 -0500
Subject: [PATCH 3/3] Add test description

---
 llvm/test/Transforms/Attributor/recursive_globals.ll | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/llvm/test/Transforms/Attributor/recursive_globals.ll b/llvm/test/Transforms/Attributor/recursive_globals.ll
index 1af6383b063f4..d6be6d6aaedcc 100644
--- a/llvm/test/Transforms/Attributor/recursive_globals.ll
+++ b/llvm/test/Transforms/Attributor/recursive_globals.ll
@@ -1,5 +1,9 @@
 ; RUN: opt < %s -passes=attributor
 
+; Global variables that reference themselves alongside a function that is called indirectly
+; used to cause an infinite loop in the attributor. The recursive reference was continually
+; pushed back into the workload, causing the attributor to hang indefinitely.
+
 @glob1 = global { ptr, ptr } { ptr @glob1, ptr @fnc1 }
 @glob2 = global { ptr, ptr } { ptr @glob3, ptr @fnc2 }
 @glob3 = global { ptr, ptr } { ptr @glob2, ptr @fnc2 }