[RISCV][WIP] Let RA do the CSR saves.

We turn the problem of saving and restoring callee-saved registers efficiently into a register allocation problem. This has the advantage that the register allocator can essentialy do shrink-wrapping on per register basis. Currently, shrink-wrapping pass saves all CSR in the same place which may be suboptimal. Also, improvements to register allocation / coalescing will translate to improvements in shrink-wrapping.

In finalizeLowering() we copy all callee-saved registers from a physical register to a virtual one. In all return blocks we copy do the reverse.

Author: Mikhail Gudim

llvm/include/llvm/CodeGen/ReachingDefAnalysis.h

@@ -114,8 +114,11 @@ class ReachingDefAnalysis : public MachineFunctionPass {
 private:
   MachineFunction *MF = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
+  const TargetInstrInfo *TII = nullptr;
   LoopTraversal::TraversalOrder TraversedMBBOrder;
   unsigned NumRegUnits = 0;
+  unsigned NumStackObjects = 0;
+  int ObjectIndexBegin = 0;
   /// Instruction that defined each register, relative to the beginning of the
   /// current basic block.  When a LiveRegsDefInfo is used to represent a
   /// live-out register, this value is relative to the end of the basic block,
@@ -138,6 +141,9 @@ class ReachingDefAnalysis : public MachineFunctionPass {
   DenseMap<MachineInstr *, int> InstIds;
 
   MBBReachingDefsInfo MBBReachingDefs;
+  using MBBFrameObjsReachingDefsInfo =
+      std::vector<std::vector<std::vector<int>>>;
+  MBBFrameObjsReachingDefsInfo MBBFrameObjsReachingDefs;
 
   /// Default values are 'nothing happened a long time ago'.
   const int ReachingDefDefaultVal = -(1 << 21);

llvm/include/llvm/CodeGen/TargetFrameLowering.h

@@ -24,15 +24,16 @@ namespace llvm {
   class CalleeSavedInfo;
   class MachineFunction;
   class RegScavenger;
-
-namespace TargetStackID {
-enum Value {
-  Default = 0,
-  SGPRSpill = 1,
-  ScalableVector = 2,
-  WasmLocal = 3,
-  NoAlloc = 255
-};
+  class ReachingDefAnalysis;
+
+  namespace TargetStackID {
+  enum Value {
+    Default = 0,
+    SGPRSpill = 1,
+    ScalableVector = 2,
+    WasmLocal = 3,
+    NoAlloc = 255
+  };
 }
 
 /// Information about stack frame layout on the target.  It holds the direction
@@ -210,6 +211,11 @@ class TargetFrameLowering {
   /// for noreturn nounwind functions.
   virtual bool enableCalleeSaveSkip(const MachineFunction &MF) const;
 
+  virtual void emitCFIsForCSRsHandledByRA(MachineFunction &MF,
+                                          ReachingDefAnalysis *RDA) const {
+    return;
+  }
+
   /// emitProlog/emitEpilog - These methods insert prolog and epilog code into
   /// the function.
   virtual void emitPrologue(MachineFunction &MF,

llvm/include/llvm/CodeGen/TargetSubtargetInfo.h

@@ -328,6 +328,8 @@ class TargetSubtargetInfo : public MCSubtargetInfo {
     return false;
   }
 
+  virtual bool doCSRSavesInRA() const;
+
   /// Classify a global function reference. This mainly used to fetch target
   /// special flags for lowering a function address. For example mark a function
   /// call should be plt or pc-related addressing.

llvm/include/llvm/DebugInfo/DWARF/DWARFDebugFrame.h

@@ -74,8 +74,7 @@ class UnwindLocation {
   bool Dereference; /// If true, the resulting location must be dereferenced
                     /// after the location value is computed.
 
-  // Constructors are private to force people to use the create static
-  // functions.
+public:
   UnwindLocation(Location K)
       : Kind(K), RegNum(InvalidRegisterNumber), Offset(0),
         AddrSpace(std::nullopt), Dereference(false) {}
@@ -88,7 +87,6 @@ class UnwindLocation {
       : Kind(DWARFExpr), RegNum(InvalidRegisterNumber), Offset(0), Expr(E),
         Dereference(Deref) {}
 
-public:
   /// Create a location whose rule is set to Unspecified. This means the
   /// register value might be in the same register but it wasn't specified in
   /// the unwind opcodes.
@@ -135,6 +133,7 @@ class UnwindLocation {
     assert(Kind == RegPlusOffset && AddrSpace);
     return *AddrSpace;
   }
+  bool getDeref() const { return Dereference; }
   int32_t getConstant() const { return Offset; }
   /// Some opcodes will modify the CFA location's register only, so we need
   /// to be able to modify the CFA register when evaluating DWARF Call Frame
@@ -148,6 +147,11 @@ class UnwindLocation {
   /// the constant value (DW_CFA_GNU_window_save which is also known as
   // DW_CFA_AARCH64_negate_ra_state).
   void setConstant(int32_t Value) { Offset = Value; }
+  void setDeref(bool NewDeref) { Dereference = NewDeref; }
+  void setKind(Location NewKind) { Kind = NewKind; }
+  bool isRegister() const {
+    return ((Kind == RegPlusOffset) && !Dereference && (Offset == 0));
+  }
 
   std::optional<DWARFExpression> getDWARFExpressionBytes() const {
     return Expr;

llvm/include/llvm/MC/MCDwarf.h

@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/MC/StringTableBuilder.h"
@@ -504,6 +505,7 @@ class MCCFIInstruction {
     OpRestoreState,
     OpOffset,
     OpLLVMDefAspaceCfa,
+    OpLLVMRegOffset,
     OpDefCfaRegister,
     OpDefCfaOffset,
     OpDefCfa,
@@ -518,7 +520,7 @@ class MCCFIInstruction {
     OpNegateRAStateWithPC,
     OpGnuArgsSize,
     OpLabel,
-    OpValOffset,
+    OpValOffset
   };
 
 private:
@@ -537,6 +539,11 @@ class MCCFIInstruction {
       unsigned Register;
       unsigned Register2;
     } RR;
+    struct {
+      unsigned Register;
+      unsigned Register2;
+      int64_t Offset;
+    } RRO;
     MCSymbol *CfiLabel;
   } U;
   OpType Operation;
@@ -569,6 +576,13 @@ class MCCFIInstruction {
     U.CfiLabel = CfiLabel;
   }
 
+  MCCFIInstruction(OpType Op, MCSymbol *L, unsigned R, unsigned R2, int64_t O,
+                   SMLoc Loc, StringRef V, StringRef Comment = "")
+      : Label(L), Operation(Op), Loc(Loc), Values(V.begin(), V.end()), Comment(Comment) {
+    assert(Op == OpLLVMRegOffset);
+    U.RRO = {R, R2, O};
+  }
+
 public:
   /// .cfi_def_cfa defines a rule for computing CFA as: take address from
   /// Register and add Offset to it.
@@ -707,6 +721,15 @@ class MCCFIInstruction {
     return MCCFIInstruction(OpValOffset, L, Register, Offset, Loc);
   }
 
+  static void createRegOffsetExpression(unsigned Reg, unsigned FrameReg, int64_t Offset, SmallString<64>& CFAExpr);
+  static MCCFIInstruction createLLVMRegOffset(MCSymbol *L, unsigned Reg, unsigned FrameReg,
+                                          int64_t Offset, SMLoc Loc = {}, StringRef Comment = "") {
+    // Build up the expression (FrameRegister + Offset)
+    SmallString<64> CFAExpr;
+    createRegOffsetExpression(Reg, FrameReg, Offset, CFAExpr);
+    return MCCFIInstruction(OpLLVMRegOffset, L, Reg, FrameReg, Offset, Loc, CFAExpr, Comment);
+  }
+
   OpType getOperation() const { return Operation; }
   MCSymbol *getLabel() const { return Label; }
 
@@ -715,6 +738,8 @@ class MCCFIInstruction {
       return U.RR.Register;
     if (Operation == OpLLVMDefAspaceCfa)
       return U.RIA.Register;
+    if (Operation == OpLLVMRegOffset)
+      return U.RRO.Register;
     assert(Operation == OpDefCfa || Operation == OpOffset ||
            Operation == OpRestore || Operation == OpUndefined ||
            Operation == OpSameValue || Operation == OpDefCfaRegister ||
@@ -723,8 +748,10 @@ class MCCFIInstruction {
   }
 
   unsigned getRegister2() const {
-    assert(Operation == OpRegister);
-    return U.RR.Register2;
+    if (Operation == OpRegister)
+      return U.RR.Register2;
+    assert (Operation == OpLLVMRegOffset);
+    return U.RRO.Register2;
   }
 
   unsigned getAddressSpace() const {
@@ -735,6 +762,8 @@ class MCCFIInstruction {
   int64_t getOffset() const {
     if (Operation == OpLLVMDefAspaceCfa)
       return U.RIA.Offset;
+    if (Operation == OpLLVMRegOffset)
+      return U.RRO.Offset;
     assert(Operation == OpDefCfa || Operation == OpOffset ||
            Operation == OpRelOffset || Operation == OpDefCfaOffset ||
            Operation == OpAdjustCfaOffset || Operation == OpGnuArgsSize ||
@@ -748,7 +777,7 @@ class MCCFIInstruction {
   }
 
   StringRef getValues() const {
-    assert(Operation == OpEscape);
+    assert(Operation == OpEscape || Operation == OpLLVMRegOffset);
     return StringRef(&Values[0], Values.size());
   }
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinterDwarf.cpp

@@ -223,6 +223,10 @@ void AsmPrinter::emitCFIInstruction(const MCCFIInstruction &Inst) const {
     OutStreamer->emitCFILLVMDefAspaceCfa(Inst.getRegister(), Inst.getOffset(),
                                          Inst.getAddressSpace(), Loc);
     break;
+  case MCCFIInstruction::OpLLVMRegOffset:
+    OutStreamer->AddComment(Inst.getComment());
+    OutStreamer->emitCFIEscape(Inst.getValues(), Loc);
+    break;
   case MCCFIInstruction::OpOffset:
     OutStreamer->emitCFIOffset(Inst.getRegister(), Inst.getOffset(), Loc);
     break;