[CodeGen] Store SP adjustment in MachineBasicBlock. NFCI.

Record the SP adjustment on entry to each basic block. This is almost
always zero except on targets like ARM which can split a basic block in
the middle of a call sequence.

This simplifies PEI::replaceFrameIndices which previously had to visit
basic blocks in a specific order and had special handling for
unreachable blocks. More importantly it paves the way for an equally
simple implementation of a backwards version of replaceFrameIndices,
which is required to fully convert PrologEpilogInserter to backwards
register scavenging, which is preferred because it does not rely on
accurate kill flags.

Differential Revision: https://reviews.llvm.org/D154281

Author: jayfoad

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -111,6 +111,12 @@ class MachineBasicBlock
 
   const BasicBlock *BB;
   int Number;
+
+  /// The SP adjustment on entry to this basic block due to call frame setup
+  /// instructions in a predecessor. This is almost always zero, unless basic
+  /// blocks are split in the middle of a call sequence.
+  int SPAdjustment = 0;
+
   MachineFunction *xParent;
   Instructions Insts;
 
@@ -1145,6 +1151,11 @@ class MachineBasicBlock
   int getNumber() const { return Number; }
   void setNumber(int N) { Number = N; }
 
+  /// Return the SP adjustment on entry to this basic block.
+  int getSPAdjustment() const { return SPAdjustment; }
+  /// Set the SP adjustment on entry to this basic block.
+  void setSPAdjustment(int N) { SPAdjustment = N; }
+
   /// Return the MCSymbol for this basic block.
   MCSymbol *getSymbol() const;
 

llvm/lib/CodeGen/MIRParser/MILexer.cpp

