1 //===----- X86DynAllocaExpander.cpp - Expand DynAlloca pseudo instruction -===// 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 file defines a pass that expands DynAlloca pseudo-instructions. 10 // 11 // It performs a conservative analysis to determine whether each allocation 12 // falls within a region of the stack that is safe to use, or whether stack 13 // probes must be emitted. 14 // 15 //===----------------------------------------------------------------------===// 16 17 #include "X86.h" 18 #include "X86InstrBuilder.h" 19 #include "X86InstrInfo.h" 20 #include "X86MachineFunctionInfo.h" 21 #include "X86Subtarget.h" 22 #include "llvm/ADT/MapVector.h" 23 #include "llvm/ADT/PostOrderIterator.h" 24 #include "llvm/CodeGen/MachineFunctionPass.h" 25 #include "llvm/CodeGen/MachineInstrBuilder.h" 26 #include "llvm/CodeGen/MachineRegisterInfo.h" 27 #include "llvm/CodeGen/Passes.h" 28 #include "llvm/CodeGen/TargetInstrInfo.h" 29 #include "llvm/IR/Function.h" 30 #include "llvm/Support/raw_ostream.h" 31 32 using namespace llvm; 33 34 namespace { 35 36 class X86DynAllocaExpander : public MachineFunctionPass { 37 public: 38 X86DynAllocaExpander() : MachineFunctionPass(ID) {} 39 40 bool runOnMachineFunction(MachineFunction &MF) override; 41 42 private: 43 /// Strategies for lowering a DynAlloca. 44 enum Lowering { TouchAndSub, Sub, Probe }; 45 46 /// Deterministic-order map from DynAlloca instruction to desired lowering. 47 typedef MapVector<MachineInstr*, Lowering> LoweringMap; 48 49 /// Compute which lowering to use for each DynAlloca instruction. 50 void computeLowerings(MachineFunction &MF, LoweringMap& Lowerings); 51 52 /// Get the appropriate lowering based on current offset and amount. 53 Lowering getLowering(int64_t CurrentOffset, int64_t AllocaAmount); 54 55 /// Lower a DynAlloca instruction. 56 void lower(MachineInstr* MI, Lowering L); 57 58 MachineRegisterInfo *MRI = nullptr; 59 const X86Subtarget *STI = nullptr; 60 const TargetInstrInfo *TII = nullptr; 61 const X86RegisterInfo *TRI = nullptr; 62 unsigned StackPtr = 0; 63 unsigned SlotSize = 0; 64 int64_t StackProbeSize = 0; 65 bool NoStackArgProbe = false; 66 67 StringRef getPassName() const override { return "X86 DynAlloca Expander"; } 68 static char ID; 69 }; 70 71 char X86DynAllocaExpander::ID = 0; 72 73 } // end anonymous namespace 74 75 FunctionPass *llvm::createX86DynAllocaExpander() { 76 return new X86DynAllocaExpander(); 77 } 78 79 /// Return the allocation amount for a DynAlloca instruction, or -1 if unknown. 80 static int64_t getDynAllocaAmount(MachineInstr *MI, MachineRegisterInfo *MRI) { 81 assert(MI->getOpcode() == X86::DYN_ALLOCA_32 || 82 MI->getOpcode() == X86::DYN_ALLOCA_64); 83 assert(MI->getOperand(0).isReg()); 84 85 Register AmountReg = MI->getOperand(0).getReg(); 86 MachineInstr *Def = MRI->getUniqueVRegDef(AmountReg); 87 88 if (!Def || 89 (Def->getOpcode() != X86::MOV32ri && Def->getOpcode() != X86::MOV64ri) || 90 !Def->getOperand(1).isImm()) 91 return -1; 92 93 return Def->getOperand(1).getImm(); 94 } 95 96 X86DynAllocaExpander::Lowering 97 X86DynAllocaExpander::getLowering(int64_t CurrentOffset, 98 int64_t AllocaAmount) { 99 // For a non-constant amount or a large amount, we have to probe. 100 if (AllocaAmount < 0 || AllocaAmount > StackProbeSize) 101 return Probe; 102 103 // If it fits within the safe region of the stack, just subtract. 104 if (CurrentOffset + AllocaAmount <= StackProbeSize) 105 return Sub; 106 107 // Otherwise, touch the current tip of the stack, then subtract. 108 return TouchAndSub; 109 } 110 111 static bool isPushPop(const MachineInstr &MI) { 112 switch (MI.getOpcode()) { 113 case X86::PUSH32i8: 114 case X86::PUSH32r: 115 case X86::PUSH32rmm: 116 case X86::PUSH32rmr: 117 case X86::PUSHi32: 118 case X86::PUSH64i8: 119 case X86::PUSH64r: 120 case X86::PUSH64rmm: 121 case X86::PUSH64rmr: 122 case X86::PUSH64i32: 123 case X86::POP32r: 124 case X86::POP64r: 125 return true; 126 default: 127 return false; 128 } 129 } 130 131 void X86DynAllocaExpander::computeLowerings(MachineFunction &MF, 132 LoweringMap &Lowerings) { 133 // Do a one-pass reverse post-order walk of the CFG to conservatively estimate 134 // the offset between the stack pointer and the lowest touched part of the 135 // stack, and use that to decide how to lower each DynAlloca instruction. 136 137 // Initialize OutOffset[B], the stack offset at exit from B, to something big. 138 DenseMap<MachineBasicBlock *, int64_t> OutOffset; 139 for (MachineBasicBlock &MBB : MF) 140 OutOffset[&MBB] = INT32_MAX; 141 142 // Note: we don't know the offset at the start of the entry block since the 143 // prologue hasn't been inserted yet, and how much that will adjust the stack 144 // pointer depends on register spills, which have not been computed yet. 145 146 // Compute the reverse post-order. 147 ReversePostOrderTraversal<MachineFunction*> RPO(&MF); 148 149 for (MachineBasicBlock *MBB : RPO) { 150 int64_t Offset = -1; 151 for (MachineBasicBlock *Pred : MBB->predecessors()) 152 Offset = std::max(Offset, OutOffset[Pred]); 153 if (Offset == -1) Offset = INT32_MAX; 154 155 for (MachineInstr &MI : *MBB) { 156 if (MI.getOpcode() == X86::DYN_ALLOCA_32 || 157 MI.getOpcode() == X86::DYN_ALLOCA_64) { 158 // A DynAlloca moves StackPtr, and potentially touches it. 159 int64_t Amount = getDynAllocaAmount(&MI, MRI); 160 Lowering L = getLowering(Offset, Amount); 161 Lowerings[&MI] = L; 162 switch (L) { 163 case Sub: 164 Offset += Amount; 165 break; 166 case TouchAndSub: 167 Offset = Amount; 168 break; 169 case Probe: 170 Offset = 0; 171 break; 172 } 173 } else if (MI.isCall() || isPushPop(MI)) { 174 // Calls, pushes and pops touch the tip of the stack. 175 Offset = 0; 176 } else if (MI.getOpcode() == X86::ADJCALLSTACKUP32 || 177 MI.getOpcode() == X86::ADJCALLSTACKUP64) { 178 Offset -= MI.getOperand(0).getImm(); 179 } else if (MI.getOpcode() == X86::ADJCALLSTACKDOWN32 || 180 MI.getOpcode() == X86::ADJCALLSTACKDOWN64) { 181 Offset += MI.getOperand(0).getImm(); 182 } else if (MI.modifiesRegister(StackPtr, TRI)) { 183 // Any other modification of SP means we've lost track of it. 184 Offset = INT32_MAX; 185 } 186 } 187 188 OutOffset[MBB] = Offset; 189 } 190 } 191 192 static unsigned getSubOpcode(bool Is64Bit, int64_t Amount) { 193 if (Is64Bit) 194 return isInt<8>(Amount) ? X86::SUB64ri8 : X86::SUB64ri32; 195 return isInt<8>(Amount) ? X86::SUB32ri8 : X86::SUB32ri; 196 } 197 198 void X86DynAllocaExpander::lower(MachineInstr *MI, Lowering L) { 199 const DebugLoc &DL = MI->getDebugLoc(); 200 MachineBasicBlock *MBB = MI->getParent(); 201 MachineBasicBlock::iterator I = *MI; 202 203 int64_t Amount = getDynAllocaAmount(MI, MRI); 204 if (Amount == 0) { 205 MI->eraseFromParent(); 206 return; 207 } 208 209 // These two variables differ on x32, which is a 64-bit target with a 210 // 32-bit alloca. 211 bool Is64Bit = STI->is64Bit(); 212 bool Is64BitAlloca = MI->getOpcode() == X86::DYN_ALLOCA_64; 213 assert(SlotSize == 4 || SlotSize == 8); 214 215 std::optional<MachineFunction::DebugInstrOperandPair> InstrNum; 216 if (unsigned Num = MI->peekDebugInstrNum()) { 217 // Operand 2 of DYN_ALLOCAs contains the stack def. 218 InstrNum = {Num, 2}; 219 } 220 221 switch (L) { 222 case TouchAndSub: { 223 assert(Amount >= SlotSize); 224 225 // Use a push to touch the top of the stack. 226 unsigned RegA = Is64Bit ? X86::RAX : X86::EAX; 227 BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::PUSH64r : X86::PUSH32r)) 228 .addReg(RegA, RegState::Undef); 229 Amount -= SlotSize; 230 if (!Amount) 231 break; 232 233 // Fall through to make any remaining adjustment. 234 [[fallthrough]]; 235 } 236 case Sub: 237 assert(Amount > 0); 238 if (Amount == SlotSize) { 239 // Use push to save size. 240 unsigned RegA = Is64Bit ? X86::RAX : X86::EAX; 241 BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::PUSH64r : X86::PUSH32r)) 242 .addReg(RegA, RegState::Undef); 243 } else { 244 // Sub. 245 BuildMI(*MBB, I, DL, 246 TII->get(getSubOpcode(Is64BitAlloca, Amount)), StackPtr) 247 .addReg(StackPtr) 248 .addImm(Amount); 249 } 250 break; 251 case Probe: 252 if (!NoStackArgProbe) { 253 // The probe lowering expects the amount in RAX/EAX. 254 unsigned RegA = Is64BitAlloca ? X86::RAX : X86::EAX; 255 BuildMI(*MBB, MI, DL, TII->get(TargetOpcode::COPY), RegA) 256 .addReg(MI->getOperand(0).getReg()); 257 258 // Do the probe. 259 STI->getFrameLowering()->emitStackProbe(*MBB->getParent(), *MBB, MI, DL, 260 /*InProlog=*/false, InstrNum); 261 } else { 262 // Sub 263 BuildMI(*MBB, I, DL, 264 TII->get(Is64BitAlloca ? X86::SUB64rr : X86::SUB32rr), StackPtr) 265 .addReg(StackPtr) 266 .addReg(MI->getOperand(0).getReg()); 267 } 268 break; 269 } 270 271 Register AmountReg = MI->getOperand(0).getReg(); 272 MI->eraseFromParent(); 273 274 // Delete the definition of AmountReg. 275 if (MRI->use_empty(AmountReg)) 276 if (MachineInstr *AmountDef = MRI->getUniqueVRegDef(AmountReg)) 277 AmountDef->eraseFromParent(); 278 } 279 280 bool X86DynAllocaExpander::runOnMachineFunction(MachineFunction &MF) { 281 if (!MF.getInfo<X86MachineFunctionInfo>()->hasDynAlloca()) 282 return false; 283 284 MRI = &MF.getRegInfo(); 285 STI = &MF.getSubtarget<X86Subtarget>(); 286 TII = STI->getInstrInfo(); 287 TRI = STI->getRegisterInfo(); 288 StackPtr = TRI->getStackRegister(); 289 SlotSize = TRI->getSlotSize(); 290 StackProbeSize = STI->getTargetLowering()->getStackProbeSize(MF); 291 NoStackArgProbe = MF.getFunction().hasFnAttribute("no-stack-arg-probe"); 292 if (NoStackArgProbe) 293 StackProbeSize = INT64_MAX; 294 295 LoweringMap Lowerings; 296 computeLowerings(MF, Lowerings); 297 for (auto &P : Lowerings) 298 lower(P.first, P.second); 299 300 return true; 301 } 302