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