Index: lib/CodeGen/ShrinkWrap.cpp
===================================================================
--- lib/CodeGen/ShrinkWrap.cpp
+++ lib/CodeGen/ShrinkWrap.cpp
@@ -175,6 +175,17 @@
   /// this call.
   void updateSaveRestorePoints(MachineBasicBlock &MBB, RegScavenger *RS);
 
+  /// Collect reachable blocks by use or def of CSRs/FI.
+  void collectReachableBlocksByCSRFI(MachineFunction &MF, RegScavenger *RS);
+
+  /// Try to split a restore point if doing so can isolate the restore point
+  /// only for blocks reachable by use or def of CSRs/FI.
+  MachineBasicBlock *tryToSplitRestore(MachineBasicBlock *MBB,
+                                       RegScavenger *RS);
+
+  /// Keep blocks reachable from the use of def of CSRs/FI.
+  DenseSet<const MachineBasicBlock *> ReachableByCSRFI;
+
   /// \brief Initialize the pass for \p MF.
   void init(MachineFunction &MF) {
     RCI.runOnMachineFunction(MF);
@@ -291,6 +302,91 @@
   return IDom;
 }
 
+static void markAllReachable(DenseSet<const MachineBasicBlock *> &Visited,
+                             MachineBasicBlock &MBB) {
+
+  SmallVector<MachineBasicBlock *, 4> Worklist(MBB.succ_begin(),
+                                               MBB.succ_end());
+  Visited.insert(&MBB);
+  while (!Worklist.empty()) {
+    MachineBasicBlock *SuccMBB = Worklist.pop_back_val();
+    if (!Visited.insert(SuccMBB).second)
+      continue;
+    Worklist.append(SuccMBB->succ_begin(), SuccMBB->succ_end());
+  }
+}
+
+static bool isAnalyzableBB(const TargetInstrInfo &TII,
+                           MachineBasicBlock &Entry) {
+  // Check if the block is analyzable.
+  MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
+  SmallVector<MachineOperand, 4> Cond;
+  if (TII.analyzeBranch(Entry, TBB, FBB, Cond))
+    return false;
+  return true;
+}
+
+MachineBasicBlock *ShrinkWrap::tryToSplitRestore(MachineBasicBlock *MBB,
+                                                 RegScavenger *RS) {
+  MachineFunction *MF = MBB->getParent();
+  const TargetInstrInfo *TII = MF->getSubtarget().getInstrInfo();
+  SmallVector<MachineBasicBlock *, 2> DirtySuccs;
+  SmallVector<MachineBasicBlock *, 2> CleanSuccs;
+  for (auto SuccBB : MBB->predecessors()) {
+    if (!isAnalyzableBB(*TII, *SuccBB))
+      return MBB;
+
+    if (ReachableByCSRFI.count(SuccBB))
+      DirtySuccs.push_back(SuccBB);
+    else
+      CleanSuccs.push_back(SuccBB);
+  }
+
+  if ((CleanSuccs.empty() || DirtySuccs.empty()))
+    return MBB;
+
+  for (const MachineInstr &MI : *MBB) {
+    if (useOrDefCSROrFI(MI, RS))
+      return MBB;
+  }
+
+  MachineBasicBlock *NMBB = MF->CreateMachineBasicBlock();
+  MF->insert(MachineFunction::iterator(MBB), NMBB);
+
+  for (auto SuccBB : DirtySuccs) {
+    SuccBB->ReplaceUsesOfBlockWith(MBB, NMBB);
+    SuccBB->updateTerminator();
+  }
+
+  for (auto SuccBB : CleanSuccs)
+    SuccBB->updateTerminator();
+
+  TII->insertUnconditionalBranch(*NMBB, MBB, DebugLoc());
+  NMBB->addSuccessor(MBB);
+
+  ReachableByCSRFI.insert(NMBB);
+  MDT->runOnMachineFunction(*MF);
+  MPDT->runOnMachineFunction(*MF);
+  return NMBB;
+}
+
+void ShrinkWrap::collectReachableBlocksByCSRFI(MachineFunction &MF,
+                                               RegScavenger *RS) {
+  ReachableByCSRFI.clear();
+  DenseSet<const MachineBasicBlock *> Visited;
+  for (MachineBasicBlock &MBB : MF) {
+    if (Visited.count(&MBB) || ReachableByCSRFI.count(&MBB))
+      continue;
+    Visited.insert(&MBB);
+    for (const MachineInstr &MI : MBB) {
+      if (!useOrDefCSROrFI(MI, RS))
+        continue;
+      // Mark all offsprings as reachable by use or def of CSRs or FI.
+      markAllReachable(ReachableByCSRFI, MBB);
+    }
+  }
+}
+
 void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB,
                                          RegScavenger *RS) {
   // Get rid of the easy cases first.
@@ -338,6 +434,8 @@
     return;
   }
 
+  Restore = tryToSplitRestore(Restore, RS);
+
   // 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.
@@ -471,6 +569,8 @@
   std::unique_ptr<RegScavenger> RS(
       TRI->requiresRegisterScavenging(MF) ? new RegScavenger() : nullptr);
 
+  collectReachableBlocksByCSRFI(MF, RS.get());
+
   for (MachineBasicBlock &MBB : MF) {
     DEBUG(dbgs() << "Look into: " << MBB.getNumber() << ' ' << MBB.getName()
                  << '\n');
Index: test/CodeGen/AArch64/srink-wrap-split-restore.ll
===================================================================
--- /dev/null
+++ test/CodeGen/AArch64/srink-wrap-split-restore.ll
@@ -0,0 +1,30 @@
+; RUN: llc < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64"
+
+; CHECK-LABEL: shrinkwrapme
+; CHECK-LABEL: %bb.0:
+; CHECK-NOT: str {{.*}} [sp{{.*}}]
+; CHECK-LABEL: %bb.1:
+; CHECK: str {{.*}} [sp{{.*}}]
+
+define void @shrinkwrapme(i32 %a) {
+entry:
+  %cmp5 = icmp sgt i32 %a, 0
+  br i1 %cmp5, label %BB0, label %Exit
+
+BB0:
+  %call = call i32 @fun()
+  %c = icmp eq i32 %call, 0
+  br i1 %c, label %BB1, label %Exit
+
+BB1:
+  %call2 = call i32 @fun()
+  br label %Exit
+
+Exit:
+  ret void
+}
+
+declare i32 @fun()