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