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

Previously, we would need to update stored call frame sizes whenever a
MachineBB with a call sequence is split or call frame instructions are
moved. By recomputing them instead, this change avoids the need for
keeping track of the call frame sizes throughout transformations. It
also fixes several cases where updates of call frame sizes were missing.

The overhead according to the compile time tracker is < +0.03% executed
instructions on all CT benchmarks.

This also allows distinguishing a zero-sized call frame from no open
call frame, which is helpful to avoid nested CALLSEQs (which are illegal
and cause machine verifier errors).

Author: ritter-x2a

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -145,15 +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. Usually does not contain a value, unless
-  /// basic blocks are split in the middle of a call sequence. Zero-sized call
-  /// frames are possible.
-  ///
-  /// This information is only maintained until PrologEpilogInserter eliminates
-  /// call frame pseudos.
-  std::optional<unsigned> CallFrameSize;
-
   MachineFunction *xParent;
   Instructions Insts;
 
@@ -1216,19 +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.
-  std::optional<unsigned> getCallFrameSize() const { return CallFrameSize; }
-
-  /// Return the call frame size on entry to this basic block if it is part of a
-  /// call sequence and 0 otherwise.
-  unsigned getCallFrameSizeOrZero() const { return CallFrameSize.value_or(0); }
-
-  /// Set the call frame size on entry to this basic block.
-  void setCallFrameSize(std::optional<unsigned> N) { CallFrameSize = N; }
-
-  /// Reset the stored call frame size.
-  void clearCallFrameSize() { CallFrameSize.reset(); }
-
   /// Return the MCSymbol for this basic block.
   MCSymbol *getSymbol() const;
 

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -26,6 +26,8 @@ class MachineFunction;
 class MachineBasicBlock;
 class BitVector;
 class AllocaInst;
+class MachineFrameSizeInfo;
+class TargetInstrInfo;
 
 /// The CalleeSavedInfo class tracks the information need to locate where a
 /// callee saved register is in the current frame.
@@ -282,6 +284,10 @@ class MachineFrameInfo {
   /// It is only valid during and after prolog/epilog code insertion.
   uint64_t MaxCallFrameSize = ~UINT64_C(0);
 
+  /// Call frame sizes for the MachineFunction's MachineBasicBlocks. This is set
+  /// by the MachineFrameSizeInfo constructor and cleared by its destructor.
+  MachineFrameSizeInfo *SizeInfo = nullptr;
+
   /// The number of bytes of callee saved registers that the target wants to
   /// report for the current function in the CodeView S_FRAMEPROC record.
   unsigned CVBytesOfCalleeSavedRegisters = 0;
@@ -675,6 +681,13 @@ class MachineFrameInfo {
   }
   void setMaxCallFrameSize(uint64_t S) { MaxCallFrameSize = S; }
 
+  /// Return an object that can be queried for call frame sizes at specific
+  /// locations in the MachineFunction. Constructing a MachineFrameSizeInfo
+  /// object for the MachineFunction automatically makes it available via this
+  /// field during the object's lifetime.
+  MachineFrameSizeInfo *getSizeInfo() const { return SizeInfo; }
+  void setSizeInfo(MachineFrameSizeInfo *SI) { SizeInfo = SI; }
+
   /// Returns how many bytes of callee-saved registers the target pushed in the
   /// prologue. Only used for debug info.
   unsigned getCVBytesOfCalleeSavedRegisters() const {
@@ -846,6 +859,73 @@ 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. Upon
+/// construction, a MachineFrameSizeInfo object registers itself in the
+/// MachineFunction's MachineFrameInfo (and it unregisters when destructed).
+/// While registered, it can be retrieved via MachineFrameInfo::getSizeInfo().
+///
+/// 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) {
+    assert(MF.getFrameInfo().getSizeInfo() == nullptr);
+    MF.getFrameInfo().setSizeInfo(this);
+  }
+
+  ~MachineFrameSizeInfo() { MF.getFrameInfo().setSizeInfo(nullptr); }
+
+  /// 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,16 +2265,6 @@ class TargetInstrInfo : public MCInstrInfo {
     return false;
   }
 
-  /// 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) 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/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(std::optional<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(std::optional<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;
-  std::optional<unsigned> CallFrameSize;
   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.has_value()) {
-      os << (hasAttributes ? ", " : " (");
-      os << "call-frame-size " << CallFrameSize.value();
-      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;