[CodeGen] Compute call frame sizes instead of storing them in MachineBBs

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -145,14 +145,6 @@ class MachineBasicBlock
   const BasicBlock *BB;
   int Number;
 
-  /// The call frame size on entry to this basic block due to call frame setup
-  /// instructions in a predecessor. This is usually zero, unless basic blocks
-  /// are split in the middle of a call sequence.
-  ///
-  /// This information is only maintained until PrologEpilogInserter eliminates
-  /// call frame pseudos.
-  unsigned CallFrameSize = 0;
-
   MachineFunction *xParent;
   Instructions Insts;
 
@@ -1215,11 +1207,6 @@ class MachineBasicBlock
   int getNumber() const { return Number; }
   void setNumber(int N) { Number = N; }
 
-  /// Return the call frame size on entry to this basic block.
-  unsigned getCallFrameSize() const { return CallFrameSize; }
-  /// Set the call frame size on entry to this basic block.
-  void setCallFrameSize(unsigned N) { CallFrameSize = N; }
-
   /// Return the MCSymbol for this basic block.
   MCSymbol *getSymbol() const;
 

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -26,6 +26,7 @@ class MachineFunction;
 class MachineBasicBlock;
 class BitVector;
 class AllocaInst;
+class TargetInstrInfo;
 
 /// The CalleeSavedInfo class tracks the information need to locate where a
 /// callee saved register is in the current frame.
@@ -846,6 +847,64 @@ class MachineFrameInfo {
   void dump(const MachineFunction &MF) const;
 };
 
+/// Computes and stores the call frame sizes at the entries and exits of
+/// MachineBasicBlocks of a MachineFunction based on call frame setup and
+/// destroy pseudo instructions. Usually, no call frame is open at block
+/// boundaries, except if a call sequence has been split into multiple blocks.
+/// Computing this information is deferred until it is queried.
+///
+/// This class assumes that call frame instructions are placed properly, i.e.,
+/// every program path hits a frame destroy of equal size after hitting a frame
+/// setup, and a frame setup of equal size before a frame destroy. Nested call
+/// frame sequences are not allowed.
+class MachineFrameSizeInfo {
+public:
+  MachineFrameSizeInfo(MachineFunction &MF) : MF(MF) {}
+
+  /// Get the call frame size just before MI. Contains no value if MI is not in
+  /// a call sequence. Zero-sized call frames are possible.
+  std::optional<unsigned> getCallFrameSizeAt(MachineInstr &MI);
+
+  /// 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);
+
+  /// Get the call frame size at the entry of MBB. Contains no value if the
+  /// entry of MBB is not in a call sequence. Zero-sized call frames are
+  /// possible. Prefer this over getCallFrameSizeAt(MBB, MBB.begin()).
+  std::optional<unsigned> getCallFrameSizeAtBegin(MachineBasicBlock &MBB);
+
+  /// Get the call frame size at the exit of MBB. Contains no value if the exit
+  /// of MBB is not in a call sequence. Zero-sized call frames are possible.
+  /// Prefer this over getCallFrameSizeAt(MBB, MBB.end()).
+  std::optional<unsigned> getCallFrameSizeAtEnd(MachineBasicBlock &MBB);
+
+private:
+  /// Stores the call frame sizes at the boundaries of a MachineBasicBlock.
+  struct MachineFrameSizeInfoForBB {
+    MachineFrameSizeInfoForBB() = default;
+    MachineFrameSizeInfoForBB(std::optional<unsigned> EntryVal,
+                              std::optional<unsigned> ExitVal)
+        : Entry(EntryVal), Exit(ExitVal) {}
+
+    std::optional<unsigned> Entry;
+    std::optional<unsigned> Exit;
+  };
+
+  /// Compute call frame sizes at the boundaries of each MachineBasicBlock.
+  void computeSizes();
+
+  MachineFunction &MF;
+  const TargetInstrInfo *TII;
+  unsigned FrameSetupOpcode = ~0u;
+  unsigned FrameDestroyOpcode = ~0u;
+  bool HasFrameOpcodes = false;
+  bool HasNoBrokenUpCallSeqs = false;
+  SmallVector<MachineFrameSizeInfoForBB, 8> State;
+  bool IsComputed = false;
+};
+
 } // End llvm namespace
 
 #endif

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2265,9 +2265,6 @@ class TargetInstrInfo : public MCInstrInfo {
     return false;
   }
 
-  // Get the call frame size just before MI.
-  unsigned getCallFrameSizeAt(MachineInstr &MI) 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/MIRParser/MILexer.cpp

@@ -285,7 +285,6 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("ir-block-address-taken", MIToken::kw_ir_block_address_taken)
       .Case("machine-block-address-taken",
             MIToken::kw_machine_block_address_taken)
