1 | //===- LocalStackSlotAllocation.cpp - Pre-allocate locals to stack slots --===// |
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 | // |
9 | // This pass assigns local frame indices to stack slots relative to one another |
10 | // and allocates additional base registers to access them when the target |
11 | // estimates they are likely to be out of range of stack pointer and frame |
12 | // pointer relative addressing. |
13 | // |
14 | //===----------------------------------------------------------------------===// |
15 | |
16 | #include "llvm/ADT/SetVector.h" |
17 | #include "llvm/ADT/SmallSet.h" |
18 | #include "llvm/ADT/SmallVector.h" |
19 | #include "llvm/ADT/Statistic.h" |
20 | #include "llvm/CodeGen/MachineBasicBlock.h" |
21 | #include "llvm/CodeGen/MachineFrameInfo.h" |
22 | #include "llvm/CodeGen/MachineFunction.h" |
23 | #include "llvm/CodeGen/MachineFunctionPass.h" |
24 | #include "llvm/CodeGen/MachineInstr.h" |
25 | #include "llvm/CodeGen/MachineOperand.h" |
26 | #include "llvm/CodeGen/TargetFrameLowering.h" |
27 | #include "llvm/CodeGen/TargetOpcodes.h" |
28 | #include "llvm/CodeGen/TargetRegisterInfo.h" |
29 | #include "llvm/CodeGen/TargetSubtargetInfo.h" |
30 | #include "llvm/InitializePasses.h" |
31 | #include "llvm/Pass.h" |
32 | #include "llvm/Support/Debug.h" |
33 | #include "llvm/Support/ErrorHandling.h" |
34 | #include "llvm/Support/raw_ostream.h" |
35 | #include <algorithm> |
36 | #include <cassert> |
37 | #include <cstdint> |
38 | #include <tuple> |
39 | |
40 | using namespace llvm; |
41 | |
42 | #define DEBUG_TYPE "localstackalloc" |
43 | |
44 | STATISTIC(NumAllocations, "Number of frame indices allocated into local block" ); |
45 | STATISTIC(NumBaseRegisters, "Number of virtual frame base registers allocated" ); |
46 | STATISTIC(NumReplacements, "Number of frame indices references replaced" ); |
47 | |
48 | namespace { |
49 | |
50 | class FrameRef { |
51 | MachineBasicBlock::iterator MI; // Instr referencing the frame |
52 | int64_t LocalOffset; // Local offset of the frame idx referenced |
53 | int FrameIdx; // The frame index |
54 | |
55 | // Order reference instruction appears in program. Used to ensure |
56 | // deterministic order when multiple instructions may reference the same |
57 | // location. |
58 | unsigned Order; |
59 | |
60 | public: |
61 | FrameRef(MachineInstr *I, int64_t Offset, int Idx, unsigned Ord) : |
62 | MI(I), LocalOffset(Offset), FrameIdx(Idx), Order(Ord) {} |
63 | |
64 | bool operator<(const FrameRef &RHS) const { |
65 | return std::tie(args: LocalOffset, args: FrameIdx, args: Order) < |
66 | std::tie(args: RHS.LocalOffset, args: RHS.FrameIdx, args: RHS.Order); |
67 | } |
68 | |
69 | MachineBasicBlock::iterator getMachineInstr() const { return MI; } |
70 | int64_t getLocalOffset() const { return LocalOffset; } |
71 | int getFrameIndex() const { return FrameIdx; } |
72 | }; |
73 | |
74 | class LocalStackSlotPass: public MachineFunctionPass { |
75 | SmallVector<int64_t, 16> LocalOffsets; |
76 | |
77 | /// StackObjSet - A set of stack object indexes |
78 | using StackObjSet = SmallSetVector<int, 8>; |
79 | |
80 | void AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx, int64_t &Offset, |
81 | bool StackGrowsDown, Align &MaxAlign); |
82 | void AssignProtectedObjSet(const StackObjSet &UnassignedObjs, |
83 | SmallSet<int, 16> &ProtectedObjs, |
84 | MachineFrameInfo &MFI, bool StackGrowsDown, |
85 | int64_t &Offset, Align &MaxAlign); |
86 | void calculateFrameObjectOffsets(MachineFunction &Fn); |
87 | bool insertFrameReferenceRegisters(MachineFunction &Fn); |
88 | |
89 | public: |
90 | static char ID; // Pass identification, replacement for typeid |
91 | |
92 | explicit LocalStackSlotPass() : MachineFunctionPass(ID) { |
93 | initializeLocalStackSlotPassPass(*PassRegistry::getPassRegistry()); |
94 | } |
95 | |
96 | bool runOnMachineFunction(MachineFunction &MF) override; |
97 | |
98 | void getAnalysisUsage(AnalysisUsage &AU) const override { |
99 | AU.setPreservesCFG(); |
100 | MachineFunctionPass::getAnalysisUsage(AU); |
101 | } |
102 | }; |
103 | |
104 | } // end anonymous namespace |
105 | |
106 | char LocalStackSlotPass::ID = 0; |
107 | |
108 | char &llvm::LocalStackSlotAllocationID = LocalStackSlotPass::ID; |
109 | INITIALIZE_PASS(LocalStackSlotPass, DEBUG_TYPE, |
110 | "Local Stack Slot Allocation" , false, false) |
111 | |
112 | bool LocalStackSlotPass::runOnMachineFunction(MachineFunction &MF) { |
113 | MachineFrameInfo &MFI = MF.getFrameInfo(); |
114 | const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo(); |
115 | unsigned LocalObjectCount = MFI.getObjectIndexEnd(); |
116 | |
117 | // If the target doesn't want/need this pass, or if there are no locals |
118 | // to consider, early exit. |
119 | if (LocalObjectCount == 0 || !TRI->requiresVirtualBaseRegisters(MF)) |
120 | return false; |
121 | |
122 | // Make sure we have enough space to store the local offsets. |
123 | LocalOffsets.resize(N: MFI.getObjectIndexEnd()); |
124 | |
125 | // Lay out the local blob. |
126 | calculateFrameObjectOffsets(Fn&: MF); |
127 | |
128 | // Insert virtual base registers to resolve frame index references. |
129 | bool UsedBaseRegs = insertFrameReferenceRegisters(Fn&: MF); |
130 | |
131 | // Tell MFI whether any base registers were allocated. PEI will only |
132 | // want to use the local block allocations from this pass if there were any. |
133 | // Otherwise, PEI can do a bit better job of getting the alignment right |
134 | // without a hole at the start since it knows the alignment of the stack |
135 | // at the start of local allocation, and this pass doesn't. |
136 | MFI.setUseLocalStackAllocationBlock(UsedBaseRegs); |
137 | |
138 | return true; |
139 | } |
140 | |
141 | /// AdjustStackOffset - Helper function used to adjust the stack frame offset. |
142 | void LocalStackSlotPass::AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx, |
143 | int64_t &Offset, bool StackGrowsDown, |
144 | Align &MaxAlign) { |
145 | // If the stack grows down, add the object size to find the lowest address. |
146 | if (StackGrowsDown) |
147 | Offset += MFI.getObjectSize(ObjectIdx: FrameIdx); |
148 | |
149 | Align Alignment = MFI.getObjectAlign(ObjectIdx: FrameIdx); |
150 | |
151 | // If the alignment of this object is greater than that of the stack, then |
152 | // increase the stack alignment to match. |
153 | MaxAlign = std::max(a: MaxAlign, b: Alignment); |
154 | |
155 | // Adjust to alignment boundary. |
156 | Offset = alignTo(Size: Offset, A: Alignment); |
157 | |
158 | int64_t LocalOffset = StackGrowsDown ? -Offset : Offset; |
159 | LLVM_DEBUG(dbgs() << "Allocate FI(" << FrameIdx << ") to local offset " |
160 | << LocalOffset << "\n" ); |
161 | // Keep the offset available for base register allocation |
162 | LocalOffsets[FrameIdx] = LocalOffset; |
163 | // And tell MFI about it for PEI to use later |
164 | MFI.mapLocalFrameObject(ObjectIndex: FrameIdx, Offset: LocalOffset); |
165 | |
166 | if (!StackGrowsDown) |
167 | Offset += MFI.getObjectSize(ObjectIdx: FrameIdx); |
168 | |
169 | ++NumAllocations; |
170 | } |
171 | |
172 | /// AssignProtectedObjSet - Helper function to assign large stack objects (i.e., |
173 | /// those required to be close to the Stack Protector) to stack offsets. |
174 | void LocalStackSlotPass::AssignProtectedObjSet( |
175 | const StackObjSet &UnassignedObjs, SmallSet<int, 16> &ProtectedObjs, |
176 | MachineFrameInfo &MFI, bool StackGrowsDown, int64_t &Offset, |
177 | Align &MaxAlign) { |
178 | for (int i : UnassignedObjs) { |
179 | AdjustStackOffset(MFI, FrameIdx: i, Offset, StackGrowsDown, MaxAlign); |
180 | ProtectedObjs.insert(V: i); |
181 | } |
182 | } |
183 | |
184 | /// calculateFrameObjectOffsets - Calculate actual frame offsets for all of the |
185 | /// abstract stack objects. |
186 | void LocalStackSlotPass::calculateFrameObjectOffsets(MachineFunction &Fn) { |
187 | // Loop over all of the stack objects, assigning sequential addresses... |
188 | MachineFrameInfo &MFI = Fn.getFrameInfo(); |
189 | const TargetFrameLowering &TFI = *Fn.getSubtarget().getFrameLowering(); |
190 | bool StackGrowsDown = |
191 | TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown; |
192 | int64_t Offset = 0; |
193 | Align MaxAlign; |
194 | |
195 | // Make sure that the stack protector comes before the local variables on the |
196 | // stack. |
197 | SmallSet<int, 16> ProtectedObjs; |
198 | if (MFI.hasStackProtectorIndex()) { |
199 | int StackProtectorFI = MFI.getStackProtectorIndex(); |
200 | |
201 | // We need to make sure we didn't pre-allocate the stack protector when |
202 | // doing this. |
203 | // If we already have a stack protector, this will re-assign it to a slot |
204 | // that is **not** covering the protected objects. |
205 | assert(!MFI.isObjectPreAllocated(StackProtectorFI) && |
206 | "Stack protector pre-allocated in LocalStackSlotAllocation" ); |
207 | |
208 | StackObjSet LargeArrayObjs; |
209 | StackObjSet SmallArrayObjs; |
210 | StackObjSet AddrOfObjs; |
211 | |
212 | // Only place the stack protector in the local stack area if the target |
213 | // allows it. |
214 | if (TFI.isStackIdSafeForLocalArea(StackId: MFI.getStackID(ObjectIdx: StackProtectorFI))) |
215 | AdjustStackOffset(MFI, FrameIdx: StackProtectorFI, Offset, StackGrowsDown, |
216 | MaxAlign); |
217 | |
218 | // Assign large stack objects first. |
219 | for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) { |
220 | if (MFI.isDeadObjectIndex(ObjectIdx: i)) |
221 | continue; |
222 | if (StackProtectorFI == (int)i) |
223 | continue; |
224 | if (!TFI.isStackIdSafeForLocalArea(StackId: MFI.getStackID(ObjectIdx: i))) |
225 | continue; |
226 | |
227 | switch (MFI.getObjectSSPLayout(ObjectIdx: i)) { |
228 | case MachineFrameInfo::SSPLK_None: |
229 | continue; |
230 | case MachineFrameInfo::SSPLK_SmallArray: |
231 | SmallArrayObjs.insert(X: i); |
232 | continue; |
233 | case MachineFrameInfo::SSPLK_AddrOf: |
234 | AddrOfObjs.insert(X: i); |
235 | continue; |
236 | case MachineFrameInfo::SSPLK_LargeArray: |
237 | LargeArrayObjs.insert(X: i); |
238 | continue; |
239 | } |
240 | llvm_unreachable("Unexpected SSPLayoutKind." ); |
241 | } |
242 | |
243 | AssignProtectedObjSet(UnassignedObjs: LargeArrayObjs, ProtectedObjs, MFI, StackGrowsDown, |
244 | Offset, MaxAlign); |
245 | AssignProtectedObjSet(UnassignedObjs: SmallArrayObjs, ProtectedObjs, MFI, StackGrowsDown, |
246 | Offset, MaxAlign); |
247 | AssignProtectedObjSet(UnassignedObjs: AddrOfObjs, ProtectedObjs, MFI, StackGrowsDown, |
248 | Offset, MaxAlign); |
249 | } |
250 | |
251 | // Then assign frame offsets to stack objects that are not used to spill |
252 | // callee saved registers. |
253 | for (unsigned i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) { |
254 | if (MFI.isDeadObjectIndex(ObjectIdx: i)) |
255 | continue; |
256 | if (MFI.getStackProtectorIndex() == (int)i) |
257 | continue; |
258 | if (ProtectedObjs.count(V: i)) |
259 | continue; |
260 | if (!TFI.isStackIdSafeForLocalArea(StackId: MFI.getStackID(ObjectIdx: i))) |
261 | continue; |
262 | |
263 | AdjustStackOffset(MFI, FrameIdx: i, Offset, StackGrowsDown, MaxAlign); |
264 | } |
265 | |
266 | // Remember how big this blob of stack space is |
267 | MFI.setLocalFrameSize(Offset); |
268 | MFI.setLocalFrameMaxAlign(MaxAlign); |
269 | } |
270 | |
271 | static inline bool |
272 | lookupCandidateBaseReg(unsigned BaseReg, |
273 | int64_t BaseOffset, |
274 | int64_t FrameSizeAdjust, |
275 | int64_t LocalFrameOffset, |
276 | const MachineInstr &MI, |
277 | const TargetRegisterInfo *TRI) { |
278 | // Check if the relative offset from the where the base register references |
279 | // to the target address is in range for the instruction. |
280 | int64_t Offset = FrameSizeAdjust + LocalFrameOffset - BaseOffset; |
281 | return TRI->isFrameOffsetLegal(MI: &MI, BaseReg, Offset); |
282 | } |
283 | |
284 | bool LocalStackSlotPass::insertFrameReferenceRegisters(MachineFunction &Fn) { |
285 | // Scan the function's instructions looking for frame index references. |
286 | // For each, ask the target if it wants a virtual base register for it |
287 | // based on what we can tell it about where the local will end up in the |
288 | // stack frame. If it wants one, re-use a suitable one we've previously |
289 | // allocated, or if there isn't one that fits the bill, allocate a new one |
290 | // and ask the target to create a defining instruction for it. |
291 | |
292 | MachineFrameInfo &MFI = Fn.getFrameInfo(); |
293 | const TargetRegisterInfo *TRI = Fn.getSubtarget().getRegisterInfo(); |
294 | const TargetFrameLowering &TFI = *Fn.getSubtarget().getFrameLowering(); |
295 | bool StackGrowsDown = |
296 | TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsDown; |
297 | |
298 | // Collect all of the instructions in the block that reference |
299 | // a frame index. Also store the frame index referenced to ease later |
300 | // lookup. (For any insn that has more than one FI reference, we arbitrarily |
301 | // choose the first one). |
302 | SmallVector<FrameRef, 64> FrameReferenceInsns; |
303 | |
304 | unsigned Order = 0; |
305 | |
306 | for (MachineBasicBlock &BB : Fn) { |
307 | for (MachineInstr &MI : BB) { |
308 | // Debug value, stackmap and patchpoint instructions can't be out of |
309 | // range, so they don't need any updates. |
310 | if (MI.isDebugInstr() || MI.getOpcode() == TargetOpcode::STATEPOINT || |
311 | MI.getOpcode() == TargetOpcode::STACKMAP || |
312 | MI.getOpcode() == TargetOpcode::PATCHPOINT) |
313 | continue; |
314 | |
315 | // For now, allocate the base register(s) within the basic block |
316 | // where they're used, and don't try to keep them around outside |
317 | // of that. It may be beneficial to try sharing them more broadly |
318 | // than that, but the increased register pressure makes that a |
319 | // tricky thing to balance. Investigate if re-materializing these |
320 | // becomes an issue. |
321 | for (const MachineOperand &MO : MI.operands()) { |
322 | // Consider replacing all frame index operands that reference |
323 | // an object allocated in the local block. |
324 | if (MO.isFI()) { |
325 | // Don't try this with values not in the local block. |
326 | if (!MFI.isObjectPreAllocated(ObjectIdx: MO.getIndex())) |
327 | break; |
328 | int Idx = MO.getIndex(); |
329 | int64_t LocalOffset = LocalOffsets[Idx]; |
330 | if (!TRI->needsFrameBaseReg(MI: &MI, Offset: LocalOffset)) |
331 | break; |
332 | FrameReferenceInsns.push_back(Elt: FrameRef(&MI, LocalOffset, Idx, Order++)); |
333 | break; |
334 | } |
335 | } |
336 | } |
337 | } |
338 | |
339 | // Sort the frame references by local offset. |
340 | // Use frame index as a tie-breaker in case MI's have the same offset. |
341 | llvm::sort(C&: FrameReferenceInsns); |
342 | |
343 | MachineBasicBlock *Entry = &Fn.front(); |
344 | |
345 | Register BaseReg; |
346 | int64_t BaseOffset = 0; |
347 | |
348 | // Loop through the frame references and allocate for them as necessary. |
349 | for (int ref = 0, e = FrameReferenceInsns.size(); ref < e ; ++ref) { |
350 | FrameRef &FR = FrameReferenceInsns[ref]; |
351 | MachineInstr &MI = *FR.getMachineInstr(); |
352 | int64_t LocalOffset = FR.getLocalOffset(); |
353 | int FrameIdx = FR.getFrameIndex(); |
354 | assert(MFI.isObjectPreAllocated(FrameIdx) && |
355 | "Only pre-allocated locals expected!" ); |
356 | |
357 | // We need to keep the references to the stack protector slot through frame |
358 | // index operands so that it gets resolved by PEI rather than this pass. |
359 | // This avoids accesses to the stack protector though virtual base |
360 | // registers, and forces PEI to address it using fp/sp/bp. |
361 | if (MFI.hasStackProtectorIndex() && |
362 | FrameIdx == MFI.getStackProtectorIndex()) |
363 | continue; |
364 | |
365 | LLVM_DEBUG(dbgs() << "Considering: " << MI); |
366 | |
367 | unsigned idx = 0; |
368 | for (unsigned f = MI.getNumOperands(); idx != f; ++idx) { |
369 | if (!MI.getOperand(i: idx).isFI()) |
370 | continue; |
371 | |
372 | if (FrameIdx == MI.getOperand(i: idx).getIndex()) |
373 | break; |
374 | } |
375 | |
376 | assert(idx < MI.getNumOperands() && "Cannot find FI operand" ); |
377 | |
378 | int64_t Offset = 0; |
379 | int64_t FrameSizeAdjust = StackGrowsDown ? MFI.getLocalFrameSize() : 0; |
380 | |
381 | LLVM_DEBUG(dbgs() << " Replacing FI in: " << MI); |
382 | |
383 | // If we have a suitable base register available, use it; otherwise |
384 | // create a new one. Note that any offset encoded in the |
385 | // instruction itself will be taken into account by the target, |
386 | // so we don't have to adjust for it here when reusing a base |
387 | // register. |
388 | if (BaseReg.isValid() && |
389 | lookupCandidateBaseReg(BaseReg, BaseOffset, FrameSizeAdjust, |
390 | LocalFrameOffset: LocalOffset, MI, TRI)) { |
391 | LLVM_DEBUG(dbgs() << " Reusing base register " << BaseReg << "\n" ); |
392 | // We found a register to reuse. |
393 | Offset = FrameSizeAdjust + LocalOffset - BaseOffset; |
394 | } else { |
395 | // No previously defined register was in range, so create a new one. |
396 | int64_t InstrOffset = TRI->getFrameIndexInstrOffset(MI: &MI, Idx: idx); |
397 | |
398 | int64_t CandBaseOffset = FrameSizeAdjust + LocalOffset + InstrOffset; |
399 | |
400 | // We'd like to avoid creating single-use virtual base registers. |
401 | // Because the FrameRefs are in sorted order, and we've already |
402 | // processed all FrameRefs before this one, just check whether or not |
403 | // the next FrameRef will be able to reuse this new register. If not, |
404 | // then don't bother creating it. |
405 | if (ref + 1 >= e || |
406 | !lookupCandidateBaseReg( |
407 | BaseReg, BaseOffset: CandBaseOffset, FrameSizeAdjust, |
408 | LocalFrameOffset: FrameReferenceInsns[ref + 1].getLocalOffset(), |
409 | MI: *FrameReferenceInsns[ref + 1].getMachineInstr(), TRI)) |
410 | continue; |
411 | |
412 | // Save the base offset. |
413 | BaseOffset = CandBaseOffset; |
414 | |
415 | // Tell the target to insert the instruction to initialize |
416 | // the base register. |
417 | // MachineBasicBlock::iterator InsertionPt = Entry->begin(); |
418 | BaseReg = TRI->materializeFrameBaseRegister(MBB: Entry, FrameIdx, Offset: InstrOffset); |
419 | |
420 | LLVM_DEBUG(dbgs() << " Materialized base register at frame local offset " |
421 | << LocalOffset + InstrOffset |
422 | << " into " << printReg(BaseReg, TRI) << '\n'); |
423 | |
424 | // The base register already includes any offset specified |
425 | // by the instruction, so account for that so it doesn't get |
426 | // applied twice. |
427 | Offset = -InstrOffset; |
428 | |
429 | ++NumBaseRegisters; |
430 | } |
431 | assert(BaseReg && "Unable to allocate virtual base register!" ); |
432 | |
433 | // Modify the instruction to use the new base register rather |
434 | // than the frame index operand. |
435 | TRI->resolveFrameIndex(MI, BaseReg, Offset); |
436 | LLVM_DEBUG(dbgs() << "Resolved: " << MI); |
437 | |
438 | ++NumReplacements; |
439 | } |
440 | |
441 | return BaseReg.isValid(); |
442 | } |
443 | |