[CodeGen] Provide an advanced shrink-wrapping interface

[thegameg]

Summary

In order to extend shrink-wrapping we need to make some changes to the interface.

When trying to extend shrink-wrapping I came across the following limitations in terms of interface:

• We cannot have multiple save points.
• We cannot have multiple restore points.
• We cannot choose which registers are being saved / restored for each block.
• We cannot separate the Prologue / Epilogue emission from the callee saved register storing / loading.
• We cannot reuse the same algorithm to shrink-wrap anything else than CSRs and stack usage.

The goal here would be to provide an interface that would allow us to develop shrink-wrapping extensions in-tree, and being able to switch between implementations without actually having to handle both interfaces in PEI.

In order to be able to simply re-use this, we need to know which basic blocks use which element (as an input), and which basic blocks are the best suited for saving / restoring each element (as an output).

For that, the interface is split in two parts:

ShrinkWrapInfo: it is an interface that determines the "uses" to be used by the shrink-wrapping algorithm.
Example of uses:

• For shrink-wrapping CSR save / restore points: uses/defs of CSRs
• For shrink-wrapping stack setup code: uses/defs of SP, FP, FrameIndex operands

In order to have individual shrink-wrapping for each element, we map every different element to an index and store it in a BitVector. If there is an use of the element in a basic block, the bit at the index of the element is set. The interface provides uses per basic block:

const BitVector *getUses(unsigned MBBNum) const

In order to choose what a "use" is, this class has to be sub-classed and it has to compute the uses per basic block that are consumed by getUses().

• Example 1: Callee saved registers:

0:
    elements:         r0 r1 r2 r3 r4 r5 r6 r7 r8 r9
    BitVector index:  0  1  2  3  4  5  6  7  8  9
    BitVector values: 1  0  0  0  0  0  0  0  0  1
    This means that r0 and r9 are used in %bb.0.

• Example 2: Stack + CSRs:

0:
      elements:         Stack r0 r1
      BitVector index:  0     1  2
      BitVector values: 1     0  1
4:
      elements:         Stack r0 r1
      BitVector index:  0     1  2
      BitVector values: 0     1  0

This means that r0 is used in %bb.4 and the Stack and r0 are used in %bb.0.

Providing this as an input to the shrink-wrapping algorithm allows the algorithm to be independent of the meaning of the input.

ShrinkWrapper: it is an interface for the shrink-wrapping algorithm. It makes it easy to switch between the current one (dominator-based) and other ones like Fred Chow's. This interface provides a result in the following format:

• DenseMap<unsigned, BitVector> Saves, Restores.
• It maps a MachineBasicBlock number to a BitVector containing the elements that need to be saved / restored.

• Example 1: Saves:

0:
    elements:         r0 r1 r2 r3 r4 r5 r6 r7 r8 r9
    BitVector index:  0  1  2  3  4  5  6  7  8  9
    BitVector values: 1  0  0  0  0  0  0  0  0  1

This means that r0 and r9 need to be saved in %bb.0.

• Example 2: Saves:

0:
      elements:         Stack r0 r1
      BitVector index:  0     1  2
      BitVector values: 1     0  1
4:
      elements:         Stack r0 r1
      BitVector index:  0     1  2
      BitVector values: 0     1  0

This means that we have to set up the stack and save r1 in %bb.0, and save r0 in %bb.4.

In order to implement an algorithm, this class has to be sub-classed and compute the results of Saves and Restores in the constructor.

The current dominator-based algorithm port to this interface can allow us to experiment with shrink-wrapping registers and stack uses individually, through the ShrinkWrapInfo class.

This class is meant to be used from PEI to collect better save / restore points, but for now I'll keep it in the ShrinkWrap pass to simplify the review of the patch and focus on the design of the interface.

For this review, I recommend to start by reviewing ShrinkWrapper.h and the ShrinkWrapper / ShrinkWrapInfo abstract classes. The DominatorShrinkWrapper is the port of the current shrink-wrapper to this new interface, and the CSRFIShrinkWrapInfo collects all the basic blocks that use any CSRs or the stack.

[junbuml]

Through this change, do you intend to support the partial shrink-wrapping? What other extensions are you planning on after this change?

[thegameg]

In D41765#968616, @junbuml wrote:

Through this change, do you intend to support the partial shrink-wrapping? What other extensions are you planning on after this change?

After this change I want to experiment with shrink-wrapping CSRs separately and have one save point per register. This would still use the current dominator-based algorithm and would basically allow us to fix all the issues that we encounter in terms of unwinding, debugging, PEI and FrameLowering assumptions about prologues and epilogues, etc.

Then the bigger plan is to integrate https://reviews.llvm.org/D30808 in tree and iterate on it until the results are satisfying enough to switch to it as the default algorithm, but IMO a lot of work will have to go into actually supporting "partial shrink-wrapping" than the algorithm itself.

[thegameg]

DefaultShrinkWrapper doesn't need the RegScavenger, only DefaultShrinkWrapInfo.

[qcolombet]

lib/CodeGen/DefaultShrinkWrapper.cpp

1 ↗

(On Diff #128923)

Could we use a different name that Default because given the plan is to change the default, this is just going to be confusing :)

[junbuml]

include/llvm/CodeGen/ShrinkWrapper.h

61

This means that r0 is used in %bb.4 and the Stack and r1 are used in

lib/CodeGen/DefaultShrinkWrapper.cpp

98 ↗