-      .Case("call-frame-size", MIToken::kw_call_frame_size)
       .Case("noconvergent", MIToken::kw_noconvergent)
       .Default(MIToken::Identifier);
 }

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -139,7 +139,6 @@ struct MIToken {
     kw_unknown_address,
     kw_ir_block_address_taken,
     kw_machine_block_address_taken,
-    kw_call_frame_size,
     kw_noconvergent,
 
     // Metadata types.

llvm/lib/CodeGen/MIRParser/MIParser.cpp

@@ -502,7 +502,6 @@ class MIParser {
   bool parseAddrspace(unsigned &Addrspace);
   bool parseSectionID(std::optional<MBBSectionID> &SID);
   bool parseBBID(std::optional<UniqueBBID> &BBID);
-  bool parseCallFrameSize(unsigned &CallFrameSize);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
@@ -684,18 +683,6 @@ bool MIParser::parseBBID(std::optional<UniqueBBID> &BBID) {
   return false;
 }
 
-// Parse basic block call frame size.
-bool MIParser::parseCallFrameSize(unsigned &CallFrameSize) {
-  assert(Token.is(MIToken::kw_call_frame_size));
-  lex();
-  unsigned Value = 0;
-  if (getUnsigned(Value))
-    return error("Unknown call frame size");
-  CallFrameSize = Value;
-  lex();
-  return false;
-}
-
 bool MIParser::parseBasicBlockDefinition(
     DenseMap<unsigned, MachineBasicBlock *> &MBBSlots) {
   assert(Token.is(MIToken::MachineBasicBlockLabel));
@@ -713,7 +700,6 @@ bool MIParser::parseBasicBlockDefinition(
   std::optional<MBBSectionID> SectionID;
   uint64_t Alignment = 0;
   std::optional<UniqueBBID> BBID;
-  unsigned CallFrameSize = 0;
   BasicBlock *BB = nullptr;
   if (consumeIfPresent(MIToken::lparen)) {
     do {
@@ -758,10 +744,6 @@ bool MIParser::parseBasicBlockDefinition(
         if (parseBBID(BBID))
           return true;
         break;
-      case MIToken::kw_call_frame_size:
-        if (parseCallFrameSize(CallFrameSize))
-          return true;
-        break;
       default:
         break;
       }
@@ -799,7 +781,6 @@ bool MIParser::parseBasicBlockDefinition(
     MBB->setSectionID(*SectionID);
     MF.setBBSectionsType(BasicBlockSection::List);
   }
-  MBB->setCallFrameSize(CallFrameSize);
   return false;
 }
 

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -578,11 +578,6 @@ void MachineBasicBlock::printName(raw_ostream &os, unsigned printNameFlags,
         os << " " << getBBID()->CloneID;
       hasAttributes = true;
     }
-    if (CallFrameSize != 0) {
-      os << (hasAttributes ? ", " : " (");
-      os << "call-frame-size " << CallFrameSize;
-      hasAttributes = true;
-    }
   }
 
   if (hasAttributes)
@@ -1155,7 +1150,6 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
   MachineBasicBlock *PrevFallthrough = getNextNode();
 
   MachineBasicBlock *NMBB = MF->CreateMachineBasicBlock();
-  NMBB->setCallFrameSize(Succ->getCallFrameSize());
 
   // Is there an indirect jump with jump table?
   bool ChangedIndirectJump = false;

llvm/lib/CodeGen/MachineFrameInfo.cpp

@@ -13,6 +13,7 @@
 #include "llvm/CodeGen/MachineFrameInfo.h"
 
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
@@ -23,6 +24,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
+#include <optional>
 
 #define DEBUG_TYPE "codegen"
 
@@ -246,6 +248,152 @@ void MachineFrameInfo::print(const MachineFunction &MF, raw_ostream &OS) const{
   }
 }
 
+std::optional<unsigned>
+MachineFrameSizeInfo::getCallFrameSizeAt(MachineInstr &MI) {
+  return this->getCallFrameSizeAt(*MI.getParent(), MI.getIterator());
+}
+
+std::optional<unsigned>
+MachineFrameSizeInfo::getCallFrameSizeAtBegin(MachineBasicBlock &MBB) {
+  if (!IsComputed)
+    computeSizes();
+  if (HasNoBrokenUpCallSeqs || !HasFrameOpcodes)
+    return std::nullopt;
+  return State[MBB.getNumber()].Entry;
+}
+
+std::optional<unsigned>
+MachineFrameSizeInfo::getCallFrameSizeAtEnd(MachineBasicBlock &MBB) {
+  if (!IsComputed)
+    computeSizes();
+  if (HasNoBrokenUpCallSeqs || !HasFrameOpcodes)
+    return std::nullopt;
+  return State[MBB.getNumber()].Exit;
+}
+
+std::optional<unsigned>
+MachineFrameSizeInfo::getCallFrameSizeAt(MachineBasicBlock &MBB,
+                                         MachineBasicBlock::iterator MII) {
+  if (!IsComputed)
+    computeSizes();
+
+  if (!HasFrameOpcodes)
+    return std::nullopt;
+
+  if (MII == MBB.end()) {
+    if (HasNoBrokenUpCallSeqs)
+      return std::nullopt;
+    return State[MBB.getNumber()].Exit;
+  }
+
+  if (MII == MBB.begin()) {
+    if (HasNoBrokenUpCallSeqs)
+      return std::nullopt;
+    return State[MBB.getNumber()].Entry;
+  }
+
+  // Search backwards from MI for the most recent call frame instruction.
+  for (auto &AdjI : reverse(make_range(MBB.begin(), MII))) {
+    if (AdjI.getOpcode() == FrameSetupOpcode)
+      return TII->getFrameTotalSize(AdjI);
+    if (AdjI.getOpcode() == FrameDestroyOpcode)
+      return std::nullopt;
+  }
+
+  // If none was found, use the call frame size from the start of the basic
+  // block.
+  if (HasNoBrokenUpCallSeqs)
+    return std::nullopt;
+  return State[MBB.getNumber()].Entry;
+}
+
+void MachineFrameSizeInfo::computeSizes() {
+  if (!IsComputed) {
+    // Populate fields that are only required once we compute the frame sizes.
+    TII = MF.getSubtarget().getInstrInfo();
+    FrameSetupOpcode = TII->getCallFrameSetupOpcode();
+    FrameDestroyOpcode = TII->getCallFrameDestroyOpcode();
+    HasFrameOpcodes = FrameSetupOpcode != ~0u || FrameDestroyOpcode != ~0u;
+    assert(!HasFrameOpcodes || FrameSetupOpcode != FrameDestroyOpcode);
+    IsComputed = true;
+  }
+  // If the target has no call frame pseudo instructions, don't compute
+  // anything, we always return std::nullopt if queried.
+  if (!HasFrameOpcodes)
+    return;
+
+  // Returns true if a call sequence in MF is broken up over multiple blocks.
+  auto FindBrokenUpCallSeq = [](const MachineFunction &MF,
+                                unsigned FrameSetupOpcode,
+                                unsigned FrameDestroyOpcode) {
+    for (const auto &MBB : MF) {
+      for (const auto &I : MBB) {
+        unsigned Opcode = I.getOpcode();
+        if (Opcode == FrameSetupOpcode)
+          break;
+        if (Opcode == FrameDestroyOpcode) {
+          // A FrameDestroy without a preceeding FrameSetup in the MBB. If
+          // FrameInstructions are placed correctly (which we assume), this
+          // occurs if and only if a call sequence is broken into multiple
+          // blocks.
+          return true;
+        }
+      }
+    }
+    return false;
+  };
+
+  HasNoBrokenUpCallSeqs =
+      !FindBrokenUpCallSeq(MF, FrameSetupOpcode, FrameDestroyOpcode);
+
+  // If every call sequence is limited to a single basic block, the frame sizes
+  // at entry and exit of each basic block need to be std::nullopt, so there is
+  // nothing to compute.
+  if (HasNoBrokenUpCallSeqs)
+    return;
+
+  State.resize(MF.getNumBlockIDs());
+
+  df_iterator_default_set<const MachineBasicBlock *> Reachable;
+
+  // Visit the MBBs in DFS order.
+  for (df_ext_iterator<MachineFunction *,
+                       df_iterator_default_set<const MachineBasicBlock *>>
+           DFI = df_ext_begin(&MF, Reachable),
+           DFE = df_ext_end(&MF, Reachable);
+       DFI != DFE; ++DFI) {
+    const MachineBasicBlock *MBB = *DFI;
+
+    MachineFrameSizeInfoForBB BBState;
+
+    // Use the exit state of the DFS stack predecessor as entry state for this
+    // block. With correctly placed call frame instructions, all other
+    // predecessors must have the same call frame size at exit.
+    if (DFI.getPathLength() >= 2) {
+      const MachineBasicBlock *StackPred = DFI.getPath(DFI.getPathLength() - 2);
+      assert(Reachable.count(StackPred) &&
+             "DFS stack predecessor is already visited.\n");
+      BBState.Entry = State[StackPred->getNumber()].Exit;
+      BBState.Exit = BBState.Entry;
+    }
+
+    // Search backwards for the last call frame instruction and use its implied
+    // state for the block exit. Otherwise, the exit state remains equal to the
+    // entry state.
+    for (auto &AdjI : reverse(make_range(MBB->begin(), MBB->end()))) {
+      if (AdjI.getOpcode() == FrameSetupOpcode) {
+        BBState.Exit = TII->getFrameTotalSize(AdjI);
+        break;
+      }
+      if (AdjI.getOpcode() == FrameDestroyOpcode) {
+        BBState.Exit = std::nullopt;
+        break;
+      }
+    }
+    State[MBB->getNumber()] = BBState;
+  }
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void MachineFrameInfo::dump(const MachineFunction &MF) const {
   print(MF, dbgs());