[llvm] [DeadArgElim] Renumber allocsize instead of dropping it (PR #221868)
Larry Meadows via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 8 13:56:25 PDT 2026
https://github.com/lfmeadow updated https://github.com/llvm/llvm-project/pull/221868
>From e27b87c7df29935f562f28e528d591bd8dc110c0 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Sun, 6 Sep 2026 20:39:09 -0500
Subject: [PATCH 1/3] [DeadArgElim] Renumber allocsize instead of dropping it
Deleting an argument shifts the indices of the ones after it, so
DeadArgumentElimination removes any allocsize attribute on a function whose
signature it rewrites. That is safe -- the verifier rejects an out-of-range
allocsize index -- but what it leaves behind is a function that still allocates
and can no longer report how much. Every consumer of allocation attributes is
then working from a description the function cannot honour. allockind and
alloc-family stay, so isAllocationFn() still recognizes it as an allocator,
while getAllocSize() and ObjectSizeOffsetEvaluator can no longer size it.
Keep the arguments allocsize names alive during the survey and renumber the
attribute to their new positions in the rewrite, dropping it only if they
disappeared anyway. The cost is at most one otherwise-dead argument on a
function that allocates.
In OpenMP device code this is what took apart __kmpc_alloc_shared. After
inlining, its size parameter looks dead, so the internalized copy came out with
no arguments at all and no allocsize, and heap-to-stack stopped being able to
promote the globalized locals allocated through it.
---
.../IPO/DeadArgumentElimination.cpp | 60 ++++++++++++++++---
llvm/test/Transforms/DeadArgElim/allocsize.ll | 46 +++++++++++++-
2 files changed, 97 insertions(+), 9 deletions(-)
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 277c7a11d7eea..46d9460fea158 100644
--- a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -580,6 +580,14 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Inspecting args for fn: "
<< F.getName() << "\n");
+ std::optional<std::pair<unsigned, std::optional<unsigned>>> AllocSizeArgs =
+ F.getAttributes().getFnAttrs().getAllocSizeArgs();
+ auto IsAllocSizeArg = [&](unsigned ArgI) {
+ return AllocSizeArgs &&
+ (ArgI == AllocSizeArgs->first ||
+ (AllocSizeArgs->second && ArgI == *AllocSizeArgs->second));
+ };
+
// Now, check all of our arguments.
unsigned ArgI = 0;
UseVector MaybeLiveArgUses;
@@ -593,6 +601,13 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
// register and stack HFAs very differently, and this is reflected in the
// IR which has already been generated.
Result = Live;
+ } else if (IsAllocSizeArg(ArgI)) {
+ // Removing an argument allocsize points at would force the attribute off
+ // the function, since the verifier rejects an out-of-range allocsize
+ // index. What that leaves is a function that still allocates but can no
+ // longer say how much, which costs its callers more than a dead
+ // argument does.
+ Result = Live;
} else {
// See what the effect of this use is (recording any uses that cause
// MaybeLive in MaybeLiveArgUses).
@@ -839,9 +854,42 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
AttributeSet RetAttrs = AttributeSet::get(F->getContext(), RAttrs);
- // Strip allocsize attributes. They might refer to the deleted arguments.
- AttributeSet FnAttrs =
- PAL.getFnAttrs().removeAttribute(F->getContext(), Attribute::AllocSize);
+ // allocsize names parameters by index, so deleting an argument ahead of one
+ // renumbers it. surveyFunction() keeps the arguments allocsize points at
+ // alive, so renumber the attribute rather than dropping it; drop it only if
+ // they went away regardless, which leaves the callee no way to report the
+ // size it allocates.
+ auto UpdateAllocSize = [&](AttributeSet FnAttrs) {
+ std::optional<std::pair<unsigned, std::optional<unsigned>>> Args =
+ FnAttrs.getAllocSizeArgs();
+ if (!Args)
+ return FnAttrs;
+
+ auto NewIdx = [&](unsigned Old) -> std::optional<unsigned> {
+ if (Old >= ArgAlive.size() || !ArgAlive[Old])
+ return std::nullopt;
+ return static_cast<unsigned>(
+ std::count(ArgAlive.begin(), ArgAlive.begin() + Old, true));
+ };
+
+ std::optional<unsigned> ElemSizeArg = NewIdx(Args->first);
+ std::optional<unsigned> NumElemsArg;
+ if (ElemSizeArg && Args->second) {
+ NumElemsArg = NewIdx(*Args->second);
+ if (!NumElemsArg)
+ ElemSizeArg = std::nullopt;
+ }
+
+ FnAttrs = FnAttrs.removeAttribute(F->getContext(), Attribute::AllocSize);
+ if (!ElemSizeArg)
+ return FnAttrs;
+
+ AttrBuilder B(F->getContext());
+ B.addAllocSizeAttr(*ElemSizeArg, NumElemsArg);
+ return FnAttrs.addAttributes(F->getContext(), B);
+ };
+
+ AttributeSet FnAttrs = UpdateAllocSize(PAL.getFnAttrs());
// Reconstruct the AttributesList based on the vector we constructed.
assert(ArgAttrVec.size() == Params.size());
@@ -916,10 +964,8 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
// Reconstruct the AttributesList based on the vector we constructed.
assert(ArgAttrVec.size() == Args.size());
- // Again, be sure to remove any allocsize attributes, since their indices
- // may now be incorrect.
- AttributeSet FnAttrs = CallPAL.getFnAttrs().removeAttribute(
- F->getContext(), Attribute::AllocSize);
+ // Again, renumber allocsize, since its indices may now be incorrect.
+ AttributeSet FnAttrs = UpdateAllocSize(CallPAL.getFnAttrs());
AttributeList NewCallPAL =
AttributeList::get(F->getContext(), FnAttrs, RetAttrs, ArgAttrVec);
diff --git a/llvm/test/Transforms/DeadArgElim/allocsize.ll b/llvm/test/Transforms/DeadArgElim/allocsize.ll
index 4115de218bc7c..0136e2ab6906a 100644
--- a/llvm/test/Transforms/DeadArgElim/allocsize.ll
+++ b/llvm/test/Transforms/DeadArgElim/allocsize.ll
@@ -1,18 +1,60 @@
; RUN: opt < %s -passes=deadargelim -S | FileCheck %s
; PR36867
-; CHECK-LABEL: @MagickMallocAligned
-; CHECK-NOT: allocsize
+; Deleting an argument ahead of the one allocsize names renumbers it, so the
+; attribute has to follow rather than be dropped.
+
+; CHECK-LABEL: define internal i64 @MagickMallocAligned(i64 %s)
+; CHECK-SAME: #[[ONE:[0-9]+]]
define internal i64 @MagickMallocAligned(i64 %DEADARG1, i64 %s) allocsize(1) {
ret i64 %s
}
define i64 @NeedsArg(i64 %s) {
+; CHECK-LABEL: define i64 @NeedsArg(
+; CHECK: call i64 @MagickMallocAligned(i64 %s)
%c = call i64 @MagickMallocAligned(i64 0, i64 %s)
ret i64 %c
}
define i64 @Test2(i64 %s) {
+; CHECK-LABEL: define i64 @Test2(
+; CHECK: call i64 @MagickMallocAligned(i64 %s) #[[ONE]]
%c = call i64 @MagickMallocAligned(i64 0, i64 %s) allocsize(1)
ret i64 %c
}
+
+; Both indices of a calloc-like allocsize are renumbered.
+
+; CHECK-LABEL: define internal ptr @two_args(i64 %n, i64 %sz)
+; CHECK-SAME: #[[TWO:[0-9]+]]
+define internal ptr @two_args(i64 %DEADARG2, i64 %n, i64 %sz) allocsize(1, 2) {
+ %p = call ptr @allocate(i64 %n, i64 %sz)
+ ret ptr %p
+}
+
+define ptr @calls_two_args(i64 %n, i64 %sz) {
+ %p = call ptr @two_args(i64 0, i64 %n, i64 %sz)
+ ret ptr %p
+}
+
+; The argument allocsize names is otherwise unused, but removing it would cost
+; the attribute, so it is kept alive.
+
+; CHECK-LABEL: define internal ptr @unused_size(i64 %n)
+; CHECK-SAME: #[[THREE:[0-9]+]]
+define internal ptr @unused_size(i64 %DEADARG3, i64 %n) allocsize(1) nounwind {
+ %p = call ptr @allocate(i64 0, i64 0)
+ ret ptr %p
+}
+
+define ptr @calls_unused_size(i64 %n) {
+ %p = call ptr @unused_size(i64 0, i64 %n)
+ ret ptr %p
+}
+
+declare ptr @allocate(i64, i64)
+
+; CHECK-DAG: attributes #[[ONE]] = { allocsize(0) }
+; CHECK-DAG: attributes #[[TWO]] = { allocsize(0,1) }
+; CHECK-DAG: attributes #[[THREE]] = { nounwind allocsize(0) }
>From 962059254a65d76c72371f7b1efe82594d0588c8 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Tue, 8 Sep 2026 05:45:39 -0500
Subject: [PATCH 2/3] Say why a call-site allocsize can still be dropped
---
llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 46d9460fea158..35711d4167fe1 100644
--- a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -855,10 +855,10 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
AttributeSet RetAttrs = AttributeSet::get(F->getContext(), RAttrs);
// allocsize names parameters by index, so deleting an argument ahead of one
- // renumbers it. surveyFunction() keeps the arguments allocsize points at
- // alive, so renumber the attribute rather than dropping it; drop it only if
- // they went away regardless, which leaves the callee no way to report the
- // size it allocates.
+ // renumbers it. surveyFunction() keeps the arguments the function's own
+ // allocsize names alive, so those always renumber. A call site's allocsize is
+ // independent of the callee's and gets no such treatment, so it can still
+ // name an argument that is gone, and is dropped.
auto UpdateAllocSize = [&](AttributeSet FnAttrs) {
std::optional<std::pair<unsigned, std::optional<unsigned>>> Args =
FnAttrs.getAllocSizeArgs();
>From f92af4bbc3dff733c00a07685f6adbcaf7c6b09e Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Tue, 8 Sep 2026 15:55:23 -0500
Subject: [PATCH 3/3] Keep a call site's allocsize arguments alive too
A call site can carry its own allocsize, and getAllocationSize() prefers it to
the callee's, so it is the one heap-to-stack reads there. Deleting an argument
it names forced it off and left that call an allocator that could not say how
much it allocates -- the shape the function's own allocsize was already being
kept alive to avoid.
Collect the indices from the call sites as well as from the function while
surveying, so both are kept. With no allocsize left that can name a deleted
argument, the attribute is now always renumbered and the code to drop it goes.
---
.../IPO/DeadArgumentElimination.cpp | 59 +++++++++----------
llvm/test/Transforms/DeadArgElim/allocsize.ll | 16 +++++
2 files changed, 45 insertions(+), 30 deletions(-)
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 35711d4167fe1..2819cc03ec912 100644
--- a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -17,6 +17,7 @@
//===----------------------------------------------------------------------===//
#include "llvm/Transforms/IPO/DeadArgumentElimination.h"
+#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
@@ -518,6 +519,17 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
// of them turn out to be live.
unsigned NumLiveRetVals = 0;
+ // Parameters named by an allocsize, on the function or on any call site.
+ SmallSet<unsigned, 4> AllocSizeArgIdxs;
+ auto NoteAllocSizeArgs = [&](AttributeSet FnAttrs) {
+ if (auto Args = FnAttrs.getAllocSizeArgs()) {
+ AllocSizeArgIdxs.insert(Args->first);
+ if (Args->second)
+ AllocSizeArgIdxs.insert(*Args->second);
+ }
+ };
+ NoteAllocSizeArgs(F.getAttributes().getFnAttrs());
+
// Loop all uses of the function.
for (const Use &U : F.uses()) {
// If the function is PASSED IN as an argument, its address has been
@@ -538,6 +550,8 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
// If we end up here, we are looking at a direct call to our function.
+ NoteAllocSizeArgs(CB->getAttributes().getFnAttrs());
+
// Now, check how our return value(s) is/are used in this caller. Don't
// bother checking return values if all of them are live already.
if (NumLiveRetVals == RetCount)
@@ -580,14 +594,6 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Inspecting args for fn: "
<< F.getName() << "\n");
- std::optional<std::pair<unsigned, std::optional<unsigned>>> AllocSizeArgs =
- F.getAttributes().getFnAttrs().getAllocSizeArgs();
- auto IsAllocSizeArg = [&](unsigned ArgI) {
- return AllocSizeArgs &&
- (ArgI == AllocSizeArgs->first ||
- (AllocSizeArgs->second && ArgI == *AllocSizeArgs->second));
- };
-
// Now, check all of our arguments.
unsigned ArgI = 0;
UseVector MaybeLiveArgUses;
@@ -601,12 +607,10 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
// register and stack HFAs very differently, and this is reflected in the
// IR which has already been generated.
Result = Live;
- } else if (IsAllocSizeArg(ArgI)) {
- // Removing an argument allocsize points at would force the attribute off
- // the function, since the verifier rejects an out-of-range allocsize
- // index. What that leaves is a function that still allocates but can no
- // longer say how much, which costs its callers more than a dead
- // argument does.
+ } else if (AllocSizeArgIdxs.contains(ArgI)) {
+ // Dropping this argument would take allocsize with it, since the
+ // verifier rejects an out-of-range index, and the attribute is worth
+ // more.
Result = Live;
} else {
// See what the effect of this use is (recording any uses that cause
@@ -855,37 +859,32 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
AttributeSet RetAttrs = AttributeSet::get(F->getContext(), RAttrs);
// allocsize names parameters by index, so deleting an argument ahead of one
- // renumbers it. surveyFunction() keeps the arguments the function's own
- // allocsize names alive, so those always renumber. A call site's allocsize is
- // independent of the callee's and gets no such treatment, so it can still
- // name an argument that is gone, and is dropped.
+ // shifts it. surveyFunction() keeps those arguments alive, so the attribute
+ // can be renumbered rather than dropped.
auto UpdateAllocSize = [&](AttributeSet FnAttrs) {
std::optional<std::pair<unsigned, std::optional<unsigned>>> Args =
FnAttrs.getAllocSizeArgs();
if (!Args)
return FnAttrs;
- auto NewIdx = [&](unsigned Old) -> std::optional<unsigned> {
- if (Old >= ArgAlive.size() || !ArgAlive[Old])
- return std::nullopt;
+ auto NewIdx = [&](unsigned Old) {
+ // An index past the parameters names a variadic argument; a vararg
+ // function keeps all of its parameters, so nothing ahead of it moved.
+ if (Old >= ArgAlive.size())
+ return Old;
+ assert(ArgAlive[Old] && "allocsize parameter was not kept alive");
return static_cast<unsigned>(
std::count(ArgAlive.begin(), ArgAlive.begin() + Old, true));
};
- std::optional<unsigned> ElemSizeArg = NewIdx(Args->first);
+ unsigned ElemSizeArg = NewIdx(Args->first);
std::optional<unsigned> NumElemsArg;
- if (ElemSizeArg && Args->second) {
+ if (Args->second)
NumElemsArg = NewIdx(*Args->second);
- if (!NumElemsArg)
- ElemSizeArg = std::nullopt;
- }
FnAttrs = FnAttrs.removeAttribute(F->getContext(), Attribute::AllocSize);
- if (!ElemSizeArg)
- return FnAttrs;
-
AttrBuilder B(F->getContext());
- B.addAllocSizeAttr(*ElemSizeArg, NumElemsArg);
+ B.addAllocSizeAttr(ElemSizeArg, NumElemsArg);
return FnAttrs.addAttributes(F->getContext(), B);
};
diff --git a/llvm/test/Transforms/DeadArgElim/allocsize.ll b/llvm/test/Transforms/DeadArgElim/allocsize.ll
index 0136e2ab6906a..87f3625d89002 100644
--- a/llvm/test/Transforms/DeadArgElim/allocsize.ll
+++ b/llvm/test/Transforms/DeadArgElim/allocsize.ll
@@ -53,6 +53,22 @@ define ptr @calls_unused_size(i64 %n) {
ret ptr %p
}
+; A call site's allocsize names an argument of a callee that has no allocsize of
+; its own and never uses it. Keeping the attribute means keeping the argument.
+
+; CHECK-LABEL: define internal ptr @callsite_names_size(i64 %sz)
+define internal ptr @callsite_names_size(i64 %DEADARG4, i64 %sz) {
+ %p = call ptr @allocate(i64 0, i64 0)
+ ret ptr %p
+}
+
+define ptr @calls_callsite_names_size(i64 %sz) {
+; CHECK-LABEL: define ptr @calls_callsite_names_size(
+; CHECK: call ptr @callsite_names_size(i64 %sz) #[[ONE]]
+ %p = call ptr @callsite_names_size(i64 0, i64 %sz) allocsize(1)
+ ret ptr %p
+}
+
declare ptr @allocate(i64, i64)
; CHECK-DAG: attributes #[[ONE]] = { allocsize(0) }
More information about the llvm-commits
mailing list