[llvm] [Draft] Support save/restore point splitting in shrink-wrap (PR #119359)

[dong jianqiang via llvm-commits]

================
@@ -811,9 +975,86 @@ static bool giveUpWithRemarks(MachineOptimizationRemarkEmitter *ORE,
   return false;
 }
 
+void ShrinkWrap::setupSaveRestorePoints(MachineFunction &MF) {
+  for (unsigned Reg : getTargetCSRList(MF)) {
+    auto [Save, Restore] = SavedRegs[Reg];
+    if (SavedRegs.contains(Reg) && Save && Restore)
+      continue;
+
+    SavePoints.insertReg(Reg, &MF.front(), SaveBlocks);
+    for (MachineBasicBlock &MBB : MF) {
+      if (MBB.isEHFuncletEntry())
+        SavePoints.insertReg(Reg, &MBB, SaveBlocks);
+      if (MBB.isReturnBlock())
+        RestorePoints.insertReg(Reg, &MBB, RestoreBlocks);
+    }
+  }
+
+  for (auto [Reg, SaveRestoreBlocks] : SavedRegs) {
+    auto [Save, Restore] = SaveRestoreBlocks;
+    if (Save && Restore) {
+      SavePoints.insertReg(Reg, Save, SaveBlocks);
+      if (!Restore->succ_empty() || Restore->isReturnBlock())
+        RestorePoints.insertReg(Reg, Restore, RestoreBlocks);
+      else
+        RestorePoints.insertReg(Reg, Restore, std::nullopt);
+    }
+  }
+}
+
+bool ShrinkWrap::canSplitSaveRestorePoints(
+    const ReversePostOrderTraversal<MachineBasicBlock *> &RPOT,
+    RegScavenger *RS) {
+  for (MachineBasicBlock *MBB : RPOT) {
+    if (MBB->isEHPad() || MBB->isInlineAsmBrIndirectTarget())
+      return false;
+
+    // Check if we found any stack accesses in the predecessors. We are not
+    // doing a full dataflow analysis here to keep things simple but just
+    // rely on a reverse portorder traversal (RPOT) to guarantee predecessors
+    // are already processed except for loops (and accept the conservative
+    // result for loops).
+    bool StackAddressUsed = any_of(MBB->predecessors(), [&](auto *Pred) {
+      return StackAddressUsedBlockInfo.test(Pred->getNumber());
+    });
+
+    for (const MachineInstr &MI : *MBB) {
+      if (useOrDefFI(MI, RS, StackAddressUsed))
+        return false;
+
+      if (useOrDefCSR(MI, RS, nullptr))
+        StackAddressUsed = true;
+    }
+
+    StackAddressUsedBlockInfo[MBB->getNumber()] = StackAddressUsed;
+  }
+  return true;
+}
+
+void ShrinkWrap::performSimpleShrinkWrap(RegScavenger *RS,
+                                         MachineBasicBlock &SavePoint) {
+  auto MF = SavePoint.getParent();
+  auto CSRs = getTargetCSRList(*MF);
----------------
dongjianqiang2 wrote:

Whenever a basic block touches any callee-saved register (CSR), the pass updates save/restore points for all CSRs. This can make the result overly conservative: saves are moved closer to the function entry and restores closer to the exit, even for registers that are not actually used in the current block.

https://github.com/llvm/llvm-project/pull/119359