@@ -281,6 +281,7 @@ 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("sp-adjustment", MIToken::kw_sp_adjustment)
       .Default(MIToken::Identifier);
 }
 

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -135,6 +135,7 @@ struct MIToken {
     kw_unknown_address,
     kw_ir_block_address_taken,
     kw_machine_block_address_taken,
+    kw_sp_adjustment,
 
     // Metadata types.
     kw_distinct,

llvm/lib/CodeGen/MIRParser/MIParser.cpp

@@ -501,6 +501,7 @@ class MIParser {
   bool parseAddrspace(unsigned &Addrspace);
   bool parseSectionID(std::optional<MBBSectionID> &SID);
   bool parseBBID(std::optional<unsigned> &BBID);
+  bool parseSPAdjustment(int &SPAdjustment);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
@@ -676,6 +677,18 @@ bool MIParser::parseBBID(std::optional<unsigned> &BBID) {
   return false;
 }
 
+// Parse basic block SP adjustment.
+bool MIParser::parseSPAdjustment(int &SPAdjustment) {
+  assert(Token.is(MIToken::kw_sp_adjustment));
+  lex();
+  unsigned Value = 0;
+  if (getUnsigned(Value) || !isInt<32>(Value))
+    return error("Unknown SP adjustment");
+  SPAdjustment = (int)Value;
+  lex();
+  return false;
+}
+
 bool MIParser::parseBasicBlockDefinition(
     DenseMap<unsigned, MachineBasicBlock *> &MBBSlots) {
   assert(Token.is(MIToken::MachineBasicBlockLabel));
@@ -693,6 +706,7 @@ bool MIParser::parseBasicBlockDefinition(
   std::optional<MBBSectionID> SectionID;
   uint64_t Alignment = 0;
   std::optional<unsigned> BBID;
+  int SPAdjustment = 0;
   BasicBlock *BB = nullptr;
   if (consumeIfPresent(MIToken::lparen)) {
     do {
@@ -737,6 +751,10 @@ bool MIParser::parseBasicBlockDefinition(
         if (parseBBID(BBID))
           return true;
         break;
+      case MIToken::kw_sp_adjustment:
+        if (parseSPAdjustment(SPAdjustment))
+          return true;
+        break;
       default:
         break;
       }
@@ -781,6 +799,7 @@ bool MIParser::parseBasicBlockDefinition(
       MF.setBBSectionsType(BasicBlockSection::Labels);
     MBB->setBBID(BBID.value());
   }
+  MBB->setSPAdjustment(SPAdjustment);
   return false;
 }
 

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -570,6 +570,11 @@ void MachineBasicBlock::printName(raw_ostream &os, unsigned printNameFlags,
       os << "bb_id " << *getBBID();
       hasAttributes = true;
     }
+    if (SPAdjustment != 0) {
+      os << (hasAttributes ? ", " : " (");
+      os << "sp-adjustment " << SPAdjustment;
+      hasAttributes = true;
+    }
   }
 
   if (hasAttributes)

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -17,7 +17,6 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
-#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -1341,34 +1340,17 @@ void PEI::replaceFrameIndices(MachineFunction &MF) {
   FrameIndexEliminationScavenging = (RS && !FrameIndexVirtualScavenging) ||
     TRI->requiresFrameIndexReplacementScavenging(MF);
 
-  // Store SPAdj at exit of a basic block.
-  SmallVector<int, 8> SPState;
-  SPState.resize(MF.getNumBlockIDs());
-  df_iterator_default_set<MachineBasicBlock*> Reachable;
-
-  // Iterate over the reachable blocks in DFS order.
-  for (auto DFI = df_ext_begin(&MF, Reachable), DFE = df_ext_end(&MF, Reachable);
-       DFI != DFE; ++DFI) {
-    int SPAdj = 0;
-    // Check the exit state of the DFS stack predecessor.
-    if (DFI.getPathLength() >= 2) {
-      MachineBasicBlock *StackPred = DFI.getPath(DFI.getPathLength() - 2);
-      assert(Reachable.count(StackPred) &&
-             "DFS stack predecessor is already visited.\n");
-      SPAdj = SPState[StackPred->getNumber()];
-    }
-    MachineBasicBlock *BB = *DFI;
-    replaceFrameIndices(BB, MF, SPAdj);
-    SPState[BB->getNumber()] = SPAdj;
-  }
+  for (auto &MBB : MF) {
+    int SPAdj = MBB.getSPAdjustment();
+    replaceFrameIndices(&MBB, MF, SPAdj);
 
-  // Handle the unreachable blocks.
-  for (auto &BB : MF) {
-    if (Reachable.count(&BB))
-      // Already handled in DFS traversal.
-      continue;
-    int SPAdj = 0;
-    replaceFrameIndices(&BB, MF, SPAdj);
+    // If call frame pseudo ops have already been simplified then we can no
+    // longer track the SP adjustment (but that's OK because in that case, frame
+    // index elimination does not care about the SP adjustment).
+    if (!TFI.canSimplifyCallFramePseudos(MF)) {
+      for (auto *Succ : MBB.successors())
+        assert(Succ->getSPAdjustment() == SPAdj);
+    }
   }
 }
 

llvm/lib/Target/ARM/ARMISelLowering.cpp

@@ -11369,6 +11369,21 @@ static void emitPostSt(MachineBasicBlock *BB, MachineBasicBlock::iterator Pos,
   }
 }
 
+// Get the SP adjustment just before MI.
+int ARMTargetLowering::getSPAdjustment(MachineInstr &MI) const {
+  const ARMBaseInstrInfo *TII = Subtarget->getInstrInfo();
+
+  // Search backwards from MI for the most recent call frame setup instruction.
+  MachineBasicBlock *MBB = MI.getParent();
+  for (auto &AdjI : reverse(make_range(MBB->instr_begin(), MI.getIterator()))) {
+    if (AdjI.getOpcode() == ARM::ADJCALLSTACKDOWN)
+      return TII->getSPAdjust(AdjI);
+  }
+
+  // If none was found, use the SP adjustment from the start of the basic block.
+  return MBB->getSPAdjustment();
+}
+
 MachineBasicBlock *
 ARMTargetLowering::EmitStructByval(MachineInstr &MI,
                                    MachineBasicBlock *BB) const {
@@ -11485,6 +11500,11 @@ ARMTargetLowering::EmitStructByval(MachineInstr &MI,
   MF->insert(It, loopMBB);
   MF->insert(It, exitMBB);
 
+  // Set the SP adjustment on entry to the new basic blocks.
+  int SPAdj = getSPAdjustment(MI);
+  loopMBB->setSPAdjustment(SPAdj);
+  exitMBB->setSPAdjustment(SPAdj);
+
   // Transfer the remainder of BB and its successor edges to exitMBB.
   exitMBB->splice(exitMBB->begin(), BB,
                   std::next(MachineBasicBlock::iterator(MI)), BB->end());

llvm/lib/Target/ARM/ARMISelLowering.h

@@ -970,6 +970,8 @@ class VectorType;
 
     void EmitSjLjDispatchBlock(MachineInstr &MI, MachineBasicBlock *MBB) const;
 
+    int getSPAdjustment(MachineInstr &MI) const;
+
     MachineBasicBlock *EmitStructByval(MachineInstr &MI,
                                        MachineBasicBlock *MBB) const;
 

llvm/test/CodeGen/ARM/2013-06-03-ByVal-2Kbytes.ll

@@ -1,4 +1,6 @@
 ; RUN: llc < %s -mcpu=cortex-a15 | FileCheck %s
+; RUN: llc < %s -mcpu=cortex-a15 -stop-after=finalize-isel -o %t.mir
+; RUN: llc %t.mir -mcpu=cortex-a15 -start-after=finalize-isel -o - | FileCheck %s
 ; ModuleID = 'attri_16.c'
 target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:64:128-a0:0:64-n32-S64"
 target triple = "armv4t--linux-gnueabihf"

llvm/test/CodeGen/MIR/Generic/basic-blocks.mir

@@ -18,6 +18,11 @@
     ret i32 0
   }
 
+  define i32 @spadj() {
+  start:
+    ret i32 0
+  }
+
 ...
 ---
 # CHECK-LABEL: name: foo
@@ -47,3 +52,11 @@ body: |
   bb.0.start (align 4, machine-block-address-taken):
   bb.1 (machine-block-address-taken, align 4):
 ...
+---
+# CHECK-LABEL: name: spadj
+# CHECK: body:
+# CHECK-NEXT: bb.0 (sp-adjustment 100):
+name: spadj
+body: |
+  bb.0 (sp-adjustment 100):
+...

llvm/tools/llvm-reduce/ReducerWorkItem.cpp

@@ -225,6 +225,8 @@ static std::unique_ptr<MachineFunction> cloneMF(MachineFunction *SrcMF,
         DstMF->CreateMachineBasicBlock(SrcMBB.getBasicBlock());
     Src2DstMBB[&SrcMBB] = DstMBB;
 
+    DstMBB->setSPAdjustment(SrcMBB.getSPAdjustment());
+
     if (SrcMBB.isIRBlockAddressTaken())
       DstMBB->setAddressTakenIRBlock(SrcMBB.getAddressTakenIRBlock());
     if (SrcMBB.isMachineBlockAddressTaken())