Use Pseudo Instruction to carry stack probing information

[serge-sans-paille]

Summary

Instead of using a fake call and metadata to temporarily represent a probed
static alloca, use a pseudo instruction.

This is inspired by the SystemZ approach proposed in https://reviews.llvm.org/D78717.

[craig.topper]

llvm/lib/Target/X86/X86FrameLowering.cpp

531

Can we use llvm::find_if here which automatically calls begin/end?

[serge-sans-paille]

Use llvm::find_if

[craig.topper]

LGTM

llvm/lib/Target/X86/X86FrameLowering.cpp

void X86FrameLowering::emitSPUpdate(MachineBasicBlock &MBB,
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   const X86TargetLowering &TLI = *STI.getTargetLowering();
   const bool EmitInlineStackProbe = TLI.hasInlineStackProbe(MF);
 
   // It's ok to not take into account large chunks when probing, as the
   // allocation is split in smaller chunks anyway.
   if (EmitInlineStackProbe && !InEpilogue) {
 
-    // Delegate stack probing to the `inlineStackProbe` mechanism to avoid
-    // complications.
-    MachineInstr *Stub = emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
-
-    // Encode the static offset as a metadata attached to the stub.
-    LLVMContext &Context = MF.getFunction().getContext();
-    MachineInstrBuilder(MF, Stub).addMetadata(
-        MDTuple::get(Context, {ConstantAsMetadata::get(ConstantInt::get(
-                                  IntegerType::get(Context, 64), Offset))}));
+    // This pseudo-instruction is going to be expanded, potentially using a
+    // loop, by inlineStackProbe().
+    BuildMI(MBB, MBBI, DL, TII.get(X86::STACKALLOC_W_PROBING)).addImm(Offset);
     return;
   } else if (Offset > Chunk) {
     // Rather than emit a long series of instructions for large offsets,
     // load the offset into a register and do one sub/add
     unsigned Reg = 0;
     unsigned Rax = (unsigned)(Is64Bit ? X86::RAX : X86::EAX);
 
     if (isSub && !isEAXLiveIn(MBB))

 
 void X86FrameLowering::emitStackProbe(MachineFunction &MF,
                                       MachineBasicBlock &MBB,
                                       MachineBasicBlock::iterator MBBI,
                                       const DebugLoc &DL, bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   if (STI.isTargetWindowsCoreCLR()) {
     if (InProlog) {
-      emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
+      BuildMI(MBB, MBBI, DL, TII.get(X86::STACKALLOC_W_PROBING))
+          .addImm(0 /* no explicit stack size */);
     } else {
       emitStackProbeInline(MF, MBB, MBBI, DL, false);
     }
   } else {
     emitStackProbeCall(MF, MBB, MBBI, DL, InProlog);
   }
 }
 
 void X86FrameLowering::inlineStackProbe(MachineFunction &MF,
                                         MachineBasicBlock &PrologMBB) const {
-  const StringRef ChkStkStubSymbol = "__chkstk_stub";
-  MachineInstr *ChkStkStub = nullptr;
-
-  for (MachineInstr &MI : PrologMBB) {
-    if (MI.isCall() && MI.getOperand(0).isSymbol() &&
-        ChkStkStubSymbol == MI.getOperand(0).getSymbolName()) {
-      ChkStkStub = &MI;
-      break;
-    }
-  }
-
-  if (ChkStkStub != nullptr) {
-    assert(!ChkStkStub->isBundled() &&
-           "Not expecting bundled instructions here");
-    MachineBasicBlock::iterator MBBI = std::next(ChkStkStub->getIterator());
-    assert(std::prev(MBBI) == ChkStkStub &&
-           "MBBI expected after __chkstk_stub.");
-    DebugLoc DL = PrologMBB.findDebugLoc(MBBI);
-    emitStackProbeInline(MF, PrologMBB, MBBI, DL, true);
-    ChkStkStub->eraseFromParent();
+  auto Where = llvm::find_if(PrologMBB, [](MachineInstr &MI) {
+    return MI.getOpcode() == X86::STACKALLOC_W_PROBING;
+  });
+  if (Where != PrologMBB.end()) {
+    DebugLoc DL = PrologMBB.findDebugLoc(Where);
+    emitStackProbeInline(MF, PrologMBB, Where, DL, true);
+    Where->eraseFromParent();
   }
 }
 
 void X86FrameLowering::emitStackProbeInline(MachineFunction &MF,
                                             MachineBasicBlock &MBB,
                                             MachineBasicBlock::iterator MBBI,
                                             const DebugLoc &DL,
                                             bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   if (STI.isTargetWindowsCoreCLR() && STI.is64Bit())
     emitStackProbeInlineWindowsCoreCLR64(MF, MBB, MBBI, DL, InProlog);
   else
     emitStackProbeInlineGeneric(MF, MBB, MBBI, DL, InProlog);
 }
 
 void X86FrameLowering::emitStackProbeInlineGeneric(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator MBBI, const DebugLoc &DL, bool InProlog) const {
-  MachineInstr &CallToInline = *std::prev(MBBI);
-  assert(CallToInline.getOperand(1).isMetadata() &&
-         "no metadata attached to that probe");
-  uint64_t Offset =
-      cast<ConstantInt>(
-          cast<ConstantAsMetadata>(
-              cast<MDTuple>(CallToInline.getOperand(1).getMetadata())
-                  ->getOperand(0))
-              ->getValue())
-          ->getZExtValue();
+  MachineInstr &AllocWithProbe = *MBBI;
+  uint64_t Offset = AllocWithProbe.getOperand(0).getImm();
 
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   const X86TargetLowering &TLI = *STI.getTargetLowering();
   assert(!(STI.is64Bit() && STI.isTargetWindowsCoreCLR()) &&
          "different expansion expected for CoreCLR 64 bit");
 
   const uint64_t StackProbeSize = TLI.getStackProbeSize(MF);
   uint64_t ProbeChunk = StackProbeSize * 8;
void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
 
   if (InProlog) {
     // Apply the frame setup flag to all inserted instrs.
     for (++ExpansionMBBI; ExpansionMBBI != MBBI; ++ExpansionMBBI)
       ExpansionMBBI->setFlag(MachineInstr::FrameSetup);
   }
 }
 
-MachineInstr *X86FrameLowering::emitStackProbeInlineStub(
-    MachineFunction &MF, MachineBasicBlock &MBB,
-    MachineBasicBlock::iterator MBBI, const DebugLoc &DL, bool InProlog) const {
-
-  assert(InProlog && "ChkStkStub called outside prolog!");
-
-  return BuildMI(MBB, MBBI, DL, TII.get(X86::CALLpcrel32))
-      .addExternalSymbol("__chkstk_stub");
-}
-
 static unsigned calculateSetFPREG(uint64_t SPAdjust) {
   // Win64 ABI has a less restrictive limitation of 240; 128 works equally well
   // and might require smaller successive adjustments.
   const uint64_t Win64MaxSEHOffset = 128;
   uint64_t SEHFrameOffset = std::min(SPAdjust, Win64MaxSEHOffset);
   // Win64 ABI requires 16-byte alignment for the UWOP_SET_FPREG opcode.
   return SEHFrameOffset & -16;
 }

llvm/lib/Target/X86/X86InstrCompiler.td

 let Defs = [RAX, RSP, EFLAGS], Uses = [RSP] in
 def PROBED_ALLOCA_64 : I<0, Pseudo, (outs GR64:$dst), (ins GR64:$size),
                       "# variable sized alloca with probing",
                       [(set GR64:$dst,
                          (X86ProbedAlloca GR64:$size))]>,
                     Requires<[In64BitMode]>;
 }
 
+let hasNoSchedulingInfo = 1 in
+def STACKALLOC_W_PROBING : I<0, Pseudo, (outs), (ins i64imm:$stacksize),
+                             "# fixed size alloca with probing",
+                             []>;
+
 // Dynamic stack allocation yields a _chkstk or _alloca call for all Windows
 // targets.  These calls are needed to probe the stack when allocating more than
 // 4k bytes in one go. Touching the stack at 4K increments is necessary to
 // ensure that the guard pages used by the OS virtual memory manager are
 // allocated in correct sequence.
 // The main point of having separate instruction are extra unmodelled effects
 // (compared to ordinary calls) like stack pointer change.