[llvm] llvm-reduce: Filter function based on uses before removing arguments (PR #133412)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/133412

>From e1fe44b9d039fac59d60102eae58127f9a0e4c25 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Fri, 28 Mar 2025 17:09:49 +0700
Subject: [PATCH] llvm-reduce: Filter function based on uses before removing
 arguments

Invokes and others are not handled, so this was leaving broken callsites
behind for anything other than CallInst
---
 .../llvm-reduce/reduce-arguments-invoke.ll    | 41 +++++++++++++++++++
 .../reduce-arguments-non-callee-use.ll        | 32 +++++++++++++++
 .../llvm-reduce/deltas/ReduceArguments.cpp    | 17 +++++++-
 3 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/tools/llvm-reduce/reduce-arguments-invoke.ll
 create mode 100644 llvm/test/tools/llvm-reduce/reduce-arguments-non-callee-use.ll

diff --git a/llvm/test/tools/llvm-reduce/reduce-arguments-invoke.ll b/llvm/test/tools/llvm-reduce/reduce-arguments-invoke.ll
new file mode 100644
index 0000000000000..fb8a4d6367b3b
--- /dev/null
+++ b/llvm/test/tools/llvm-reduce/reduce-arguments-invoke.ll
@@ -0,0 +1,41 @@
+; RUN: llvm-reduce %s -o %t --abort-on-invalid-reduction --delta-passes=arguments --test FileCheck --test-arg %s --test-arg --check-prefix=INTERESTING --test-arg --input-file
+; RUN: FileCheck %s --input-file %t --check-prefix=REDUCED
+
+; Test that we don't break the callsite for an unhandled invoke
+
+declare void @did_not_throw(i32)
+declare void @thrown()
+
+; INTERESTING-LABEL: define i32 @maybe_throwing_callee(
+
+; REDUCED-LABEL: define i32 @maybe_throwing_callee(i32 %arg0, i32 %arg1) {
+define i32 @maybe_throwing_callee(i32 %arg0, i32 %arg1) {
+  call void @thrown()
+  ret i32 %arg1
+}
+
+ at initializer_user = global [1 x ptr] [ptr @maybe_throwing_callee ]
+
+; REDUCED-LABEL: define void @caller()
+; REDUCED: %i0 = invoke i32 @maybe_throwing_callee(i32 0, i32 0) #0
+define void @caller(i32 %arg0, ptr %arg1) personality ptr @__gxx_personality_v0 {
+bb:
+  %val = load i32, ptr %arg1
+  %i0 = invoke i32 @maybe_throwing_callee(i32 0, i32 0) nofree
+          to label %bb3 unwind label %bb1
+
+bb1:
+  landingpad { ptr, i32 }
+  catch ptr null
+  call void @thrown()
+  br label %bb4
+
+bb3:
+  call void @did_not_throw(i32 %i0)
+  br label %bb4
+
+bb4:
+  ret void
+}
+
+declare i32 @__gxx_personality_v0(...)
diff --git a/llvm/test/tools/llvm-reduce/reduce-arguments-non-callee-use.ll b/llvm/test/tools/llvm-reduce/reduce-arguments-non-callee-use.ll
new file mode 100644
index 0000000000000..63b2909404b7c
--- /dev/null
+++ b/llvm/test/tools/llvm-reduce/reduce-arguments-non-callee-use.ll
@@ -0,0 +1,32 @@
+; RUN: llvm-reduce %s -o %t --abort-on-invalid-reduction --delta-passes=arguments --test FileCheck --test-arg %s --test-arg --check-prefix=INTERESTING --test-arg --input-file
+; RUN: FileCheck %s --input-file %t --check-prefix=REDUCED
+
+; INTERESTING: @initializer_user
+; REDUCED: @initializer_user = global [1 x ptr] [ptr @captured_func]
+ at initializer_user = global [1 x ptr] [ptr @captured_func ]
+
+; INTERESTING-LABEL: define i32 @captured_func(
+
+; REDUCED-LABEL: define i32 @captured_func() {
+define i32 @captured_func(i32 %a, i32 %b) {
+  %mul = mul i32 %a, %b
+  ret i32 %mul
+}
+
+; INTERESTING-LABEL: declare void @captures(
+declare void @captures(i32, ptr, i32)
+
+
+; INTERESTING-LABEL: define i32 @caller(
+; INTERESTING: = call
+; INTERESTING: = call
+
+; REDUCED-LABEL: define i32 @caller(i32 %a, i32 %b) {
+; REDUCED: %call0 = call i32 @captures(i32 %a, ptr @captured_func, i32 %b)
+; REDUCED: %call1 = call i32 @captured_func()
+define i32 @caller(i32 %a, i32 %b) {
+  %call0 = call i32 @captures(i32 %a, ptr @captured_func, i32 %b)
+  %call1 = call i32 @captured_func(i32 %a, i32 %b)
+  %add = add i32 %call0, %call1
+  ret i32 %add
+}
diff --git a/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp b/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
index f1b72a75564de..bf9d6ba145f6f 100644
--- a/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
+++ b/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
@@ -96,6 +96,20 @@ static bool shouldRemoveArguments(const Function &F) {
   return !F.arg_empty() && !F.isIntrinsic();
 }
 
+static bool allFuncUsersRewritable(const Function &F) {
+  for (const Use &U : F.uses()) {
+    const CallBase *CB = dyn_cast<CallBase>(U.getUser());
+    if (!CB || !CB->isCallee(&U))
+      continue;
+
+    // TODO: Handle all CallBase cases.
+    if (!isa<CallInst>(CB))
+      return false;
+  }
+
+  return true;
+}
+
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractArgumentsFromModule(Oracle &O, ReducerWorkItem &WorkItem) {
@@ -107,7 +121,8 @@ static void extractArgumentsFromModule(Oracle &O, ReducerWorkItem &WorkItem) {
   for (auto &F : Program) {
     if (!shouldRemoveArguments(F))
       continue;
-
+    if (!allFuncUsersRewritable(F))
+      continue;
     Funcs.push_back(&F);
     for (auto &A : F.args()) {
       if (callingConvRequiresArgument(F, A) || O.shouldKeep())