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