1//===-- AArch64StackTaggingPreRA.cpp --- Stack Tagging for AArch64 -----===//
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#include "AArch64.h"
10#include "AArch64InstrInfo.h"
11#include "AArch64MachineFunctionInfo.h"
12#include "llvm/ADT/SetVector.h"
13#include "llvm/ADT/Statistic.h"
14#include "llvm/CodeGen/MachineFrameInfo.h"
15#include "llvm/CodeGen/MachineFunction.h"
16#include "llvm/CodeGen/MachineFunctionPass.h"
17#include "llvm/CodeGen/MachineInstrBuilder.h"
18#include "llvm/CodeGen/MachineRegisterInfo.h"
19#include "llvm/CodeGen/MachineTraceMetrics.h"
20#include "llvm/CodeGen/Passes.h"
21#include "llvm/CodeGen/TargetInstrInfo.h"
22#include "llvm/CodeGen/TargetRegisterInfo.h"
23#include "llvm/CodeGen/TargetSubtargetInfo.h"
24#include "llvm/Support/Debug.h"
25#include "llvm/Support/raw_ostream.h"
26
27using namespace llvm;
28
29#define DEBUG_TYPE "aarch64-stack-tagging-pre-ra"
30
31namespace {
32
33class AArch64StackTaggingPreRAImpl {
34 MachineFunction *MF;
35 AArch64FunctionInfo *AFI;
36 MachineFrameInfo *MFI;
37 MachineRegisterInfo *MRI;
38 const AArch64RegisterInfo *TRI;
39 const AArch64InstrInfo *TII;
40
41 SmallVector<MachineInstr*, 16> ReTags;
42
43public:
44 bool run(MachineFunction &Func);
45
46private:
47 bool mayUseUncheckedLoadStore();
48 void uncheckUsesOf(unsigned TaggedReg, int FI);
49 void uncheckLoadsAndStores();
50 std::optional<int> findFirstSlotCandidate();
51};
52
53class AArch64StackTaggingPreRALegacy : public MachineFunctionPass {
54public:
55 static char ID;
56 AArch64StackTaggingPreRALegacy() : MachineFunctionPass(ID) {}
57
58 bool runOnMachineFunction(MachineFunction &MF) override {
59 if (skipFunction(F: MF.getFunction()))
60 return false;
61 return AArch64StackTaggingPreRAImpl().run(Func&: MF);
62 }
63
64 StringRef getPassName() const override {
65 return "AArch64 Stack Tagging PreRA";
66 }
67
68 void getAnalysisUsage(AnalysisUsage &AU) const override {
69 AU.setPreservesCFG();
70 MachineFunctionPass::getAnalysisUsage(AU);
71 }
72};
73} // end anonymous namespace
74
75char AArch64StackTaggingPreRALegacy::ID = 0;
76
77INITIALIZE_PASS_BEGIN(AArch64StackTaggingPreRALegacy,
78 "aarch64-stack-tagging-pre-ra",
79 "AArch64 Stack Tagging PreRA Pass", false, false)
80INITIALIZE_PASS_END(AArch64StackTaggingPreRALegacy,
81 "aarch64-stack-tagging-pre-ra",
82 "AArch64 Stack Tagging PreRA Pass", false, false)
83
84FunctionPass *llvm::createAArch64StackTaggingPreRALegacyPass() {
85 return new AArch64StackTaggingPreRALegacy();
86}
87
88PreservedAnalyses
89AArch64StackTaggingPreRAPass::run(MachineFunction &MF,
90 MachineFunctionAnalysisManager &MFAM) {
91 if (AArch64StackTaggingPreRAImpl().run(Func&: MF)) {
92 PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
93 PA.preserveSet<CFGAnalyses>();
94 return PA;
95 }
96 return PreservedAnalyses::all();
97}
98
99static bool isUncheckedLoadOrStoreOpcode(unsigned Opcode) {
100 switch (Opcode) {
101 case AArch64::LDRBBui:
102 case AArch64::LDRHHui:
103 case AArch64::LDRWui:
104 case AArch64::LDRXui:
105
106 case AArch64::LDRBui:
107 case AArch64::LDRHui:
108 case AArch64::LDRSui:
109 case AArch64::LDRDui:
110 case AArch64::LDRQui:
111
112 case AArch64::LDRSHWui:
113 case AArch64::LDRSHXui:
114
115 case AArch64::LDRSBWui:
116 case AArch64::LDRSBXui:
117
118 case AArch64::LDRSWui:
119
120 case AArch64::STRBBui:
121 case AArch64::STRHHui:
122 case AArch64::STRWui:
123 case AArch64::STRXui:
124
125 case AArch64::STRBui:
126 case AArch64::STRHui:
127 case AArch64::STRSui:
128 case AArch64::STRDui:
129 case AArch64::STRQui:
130
131 case AArch64::LDPWi:
132 case AArch64::LDPXi:
133 case AArch64::LDPSi:
134 case AArch64::LDPDi:
135 case AArch64::LDPQi:
136
137 case AArch64::LDPSWi:
138
139 case AArch64::STPWi:
140 case AArch64::STPXi:
141 case AArch64::STPSi:
142 case AArch64::STPDi:
143 case AArch64::STPQi:
144 return true;
145 default:
146 return false;
147 }
148}
149
150bool AArch64StackTaggingPreRAImpl::mayUseUncheckedLoadStore() {
151 AArch64::UncheckedLdStMode Mode = MF->getSubtarget<AArch64Subtarget>()
152 .getCLOpts()
153 .stack_tagging_unchecked_ld_st;
154 if (Mode == AArch64::UncheckedLdStMode::Never)
155 return false;
156 else if (Mode == AArch64::UncheckedLdStMode::Always)
157 return true;
158
159 // This estimate can be improved if we had harder guarantees about stack frame
160 // layout. With LocalStackAllocation we can estimate SP offset to any
161 // preallocated slot. AArch64FrameLowering::orderFrameObjects could put tagged
162 // objects ahead of non-tagged ones, but that's not always desirable.
163 //
164 // Underestimating SP offset here may require the use of LDG to materialize
165 // the tagged address of the stack slot, along with a scratch register
166 // allocation (post-regalloc!).
167 //
168 // For now we do the safe thing here and require that the entire stack frame
169 // is within range of the shortest of the unchecked instructions.
170 unsigned FrameSize = 0;
171 for (unsigned i = 0, e = MFI->getObjectIndexEnd(); i != e; ++i)
172 FrameSize += MFI->getObjectSize(ObjectIdx: i);
173 bool EntireFrameReachableFromSP = FrameSize < 0xf00;
174 return !MFI->hasVarSizedObjects() && EntireFrameReachableFromSP;
175}
176
177void AArch64StackTaggingPreRAImpl::uncheckUsesOf(unsigned TaggedReg, int FI) {
178 for (MachineInstr &UseI :
179 llvm::make_early_inc_range(Range: MRI->use_instructions(Reg: TaggedReg))) {
180 if (isUncheckedLoadOrStoreOpcode(Opcode: UseI.getOpcode())) {
181 // FI operand is always the one before the immediate offset.
182 unsigned OpIdx = TII->getLoadStoreImmIdx(Opc: UseI.getOpcode()) - 1;
183 if (UseI.getOperand(i: OpIdx).isReg() &&
184 UseI.getOperand(i: OpIdx).getReg() == TaggedReg) {
185 UseI.getOperand(i: OpIdx).ChangeToFrameIndex(Idx: FI);
186 UseI.getOperand(i: OpIdx).setTargetFlags(AArch64II::MO_TAGGED);
187 }
188 } else if (UseI.isCopy() && UseI.getOperand(i: 0).getReg().isVirtual()) {
189 uncheckUsesOf(TaggedReg: UseI.getOperand(i: 0).getReg(), FI);
190 }
191 }
192}
193
194void AArch64StackTaggingPreRAImpl::uncheckLoadsAndStores() {
195 for (auto *I : ReTags) {
196 Register TaggedReg = I->getOperand(i: 0).getReg();
197 int FI = I->getOperand(i: 1).getIndex();
198 uncheckUsesOf(TaggedReg, FI);
199 }
200}
201
202namespace {
203struct SlotWithTag {
204 int FI;
205 int Tag;
206 SlotWithTag(int FI, int Tag) : FI(FI), Tag(Tag) {}
207 explicit SlotWithTag(const MachineInstr &MI)
208 : FI(MI.getOperand(i: 1).getIndex()), Tag(MI.getOperand(i: 4).getImm()) {}
209 bool operator==(const SlotWithTag &Other) const {
210 return FI == Other.FI && Tag == Other.Tag;
211 }
212};
213} // namespace
214
215namespace llvm {
216template <> struct DenseMapInfo<SlotWithTag> {
217 static unsigned getHashValue(const SlotWithTag &V) {
218 return hash_combine(args: DenseMapInfo<int>::getHashValue(Val: V.FI),
219 args: DenseMapInfo<int>::getHashValue(Val: V.Tag));
220 }
221 static bool isEqual(const SlotWithTag &A, const SlotWithTag &B) {
222 return A == B;
223 }
224};
225} // namespace llvm
226
227static bool isSlotPreAllocated(MachineFrameInfo *MFI, int FI) {
228 return MFI->getUseLocalStackAllocationBlock() &&
229 MFI->isObjectPreAllocated(ObjectIdx: FI);
230}
231
232// Pin one of the tagged slots to offset 0 from the tagged base pointer.
233// This would make its address available in a virtual register (IRG's def), as
234// opposed to requiring an ADDG instruction to materialize. This effectively
235// eliminates a vreg (by replacing it with direct uses of IRG, which is usually
236// live almost everywhere anyway), and therefore needs to happen before
237// regalloc.
238std::optional<int> AArch64StackTaggingPreRAImpl::findFirstSlotCandidate() {
239 // Find the best (FI, Tag) pair to pin to offset 0.
240 // Looking at the possible uses of a tagged address, the advantage of pinning
241 // is:
242 // - COPY to physical register.
243 // Does not matter, this would trade a MOV instruction for an ADDG.
244 // - ST*G matter, but those mostly appear near the function prologue where all
245 // the tagged addresses need to be materialized anyway; also, counting ST*G
246 // uses would overweight large allocas that require more than one ST*G
247 // instruction.
248 // - Load/Store instructions in the address operand do not require a tagged
249 // pointer, so they also do not benefit. These operands have already been
250 // eliminated (see uncheckLoadsAndStores) so all remaining load/store
251 // instructions count.
252 // - Any other instruction may benefit from being pinned to offset 0.
253 LLVM_DEBUG(
254 dbgs() << "AArch64StackTaggingPreRAImpl::findFirstSlotCandidate\n");
255 if (!MF->getSubtarget<AArch64Subtarget>()
256 .getCLOpts()
257 .stack_tagging_first_slot_opt)
258 return std::nullopt;
259
260 DenseMap<SlotWithTag, int> RetagScore;
261 SlotWithTag MaxScoreST{-1, -1};
262 int MaxScore = -1;
263 for (auto *I : ReTags) {
264 SlotWithTag ST{*I};
265 if (isSlotPreAllocated(MFI, FI: ST.FI))
266 continue;
267
268 Register RetagReg = I->getOperand(i: 0).getReg();
269 if (!RetagReg.isVirtual())
270 continue;
271
272 int Score = 0;
273 SmallVector<Register, 8> WorkList;
274 WorkList.push_back(Elt: RetagReg);
275
276 while (!WorkList.empty()) {
277 Register UseReg = WorkList.pop_back_val();
278 for (auto &UseI : MRI->use_instructions(Reg: UseReg)) {
279 unsigned Opcode = UseI.getOpcode();
280 if (Opcode == AArch64::STGi || Opcode == AArch64::ST2Gi ||
281 Opcode == AArch64::STZGi || Opcode == AArch64::STZ2Gi ||
282 Opcode == AArch64::STGPi || Opcode == AArch64::STGloop ||
283 Opcode == AArch64::STZGloop || Opcode == AArch64::STGloop_wback ||
284 Opcode == AArch64::STZGloop_wback)
285 continue;
286 if (UseI.isCopy()) {
287 Register DstReg = UseI.getOperand(i: 0).getReg();
288 if (DstReg.isVirtual())
289 WorkList.push_back(Elt: DstReg);
290 continue;
291 }
292 LLVM_DEBUG(dbgs() << "[" << ST.FI << ":" << ST.Tag << "] use of "
293 << printReg(UseReg) << " in " << UseI << "\n");
294 Score++;
295 }
296 }
297
298 int TotalScore = RetagScore[ST] += Score;
299 if (TotalScore > MaxScore ||
300 (TotalScore == MaxScore && ST.FI > MaxScoreST.FI)) {
301 MaxScore = TotalScore;
302 MaxScoreST = ST;
303 }
304 }
305
306 if (MaxScoreST.FI < 0)
307 return std::nullopt;
308
309 // If FI's tag is already 0, we are done.
310 if (MaxScoreST.Tag == 0)
311 return MaxScoreST.FI;
312
313 // Otherwise, find a random victim pair (FI, Tag) where Tag == 0.
314 SlotWithTag SwapST{-1, -1};
315 for (auto *I : ReTags) {
316 SlotWithTag ST{*I};
317 if (ST.Tag == 0) {
318 SwapST = ST;
319 break;
320 }
321 }
322
323 // Swap tags between the victim and the highest scoring pair.
324 // If SwapWith is still (-1, -1), that's fine, too - we'll simply take tag for
325 // the highest score slot without changing anything else.
326 for (auto *&I : ReTags) {
327 SlotWithTag ST{*I};
328 MachineOperand &TagOp = I->getOperand(i: 4);
329 if (ST == MaxScoreST) {
330 TagOp.setImm(0);
331 } else if (ST == SwapST) {
332 TagOp.setImm(MaxScoreST.Tag);
333 }
334 }
335 return MaxScoreST.FI;
336}
337
338bool AArch64StackTaggingPreRAImpl::run(MachineFunction &Func) {
339 MF = &Func;
340 MRI = &MF->getRegInfo();
341 AFI = MF->getInfo<AArch64FunctionInfo>();
342 TII = static_cast<const AArch64InstrInfo *>(MF->getSubtarget().getInstrInfo());
343 TRI = static_cast<const AArch64RegisterInfo *>(
344 MF->getSubtarget().getRegisterInfo());
345 MFI = &MF->getFrameInfo();
346 ReTags.clear();
347
348 assert(MRI->isSSA());
349
350 LLVM_DEBUG(dbgs() << "********** AArch64 Stack Tagging PreRA **********\n"
351 << "********** Function: " << MF->getName() << '\n');
352
353 SmallSetVector<int, 8> TaggedSlots;
354 for (auto &BB : *MF) {
355 for (auto &I : BB) {
356 if (I.getOpcode() == AArch64::TAGPstack) {
357 ReTags.push_back(Elt: &I);
358 int FI = I.getOperand(i: 1).getIndex();
359 TaggedSlots.insert(X: FI);
360 // There should be no offsets in TAGP yet.
361 assert(I.getOperand(2).getImm() == 0);
362 }
363 }
364 }
365
366 // Take over from SSP. It does nothing for tagged slots, and should not really
367 // have been enabled in the first place.
368 for (int FI : TaggedSlots)
369 MFI->setObjectSSPLayout(ObjectIdx: FI, Kind: MachineFrameInfo::SSPLK_None);
370
371 if (ReTags.empty())
372 return false;
373
374 if (mayUseUncheckedLoadStore())
375 uncheckLoadsAndStores();
376
377 // Find a slot that is used with zero tag offset, like ADDG #fi, 0.
378 // If the base tagged pointer is set up to the address of this slot,
379 // the ADDG instruction can be eliminated.
380 std::optional<int> BaseSlot = findFirstSlotCandidate();
381 if (BaseSlot)
382 AFI->setTaggedBasePointerIndex(*BaseSlot);
383
384 for (auto *I : ReTags) {
385 int FI = I->getOperand(i: 1).getIndex();
386 int Tag = I->getOperand(i: 4).getImm();
387 Register Base = I->getOperand(i: 3).getReg();
388 if (Tag == 0 && FI == BaseSlot) {
389 BuildMI(BB&: *I->getParent(), I, MIMD: {}, MCID: TII->get(Opcode: AArch64::COPY),
390 DestReg: I->getOperand(i: 0).getReg())
391 .addReg(RegNo: Base);
392 I->eraseFromParent();
393 }
394 }
395
396 return true;
397}
398