[WIP][CodeGen] Enable early CFI

llvm/include/llvm/CodeGen/TargetFrameLowering.h

@@ -25,16 +25,17 @@ namespace llvm {
   class CalleeSavedInfo;
   class MachineFunction;
   class RegScavenger;
-
-namespace TargetStackID {
-enum Value {
-  Default = 0,
-  SGPRSpill = 1,
-  ScalableVector = 2,
-  WasmLocal = 3,
-  ScalablePredicateVector = 4,
-  NoAlloc = 255
-};
+  class ReachingDefInfo;
+
+  namespace TargetStackID {
+  enum Value {
+    Default = 0,
+    SGPRSpill = 1,
+    ScalableVector = 2,
+    WasmLocal = 3,
+    ScalablePredicateVector = 4,
+    NoAlloc = 255
+  };
 }
 
 /// Information about stack frame layout on the target.  It holds the direction
@@ -212,6 +213,12 @@ class LLVM_ABI TargetFrameLowering {
   /// for noreturn nounwind functions.
   virtual bool enableCalleeSaveSkip(const MachineFunction &MF) const;
 
+  /// If savesCSRsEarly is true, we don't really know where the CSRs are
+  /// saved. This function calculates where each CSR is at every point in the
+  /// function and inserts necessary CFIs. It has to run before frame indicies
+  /// are resolved.
+  virtual void emitCFIsEarly(MachineFunction &MF, ReachingDefInfo &RDA) const {}
+
   /// emitProlog/emitEpilog - These methods insert prolog and epilog code into
   /// the function.
   virtual void emitPrologue(MachineFunction &MF,

llvm/lib/CodeGen/PrologEpilogInserter.cpp

@@ -36,6 +36,7 @@
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/PEI.h"
+#include "llvm/CodeGen/ReachingDefAnalysis.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
@@ -79,6 +80,7 @@ namespace {
 
 class PEIImpl {
   RegScavenger *RS = nullptr;
+  ReachingDefInfo &RDI;
 
   // MinCSFrameIndex, MaxCSFrameIndex - Keeps the range of callee saved
   // stack frame indexes.
@@ -125,7 +127,8 @@ class PEIImpl {
   void insertZeroCallUsedRegs(MachineFunction &MF);
 
 public:
-  PEIImpl(MachineOptimizationRemarkEmitter *ORE) : ORE(ORE) {}
+  PEIImpl(ReachingDefInfo &RDI, MachineOptimizationRemarkEmitter *ORE)
+      : RDI(RDI), ORE(ORE) {}
   bool run(MachineFunction &MF);
 };
 
@@ -155,6 +158,7 @@ INITIALIZE_PASS_BEGIN(PEILegacy, DEBUG_TYPE, "Prologue/Epilogue Insertion",
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
+INITIALIZE_PASS_DEPENDENCY(ReachingDefInfoWrapperPass)
 INITIALIZE_PASS_END(PEILegacy, DEBUG_TYPE,
                     "Prologue/Epilogue Insertion & Frame Finalization", false,
                     false)
@@ -171,6 +175,7 @@ void PEILegacy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addPreserved<MachineLoopInfoWrapperPass>();
   AU.addPreserved<MachineDominatorTreeWrapperPass>();
   AU.addRequired<MachineOptimizationRemarkEmitterPass>();
+  AU.addRequired<ReachingDefInfoWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
@@ -361,15 +366,17 @@ bool PEIImpl::run(MachineFunction &MF) {
 bool PEILegacy::runOnMachineFunction(MachineFunction &MF) {
   MachineOptimizationRemarkEmitter *ORE =
       &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
-  return PEIImpl(ORE).run(MF);
+  auto &RDI = getAnalysis<ReachingDefInfoWrapperPass>().getRDI();
+  return PEIImpl(RDI, ORE).run(MF);
 }
 
 PreservedAnalyses
 PrologEpilogInserterPass::run(MachineFunction &MF,
                               MachineFunctionAnalysisManager &MFAM) {
   MachineOptimizationRemarkEmitter &ORE =
       MFAM.getResult<MachineOptimizationRemarkEmitterAnalysis>(MF);
-  if (!PEIImpl(&ORE).run(MF))
+  auto &RDI = MFAM.getResult<ReachingDefAnalysis>(MF);
+  if (!PEIImpl(RDI, &ORE).run(MF))
     return PreservedAnalyses::all();
 
   return getMachineFunctionPassPreservedAnalyses()
@@ -1188,6 +1195,11 @@ void PEIImpl::calculateFrameObjectOffsets(MachineFunction &MF) {
 void PEIImpl::insertPrologEpilogCode(MachineFunction &MF) {
   const TargetFrameLowering &TFI = *MF.getSubtarget().getFrameLowering();
 
+  if (MF.getSubtarget().savesCSRsEarly()) {
+    RDI.reset();
+    TFI.emitCFIsEarly(MF, RDI);
+  }
+
   // Add prologue to the function...
   for (MachineBasicBlock *SaveBlock : SaveBlocks)
     TFI.emitPrologue(MF, *SaveBlock);

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

@@ -14,13 +14,15 @@
 #include "MCTargetDesc/RISCVBaseInfo.h"
 #include "RISCVMachineFunctionInfo.h"
 #include "RISCVSubtarget.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/CFIInstBuilder.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/ReachingDefAnalysis.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/MC/MCDwarf.h"
@@ -2379,6 +2381,153 @@ bool RISCVFrameLowering::canUseAsEpilogue(const MachineBasicBlock &MBB) const {
   return SuccMBB->isReturnBlock() && SuccMBB->size() == 1;
 }
 
+struct CFIBuildInfo {
+  MachineBasicBlock *MBB;
+  MachineInstr *InsertAfterMI; // nullptr means insert at MBB.begin()
+  DebugLoc DL;
+  unsigned CFIIndex;
+};
+
+class EarlyCFIEmitter {
+public:
+  EarlyCFIEmitter(MachineFunction &MF_, const ReachingDefInfo &RDI_,
+                  const RISCVInstrInfo &TII_, const RISCVRegisterInfo &TRI_,
+                  const RISCVFrameLowering &TFI_)
+      : MF{MF_}, RDI{RDI_}, TII{TII_}, TRI{TRI_}, TFI{TFI_} {};
+  void trackRegisterAndEmitCFIs(
+      Register Reg, int64_t DwarfEHRegNum, MachineInstr &MI,
+      SmallVectorImpl<CFIBuildInfo> &CFIBuildInfos,
+      SmallPtrSetImpl<MachineInstr *> &VisitedRestorePoints,
+      SmallPtrSetImpl<MachineInstr *> &VisitedDefs);
+
+private:
+  MachineFunction &MF;
+  const ReachingDefInfo &RDI;
+  const RISCVInstrInfo &TII;
+  const RISCVRegisterInfo &TRI;
+  const RISCVFrameLowering &TFI;
+};
+
+void EarlyCFIEmitter::trackRegisterAndEmitCFIs(
+    Register Reg, int64_t DwarfEHRegNum, MachineInstr &MI,
+    SmallVectorImpl<CFIBuildInfo> &CFIBuildInfos,
+    SmallPtrSetImpl<MachineInstr *> &VisitedRestorePoints,
+    SmallPtrSetImpl<MachineInstr *> &VisitedDefs) {
+  if (VisitedRestorePoints.find(&MI) != VisitedRestorePoints.end())
+    return;
+  VisitedRestorePoints.insert(&MI);
+
+  SmallPtrSet<MachineInstr *, 2> Defs;
+  RDI.getGlobalReachingDefs(&MI, Reg, Defs);
+  if (Defs.empty())
+    // it's a live-in register at the entry block.
+    return;
+
+  int FrameIndex = std::numeric_limits<int>::min();
+  for (MachineInstr *Def : Defs) {
+    if (VisitedDefs.find(Def) != VisitedDefs.end())
+      continue;
+    VisitedDefs.insert(Def);
+
+    MachineBasicBlock &MBB = *Def->getParent();
+    const DebugLoc &DL = Def->getDebugLoc();
+
+    if (Register StoredReg = TII.isStoreToStackSlot(*Def, FrameIndex)) {
+      assert(FrameIndex == Reg.stackSlotIndex());
+
+      Register FrameReg;
+      StackOffset Offset = TFI.getFrameIndexReference(MF, FrameIndex, FrameReg);
+      int64_t FixedOffset = Offset.getFixed();
+      int64_t ScalableOffset = Offset.getScalable();
+
+      std::string CommentBuffer;
+      llvm::raw_string_ostream Comment(CommentBuffer);
+      int DwarfEHFrameReg = TRI.getDwarfRegNum(FrameReg, true);
+      Register LLVMReg = *TRI.getLLVMRegNum(DwarfEHRegNum, true);
+      Comment << printReg(LLVMReg, &TRI) << " @";
+      Comment << printReg(FrameReg, &TRI) << " + ";
+      Comment << "vlenb * " << ScalableOffset << " + ";
+      Comment << FixedOffset;
+      unsigned CFIIndex = MF.addFrameInst(
+          MCCFIInstruction::createLLVMRegAtScalableOffsetFromReg(
+              nullptr, DwarfEHRegNum, DwarfEHFrameReg, ScalableOffset,
+              FixedOffset, SMLoc(), Comment.str()));
+
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(StoredReg, DwarfEHRegNum, *Def, CFIBuildInfos,
+                               VisitedRestorePoints, VisitedDefs);
+    } else if (Register LoadedReg = TII.isLoadFromStackSlot(*Def, FrameIndex)) {
+      assert(LoadedReg == Reg);
+
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRegister(
+          nullptr, DwarfEHRegNum, TRI.getDwarfRegNum(LoadedReg, true)));
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(Register::index2StackSlot(FrameIndex),
+                               DwarfEHRegNum, *Def, CFIBuildInfos,
+                               VisitedRestorePoints, VisitedDefs);
+    } else if (auto DstSrc = TII.isCopyInstr(*Def)) {
+      Register DstReg = DstSrc->Destination->getReg();
+      Register SrcReg = DstSrc->Source->getReg();
+      assert(DstReg == Reg);
+
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRegister(
+          nullptr, DwarfEHRegNum, TRI.getDwarfRegNum(DstReg, true)));
+      CFIBuildInfos.push_back({&MBB, Def, DL, CFIIndex});
+      trackRegisterAndEmitCFIs(SrcReg, DwarfEHRegNum, *Def, CFIBuildInfos,
+                               VisitedRestorePoints, VisitedDefs);
+    } else
+      llvm_unreachable("Unexpected instruction");
+  }
+  return;
+}
+
+void RISCVFrameLowering::emitCFIsEarly(MachineFunction &MF,
+                                       ReachingDefInfo &RDI) const {
+  BitVector EarlyCSRs;
+  determineEarlyCalleeSaves(MF, EarlyCSRs);
+
+  SmallVector<MachineInstr *, 4> RestorePoints;
+  for (MachineBasicBlock &MBB : MF) {
+    if (MBB.isReturnBlock())
+      RestorePoints.push_back(&MBB.back());
+  }
+  SmallVector<CFIBuildInfo, 32> CFIBuildInfos;
+  const RISCVInstrInfo &TII = *STI.getInstrInfo();
+  const RISCVRegisterInfo &TRI = *STI.getRegisterInfo();
+  EarlyCFIEmitter EarlyCFIE(MF, RDI, TII, TRI, *STI.getFrameLowering());
+  const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
+  for (unsigned i = 0; CSRegs[i]; ++i) {
+    unsigned Reg = CSRegs[i];
+    if (!EarlyCSRs[Reg])
+      continue;
+    SmallPtrSet<MachineInstr *, 32> VisitedDefs;
+    for (MachineInstr *RestorePoint : RestorePoints) {
+      SmallPtrSet<MachineInstr *, 32> VisitedRestorePoints;
+      EarlyCFIE.trackRegisterAndEmitCFIs(Reg, TRI.getDwarfRegNum(Reg, true),
+                                         *RestorePoint, CFIBuildInfos,
+                                         VisitedRestorePoints, VisitedDefs);
+    }
+  }
+  for (CFIBuildInfo &Info : CFIBuildInfos) {
+    MachineBasicBlock *MBB = Info.MBB;
+    if (Info.InsertAfterMI) {
+      auto Bundler =
+          MIBundleBuilder(*MBB, ++(Info.InsertAfterMI->getIterator()));
+      Bundler.append(Info.InsertAfterMI->removeFromParent());
+      Bundler.append(
+          BuildMI(MBB, Info.DL, TII.get(TargetOpcode::CFI_INSTRUCTION))
+              .addCFIIndex(Info.CFIIndex)
+              ->removeFromParent());
+      finalizeBundle(*MBB, Bundler.begin(), Bundler.end());
+    } else {
+      BuildMI(*MBB, MBB->begin(), Info.DL,
+              TII.get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(Info.CFIIndex);
+    }
+  }
+  return;
+}
+
 bool RISCVFrameLowering::isSupportedStackID(TargetStackID::Value ID) const {
   switch (ID) {
   case TargetStackID::Default:

llvm/lib/Target/RISCV/RISCVFrameLowering.h

@@ -74,6 +74,8 @@ class RISCVFrameLowering : public TargetFrameLowering {
 
   bool enableShrinkWrapping(const MachineFunction &MF) const override;
 
+  void emitCFIsEarly(MachineFunction &MF, ReachingDefInfo &RDI) const override;
+
   bool isSupportedStackID(TargetStackID::Value ID) const override;
   TargetStackID::Value getStackIDForScalableVectors() const override;
 

llvm/lib/Target/RISCV/RISCVSubtarget.cpp

@@ -75,9 +75,10 @@ static cl::opt<bool> EnablePExtCodeGen(
              "only partial codegen is currently supported)"),
     cl::init(false), cl::Hidden);
 
-static cl::opt<bool> SaveCSREarly("riscv-save-csrs-early",
-                                  cl::desc("Save CSRs early"),
-                                  cl::init(false), cl::Hidden);
+static cl::opt<bool> RISCVSaveCSRsEarly(
+    "riscv-save-csrs-early",
+    cl::desc("Let register alloctor do csr saves/restores"), cl::init(false),
+    cl::Hidden);
 
 void RISCVSubtarget::anchor() {}
 
@@ -221,8 +222,8 @@ bool RISCVSubtarget::enableMachinePipeliner() const {
   return getSchedModel().hasInstrSchedModel();
 }
 
-  /// Enable use of alias analysis during code generation (during MI
-  /// scheduling, DAGCombine, etc.).
+/// Enable use of alias analysis during code generation (during MI
+/// scheduling, DAGCombine, etc.).
 bool RISCVSubtarget::useAA() const { return UseAA; }
 
 unsigned RISCVSubtarget::getMinimumJumpTableEntries() const {
@@ -270,6 +271,4 @@ bool RISCVSubtarget::useMIPSCCMovInsn() const {
   return UseMIPSCCMovInsn && HasVendorXMIPSCMov;
 }
 
-bool RISCVSubtarget::savesCSRsEarly() const {
-  return SaveCSREarly;
-}
+bool RISCVSubtarget::savesCSRsEarly() const { return RISCVSaveCSRsEarly; }

llvm/test/CodeGen/RISCV/O0-pipeline.ll

@@ -54,6 +54,7 @@
 ; CHECK-NEXT:       Fixup Statepoint Caller Saved
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
+; CHECK-NEXT:       Reaching Definitions Analysis
 ; CHECK-NEXT:       Prologue/Epilogue Insertion & Frame Finalization
 ; CHECK-NEXT:       Post-RA pseudo instruction expansion pass
 ; CHECK-NEXT:       RISC-V post-regalloc pseudo instruction expansion pass

llvm/test/CodeGen/RISCV/O3-pipeline.ll

@@ -175,6 +175,7 @@
 ; CHECK-NEXT:       Lazy Machine Block Frequency Analysis
 ; CHECK-NEXT:       Machine Optimization Remark Emitter
 ; CHECK-NEXT:       Shrink Wrapping analysis
+; CHECK-NEXT:       Reaching Definitions Analysis
 ; CHECK-NEXT:       Prologue/Epilogue Insertion & Frame Finalization
 ; CHECK-NEXT:       Machine Late Instructions Cleanup Pass
 ; CHECK-NEXT:       Control Flow Optimizer