[llvm] [llvm][SplitModuleByCategory] Fix infinite loop on cyclic global uses (PR #206862)

[Yury Plyakhin via llvm-commits]

https://github.com/YuriPlyakhin created https://github.com/llvm/llvm-project/pull/206862

The dependency graph construction in `addUserToGraphRecursively` walked the users of globals without visited set. A cycle in the use graph caused it to push the same users forever, hanging the splitter.

Track visited users so each is processed once.

Co-Authored-By: Claude


>From ece56cfde00f78bafcff064331128ff22a24d1c9 Mon Sep 17 00:00:00 2001
From: "Plyakhin, Yury" <yury.plyakhin at intel.com>
Date: Wed, 1 Jul 2026 01:55:13 +0200
Subject: [PATCH] [llvm][SplitModuleByCategory] Fix infinite loop on cyclic
 global uses

The dependency graph construction in addUserToGraphRecursively walked
the users of globals without visited set. A cycle in the use graph caused
it to push the same users forever, hanging the splitter.

Track visited users so each is processed once.

Co-Authored-By: Claude
---
 .../Transforms/Utils/SplitModuleByCategory.h  |  4 +-
 .../Utils/SplitModuleByCategory.cpp           |  5 +-
 .../SplitByCategory/recursion-and-cycles.ll   | 47 +++++++++++++++++++
 3 files changed, 53 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/tools/llvm-split/SplitByCategory/recursion-and-cycles.ll

diff --git a/llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h b/llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h
index 0ab2f455052f9..545899e7290ef 100644
--- a/llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h
+++ b/llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h
@@ -53,8 +53,8 @@ class Function;
 ///    there is a function B with signature S. An "A" -> "B" edge will be added
 ///    to the graph;
 ///
-/// FIXME: For now, the algorithm assumes no recursion in the input Module. This
-/// will be addressed in the near future.
+/// Recursion and other cycles in the input Module (e.g. directly or mutually
+/// recursive functions, or self-referencing global variables) are supported.
 LLVM_ABI Error splitModuleTransitiveFromEntryPoints(
     std::unique_ptr<Module> M,
     function_ref<std::optional<int>(const Function &F)> EntryPointCategorizer,
diff --git a/llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp b/llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp
index a8b705ac8148d..b3d69d3affc06 100644
--- a/llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp
+++ b/llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp
@@ -159,7 +159,9 @@ class DependencyGraph {
 private:
   void addUserToGraphRecursively(const User *Root, const GlobalValue *V) {
     SmallVector<const User *, 8> WorkList;
+    SmallPtrSet<const User *, 8> Visited;
     WorkList.push_back(Root);
+    Visited.insert(Root);
 
     while (!WorkList.empty()) {
       const User *U = WorkList.pop_back_val();
@@ -173,7 +175,8 @@ class DependencyGraph {
         // bitcast or gep). We trace users of this constant further to reach
         // global objects they are used by and add them to the graph.
         for (const User *UU : U->users())
-          WorkList.push_back(UU);
+          if (Visited.insert(UU).second)
+            WorkList.push_back(UU);
       } else {
         llvm_unreachable("Unhandled type of function user");
       }
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/recursion-and-cycles.ll b/llvm/test/tools/llvm-split/SplitByCategory/recursion-and-cycles.ll
new file mode 100644
index 0000000000000..52dd6bf6c70ad
--- /dev/null
+++ b/llvm/test/tools/llvm-split/SplitByCategory/recursion-and-cycles.ll
@@ -0,0 +1,47 @@
+; Check that Module splitting produces correct output in the presence of cycles
+; in the dependency graph.
+
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
+; RUN:     --implicit-check-not @gptr --implicit-check-not @kernel_A
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
+; RUN:     --implicit-check-not @self --implicit-check-not @kernel_B
+
+; CHECK0-DAG: define spir_func void @ping()
+; CHECK0-DAG: define spir_func void @pong()
+; CHECK0-DAG: define spir_func void @self()
+; CHECK0-DAG: define spir_kernel void @kernel_B()
+
+; CHECK1-DAG: @gptr = private global ptr @gptr
+; CHECK1-DAG: define spir_func void @ping()
+; CHECK1-DAG: define spir_func void @pong()
+; CHECK1-DAG: define spir_kernel void @kernel_A()
+
+ at gptr = private global ptr @gptr
+
+define spir_func void @ping() {
+  call spir_func void @pong()
+  ret void
+}
+
+define spir_func void @pong() {
+  call spir_func void @ping()
+  ret void
+}
+
+define spir_func void @self() {
+  call spir_func void @self()
+  ret void
+}
+
+define spir_kernel void @kernel_A() {
+  call spir_func void @ping()
+  %v = load ptr, ptr @gptr
+  ret void
+}
+
+define spir_kernel void @kernel_B() {
+  call spir_func void @ping()
+  call spir_func void @self()
+  ret void
+}