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