Skip to content

Commit c7aacbb

Browse files
nikicnikic
committed
[ArgPromotion] Update allocsize indices after promotion
Promotion can add/remove arguments. We need to update the indices in the allocsize attribute accordingly. Fixes #66103.
1 parent d8d0588 commit c7aacbb

File tree

2 files changed

+115
-1
lines changed

2 files changed

+115
-1
lines changed

‎llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,26 +121,33 @@ doPromotion(Function *F, FunctionAnalysisManager &FAM,
121121
// that we are *not* promoting. For the ones that we do promote, the parameter
122122
// attributes are lost
123123
SmallVector<AttributeSet, 8> ArgAttrVec;
124+
// Mapping from old to new argument indices. -1 for promoted or removed
125+
// arguments.
126+
SmallVector<unsigned> NewArgIndices;
124127
AttributeList PAL = F->getAttributes();
125128

126129
// First, determine the new argument list
127-
unsigned ArgNo = 0;
130+
unsigned ArgNo = 0, NewArgNo = 0;
128131
for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
129132
++I, ++ArgNo) {
130133
if (!ArgsToPromote.count(&*I)) {
131134
// Unchanged argument
132135
Params.push_back(I->getType());
133136
ArgAttrVec.push_back(PAL.getParamAttrs(ArgNo));
137+
NewArgIndices.push_back(NewArgNo++);
134138
} else if (I->use_empty()) {
135139
// Dead argument (which are always marked as promotable)
136140
++NumArgumentsDead;
141+
NewArgIndices.push_back((unsigned)-1);
137142
} else {
138143
const auto &ArgParts = ArgsToPromote.find(&*I)->second;
139144
for (const auto &Pair : ArgParts) {
140145
Params.push_back(Pair.second.Ty);
141146
ArgAttrVec.push_back(AttributeSet());
142147
}
143148
++NumArgumentsPromoted;
149+
NewArgIndices.push_back((unsigned)-1);
150+
NewArgNo += ArgParts.size();
144151
}
145152
}
146153

