[llvm] llvm-reduce: Fix losing callsite attributes when removing arguments (PR #133407)

[Matt Arsenault via llvm-commits]

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

>From 52eb2ed51e48d7750d5c5ecb48a5d770c6f8408a Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Fri, 28 Mar 2025 15:47:46 +0700
Subject: [PATCH] llvm-reduce: Fix losing callsite attributes when removing
 arguments

The attribute APIs make this cumbersome. There seem to be missing
overloads using AttrBuilder for the function attrs. Plus there doesn't
seem to be a direct way to set the function attrs on the call.
---
 .../reduce-arguments-x86_intrcc.ll            |  2 +-
 .../remove-argument-preserve-attributes.ll    | 55 +++++++++++++++++++
 .../llvm-reduce/deltas/ReduceArguments.cpp    | 24 +++++++-
 3 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/tools/llvm-reduce/remove-argument-preserve-attributes.ll

diff --git a/llvm/test/tools/llvm-reduce/reduce-arguments-x86_intrcc.ll b/llvm/test/tools/llvm-reduce/reduce-arguments-x86_intrcc.ll
index 50aac3cce32fa..7589835f0c0fe 100644
--- a/llvm/test/tools/llvm-reduce/reduce-arguments-x86_intrcc.ll
+++ b/llvm/test/tools/llvm-reduce/reduce-arguments-x86_intrcc.ll
@@ -20,7 +20,7 @@ declare x86_intrcc void @extern_decl(ptr byval(i32), i32, i32)
 
 ; INTERESTING-LABEL: void @callsite(
 ; INTERESTING: call
-; REDUCED: call x86_intrcc void @func(ptr %k, i32 %other.keep)
+; REDUCED: call x86_intrcc void @func(ptr byval(i32) %k, i32 %other.keep)
 define void @callsite(ptr %k, i32 %other.keep, i32 %other.drop) {
   call x86_intrcc void @func(ptr byval(i32) %k, i32 %other.keep, i32 %other.drop)
   ret void
diff --git a/llvm/test/tools/llvm-reduce/remove-argument-preserve-attributes.ll b/llvm/test/tools/llvm-reduce/remove-argument-preserve-attributes.ll
new file mode 100644
index 0000000000000..0491dec983438
--- /dev/null
+++ b/llvm/test/tools/llvm-reduce/remove-argument-preserve-attributes.ll
@@ -0,0 +1,55 @@
+; Check that when removing arguments, existing callsite attributes are preserved
+
+; RUN: llvm-reduce --abort-on-invalid-reduction --delta-passes=arguments --test FileCheck --test-arg --check-prefixes=INTERESTING --test-arg %s --test-arg --input-file %s -o %t
+; RUN: FileCheck --check-prefixes=RESULT %s < %t
+
+; INTERESTING-LABEL: define void @callee0(
+define void @callee0(ptr %interesting0, ptr %interesting1, i32 %uninteresting2) {
+  ret void
+}
+
+; INTERESTING-LABEL: define void @callee1(
+define void @callee1(ptr byval(i64) %interesting0, ptr %interesting1, i32 %uninteresting2) {
+  ret void
+}
+
+; INTERESTING-LABEL: define void @caller0(
+
+; INTERESTING: byval
+; INTERESTING-SAME: "some-attr"
+
+; INTERESTING: byval
+; INTERESTING-SAME: "more-attr"
+
+; RESULT-LABEL: define void @caller0(ptr %val0) {
+; RESULT: call void @callee0(ptr byval(i32) %val0, ptr "some-attr" %alloca0) #0
+; RESULT: call void @callee1(ptr byval(i64) %alloca1, ptr "more-attr" %alloca1) #1
+define void @caller0(ptr %val0, i32 %val1) {
+  %alloca0 = alloca i32
+  %alloca1 = alloca i64
+  call void @callee0(ptr byval(i32) %val0, ptr "some-attr" %alloca0, i32 %val1) nounwind memory(none) "a-func-attr"
+  call void @callee1(ptr byval(i64) %alloca1, ptr "more-attr" %alloca1, i32 9) "val-func-attr="="something"
+  ret void
+}
+
+; RESULT-LABEL: define ptr @callee2() {
+; RESULT-NEXT: ret ptr null
+define ptr @callee2(ptr %val0, i32 %b) {
+  store i32 %b, ptr %val0
+  ret ptr %val0
+}
+
+; Make sure ret attributes are preserved
+; INTERESTING: define ptr @caller1(
+; INTERESTING: call
+
+; RESULT-LABEL: define ptr @caller1() {
+; RESULT: %ret = call align 4 "ret-attr" ptr @callee2()
+
+define ptr @caller1(ptr %val0, i32 %val1) {
+  %ret = call align 4 "ret-attr" ptr @callee2(ptr %val0, i32 %val1)
+  ret ptr %ret
+}
+
+; RESULT: attributes #0 = { nounwind memory(none) "a-func-attr" }
+; RESULT: attributes #1 = { "val-func-attr="="something" }
diff --git a/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp b/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
index 134801fef020b..ec2699a990311 100644
--- a/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
+++ b/llvm/tools/llvm-reduce/deltas/ReduceArguments.cpp
@@ -39,6 +39,8 @@ static bool callingConvRequiresArgument(const Function &F,
 /// Goes over OldF calls and replaces them with a call to NewF
 static void replaceFunctionCalls(Function &OldF, Function &NewF,
                                  const std::set<int> &ArgIndexesToKeep) {
+  LLVMContext &Ctx = OldF.getContext();
+
   const auto &Users = OldF.users();
   for (auto I = Users.begin(), E = Users.end(); I != E; )
     if (auto *CI = dyn_cast<CallInst>(*I++)) {
@@ -47,12 +49,30 @@ static void replaceFunctionCalls(Function &OldF, Function &NewF,
       if (CI->getCalledFunction() != &OldF)
         continue;
       SmallVector<Value *, 8> Args;
-      for (auto ArgI = CI->arg_begin(), E = CI->arg_end(); ArgI != E; ++ArgI)
-        if (ArgIndexesToKeep.count(ArgI - CI->arg_begin()))
+      SmallVector<AttrBuilder, 8> ArgAttrs;
+
+      for (auto ArgI = CI->arg_begin(), E = CI->arg_end(); ArgI != E; ++ArgI) {
+        unsigned ArgIdx = ArgI - CI->arg_begin();
+        if (ArgIndexesToKeep.count(ArgIdx)) {
           Args.push_back(*ArgI);
+          ArgAttrs.emplace_back(Ctx, CI->getParamAttributes(ArgIdx));
+        }
+      }
 
+      // FIXME: Losing bundles, fast math flags and metadata
       CallInst *NewCI = CallInst::Create(&NewF, Args);
       NewCI->setCallingConv(NewF.getCallingConv());
+
+      AttrBuilder CallSiteAttrs(Ctx, CI->getAttributes().getFnAttrs());
+      NewCI->setAttributes(
+          AttributeList::get(Ctx, AttributeList::FunctionIndex, CallSiteAttrs));
+      NewCI->addRetAttrs(AttrBuilder(Ctx, CI->getRetAttributes()));
+
+      unsigned AttrIdx = 0;
+      for (auto ArgI = NewCI->arg_begin(), E = NewCI->arg_end(); ArgI != E;
+           ++ArgI, ++AttrIdx)
+        NewCI->addParamAttrs(AttrIdx, ArgAttrs[AttrIdx]);
+
       if (!CI->use_empty())
         CI->replaceAllUsesWith(NewCI);
       ReplaceInstWithInst(CI, NewCI);