(On Diff #128923)

Just to remind you that your previous commit (r322584) should be applied in this function.

178 ↗

(On Diff #128923)

As we handle CSRs/FI all together in the current implementation, it will be good to mention that '0' implicitly means that.

393 ↗

(On Diff #128923)

I think Elt should be printed together here.

407 ↗

(On Diff #128923)

Don't you need do increase NumCandidates here ?

[thegameg]

Split DefaultShrinkWrapper and DefaultShrinkWrapInfo into:

• DominatorShrinkWrapper
• CSRFIShrinkWrapInfo

I'm not very happy with these names, but I couldn't think of anything else.

Also,

• Rebase
• Fix NumCandidates increment
• Add enum to give more meaning to the bit 0

[thegameg]

• Stop looping if none of the elements are tracked anymore.
• Let the caller decide what to do with untracked elements.
• Comments / style updates.

[junbuml]

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

103–104

Little confusing about this comment. Should it be like : since we cannot restore before a terminator using CSR?

107

Don't we need to check all terminators in MBB, instead of checking if MI is a terminator?

lib/CodeGen/CSRFIShrinkWrapInfo.h

57

typo : ressources

lib/CodeGen/DominatorShrinkWrapper.cpp

213

what about something like hasElementToTrack() ?

lib/CodeGen/ShrinkWrap.cpp

107–126

Why don't we make this as a method in ShringWrapper?

[thegameg]

Address review concerns.

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

107

Is that really necessary? I couldn't find any way in which this approach (marking the successors) doesn't work.

lib/CodeGen/CSRFIShrinkWrapInfo.h

57

Sorry, I'm not sure what the issue is here?

lib/CodeGen/ShrinkWrap.cpp

107–126

My plan is to separate the ShrinkWrapper and ShrinkWrapInfo from MachineFrameInfo, so that it can be reused easily.

One of my next steps is to move the invocation from the ShrinkWrap pass to PEI. That way, we can easily switch between CSRFIShrinkWrapInfo and something like PartialShrinkWrapInfo where we would track more than one element, so we would have more than one element to save/restore.

For something like PartialShrinkWrapInfo, we would need more than MFI.SavePoint and MFI.RestorePoint. We would need to know which registers need to be saved in which basic block, and separately, where to emit the prologue/epilogue.

So I would like to keep the "interpretation" of the results of the ShrinkWrapper separate from the algorithm itself. We could think of some way to put this in CSRFIShrinkWrapInfo, but I'm not sure how to do it in a reusable way.

[junbuml]

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

107

As far as I understand, if a block has an instruction using csr and also has a terminator using csr. Then, we have to mark its successors as used as well so that we can safely place the restore point after the terminator.
Assuming what I understand is correct, if MI using csr is a non-terminator, then current code may not mark its successor even when MBB's terminator use csr. When MBB is marked, we should also check if its terminators use csr before break in line 110.

lib/CodeGen/CSRFIShrinkWrapInfo.h

57

Sorry I meant typo in line 56 : ressources -> resources.

[thegameg]

• Fix typo
• Iterate on MBB in reverse to catch terminator uses first

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

107

You're totally right, sorry. I was confused between CSRFIShrinkWrapInfo and a version which tracks multiple registers so it has to go through all the instructions.

How about iterating on the instructions in reverse order? We will reach terminators first, and if they use any CSR/FI, we can safely assume that if any non-terminator instructions use CSR/FI, the terminators have been checked first.

[junbuml]

LGTM.
It seems to be a good start to extend shrink-wrapping.
I want to encourage other reviews to take a look.

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

107

Looks good to me. Thanks!

include/llvm/CodeGen/ShrinkWrapper.h

+//===- llvm/CodeGen/ShrinkWrapper.h - Shrink Wrapping Utility ---*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the main interface to use different shrink-wrapping algorithms with
+// different kinds of inputs.
+//
+// For that, the interface is split in two parts:
+//
+// * ShrinkWrapInfo: it is an interface that determines the "uses" to be used
+// by the shrink-wrapping algorithm.
+//
+// * ShrinkWrapper: it is an interface for the shrink-wrapping algorithm. It
+// makes it easy to switch between the current one (dominator-based) and other
+// ones like Fred Chow's.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_SHRINKWRAPPER_H
+#define LLVM_CODEGEN_SHRINKWRAPPER_H
+
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
+#include <memory>
+
+namespace llvm {
+
+class MachineFunction;
+class MachineBasicBlock;
+class MachineBlockFrequencyInfo;
+class MachineOptimizationRemarkEmitter;
+class raw_ostream;
+
+/// Information about the requirements on shrink-wrapping. This should describe
+/// what does "used" mean, and it should be the main interface to work with the
+/// targets and other shrink-wrappable inputs.
+/// It is meant to be sub-classed.
+class ShrinkWrapInfo {
+  /// Example 1: Callee saved registers:
+  ///  0:
+  ///    elements:         r0 r1 r2 r3 r4 r5 r6 r7 r8 r9
+  ///    BitVector index:  0  1  2  3  4  5  6  7  8  9
+  ///    BitVector values: 1  0  0  0  0  0  0  0  0  1
+  ///  This means that r0 and r9 are used in %bb.0.
+  ///
+  /// Example 2: Stack + CSRs:
+  ///  0:
+  ///    elements:         Stack r0 r1
+  ///    BitVector index:  0     1  2
+  ///    BitVector values: 1     0  1
+  ///  4:
+  ///    elements:         Stack r0 r1
+  ///    BitVector index:  0     1  2
+  ///    BitVector values: 0     1  0
+  ///  This means that r0 is used in %bb.4 and the Stack and r1 are used in
+  ///  %bb.0.
+public:
+  ShrinkWrapInfo() = default;
+
+  /// Get the number of elements that are tracked.
+  unsigned getNumElts() const { return getAllUsedElts().size(); }
+
+  /// Return a BitVector containing all the used elements.
+  virtual const BitVector &getAllUsedElts() const = 0;
+
+  /// Get the elements that are used for a particular basic block. The result is
+  /// `nullptr` if there are no uses.
+  virtual const BitVector *getUses(unsigned MBBNum) const = 0;
+
+  virtual ~ShrinkWrapInfo() = default;
+};
+
+/// The interface of a shrink-wrapping algorithm. This class should implement a
+/// shrink-wrapping algorithm. It should get the "uses" from the ShrinkWrapInfo
+/// object and return the save / restore points in the Saves Restores maps. This
+/// allows us to switch shrink-wrapping implementations easily.
+class ShrinkWrapper {
+public:
+  /// The result is a map that maps a MachineBasicBlock number to a BitVector
+  /// containing the elements that need to be saved / restored.
+  ///
+  /// Example 1: Saves:
+  /// 0:
+  ///   elements:         r0 r1 r2 r3 r4 r5 r6 r7 r8 r9
+  ///   BitVector index:  0  1  2  3  4  5  6  7  8  9
+  ///   BitVector values: 1  0  0  0  0  0  0  0  0  1
+  /// This means that r0 and r9 need to be saved in %bb.0.
+  ///
+  /// Example 2: Saves:
+  /// 0:
+  ///   elements:         Stack r0 r1
+  ///   BitVector index:  0     1  2
+  ///   BitVector values: 1     0  1
+  /// 4:
+  ///   elements:         Stack r0 r1
+  ///   BitVector index:  0     1  2
+  ///   BitVector values: 0     1  0
+  /// This means that we have to set up the stack and save r1 in %bb.0, and
+  /// save r0 in %bb.4
+  typedef DenseMap<unsigned, BitVector> BBResultSetMap;
+
+protected:
+  /// Final results.
+  BBResultSetMap SavesResult;
+  BBResultSetMap RestoresResult;
+
+public:
+  /// Run the shrink-wrapper on the function. If there are no uses, there will
+  /// be no saves / restores.
+  ShrinkWrapper() = default;
+
+  virtual ~ShrinkWrapper() = default;
+
+  /// Get the final results.
+  /// If any of the expected elements doesn't show up at all in the Save /
+  /// Restore map, it means there is no better placement for saving / restoring
+  /// that element than the entry / return blocks. The caller has the liberty
+  /// to decide what to do with these elements.
+  const BBResultSetMap &getSaves() { return SavesResult; }
+  const BBResultSetMap &getRestores() { return RestoresResult; }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_SHRINKWRAPPER_H

lib/CodeGen/CMakeLists.txt

 add_llvm_library(LLVMCodeGen
   AggressiveAntiDepBreaker.cpp
   AllocationOrder.cpp
   Analysis.cpp
   AtomicExpandPass.cpp
   BasicTargetTransformInfo.cpp
   BranchFolding.cpp
   BranchRelaxation.cpp
   BreakFalseDeps.cpp
   BuiltinGCs.cpp
   CalcSpillWeights.cpp
   CallingConvLower.cpp
+  CSRFIShrinkWrapInfo.cpp
   CodeGen.cpp
   CodeGenPrepare.cpp
   CriticalAntiDepBreaker.cpp
   DeadMachineInstructionElim.cpp
   DetectDeadLanes.cpp
   DFAPacketizer.cpp
+  DominatorShrinkWrapper.cpp
   DwarfEHPrepare.cpp
   EarlyIfConversion.cpp
   EdgeBundles.cpp
   ExecutionDomainFix.cpp
   ExpandISelPseudos.cpp
   ExpandMemCmp.cpp
   ExpandPostRAPseudos.cpp
   ExpandReductions.cpp

lib/CodeGen/CSRFIShrinkWrapInfo.h

+//===- CSRFIShrinkWrapInfo.h - Track CSR or FI shrink-wrap uses  *- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This ShrinkWrapInfo determines blocks that use either callee-saved registers
+// or contain any stack operations and marks them as used.
+// Since this searches for blocks containing either CSRs or stack operations,
+// we can use a single bit to encode this information.
+//
+// The result should be interpreted in the following way:
+// 4:
+//   elements:         CSRFI
+//   BitVector index:  0
+//   BitVector values: 1
+//
+// Means that %bb.4 uses either CSRs or stack operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_CSRFI_SHRINKWRAPINFO_H
+#define LLVM_CODEGEN_CSRFI_SHRINKWRAPINFO_H
+
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/RegisterClassInfo.h"
+#include "llvm/CodeGen/ShrinkWrapper.h"
+
+namespace llvm {
+
+class RegScavenger;
+
+class CSRFIShrinkWrapInfo : public ShrinkWrapInfo {
+  /// Track the usage per basic block.
+  BitVector Uses;
+
+  using SetOfRegs = SmallSetVector<unsigned, 16>;
+  /// Hold callee-saved information.
+  mutable SetOfRegs CurrentCSRs;
+
+  RegisterClassInfo RCI;
+
+  /// The machine function we're shrink-wrapping.
+  const MachineFunction &MF;
+
+  /// Cache the target's frame setup / destroy opcodes.
+  unsigned FrameSetupOpcode = 0;
+  unsigned FrameDestroyOpcode = 0;
+
+  /// Determine CSRs and cache the result.
+  const SetOfRegs &getCurrentCSRs(RegScavenger *RS) const;
+  /// Return true if MI uses any of the resources that we are interested in.
+  bool useOrDefCSROrFI(const MachineInstr &MI, RegScavenger *RS) const;
+
+public:
+  enum { CSRFIUsedBit = 0, CSRFIShrinkWrapInfoSize };
+
+  /// Run the analysis.
+  CSRFIShrinkWrapInfo(const MachineFunction &MF, RegScavenger *RS);
+
+  const BitVector *getUses(unsigned MBBNum) const override;
+
+  const BitVector &getAllUsedElts() const override;
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_CSRFI_SHRINKWRAPINFO_H

lib/CodeGen/CSRFIShrinkWrapInfo.cpp

+//===- CSRFIShrinkWrapInfo.cpp - Shrink Wrapping Utility -------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+// Implementation of the CSRFIShrinkWrapInfo.
+//===----------------------------------------------------------------------===//
+
+#include "CSRFIShrinkWrapInfo.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "shrink-wrap"
+
+STATISTIC(NumFunc, "Number of functions");
+
+const CSRFIShrinkWrapInfo::SetOfRegs &
+CSRFIShrinkWrapInfo::getCurrentCSRs(RegScavenger *RS) const {
+  if (CurrentCSRs.empty()) {
+    BitVector SavedRegs;
+    const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
+
+    TFI->determineCalleeSaves(const_cast<MachineFunction &>(MF), SavedRegs, RS);
+
+    for (int Reg = SavedRegs.find_first(); Reg != -1;
+         Reg = SavedRegs.find_next(Reg))
+      CurrentCSRs.insert((unsigned)Reg);
+  }
+  return CurrentCSRs;
+}
+
+bool CSRFIShrinkWrapInfo::useOrDefCSROrFI(const MachineInstr &MI,
+                                          RegScavenger *RS) const {
+  if (MI.getOpcode() == FrameSetupOpcode ||
+      MI.getOpcode() == FrameDestroyOpcode) {
+    DEBUG(dbgs() << "Frame instruction: " << MI << '\n');
+    return true;
+  }
+  for (const MachineOperand &MO : MI.operands()) {
+    bool UseOrDefCSR = false;
+    if (MO.isReg()) {
+      // Ignore instructions like DBG_VALUE which don't read/def the register.
+      if (!MO.isDef() && !MO.readsReg())
+        continue;
+      unsigned PhysReg = MO.getReg();
+      if (!PhysReg)
+        continue;
+      assert(TargetRegisterInfo::isPhysicalRegister(PhysReg) &&
+             "Unallocated register?!");
+      UseOrDefCSR = RCI.getLastCalleeSavedAlias(PhysReg);
+    } else if (MO.isRegMask()) {
+      // Check if this regmask clobbers any of the CSRs.
+      for (unsigned Reg : getCurrentCSRs(RS)) {
+        if (MO.clobbersPhysReg(Reg)) {
+          UseOrDefCSR = true;
+          break;
+        }
+      }
+    }
+    // Skip FrameIndex operands in DBG_VALUE instructions.
+    if (UseOrDefCSR || (MO.isFI() && !MI.isDebugValue())) {
+      DEBUG(dbgs() << "Use or define CSR(" << UseOrDefCSR << ") or FI("
+                   << MO.isFI() << "): " << MI << '\n');
+      return true;
+    }
+  }
+  return false;
+}
+
+CSRFIShrinkWrapInfo::CSRFIShrinkWrapInfo(const MachineFunction &MF,
+                                         RegScavenger *RS)
+    : ShrinkWrapInfo(), Uses(MF.getNumBlockIDs()), MF(MF),
+      FrameSetupOpcode(
+          MF.getSubtarget().getInstrInfo()->getCallFrameSetupOpcode()),
+      FrameDestroyOpcode(
+          MF.getSubtarget().getInstrInfo()->getCallFrameDestroyOpcode()) {
+  ++NumFunc;
+
+  RCI.runOnMachineFunction(MF);
+  for (const MachineBasicBlock &MBB : MF) {
+    DEBUG(dbgs() << "Check for uses: " << printMBBReference(MBB) << ' '
+                 << MBB.getName() << '\n');
+
+    if (MBB.isEHFuncletEntry()) {
+      DEBUG(dbgs() << "EH Funclets are not supported yet.\n");
+      Uses.clear();
+      return;
+    }
+    for (const MachineInstr &MI :
+         make_range(MBB.rbegin(), MBB.rend())) {
+      if (useOrDefCSROrFI(MI, RS)) {
+        Uses.set(MBB.getNumber());
+        // Make sure we would be able to insert the restore code before the
+        // terminator. Mark the successors as used as well, since we can't
+        // restore after a terminator (i.e. cbz w23, #10).
+        // Note: This doesn't handle the case when a return instruction uses a
+        // CSR/FI. If that happens, the default placement in all the return
+        // blocks won't work either.
+        if (MI.isTerminator())
+          for (const MachineBasicBlock *Succ : MBB.successors())
+            Uses.set(Succ->getNumber());
+        break;
+      }
+    }
+  }
+}
+
+const BitVector *CSRFIShrinkWrapInfo::getUses(unsigned MBBNum) const {
+  if (!Uses.empty() && Uses.test(MBBNum)) {
+    // Create a dummy BitVector having only one bit which is set.
+    static BitVector HasUses;
+    if (HasUses.empty()) {
+      HasUses.resize(CSRFIShrinkWrapInfoSize);
+      HasUses.set(CSRFIUsedBit);
+    }
+    return &HasUses;
+  }
+  return nullptr;
+}
+
+const BitVector &CSRFIShrinkWrapInfo::getAllUsedElts() const {
+  if (!Uses.empty() && Uses.any()) {
+    // Create a dummy BitVector having only one bit which is set.
+    static BitVector Set;
+    if (Set.empty())
+      Set.resize(CSRFIShrinkWrapInfoSize, true);
+    return Set;
+  }
+  // Create a dummy BitVector having only one bit which is unset.
+  static BitVector Unset;
+  if (Unset.empty())
+    Unset.resize(CSRFIShrinkWrapInfoSize, false);
+  return Unset;
+}

lib/CodeGen/DominatorShrinkWrapper.h

+//===- DominatorShrinkWrapper.h - Dominator-based shrink-wrapping *- C++ *-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass looks for safe point where the prologue and epilogue can be
+// inserted.
+// The safe point for the prologue (resp. epilogue) is called Save
+// (resp. Restore).
+// A point is safe for prologue (resp. epilogue) if and only if
+// it 1) dominates (resp. post-dominates) all the frame related operations and
+// between 2) two executions of the Save (resp. Restore) point there is an
+// execution of the Restore (resp. Save) point.
+//
+// For instance, the following points are safe:
+// for (int i = 0; i < 10; ++i) {
+//   Save
+//   ...
+//   Restore
+// }
+// Indeed, the execution looks like Save -> Restore -> Save -> Restore ...
+// And the following points are not:
+// for (int i = 0; i < 10; ++i) {
+//   Save
+//   ...
+// }
+// for (int i = 0; i < 10; ++i) {
+//   ...
+//   Restore
+// }
+// Indeed, the execution looks like Save -> Save -> ... -> Restore -> Restore.
+//
+// This pass also ensures that the safe points are 3) cheaper than the regular
+// entry and exits blocks.
+//
+// Property #1 is ensured via the use of MachineDominatorTree and
+// MachinePostDominatorTree.
+// Property #2 is ensured via property #1 and MachineLoopInfo, i.e., both
+// points must be in the same loop.
+// Property #3 is ensured via the MachineBlockFrequencyInfo.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_DOMINATOR_SHRINKWRAPPER_H
+#define LLVM_CODEGEN_DOMINATOR_SHRINKWRAPPER_H
+
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/RegisterClassInfo.h"
+#include "llvm/CodeGen/ShrinkWrapper.h"
+#include "llvm/Support/Debug.h"
+
+namespace llvm {
+
+class MachineDominatorTree;
+class MachineLoopInfo;
+struct MachinePostDominatorTree;
+
+/// \brief Class to determine where the safe point to insert the
+/// prologue and epilogue are.
+/// Unlike the paper from Fred C. Chow, PLDI'88, that introduces the
+/// shrink-wrapping term for prologue/epilogue placement, this pass
+/// does not rely on expensive data-flow analysis. Instead we use the
+/// dominance properties and loop information to decide which point
+/// are safe for such insertion.
+class DominatorShrinkWrapper : public ShrinkWrapper {
+  /// This algorithm is based on common dominators.
+  MachineDominatorTree &MDT;
+  MachinePostDominatorTree &MPDT;
+
+  /// Hold the loop information. Used to determine if Save and Restore
+  /// are in the same loop.
+  MachineLoopInfo &MLI;
+
+  /// Current safe point found for saving each element.
+  /// Example:
+  /// The prologue will be inserted before the first instruction
+  /// in this basic block.
+  SmallVector<MachineBasicBlock *, 1> Saves;
+
+  /// Current safe point found for restoring each element.
+  /// Example:
+  /// The epilogue will be inserted before the first terminator instruction
+  /// in this basic block.
+  SmallVector<MachineBasicBlock *, 1> Restores;
+
+  // Still tracking.
+  BitVector Tracking;
+
+  /// Entry block.
+  const MachineBasicBlock &Entry;
+
+  /// Check whether or not Save and Restore points are still interesting for
+  /// shrink-wrapping.
+  bool arePointsInteresting(unsigned Elt) const {
+    return Saves[Elt] != &Entry && Saves[Elt] && Restores[Elt];
+  }
+
+  /// Check if we're still tracking \p Elt or we already gave up on it.
+  bool stillTrackingElt(unsigned Elt) const { return Tracking.test(Elt); }
+  /// Check if we still have elements to track.
+  bool hasElementToTrack() const { return Tracking.any(); }
+
+  /// Mark this element as untracked.
+  void stopTrackingElt(unsigned Elt) {
+    Tracking.reset(Elt);
+    Saves[Elt] = nullptr;
+    Restores[Elt] = nullptr;
+  }
+  /// Mark all elements as untracked.
+  void stopTrackingAllElts() {
+    Tracking.reset();
+  }
+
+  /// \brief Update the Save and Restore points such that \p MBB is in
+  /// the region that is dominated by Save and post-dominated by Restore
+  /// and Save and Restore still match the safe point definition.
+  /// Such point may not exist and Save and/or Restore may be null after
+  /// this call.
+  void updateSaveRestorePoints(MachineBasicBlock &MBB, unsigned Elt);
+
+public:
+  DominatorShrinkWrapper(MachineFunction &MF, const ShrinkWrapInfo &SWI,
+                         MachineDominatorTree &MDT,
+                         MachinePostDominatorTree &MPDT,
+                         MachineBlockFrequencyInfo &MBFI, MachineLoopInfo &MLI);
+
+  /// Return the elements that are tracked in the final result.
+  /// Untracked elements won't show up in the maps returned by `getSaves()` and
+  /// `getRestores()`. The usual behavior is to add them to the entry and
+  /// return blocks, but that decision is left to the caller.
+  const BitVector &getTrackedElts() { return Tracking; }
+};
+
+} // end namespace llvm
+#endif // LLVM_CODEGEN_DOMINATOR_SHRINKWRAPPER_H

lib/CodeGen/DominatorShrinkWrapper.cpp

+//===- DominatorShrinkWrapper.cpp - Shrink Wrapping Utility -----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+// Implementation of the DominatorShrinkWrap(per / Info).
+//===----------------------------------------------------------------------===//
+
+#include "DominatorShrinkWrapper.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineDominators.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "shrink-wrap"
+
+STATISTIC(NumCandidates, "Number of shrink-wrapping candidates");
+STATISTIC(NumCandidatesDropped,
+          "Number of shrink-wrapping candidates dropped because of frequency");
+
+/// Check whether the edge (\p SrcBB, \p DestBB) is a backedge according to MLI.
+/// I.e., check if it exists a loop that contains SrcBB and where DestBB is the
+/// loop header.
+static bool isProperBackedge(const MachineLoopInfo &MLI,
+                             const MachineBasicBlock *SrcBB,
+                             const MachineBasicBlock *DestBB) {
+  for (const MachineLoop *Loop = MLI.getLoopFor(SrcBB); Loop;
+       Loop = Loop->getParentLoop()) {
+    if (Loop->getHeader() == DestBB)
+      return true;
+  }
+  return false;
+}
+
+/// Check if the CFG of \p MF is irreducible.
+static bool isIrreducibleCFG(const MachineFunction &MF,
+                             const MachineLoopInfo &MLI) {
+  const MachineBasicBlock *Entry = &*MF.begin();
+  ReversePostOrderTraversal<const MachineBasicBlock *> RPOT(Entry);
+  BitVector VisitedBB(MF.getNumBlockIDs());
+  for (const MachineBasicBlock *MBB : RPOT) {
+    VisitedBB.set(MBB->getNumber());
+    for (const MachineBasicBlock *SuccBB : MBB->successors()) {
+      if (!VisitedBB.test(SuccBB->getNumber()))
+        continue;
+      // We already visited SuccBB, thus MBB->SuccBB must be a backedge.
+      // Check that the head matches what we have in the loop information.
+      // Otherwise, we have an irreducible graph.
+      if (!isProperBackedge(MLI, MBB, SuccBB))
+        return true;
+    }
+  }
+  return false;
+}
+
+/// \brief Helper function to find the immediate (post) dominator.
+template <typename ListOfBBs, typename DominanceAnalysis>
+static MachineBasicBlock *FindIDom(MachineBasicBlock &Block, ListOfBBs BBs,
+                                   DominanceAnalysis &Dom) {
+  MachineBasicBlock *IDom = &Block;
+  for (MachineBasicBlock *BB : BBs) {
+    IDom = Dom.findNearestCommonDominator(IDom, BB);
+    if (!IDom)
+      break;
+  }
+  if (IDom == &Block)
+    return nullptr;
+  return IDom;
+}
+
+void DominatorShrinkWrapper::updateSaveRestorePoints(MachineBasicBlock &MBB,
+                                                     unsigned Elt) {
+
+  MachineBasicBlock *&Save = Saves[Elt];
+  MachineBasicBlock *&Restore = Restores[Elt];
+
+  // Get rid of the easy cases first.
+  if (!Save)
+    Save = &MBB;
+  else
+    Save = MDT.findNearestCommonDominator(Save, &MBB);
+
+  if (!Save) {
+    DEBUG(dbgs() << "Found a block that is not reachable from Entry\n");
+    return;
+  }
+
+  if (!Restore)
+    Restore = &MBB;
+  else if (MPDT.getNode(&MBB)) // If the block is not in the post dom tree, it
+                               // means the block never returns. If that's the
+                               // case, we don't want to call
+                               // `findNearestCommonDominator`, which will
+                               // return `Restore`.
+    Restore = MPDT.findNearestCommonDominator(Restore, &MBB);
+  else
+    stopTrackingElt(Elt); // Abort, we can't find a restore point in this case.
+
+  if (!Restore) {
+    DEBUG(dbgs() << "Restore point needs to be spanned on several blocks\n");
+    return;
+  }
+
+  // Make sure Save and Restore are suitable for shrink-wrapping:
+  // 1. all path from Save needs to lead to Restore before exiting.
+  // 2. all path to Restore needs to go through Save from Entry.
+  // We achieve that by making sure that:
+  // A. Save dominates Restore.
+  // B. Restore post-dominates Save.
+  // C. Save and Restore are in the same loop.
+  bool SaveDominatesRestore = false;
+  bool RestorePostDominatesSave = false;
+  while (Save && Restore &&
+         (!(SaveDominatesRestore = MDT.dominates(Save, Restore)) ||
+          !(RestorePostDominatesSave = MPDT.dominates(Restore, Save)) ||
+          // Post-dominance is not enough in loops to ensure that all uses/defs
+          // are after the prologue and before the epilogue at runtime.
+          // E.g.,
+          // while(1) {
+          //  Save
+          //  Restore
+          //   if (...)
+          //     break;
+          //  use/def CSRs
+          // }
+          // All the uses/defs of CSRs are dominated by Save and post-dominated
+          // by Restore. However, the CSRs uses are still reachable after
+          // Restore and before Save are executed.
+          //
+          // For now, just push the restore/save points outside of loops.
+          // FIXME: Refine the criteria to still find interesting cases
+          // for loops.
+          MLI.getLoopFor(Save) || MLI.getLoopFor(Restore))) {
+    // Fix (A).
+    if (!SaveDominatesRestore) {
+      Save = MDT.findNearestCommonDominator(Save, Restore);
+      continue;
+    }
+    // Fix (B).
+    if (!RestorePostDominatesSave)
+      Restore = MPDT.findNearestCommonDominator(Restore, Save);
+
+    // Fix (C).
+    if (Save && Restore && (MLI.getLoopFor(Save) || MLI.getLoopFor(Restore))) {
+      if (MLI.getLoopDepth(Save) > MLI.getLoopDepth(Restore)) {
+        // Push Save outside of this loop if immediate dominator is
+        // different from save block. If immediate dominator is not
+        // different, bail out.
+        Save = FindIDom<>(*Save, Save->predecessors(), MDT);
+        if (!Save)
+          break;
+      } else {
+        // If the loop does not exit, there is no point in looking
+        // for a post-dominator outside the loop.
+        SmallVector<MachineBasicBlock *, 4> ExitBlocks;
+        MLI.getLoopFor(Restore)->getExitingBlocks(ExitBlocks);
+        // Push Restore outside of this loop.
+        // Look for the immediate post-dominator of the loop exits.
+        MachineBasicBlock *IPdom = Restore;
+        for (MachineBasicBlock *LoopExitBB : ExitBlocks) {
+          IPdom = FindIDom<>(*IPdom, LoopExitBB->successors(), MPDT);
+          if (!IPdom)
+            break;
+        }
+        // If the immediate post-dominator is not in a less nested loop,
+        // then we are stuck in a program with an infinite loop.
+        // In that case, we will not find a safe point, hence, bail out.
+        if (IPdom && MLI.getLoopDepth(IPdom) < MLI.getLoopDepth(Restore))
+          Restore = IPdom;
+        else {
+          stopTrackingElt(Elt);
+          break;
+        }
+      }
+    }
+  }
+}
+
+DominatorShrinkWrapper::DominatorShrinkWrapper(MachineFunction &MF,
+                                               const ShrinkWrapInfo &SWI,
+                                               MachineDominatorTree &MDT,
+                                               MachinePostDominatorTree &MPDT,
+                                               MachineBlockFrequencyInfo &MBFI,
+                                               MachineLoopInfo &MLI)
+    : ShrinkWrapper(), MDT(MDT), MPDT(MPDT), MLI(MLI), Saves(SWI.getNumElts()),
+      Restores(SWI.getNumElts()), Tracking(SWI.getAllUsedElts()),
+      Entry(MF.front()) {
+  if (isIrreducibleCFG(MF, MLI)) {
+    // If MF is irreducible, a block may be in a loop without
+    // MachineLoopInfo reporting it. I.e., we may use the
+    // post-dominance property in loops, which lead to incorrect
+    // results. Moreover, we may miss that the prologue and
+    // epilogue are not in the same loop, leading to unbalanced
+    // construction/deconstruction of the stack frame.
+    DEBUG(dbgs() << "Irreducible CFGs are not supported yet\n");
+    stopTrackingAllElts();
+    return;
+  }
+
+  for (MachineBasicBlock &MBB : MF) {
+    if (!hasElementToTrack())
+      return;
+    DEBUG(dbgs() << "Look into: " << printMBBReference(MBB) << '\n');
+    if (const BitVector *Uses = SWI.getUses(MBB.getNumber())) {
+      for (unsigned Elt : Uses->set_bits()) {
+        if (!stillTrackingElt(Elt))
+          continue;
+        // FIXME: Add custom printing through SWI.
+        DEBUG(dbgs() << "For elt " << Elt << ":\n");
+        // Save (resp. restore) point must dominate (resp. post dominate)
+        // MI. Look for the proper basic block for those.
+        updateSaveRestorePoints(MBB, Elt);
+        // If we are at a point where we cannot improve the placement of
+        // save/restore instructions, just give up.
+        if (!arePointsInteresting(Elt)) {
+          DEBUG(dbgs() << "No Shrink wrap candidate found\n");
+          stopTrackingElt(Elt);
+          continue;
+        }
+      }
+    }
+  }
+
+  for (unsigned Elt : Tracking.set_bits()) {
+    MachineBasicBlock *&Save = Saves[Elt];
+    MachineBasicBlock *&Restore = Restores[Elt];
+
+    DEBUG(dbgs() << "Post-process elt " << Elt << ":\n");
+    if (!arePointsInteresting(Elt)) {
+      // If the points are not interesting at this point, then they must be
+      // null because it means we did not encounter any frame/CSR related
+      // code. Otherwise, we would have returned from the previous loop.
+      assert(!Save && !Restore && "We miss a shrink-wrap opportunity?!");
+      DEBUG(dbgs() << "Nothing to shrink-wrap\n");
+      continue;
+    }
+
+    uint64_t EntryFreq = MBFI.getEntryFreq();
+    DEBUG(dbgs() << "\n ** Results **\nFrequency of the Entry: " << EntryFreq
+                 << '\n');
+
+    const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
+    do {
+      DEBUG(dbgs() << "Shrink wrap candidates (#, Name, Freq):\nSave: "
+                   << Save->getNumber() << ' ' << Save->getName() << ' '
+                   << MBFI.getBlockFreq(Save).getFrequency() << "\nRestore: "
+                   << Restore->getNumber() << ' ' << Restore->getName() << ' '
+                   << MBFI.getBlockFreq(Restore).getFrequency() << '\n');
+
+      bool IsSaveCheap, TargetCanUseSaveAsPrologue = false;
+      if (((IsSaveCheap =
+                EntryFreq >= MBFI.getBlockFreq(Save).getFrequency()) &&
+           EntryFreq >= MBFI.getBlockFreq(Restore).getFrequency()) &&
+          ((TargetCanUseSaveAsPrologue = TFI->canUseAsPrologue(*Save)) &&
+           TFI->canUseAsEpilogue(*Restore)))
+        break;
+      DEBUG(
+          dbgs() << "New points are too expensive or invalid for the target\n");
+      MachineBasicBlock *NewBB;
+      if (!IsSaveCheap || !TargetCanUseSaveAsPrologue) {
+        Save = FindIDom<>(*Save, Save->predecessors(), MDT);
+        if (!Save)
+          break;
+        NewBB = Save;
+      } else {
+        // Restore is expensive.
+        Restore = FindIDom<>(*Restore, Restore->successors(), MPDT);
+        if (!Restore)
+          break;
+        NewBB = Restore;
+      }
+      updateSaveRestorePoints(*NewBB, Elt);
+    } while (Save && Restore);
+
+    if (!arePointsInteresting(Elt)) {
+      ++NumCandidatesDropped;
+      stopTrackingElt(Elt);
+      continue;
+    }
+
+    DEBUG(dbgs() << "Final shrink wrap candidates:\nSave: " << Save->getNumber()
+                 << ' ' << Save->getName() << "\nRestore: "
+                 << Restore->getNumber() << ' ' << Restore->getName() << '\n');
+
+    /// "return" the blocks where we need to save / restore this elt.
+    /// In this case, only one save and one restore point is returned.
+    BitVector &SaveResult = SavesResult[Save->getNumber()];
+    if (SaveResult.empty())
+      SaveResult.resize(SWI.getNumElts());
+    SaveResult.set(Elt);
+
+    BitVector &RestoreResult = RestoresResult[Restore->getNumber()];
+    if (RestoreResult.empty())
+      RestoreResult.resize(SWI.getNumElts());
+    RestoreResult.set(Elt);
+
+    ++NumCandidates;
+  }
+}

lib/CodeGen/ShrinkWrap.cpp

 //===- ShrinkWrap.cpp - Compute safe point for prolog/epilog insertion ----===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
-// This pass looks for safe point where the prologue and epilogue can be
-// inserted.
-// The safe point for the prologue (resp. epilogue) is called Save
-// (resp. Restore).
-// A point is safe for prologue (resp. epilogue) if and only if
-// it 1) dominates (resp. post-dominates) all the frame related operations and
-// between 2) two executions of the Save (resp. Restore) point there is an
-// execution of the Restore (resp. Save) point.
-//
-// For instance, the following points are safe:
-// for (int i = 0; i < 10; ++i) {
-//   Save
-//   ...
-//   Restore
-// }
-// Indeed, the execution looks like Save -> Restore -> Save -> Restore ...
-// And the following points are not:
-// for (int i = 0; i < 10; ++i) {
-//   Save
-//   ...
-// }
-// for (int i = 0; i < 10; ++i) {
-//   ...
-//   Restore
-// }
-// Indeed, the execution looks like Save -> Save -> ... -> Restore -> Restore.
-//
-// This pass also ensures that the safe points are 3) cheaper than the regular
-// entry and exits blocks.
-//
-// Property #1 is ensured via the use of MachineDominatorTree and
-// MachinePostDominatorTree.
-// Property #2 is ensured via property #1 and MachineLoopInfo, i.e., both
-// points must be in the same loop.
-// Property #3 is ensured via the MachineBlockFrequencyInfo.
-//
-// If this pass found points matching all these properties, then
-// MachineFrameInfo is updated with this information.
+// This pass takes care of adding the correct dependencies, running the
+// shrink-wrapping algorithm and update MachineFrameInfo with the results if
+// they are available.
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ADT/BitVector.h"
-#include "llvm/ADT/PostOrderIterator.h"
-#include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/Statistic.h"
-#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "DominatorShrinkWrapper.h"
+#include "CSRFIShrinkWrapInfo.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
-#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
-#include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachinePostDominators.h"
-#include "llvm/CodeGen/RegisterClassInfo.h"
+#include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
-#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
-#include "llvm/IR/Attributes.h"
-#include "llvm/IR/Function.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/Pass.h"
-#include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
-#include <cassert>
-#include <cstdint>
-#include <memory>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "shrink-wrap"
 
-STATISTIC(NumFunc, "Number of functions");
-STATISTIC(NumCandidates, "Number of shrink-wrapping candidates");
-STATISTIC(NumCandidatesDropped,
-          "Number of shrink-wrapping candidates dropped because of frequency");
-
 static cl::opt<cl::boolOrDefault>
 EnableShrinkWrapOpt("enable-shrink-wrap", cl::Hidden,
                     cl::desc("enable the shrink-wrapping pass"));
 
 namespace {
 
-/// \brief Class to determine where the safe point to insert the
-/// prologue and epilogue are.
-/// Unlike the paper from Fred C. Chow, PLDI'88, that introduces the
-/// shrink-wrapping term for prologue/epilogue placement, this pass
-/// does not rely on expensive data-flow analysis. Instead we use the
-/// dominance properties and loop information to decide which point
-/// are safe for such insertion.
 class ShrinkWrap : public MachineFunctionPass {
-  /// Hold callee-saved information.
-  RegisterClassInfo RCI;
-  MachineDominatorTree *MDT;
-  MachinePostDominatorTree *MPDT;
-
-  /// Current safe point found for the prologue.
-  /// The prologue will be inserted before the first instruction
-  /// in this basic block.
-  MachineBasicBlock *Save;
-
-  /// Current safe point found for the epilogue.
-  /// The epilogue will be inserted before the first terminator instruction
-  /// in this basic block.
-  MachineBasicBlock *Restore;
-
-  /// Hold the information of the basic block frequency.
-  /// Use to check the profitability of the new points.
-  MachineBlockFrequencyInfo *MBFI;
-
-  /// Hold the loop information. Used to determine if Save and Restore
-  /// are in the same loop.
-  MachineLoopInfo *MLI;
-
-  /// Frequency of the Entry block.
-  uint64_t EntryFreq;
-
-  /// Current opcode for frame setup.
-  unsigned FrameSetupOpcode;
-
-  /// Current opcode for frame destroy.
-  unsigned FrameDestroyOpcode;
-
-  /// Entry block.
-  const MachineBasicBlock *Entry;
-
-  using SetOfRegs = SmallSetVector<unsigned, 16>;
-
-  /// Registers that need to be saved for the current function.
-  mutable SetOfRegs CurrentCSRs;
-
-  /// Current MachineFunction.
-  MachineFunction *MachineFunc;
-
-  /// \brief Check if \p MI uses or defines a callee-saved register or
-  /// a frame index. If this is the case, this means \p MI must happen
-  /// after Save and before Restore.
-  bool useOrDefCSROrFI(const MachineInstr &MI, RegScavenger *RS) const;
-
-  const SetOfRegs &getCurrentCSRs(RegScavenger *RS) const {
-    if (CurrentCSRs.empty()) {
-      BitVector SavedRegs;
-      const TargetFrameLowering *TFI =
-          MachineFunc->getSubtarget().getFrameLowering();
-
-      TFI->determineCalleeSaves(*MachineFunc, SavedRegs, RS);
-
-      for (int Reg = SavedRegs.find_first(); Reg != -1;
-           Reg = SavedRegs.find_next(Reg))
-        CurrentCSRs.insert((unsigned)Reg);
-    }
-    return CurrentCSRs;
-  }
-
-  /// \brief Update the Save and Restore points such that \p MBB is in
-  /// the region that is dominated by Save and post-dominated by Restore
-  /// and Save and Restore still match the safe point definition.
-  /// Such point may not exist and Save and/or Restore may be null after
-  /// this call.
-  void updateSaveRestorePoints(MachineBasicBlock &MBB, RegScavenger *RS);
-
-  /// \brief Initialize the pass for \p MF.
-  void init(MachineFunction &MF) {
-    RCI.runOnMachineFunction(MF);
-    MDT = &getAnalysis<MachineDominatorTree>();
-    MPDT = &getAnalysis<MachinePostDominatorTree>();
-    Save = nullptr;
-    Restore = nullptr;
-    MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
-    MLI = &getAnalysis<MachineLoopInfo>();
-    EntryFreq = MBFI->getEntryFreq();
-    const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
-    FrameSetupOpcode = TII.getCallFrameSetupOpcode();
-    FrameDestroyOpcode = TII.getCallFrameDestroyOpcode();
-    Entry = &MF.front();
-    CurrentCSRs.clear();
-    MachineFunc = &MF;
-
-    ++NumFunc;
-  }
-
-  /// Check whether or not Save and Restore points are still interesting for
-  /// shrink-wrapping.
-  bool ArePointsInteresting() const { return Save != Entry && Save && Restore; }
-
   /// \brief Check if shrink wrapping is enabled for this target and function.
   static bool isShrinkWrapEnabled(const MachineFunction &MF);
 
 public:
   static char ID;
 
   ShrinkWrap() : MachineFunctionPass(ID) {
     initializeShrinkWrapPass(*PassRegistry::getPassRegistry());

 
 INITIALIZE_PASS_BEGIN(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineBlockFrequencyInfo)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(MachinePostDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_END(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
 
-bool ShrinkWrap::useOrDefCSROrFI(const MachineInstr &MI,
-                                 RegScavenger *RS) const {
-  if (MI.getOpcode() == FrameSetupOpcode ||
-      MI.getOpcode() == FrameDestroyOpcode) {
-    DEBUG(dbgs() << "Frame instruction: " << MI << '\n');
-    return true;
-  }
-  for (const MachineOperand &MO : MI.operands()) {
-    bool UseOrDefCSR = false;
-    if (MO.isReg()) {
-      // Ignore instructions like DBG_VALUE which don't read/def the register.
-      if (!MO.isDef() && !MO.readsReg())
-        continue;
-      unsigned PhysReg = MO.getReg();
-      if (!PhysReg)
-        continue;
-      assert(TargetRegisterInfo::isPhysicalRegister(PhysReg) &&
-             "Unallocated register?!");
-      UseOrDefCSR = RCI.getLastCalleeSavedAlias(PhysReg);
-    } else if (MO.isRegMask()) {
-      // Check if this regmask clobbers any of the CSRs.
-      for (unsigned Reg : getCurrentCSRs(RS)) {
-        if (MO.clobbersPhysReg(Reg)) {
-          UseOrDefCSR = true;
-          break;
-        }
-      }
-    }
-    // Skip FrameIndex operands in DBG_VALUE instructions.
-    if (UseOrDefCSR || (MO.isFI() && !MI.isDebugValue())) {
-      DEBUG(dbgs() << "Use or define CSR(" << UseOrDefCSR << ") or FI("
-                   << MO.isFI() << "): " << MI << '\n');
-      return true;
-    }
-  }
-  return false;
-}
-
-/// \brief Helper function to find the immediate (post) dominator.
-template <typename ListOfBBs, typename DominanceAnalysis>
-static MachineBasicBlock *FindIDom(MachineBasicBlock &Block, ListOfBBs BBs,
-                                   DominanceAnalysis &Dom) {
-  MachineBasicBlock *IDom = &Block;
-  for (MachineBasicBlock *BB : BBs) {
-    IDom = Dom.findNearestCommonDominator(IDom, BB);
-    if (!IDom)
-      break;
-  }
-  if (IDom == &Block)
-    return nullptr;
-  return IDom;
-}
-
-void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB,
-                                         RegScavenger *RS) {
-  // Get rid of the easy cases first.
-  if (!Save)
-    Save = &MBB;
-  else
-    Save = MDT->findNearestCommonDominator(Save, &MBB);
-
-  if (!Save) {
-    DEBUG(dbgs() << "Found a block that is not reachable from Entry\n");
-    return;
-  }
-
-  if (!Restore)
-    Restore = &MBB;
-  else if (MPDT->getNode(&MBB)) // If the block is not in the post dom tree, it
-                                // means the block never returns. If that's the
-                                // case, we don't want to call
-                                // `findNearestCommonDominator`, which will
-                                // return `Restore`.
-    Restore = MPDT->findNearestCommonDominator(Restore, &MBB);
-  else
-    Restore = nullptr; // Abort, we can't find a restore point in this case.
-
-  // Make sure we would be able to insert the restore code before the
-  // terminator.
-  if (Restore == &MBB) {
-    for (const MachineInstr &Terminator : MBB.terminators()) {
-      if (!useOrDefCSROrFI(Terminator, RS))
-        continue;
-      // One of the terminator needs to happen before the restore point.
-      if (MBB.succ_empty()) {
-        Restore = nullptr; // Abort, we can't find a restore point in this case.
-        break;
-      }
-      // Look for a restore point that post-dominates all the successors.
-      // The immediate post-dominator is what we are looking for.
-      Restore = FindIDom<>(*Restore, Restore->successors(), *MPDT);
-      break;
-    }
-  }
-
-  if (!Restore) {
-    DEBUG(dbgs() << "Restore point needs to be spanned on several blocks\n");
-    return;
-  }
-
-  // Make sure Save and Restore are suitable for shrink-wrapping:
-  // 1. all path from Save needs to lead to Restore before exiting.
-  // 2. all path to Restore needs to go through Save from Entry.
-  // We achieve that by making sure that:
-  // A. Save dominates Restore.
-  // B. Restore post-dominates Save.
-  // C. Save and Restore are in the same loop.
-  bool SaveDominatesRestore = false;
-  bool RestorePostDominatesSave = false;
-  while (Save && Restore &&
-         (!(SaveDominatesRestore = MDT->dominates(Save, Restore)) ||
-          !(RestorePostDominatesSave = MPDT->dominates(Restore, Save)) ||
-          // Post-dominance is not enough in loops to ensure that all uses/defs
-          // are after the prologue and before the epilogue at runtime.
-          // E.g.,
-          // while(1) {
-          //  Save
-          //  Restore
-          //   if (...)
-          //     break;
-          //  use/def CSRs
-          // }
-          // All the uses/defs of CSRs are dominated by Save and post-dominated
-          // by Restore. However, the CSRs uses are still reachable after
-          // Restore and before Save are executed.
-          //
-          // For now, just push the restore/save points outside of loops.
-          // FIXME: Refine the criteria to still find interesting cases
-          // for loops.
-          MLI->getLoopFor(Save) || MLI->getLoopFor(Restore))) {
-    // Fix (A).
-    if (!SaveDominatesRestore) {
-      Save = MDT->findNearestCommonDominator(Save, Restore);
-      continue;
-    }
-    // Fix (B).
-    if (!RestorePostDominatesSave)
-      Restore = MPDT->findNearestCommonDominator(Restore, Save);
-
-    // Fix (C).
-    if (Save && Restore &&
-        (MLI->getLoopFor(Save) || MLI->getLoopFor(Restore))) {
-      if (MLI->getLoopDepth(Save) > MLI->getLoopDepth(Restore)) {
-        // Push Save outside of this loop if immediate dominator is different
-        // from save block. If immediate dominator is not different, bail out.
-        Save = FindIDom<>(*Save, Save->predecessors(), *MDT);
-        if (!Save)
-          break;
-      } else {
-        // If the loop does not exit, there is no point in looking
-        // for a post-dominator outside the loop.
-        SmallVector<MachineBasicBlock*, 4> ExitBlocks;
-        MLI->getLoopFor(Restore)->getExitingBlocks(ExitBlocks);
-        // Push Restore outside of this loop.
-        // Look for the immediate post-dominator of the loop exits.
-        MachineBasicBlock *IPdom = Restore;
-        for (MachineBasicBlock *LoopExitBB: ExitBlocks) {
-          IPdom = FindIDom<>(*IPdom, LoopExitBB->successors(), *MPDT);
-          if (!IPdom)
-            break;
-        }
-        // If the immediate post-dominator is not in a less nested loop,
-        // then we are stuck in a program with an infinite loop.
-        // In that case, we will not find a safe point, hence, bail out.
-        if (IPdom && MLI->getLoopDepth(IPdom) < MLI->getLoopDepth(Restore))
-          Restore = IPdom;
-        else {
-          Restore = nullptr;
-          break;
-        }
-      }
-    }
-  }
-}
-
-/// Check whether the edge (\p SrcBB, \p DestBB) is a backedge according to MLI.
-/// I.e., check if it exists a loop that contains SrcBB and where DestBB is the
-/// loop header.
-static bool isProperBackedge(const MachineLoopInfo &MLI,
-                             const MachineBasicBlock *SrcBB,
-                             const MachineBasicBlock *DestBB) {
-  for (const MachineLoop *Loop = MLI.getLoopFor(SrcBB); Loop;
-       Loop = Loop->getParentLoop()) {
-    if (Loop->getHeader() == DestBB)
-      return true;
-  }
-  return false;
-}
-
-/// Check if the CFG of \p MF is irreducible.
-static bool isIrreducibleCFG(const MachineFunction &MF,
-                             const MachineLoopInfo &MLI) {
-  const MachineBasicBlock *Entry = &*MF.begin();
-  ReversePostOrderTraversal<const MachineBasicBlock *> RPOT(Entry);
-  BitVector VisitedBB(MF.getNumBlockIDs());
-  for (const MachineBasicBlock *MBB : RPOT) {
-    VisitedBB.set(MBB->getNumber());
-    for (const MachineBasicBlock *SuccBB : MBB->successors()) {
-      if (!VisitedBB.test(SuccBB->getNumber()))
-        continue;
-      // We already visited SuccBB, thus MBB->SuccBB must be a backedge.
-      // Check that the head matches what we have in the loop information.
-      // Otherwise, we have an irreducible graph.
-      if (!isProperBackedge(MLI, MBB, SuccBB))
-        return true;
-    }
-  }
-  return false;
-}
-
 bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()) || MF.empty() || !isShrinkWrapEnabled(MF))
     return false;
 
-  DEBUG(dbgs() << "**** Analysing " << MF.getName() << '\n');
-
-  init(MF);
-
-  if (isIrreducibleCFG(MF, *MLI)) {
-    // If MF is irreducible, a block may be in a loop without
-    // MachineLoopInfo reporting it. I.e., we may use the
-    // post-dominance property in loops, which lead to incorrect
-    // results. Moreover, we may miss that the prologue and
-    // epilogue are not in the same loop, leading to unbalanced
-    // construction/deconstruction of the stack frame.
-    DEBUG(dbgs() << "Irreducible CFGs are not supported yet\n");
-    return false;
-  }
+  // The algorithm is based on (post-)dominator trees.
+  auto *MDT = &getAnalysis<MachineDominatorTree>();
+  auto *MPDT = &getAnalysis<MachinePostDominatorTree>();
+  // Use the information of the basic block frequency to check the
+  // profitability of the new points.
+  auto *MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
+  // Hold the loop information. Used to determine if Save and Restore
+  // are in the same loop.
+  auto *MLI = &getAnalysis<MachineLoopInfo>();
 
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   std::unique_ptr<RegScavenger> RS(
       TRI->requiresRegisterScavenging(MF) ? new RegScavenger() : nullptr);
 
-  for (MachineBasicBlock &MBB : MF) {
-    DEBUG(dbgs() << "Look into: " << MBB.getNumber() << ' ' << MBB.getName()
-                 << '\n');
-
-    if (MBB.isEHFuncletEntry()) {
-      DEBUG(dbgs() << "EH Funclets are not supported yet.\n");
-      return false;
-    }
+  DEBUG(dbgs() << "**** Analysing " << MF.getName() << '\n');
 
-    for (const MachineInstr &MI : MBB) {
-      if (!useOrDefCSROrFI(MI, RS.get()))
-        continue;
-      // Save (resp. restore) point must dominate (resp. post dominate)
-      // MI. Look for the proper basic block for those.
-      updateSaveRestorePoints(MBB, RS.get());
-      // If we are at a point where we cannot improve the placement of
-      // save/restore instructions, just give up.
-      if (!ArePointsInteresting()) {
-        DEBUG(dbgs() << "No Shrink wrap candidate found\n");
-        return false;
-      }
-      // No need to look for other instructions, this basic block
-      // will already be part of the handled region.
-      break;
-    }
-  }
-  if (!ArePointsInteresting()) {
-    // If the points are not interesting at this point, then they must be null
-    // because it means we did not encounter any frame/CSR related code.
-    // Otherwise, we would have returned from the previous loop.
-    assert(!Save && !Restore && "We miss a shrink-wrap opportunity?!");
-    DEBUG(dbgs() << "Nothing to shrink-wrap\n");
-    return false;
-  }
+  // Gather all the basic blocks that use CSRs and/or the stack.
+  CSRFIShrinkWrapInfo SWI(MF, RS.get());
+  // Run the shrink-wrapping algorithm.
+  DominatorShrinkWrapper SW(MF, SWI, *MDT, *MPDT, *MBFI, *MLI);
+  const BitVector &TrackedElts = SW.getTrackedElts();
+  if (!TrackedElts.test(CSRFIShrinkWrapInfo::CSRFIUsedBit))
 
-  DEBUG(dbgs() << "\n ** Results **\nFrequency of the Entry: " << EntryFreq
-               << '\n');
+  // If the CSRFIUsedBit is set, we can now assume that:
+  // * All used elements are tracked (only one in this case)
+  // * There is *exactly one* save
+  // * There is *exactly one* restore
+  assert(TrackedElts == SWI.getAllUsedElts() &&
+         SW.getSaves().size() == 1 && SW.getSaves().size() == 1);
 
-  const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
-  do {
-    DEBUG(dbgs() << "Shrink wrap candidates (#, Name, Freq):\nSave: "
-                 << Save->getNumber() << ' ' << Save->getName() << ' '
-                 << MBFI->getBlockFreq(Save).getFrequency() << "\nRestore: "
-                 << Restore->getNumber() << ' ' << Restore->getName() << ' '
-                 << MBFI->getBlockFreq(Restore).getFrequency() << '\n');
-
-    bool IsSaveCheap, TargetCanUseSaveAsPrologue = false;
-    if (((IsSaveCheap = EntryFreq >= MBFI->getBlockFreq(Save).getFrequency()) &&
-         EntryFreq >= MBFI->getBlockFreq(Restore).getFrequency()) &&
-        ((TargetCanUseSaveAsPrologue = TFI->canUseAsPrologue(*Save)) &&
-         TFI->canUseAsEpilogue(*Restore)))
-      break;
-    DEBUG(dbgs() << "New points are too expensive or invalid for the target\n");
-    MachineBasicBlock *NewBB;
-    if (!IsSaveCheap || !TargetCanUseSaveAsPrologue) {
-      Save = FindIDom<>(*Save, Save->predecessors(), *MDT);
-      if (!Save)
-        break;
-      NewBB = Save;
-    } else {
-      // Restore is expensive.
-      Restore = FindIDom<>(*Restore, Restore->successors(), *MPDT);
-      if (!Restore)
-        break;
-      NewBB = Restore;
-    }
-    updateSaveRestorePoints(*NewBB, RS.get());
-  } while (Save && Restore);
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+  unsigned SaveMBBNum = SW.getSaves().begin()->first;
 
-  if (!ArePointsInteresting()) {
-    ++NumCandidatesDropped;
-    return false;
-  }
+  MachineBasicBlock *SaveBlock = MF.getBlockNumbered(SaveMBBNum);
+  MFI.setSavePoint(SaveBlock);
 
-  DEBUG(dbgs() << "Final shrink wrap candidates:\nSave: " << Save->getNumber()
-               << ' ' << Save->getName() << "\nRestore: "
-               << Restore->getNumber() << ' ' << Restore->getName() << '\n');
+  unsigned RestoreMBBNum = SW.getRestores().begin()->first;
+  MachineBasicBlock *RestoreBlock = MF.getBlockNumbered(RestoreMBBNum);
+  MFI.setRestorePoint(RestoreBlock);
 
-  MachineFrameInfo &MFI = MF.getFrameInfo();
-  MFI.setSavePoint(Save);
-  MFI.setRestorePoint(Restore);
-  ++NumCandidates;
   return false;
 }
 
 bool ShrinkWrap::isShrinkWrapEnabled(const MachineFunction &MF) {
   const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
 
   switch (EnableShrinkWrapOpt) {
   case cl::BOU_UNSET: