Catch some cases where the call frame size stored in MachineBasicBlocks was not set or updated.

Author: ritter-x2a

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2269,6 +2269,12 @@ class TargetInstrInfo : public MCInstrInfo {
   /// a call sequence. Zero-sized call frames are possible.
   std::optional<unsigned> getCallFrameSizeAt(MachineInstr &MI) const;
 
+  /// Get the call frame size just before MII. Contains no value if MII is not
+  /// in a call sequence. Zero-sized call frames are possible.
+  std::optional<unsigned>
+  getCallFrameSizeAt(MachineBasicBlock &MBB,
+                     MachineBasicBlock::iterator MII) const;
+
   /// Fills in the necessary MachineOperands to refer to a frame index.
   /// The best way to understand this is to print `asm(""::"m"(x));` after
   /// finalize-isel. Example:

llvm/lib/CodeGen/TargetInstrInfo.cpp

@@ -13,6 +13,7 @@
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/BinaryFormat/Dwarf.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineCombinerPattern.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -1626,9 +1627,14 @@ TargetInstrInfo::describeLoadedValue(const MachineInstr &MI,
 
 std::optional<unsigned>
 TargetInstrInfo::getCallFrameSizeAt(MachineInstr &MI) const {
+  return this->getCallFrameSizeAt(*MI.getParent(), MI.getIterator());
+}
+
+std::optional<unsigned>
+TargetInstrInfo::getCallFrameSizeAt(MachineBasicBlock &MBB,
+                                    MachineBasicBlock::iterator MII) const {
   // Search backwards from MI for the most recent call frame instruction.
-  MachineBasicBlock *MBB = MI.getParent();
-  for (auto &AdjI : reverse(make_range(MBB->instr_begin(), MI.getIterator()))) {
+  for (auto &AdjI : reverse(make_range(MBB.begin(), MII))) {
     if (AdjI.getOpcode() == getCallFrameSetupOpcode())
       return getFrameTotalSize(AdjI);
     if (AdjI.getOpcode() == getCallFrameDestroyOpcode())
@@ -1637,7 +1643,7 @@ TargetInstrInfo::getCallFrameSizeAt(MachineInstr &MI) const {
 
   // If none was found, use the call frame size from the start of the basic
   // block.
-  return MBB->getCallFrameSize();
+  return MBB.getCallFrameSize();
 }
 
 /// Both DefMI and UseMI must be valid.  By default, call directly to the

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

@@ -816,6 +816,15 @@ bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
   MachineBasicBlock *BodyBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RemainderBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RestoreExecBB = MF->CreateMachineBasicBlock();
+
+  // Set call-frame sizes in each new BB, in case we are in a call sequence.
+  std::optional<unsigned> CallFrameSizeBefore =
+      TII->getCallFrameSizeAt(MBB, Range.begin());
+  LoopBB->setCallFrameSize(CallFrameSizeBefore);
+  BodyBB->setCallFrameSize(CallFrameSizeBefore);
+  RemainderBB->setCallFrameSize(CallFrameSizeBefore);
+  RestoreExecBB->setCallFrameSize(CallFrameSizeBefore);
+
   MachineFunction::iterator MBBI(MBB);
   ++MBBI;
   MF->insert(MBBI, LoopBB);

llvm/lib/Target/X86/X86FrameLowering.cpp

@@ -790,6 +790,12 @@ void X86FrameLowering::emitStackProbeInlineGenericLoop(
   MachineBasicBlock *testMBB = MF.CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *tailMBB = MF.CreateMachineBasicBlock(LLVM_BB);
 
+  // Set call-frame sizes in each new BB, in case we are in a call sequence.
+  std::optional<unsigned> CallFrameSizeBefore =
+      TII.getCallFrameSizeAt(MBB, MBBI);
+  testMBB->setCallFrameSize(CallFrameSizeBefore);
+  tailMBB->setCallFrameSize(CallFrameSizeBefore);
+
   MachineFunction::iterator MBBIter = ++MBB.getIterator();
   MF.insert(MBBIter, testMBB);
   MF.insert(MBBIter, tailMBB);
@@ -932,6 +938,13 @@ void X86FrameLowering::emitStackProbeInlineWindowsCoreCLR64(
   MachineBasicBlock *LoopMBB = MF.CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *ContinueMBB = MF.CreateMachineBasicBlock(LLVM_BB);
 
+  // Set call-frame sizes in each new BB, in case we are in a call sequence.
+  std::optional<unsigned> CallFrameSizeBefore =
+      TII.getCallFrameSizeAt(MBB, MBBI);
+  RoundMBB->setCallFrameSize(CallFrameSizeBefore);
+  LoopMBB->setCallFrameSize(CallFrameSizeBefore);
+  ContinueMBB->setCallFrameSize(CallFrameSizeBefore);
+
   MachineFunction::iterator MBBIter = std::next(MBB.getIterator());
   MF.insert(MBBIter, RoundMBB);
   MF.insert(MBBIter, LoopMBB);
@@ -1286,6 +1299,19 @@ void X86FrameLowering::BuildStackAlignAND(MachineBasicBlock &MBB,
                                   : Is64Bit         ? X86::R11D
                                                     : X86::EAX;
 
+      std::optional<unsigned> CallFrameSizeBefore =
+          TII.getCallFrameSizeAt(MBB, MBBI);
+      // entryMBB takes the place of MBB, so give it its stored call frame size.
+      entryMBB->setCallFrameSize(MBB.getCallFrameSize());
+      // MBB contains only a tail of its original instructions, so update its
+      // stored call frame size.
+      MBB.setCallFrameSize(CallFrameSizeBefore);
+      // For the other basic blocks, take the state at the insertion point as
+      // well.
+      headMBB->setCallFrameSize(CallFrameSizeBefore);
+      bodyMBB->setCallFrameSize(CallFrameSizeBefore);
+      footMBB->setCallFrameSize(CallFrameSizeBefore);
+
       // Setup entry block
       {
 

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -35720,6 +35720,16 @@ X86TargetLowering::EmitLoweredProbedAlloca(MachineInstr &MI,
   MachineBasicBlock *tailMBB = MF->CreateMachineBasicBlock(LLVM_BB);
   MachineBasicBlock *blockMBB = MF->CreateMachineBasicBlock(LLVM_BB);
 
+  // Set call-frame sizes in each new BB, in case we are in a call sequence.
+  // MBB is split at std::next(MachineBasicBlock::iterator(MI) and these blocks
+  // are inserted in between, so they should all have the call frame size at
+  // the split point.
+  std::optional<unsigned> CallFrameSizeAtSplit =
+      TII->getCallFrameSizeAt(*MBB, std::next(MachineBasicBlock::iterator(MI)));
+  testMBB->setCallFrameSize(CallFrameSizeAtSplit);
+  blockMBB->setCallFrameSize(CallFrameSizeAtSplit);
+  tailMBB->setCallFrameSize(CallFrameSizeAtSplit);
+
   MachineFunction::iterator MBBIter = ++MBB->getIterator();
   MF->insert(MBBIter, testMBB);
   MF->insert(MBBIter, blockMBB);