1//===----------------------------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8/// \file
9/// This transformation implements the well known scalar replacement of
10/// aggregates transformation but for logical pointers.
11/// It tries to identify promotable elements of an aggregate alloca, and
12/// promote them to multiple allocas of scalar type.
13///
14/// FIXME: nested aggregates are not fully optimized (#192619).
15/// FIXME: array are not optimized (#192620).
16///
17//===----------------------------------------------------------------------===//
18
19#include "llvm/Transforms/Scalar/LogicalSROA.h"
20#include "llvm/ADT/DenseSet.h"
21#include "llvm/ADT/SmallVector.h"
22#include "llvm/IR/IRBuilder.h"
23#include "llvm/IR/IntrinsicInst.h"
24#include "llvm/IR/PassManager.h"
25#include "llvm/Pass.h"
26#include "llvm/Transforms/Scalar.h"
27
28using namespace llvm;
29
30#define DEBUG_TYPE "logical-sroa"
31
32// Return all lifetime intrinsics with the instruction I as operand.
33static SmallVector<LifetimeIntrinsic *>
34collectLifetimeIntrinsicsUsing(Instruction &I) {
35 SmallVector<LifetimeIntrinsic *> Output;
36
37 for (User *U : I.users()) {
38 if (auto *LI = dyn_cast<LifetimeIntrinsic>(Val: U))
39 Output.push_back(Elt: LI);
40 }
41
42 return Output;
43}
44
45// Returns true if all direct and indirect users of the alloca
46// allow the split.
47static bool isAllocaSplittable(StructuredAllocaInst &SAI) {
48 SmallVector<Value *> WorkList(SAI.users());
49 DenseSet<Value *> Visited;
50
51 // Helper function to enqueue all non-visited users of `I`.
52 auto enqueueAllUsers = [&](Instruction *I) {
53 for (auto *U : I->users()) {
54 if (Visited.contains(V: U))
55 continue;
56 WorkList.push_back(Elt: U);
57 }
58 };
59
60 while (!WorkList.empty()) {
61 Instruction *I = dyn_cast<Instruction>(Val: WorkList.back());
62 WorkList.pop_back();
63
64 // User is not an instruction. Not sure what it it, in
65 // doubt, don't split.
66 if (!I)
67 return false;
68
69 Visited.insert(V: I);
70
71 // Those allow the alloca split.
72 if (isa<LifetimeIntrinsic>(Val: I))
73 continue;
74
75 // If we load the whole alloca, we cannot split,
76 // otherwise, we can stop looking into derived users.
77 if (auto *LI = dyn_cast<LoadInst>(Val: I)) {
78 if (LI->getPointerOperand() == &SAI)
79 return false;
80 continue;
81 }
82
83 // If we store to whole alloca, we cannot split,
84 // otherwise, we can stop looking into derived users.
85 if (auto *SI = dyn_cast<StoreInst>(Val: I)) {
86 if (SI->getPointerOperand() == &SAI)
87 return false;
88 continue;
89 }
90
91 // PHI and Select instruction are not inherently preventing
92 // the split, but correctly handling those requires more testing,
93 // so postponing this (See #193749)
94 if (isa<PHINode>(Val: I) || isa<SelectInst>(Val: I))
95 return false;
96
97 if (auto *SGEP = dyn_cast<StructuredGEPInst>(Val: I)) {
98 // If the SGEP has no indices and is still there, this probably means the
99 // ptr is escaping or uses as-is. For now, we bail out.
100 if (SGEP->getNumIndices() == 0)
101 return false;
102
103 enqueueAllUsers(SGEP);
104 continue;
105 }
106
107 // Any other users prevents the split (call, escape, etc).
108 return false;
109 }
110
111 return true;
112}
113
114// Returns a vector with one element for each field of the struct allocated by
115// SAI. Each element is a vector of SGEP instruction referencing this field.
116// This function ignores lifetime intrinsics.
117static SmallVector<SmallVector<StructuredGEPInst *>>
118collectPerFieldSGEP(StructuredAllocaInst &SAI) {
119 StructType *ST = cast<StructType>(Val: SAI.getAllocationType());
120 SmallVector<SmallVector<StructuredGEPInst *>> Output(ST->getNumElements());
121
122 for (User *U : SAI.users()) {
123 if (isa<LifetimeIntrinsic>(Val: U))
124 continue;
125
126 auto *SGEP = cast<StructuredGEPInst>(Val: U);
127
128 // IR rule: SGEP on struct can only use constant int as indices.
129 ConstantInt *Index = cast<ConstantInt>(Val: SGEP->getIndexOperand(Index: 0));
130 assert(Index->getZExtValue() < Output.size());
131 Output[Index->getZExtValue()].push_back(Elt: SGEP);
132 }
133
134 return Output;
135}
136
137// For each lifetime intrinsic in LifetimeIntrinsics, creates a new one, but
138// uses V as operand.
139static void copyLifetimeIntrinsicFor(IRBuilder<> &B, LifetimeIntrinsic *II,
140 Value *V) {
141 B.SetInsertPoint(II);
142
143 if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
144 B.CreateLifetimeStart(Ptr: V);
145 } else if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
146 B.CreateLifetimeEnd(Ptr: V);
147 } else
148 llvm_unreachable("invalid argument: expected a lifetime intrinsic");
149}
150
151static void rewriteSGEPChain(IRBuilder<> &B, StructuredGEPInst *SGEP,
152 StructuredAllocaInst *FieldAlloca) {
153 if (SGEP->getNumIndices() == 1) {
154 SGEP->replaceAllUsesWith(V: FieldAlloca);
155 SGEP->eraseFromParent();
156 return;
157 }
158
159 SmallVector<Value *, 4> Indices(llvm::drop_begin(RangeOrContainer: SGEP->indices()));
160 B.SetInsertPoint(SGEP);
161 auto *I = B.CreateStructuredGEP(BaseType: FieldAlloca->getAllocationType(), PtrBase: FieldAlloca,
162 Indices, Name: SGEP->getName());
163 SGEP->replaceAllUsesWith(V: I);
164 SGEP->eraseFromParent();
165}
166
167static bool runOnStructuredAlloca(StructuredAllocaInst &SAI) {
168 // For now, LogicalSROA only handles SGEP on structs.
169 StructType *ST = dyn_cast<StructType>(Val: SAI.getAllocationType());
170 if (!ST)
171 return false;
172
173 if (!isAllocaSplittable(SAI))
174 return false;
175
176 auto PerFieldSGEP = collectPerFieldSGEP(SAI);
177 assert(PerFieldSGEP.size() == ST->getNumElements());
178
179 auto LifetimeIntrinsics = collectLifetimeIntrinsicsUsing(I&: SAI);
180 IRBuilder B(&SAI);
181 for (const auto &[FieldIndex, Users] : llvm::enumerate(First&: PerFieldSGEP)) {
182 if (Users.empty())
183 continue;
184
185 B.SetInsertPoint(&SAI);
186 auto *FieldAlloca = cast<StructuredAllocaInst>(
187 Val: B.CreateStructuredAlloca(BaseType: ST->getElementType(N: FieldIndex)));
188
189 for (auto II : LifetimeIntrinsics)
190 copyLifetimeIntrinsicFor(B, II, V: FieldAlloca);
191
192 for (StructuredGEPInst *SGEP : Users)
193 rewriteSGEPChain(B, SGEP, FieldAlloca);
194 }
195
196 for (auto *II : LifetimeIntrinsics)
197 II->eraseFromParent();
198 SAI.eraseFromParent();
199 return true;
200}
201
202static bool runLogicalSROA(Function &F) {
203 SmallVector<StructuredAllocaInst *> Worklist;
204 BasicBlock &EntryBB = F.getEntryBlock();
205 for (Instruction &I : EntryBB) {
206 if (StructuredAllocaInst *SAI = dyn_cast<StructuredAllocaInst>(Val: &I))
207 Worklist.push_back(Elt: SAI);
208 }
209
210 bool Changed = false;
211 for (StructuredAllocaInst *SAI : Worklist)
212 Changed |= runOnStructuredAlloca(SAI&: *SAI);
213 return Changed;
214}
215
216PreservedAnalyses LogicalSROAPass::run(Function &F,
217 FunctionAnalysisManager &AM) {
218 if (!runLogicalSROA(F))
219 return PreservedAnalyses::all();
220
221 PreservedAnalyses PA;
222 PA.preserveSet<CFGAnalyses>();
223 return PA;
224}
225
226LogicalSROAPass::LogicalSROAPass() {}
227