@@ -173,6 +180,19 @@ doPromotion(Function *F, FunctionAnalysisManager &FAM,
173180
// the function.
174181
NF->setAttributes(AttributeList::get(F->getContext(), PAL.getFnAttrs(),
175182
PAL.getRetAttrs(), ArgAttrVec));
183+
184+
// Remap argument indices in allocsize attribute.
185+
if (auto AllocSize = NF->getAttributes().getFnAttrs().getAllocSizeArgs()) {
186+
unsigned Arg1 = NewArgIndices[AllocSize->first];
187+
assert(Arg1 != (unsigned)-1 && "allocsize cannot be promoted argument");
188+
std::optional<unsigned> Arg2;
189+
if (AllocSize->second) {
190+
Arg2 = NewArgIndices[*AllocSize->second];
191+
assert(Arg2 != (unsigned)-1 && "allocsize cannot be promoted argument");
192+
}
193+
NF->addFnAttr(Attribute::getWithAllocSizeArgs(F->getContext(), Arg1, Arg2));
194+
}
195+
176196
AttributeFuncs::updateMinLegalVectorWidthAttr(*NF, LargestVectorWidth);
177197
ArgAttrVec.clear();
178198

‎llvm/test/Transforms/ArgumentPromotion/allocsize.ll

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --version 3
2+
; RUN: opt -S -passes=argpromotion < %s | FileCheck %s
3+
4+
declare ptr @malloc(i64)
5+
declare ptr @calloc(i64, i64)
6+
7+
define internal ptr @my_alloc1(i64 %unchanged, ptr %unused, i64 %size, ptr %unused2) allocsize(2) {
8+
; CHECK: Function Attrs: allocsize(1)
9+
; CHECK-LABEL: define internal ptr @my_alloc1(
10+
; CHECK-SAME: i64 [[UNCHANGED:%.*]], i64 [[SIZE:%.*]]) #[[ATTR0:[0-9]+]] {
11+
; CHECK-NEXT: [[PTR:%.*]] = call ptr @malloc(i64 [[SIZE]])
12+
; CHECK-NEXT: ret ptr [[PTR]]
13+
;
14+
%ptr = call ptr @malloc(i64 %size)
15+
ret ptr %ptr
16+
}
17+
18+
define internal ptr @my_alloc2(i64 %unchanged, ptr %unused, i64 %size, i64 %size2, ptr %unused2) allocsize(2,3) {
19+
; CHECK: Function Attrs: allocsize(1,2)
20+
; CHECK-LABEL: define internal ptr @my_alloc2(
21+
; CHECK-SAME: i64 [[UNCHANGED:%.*]], i64 [[SIZE:%.*]], i64 [[SIZE2:%.*]]) #[[ATTR1:[0-9]+]] {
22+
; CHECK-NEXT: [[PTR:%.*]] = call ptr @calloc(i64 [[SIZE]], i64 [[SIZE2]])
23+
; CHECK-NEXT: ret ptr [[PTR]]
24+
;
25+
%ptr = call ptr @calloc(i64 %size, i64 %size2)
26+
ret ptr %ptr
27+
}
28+
29+
define internal ptr @my_alloc3(i64 %unchanged, ptr %promoted, ptr %promoted2, i64 %size) allocsize(3) {
30+
; CHECK: Function Attrs: allocsize(5)
31+
; CHECK-LABEL: define internal ptr @my_alloc3(
32+
; CHECK-SAME: i64 [[UNCHANGED:%.*]], i32 [[PROMOTED_0_VAL:%.*]], i32 [[PROMOTED_4_VAL:%.*]], i32 [[PROMOTED2_0_VAL:%.*]], i32 [[PROMOTED2_4_VAL:%.*]], i64 [[SIZE:%.*]]) #[[ATTR2:[0-9]+]] {
33+
; CHECK-NEXT: [[PTR:%.*]] = call ptr @malloc(i64 [[SIZE]])
34+
; CHECK-NEXT: ret ptr [[PTR]]
35+
;
36+
load i32, ptr %promoted
37+
%promoted.gep1 = getelementptr i8, ptr %promoted, i64 4
38+
load i32, ptr %promoted.gep1
39+
40+
load i32, ptr %promoted2
41+
%promoted2.gep1 = getelementptr i8, ptr %promoted2, i64 4
42+
load i32, ptr %promoted2.gep1
43+
44+
%ptr = call ptr @malloc(i64 %size)
45+
ret ptr %ptr
46+
}
47+
48+
define internal ptr @my_alloc4(i64 %unchanged, ptr %promoted, ptr %promoted2, i64 %size, i64 %size2) allocsize(3,4) {
49+
; CHECK: Function Attrs: allocsize(5,6)
50+
; CHECK-LABEL: define internal ptr @my_alloc4(
51+
; CHECK-SAME: i64 [[UNCHANGED:%.*]], i32 [[PROMOTED_0_VAL:%.*]], i32 [[PROMOTED_4_VAL:%.*]], i32 [[PROMOTED2_0_VAL:%.*]], i32 [[PROMOTED2_4_VAL:%.*]], i64 [[SIZE:%.*]], i64 [[SIZE2:%.*]]) #[[ATTR3:[0-9]+]] {
52+
; CHECK-NEXT: [[PTR:%.*]] = call ptr @calloc(i64 [[SIZE]], i64 [[SIZE2]])
53+
; CHECK-NEXT: ret ptr [[PTR]]
54+
;
55+
load i32, ptr %promoted
56+
%promoted.gep1 = getelementptr i8, ptr %promoted, i64 4
57+
load i32, ptr %promoted.gep1
58+
59+
load i32, ptr %promoted2
60+
%promoted2.gep1 = getelementptr i8, ptr %promoted2, i64 4
61+
load i32, ptr %promoted2.gep1
62+
63+
%ptr = call ptr @calloc(i64 %size, i64 %size2)
64+
ret ptr %ptr
65+
}
66+
67+
define void @call_my_alloc(ptr %arg, ptr %arg2) {
68+
; CHECK-LABEL: define void @call_my_alloc(
69+
; CHECK-SAME: ptr [[ARG:%.*]], ptr [[ARG2:%.*]]) {
70+
; CHECK-NEXT: [[TMP1:%.*]] = call ptr @my_alloc1(i64 0, i64 2)
71+
; CHECK-NEXT: [[TMP2:%.*]] = call ptr @my_alloc2(i64 0, i64 2, i64 2)
72+
; CHECK-NEXT: [[ARG_VAL:%.*]] = load i32, ptr [[ARG]], align 4
73+
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr i8, ptr [[ARG]], i64 4
74+
; CHECK-NEXT: [[ARG_VAL1:%.*]] = load i32, ptr [[TMP3]], align 4
75+
; CHECK-NEXT: [[ARG2_VAL:%.*]] = load i32, ptr [[ARG2]], align 4
76+
; CHECK-NEXT: [[TMP4:%.*]] = getelementptr i8, ptr [[ARG2]], i64 4
77+
; CHECK-NEXT: [[ARG2_VAL2:%.*]] = load i32, ptr [[TMP4]], align 4
78+
; CHECK-NEXT: [[TMP5:%.*]] = call ptr @my_alloc3(i64 0, i32 [[ARG_VAL]], i32 [[ARG_VAL1]], i32 [[ARG2_VAL]], i32 [[ARG2_VAL2]], i64 2)
79+
; CHECK-NEXT: [[ARG_VAL3:%.*]] = load i32, ptr [[ARG]], align 4
80+
; CHECK-NEXT: [[TMP6:%.*]] = getelementptr i8, ptr [[ARG]], i64 4
81+
; CHECK-NEXT: [[ARG_VAL4:%.*]] = load i32, ptr [[TMP6]], align 4
82+
; CHECK-NEXT: [[ARG2_VAL5:%.*]] = load i32, ptr [[ARG2]], align 4
83+
; CHECK-NEXT: [[TMP7:%.*]] = getelementptr i8, ptr [[ARG2]], i64 4
84+
; CHECK-NEXT: [[ARG2_VAL6:%.*]] = load i32, ptr [[TMP7]], align 4
85+
; CHECK-NEXT: [[TMP8:%.*]] = call ptr @my_alloc4(i64 0, i32 [[ARG_VAL3]], i32 [[ARG_VAL4]], i32 [[ARG2_VAL5]], i32 [[ARG2_VAL6]], i64 2, i64 2)
86+
; CHECK-NEXT: ret void
87+
;
88+
%ptr = call ptr @my_alloc1(i64 0, ptr null, i64 2, ptr null)
89+
%ptr2 = call ptr @my_alloc2(i64 0, ptr null, i64 2, i64 2, ptr null)
90+
%ptr3 = call ptr @my_alloc3(i64 0, ptr %arg, ptr %arg2, i64 2)
91+
%ptr4 = call ptr @my_alloc4(i64 0, ptr %arg, ptr %arg2, i64 2, i64 2)
92+
ret void
93+
}
94+

0 commit comments

Comments
 (0)