Index: lib/CodeGen/MachineSink.cpp
===================================================================
--- lib/CodeGen/MachineSink.cpp
+++ lib/CodeGen/MachineSink.cpp
@@ -26,6 +26,7 @@
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.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"
@@ -78,6 +79,7 @@
 STATISTIC(NumSplit,     "Number of critical edges split");
 STATISTIC(NumCoalesces, "Number of copies coalesced");
 STATISTIC(NumPostRACopySink, "Number of copies sunk after RA");
+STATISTIC(NumPostRASpillSink, "Number of spills sunk after RA");
 
 namespace {
 
@@ -940,6 +942,8 @@
 //===----------------------------------------------------------------------===//
 namespace {
 
+typedef DenseMap<int, SmallPtrSet<MachineBasicBlock *, 4>> StackUseMapTy;
+
 class PostRAMachineSinking : public MachineFunctionPass {
 public:
   bool runOnMachineFunction(MachineFunction &MF) override;
@@ -950,25 +954,46 @@
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
+    AU.addRequired<MachineDominatorTree>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
 private:
+  const TargetRegisterInfo *TRI;
+  const TargetInstrInfo *TII;
+  const MachineFrameInfo *MFI;
+  const MachineDominatorTree *MDT;
+
   /// Track which registers have been modified and used.
   BitVector ModifiedRegs, UsedRegs;
 
+  /// Hold the uses of stack slots.
+  StackUseMapTy StackUseMap;
+
+  /// Find basic blocks where there is a load from stack slots.
+  void scanLoadFromStackSlot(MachineFunction &MF);
+
+  /// Sink register spills close to their reloads.
+  bool tryToSinkSpill(MachineInstr &MI, MachineBasicBlock &CurBB);
+
   /// Sink Copy instructions unused in the same block close to their uses in
   /// successors.
-  bool tryToSinkCopy(MachineBasicBlock &BB, MachineFunction &MF,
-                     const TargetRegisterInfo *TRI, const TargetInstrInfo *TII);
+  bool tryToSinkCopy(MachineInstr &MI, MachineBasicBlock &CurBB,
+                     ArrayRef<MachineBasicBlock *> SinkableBBs);
+
+  /// Perform sinking for Copy and store to stack.
+  bool tryToSink(MachineBasicBlock &BB);
 };
 } // namespace
 
 char PostRAMachineSinking::ID = 0;
 char &llvm::PostRAMachineSinkingID = PostRAMachineSinking::ID;
 
-INITIALIZE_PASS(PostRAMachineSinking, "postra-machine-sink",
-                "PostRA Machine Sink", false, false)
+INITIALIZE_PASS_BEGIN(PostRAMachineSinking, "postra-machine-sink",
+                      "PostRA Machine Sink", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_END(PostRAMachineSinking, "postra-machine-sink",
+                    "PostRA Machine Sink", false, false)
 
 static MachineBasicBlock *
 getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
@@ -1004,16 +1029,168 @@
   return BB;
 }
 
-bool PostRAMachineSinking::tryToSinkCopy(MachineBasicBlock &CurBB,
-                                         MachineFunction &MF,
-                                         const TargetRegisterInfo *TRI,
-                                         const TargetInstrInfo *TII) {
+// Return true if there is potential path from FromMBB to any basic block which
+// contains uses of FI.
+static bool isReachableToUseOfFI(MachineBasicBlock *FromMBB,
+                                 StackUseMapTy &StackUseMap, int FI,
+                                 const MachineDominatorTree *MDT) {
+  SmallPtrSet<MachineBasicBlock *, 4> &UseBBs = StackUseMap[FI];
+  if (UseBBs.count(FromMBB))
+    return true;
+
+  if (any_of(UseBBs, [&](MachineBasicBlock *UseBB) {
+        return MDT->dominates(FromMBB, UseBB);
+      }))
+    return true;
+
+  // Limit the number of blocks we visit for compile times.
+  // The default value (32) was chosen arbitrarily.
+  unsigned Limit = 32;
+
+  DenseSet<const MachineBasicBlock *> Visited;
+  SmallVector<MachineBasicBlock *, 4> Worklist(FromMBB->succ_begin(),
+                                               FromMBB->succ_end());
+  while (!Worklist.empty()) {
+    MachineBasicBlock *MBB = Worklist.pop_back_val();
+    if (!Visited.insert(MBB).second)
+      continue;
+    if (UseBBs.count(MBB))
+      return true;
+
+    // If it reaches to the limit, conservatively return true (there is
+    // potentially a path).
+    if (!--Limit)
+      return true;
+
+    Worklist.append(MBB->succ_begin(), MBB->succ_end());
+  }
+  return false;
+}
+
+static MachineBasicBlock *
+getSingleReloadInSuccBB(MachineBasicBlock &CurBB, StackUseMapTy &StackUseMap,
+                        int FI, const MachineDominatorTree *MDT) {
+  // TODO: We should track access to FI just like we track the def of register.
+  if (StackUseMap[FI].count(&CurBB))
+    return nullptr;
+
+  MachineBasicBlock *SingleReachableSucc = nullptr;
+
+  // Try to find a single sinkable successor from which any use of FI is
+  // reachable.
+  for (MachineBasicBlock *SI : CurBB.successors()) {
+    if (SI == &CurBB)
+      continue;
+    if (isReachableToUseOfFI(SI, StackUseMap, FI, MDT)) {
+      if (SingleReachableSucc)
+        return nullptr;
+      SingleReachableSucc = SI;
+      if (SingleReachableSucc->pred_size() != 1)
+        return nullptr;
+    }
+  }
+  return SingleReachableSucc;
+}
+
+void PostRAMachineSinking::scanLoadFromStackSlot(MachineFunction &MF) {
+  StackUseMap.clear();
+  for (auto &MBB : MF)
+    for (auto &MI : MBB) {
+      int FI;
+      const MachineMemOperand *MMO;
+      if (TII->hasLoadFromStackSlot(MI, MMO, FI))
+        StackUseMap[FI].insert(&MBB);
+    }
+}
+
+bool PostRAMachineSinking::tryToSinkSpill(MachineInstr &MI,
+                                          MachineBasicBlock &CurBB) {
+  int FI;
+  unsigned StrValReg = TII->isStoreToStackSlot(MI, FI);
+  if (!StrValReg || !MFI->isSpillSlotObjectIndex(FI))
+    return false;
+
+  if (ModifiedRegs[StrValReg])
+    return false;
+
+  MachineBasicBlock *SuccBB =
+      getSingleReloadInSuccBB(CurBB, StackUseMap, FI, MDT);
+
+  if (!SuccBB)
+    return false;
+
+  if (UsedRegs[StrValReg]) {
+    MachineBasicBlock::iterator NI = std::next(MI.getIterator());
+    for (MachineInstr &UI : make_range(NI, CurBB.end()))
+      if (UI.killsRegister(StrValReg, TRI)) {
+        UI.clearRegisterKills(StrValReg, TRI);
+        break;
+      }
+  }
+
+  MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
+  SuccBB->splice(InsertPos, &CurBB, &MI);
+
+  if (!SuccBB->isLiveIn(StrValReg))
+    SuccBB->addLiveIn(StrValReg);
+
+  ++NumPostRASpillSink;
+  return true;
+}
+
+bool PostRAMachineSinking::tryToSinkCopy(
+    MachineInstr &MI, MachineBasicBlock &CurBB,
+    ArrayRef<MachineBasicBlock *> SinkableBBs) {
+  if (!MI.getOperand(0).isRenamable())
+    return false;
+
+  unsigned DefReg = MI.getOperand(0).getReg();
+  unsigned SrcReg = MI.getOperand(1).getReg();
+
+  // Don't sink the COPY if it would violate a register dependency.
+  if (ModifiedRegs[DefReg] || ModifiedRegs[SrcReg] || UsedRegs[DefReg])
+    return false;
+
+  MachineBasicBlock *SuccBB =
+      getSingleLiveInSuccBB(CurBB, SinkableBBs, DefReg, TRI);
+  // Don't sink if we cannot find a single sinkable successor in which Reg
+  // is live-in.
+  if (!SuccBB)
+    return false;
+
+  assert((SuccBB->pred_size() == 1 && *SuccBB->pred_begin() == &CurBB) &&
+         "Unexpected predecessor");
+
+  // Clear the kill flag if SrcReg is killed between MI and the end of the
+  // block.
+  if (UsedRegs[SrcReg]) {
+    MachineBasicBlock::iterator NI = std::next(MI.getIterator());
+    for (MachineInstr &UI : make_range(NI, CurBB.end())) {
+      if (UI.killsRegister(SrcReg, TRI)) {
+        UI.clearRegisterKills(SrcReg, TRI);
+        MI.getOperand(1).setIsKill(true);
+        break;
+      }
+    }
+  }
+
+  MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
+  SuccBB->splice(InsertPos, &CurBB, &MI);
+  SuccBB->removeLiveIn(DefReg);
+  if (!SuccBB->isLiveIn(SrcReg))
+    SuccBB->addLiveIn(SrcReg);
+
+  ++NumPostRACopySink;
+  return true;
+}
+
+bool PostRAMachineSinking::tryToSink(MachineBasicBlock &CurBB) {
   SmallVector<MachineBasicBlock *, 2> SinkableBBs;
   // FIXME: For now, we sink only to a successor which has a single predecessor
-  // so that we can directly sink COPY instructions to the successor without
-  // adding any new block or branch instruction.
+  // so that we can directly sink instructions to the successor without adding
+  // any new block or branch instruction.
   for (MachineBasicBlock *SI : CurBB.successors())
-    if (!SI->livein_empty() && SI->pred_size() == 1)
+    if (SI->pred_size() == 1)
       SinkableBBs.push_back(SI);
 
   if (SinkableBBs.empty())
@@ -1027,71 +1204,43 @@
   UsedRegs.reset();
 
   for (auto I = CurBB.rbegin(), E = CurBB.rend(); I != E;) {
-    MachineInstr *MI = &*I;
+    MachineInstr &MI = *I;
     ++I;
 
     // Do not move any instruction across function call.
-    if (MI->isCall())
-      return false;
-
-    if (!MI->isCopy() || !MI->getOperand(0).isRenamable()) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
-      continue;
-    }
+    if (MI.isCall())
+      return Changed;
 
-    unsigned DefReg = MI->getOperand(0).getReg();
-    unsigned SrcReg = MI->getOperand(1).getReg();
-    // Don't sink the COPY if it would violate a register dependency.
-    if (ModifiedRegs[DefReg] || ModifiedRegs[SrcReg] || UsedRegs[DefReg]) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
-      continue;
-    }
-
-    MachineBasicBlock *SuccBB =
-        getSingleLiveInSuccBB(CurBB, SinkableBBs, DefReg, TRI);
-    // Don't sink if we cannot find a single sinkable successor in which Reg
-    // is live-in.
-    if (!SuccBB) {
-      TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
-      continue;
-    }
-    assert((SuccBB->pred_size() == 1 && *SuccBB->pred_begin() == &CurBB) &&
-           "Unexpected predecessor");
-
-    // Clear the kill flag if SrcReg is killed between MI and the end of the
-    // block.
-    if (UsedRegs[SrcReg]) {
-      MachineBasicBlock::iterator NI = std::next(MI->getIterator());
-      for (MachineInstr &UI : make_range(NI, CurBB.end())) {
-        if (UI.killsRegister(SrcReg, TRI)) {
-          UI.clearRegisterKills(SrcReg, TRI);
-          MI->getOperand(1).setIsKill(true);
-          break;
-        }
+    if (MI.mayStore()) {
+      if (tryToSinkSpill(MI, CurBB)) {
+        Changed = true;
+        continue;
+      }
+    } else if (MI.isCopy()) {
+      if (tryToSinkCopy(MI, CurBB, SinkableBBs)) {
+        Changed = true;
+        continue;
       }
     }
-
-    MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
-    SuccBB->splice(InsertPos, &CurBB, MI);
-    SuccBB->removeLiveIn(DefReg);
-    if (!SuccBB->isLiveIn(SrcReg))
-      SuccBB->addLiveIn(SrcReg);
-
-    Changed = true;
-    ++NumPostRACopySink;
+    TII->trackRegDefsUses(MI, ModifiedRegs, UsedRegs, TRI);
   }
   return Changed;
 }
 
 bool PostRAMachineSinking::runOnMachineFunction(MachineFunction &MF) {
   bool Changed = false;
-  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
-  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  TRI = MF.getSubtarget().getRegisterInfo();
+  TII = MF.getSubtarget().getInstrInfo();
+  MFI = &MF.getFrameInfo();
+  MDT = &getAnalysis<MachineDominatorTree>();
+
   ModifiedRegs.resize(TRI->getNumRegs());
   UsedRegs.resize(TRI->getNumRegs());
 
+  scanLoadFromStackSlot(MF);
+
   for (auto &BB : MF)
-    Changed |= tryToSinkCopy(BB, MF, TRI, TII);
+    Changed |= tryToSink(BB);
 
   return Changed;
 }
Index: test/CodeGen/AArch64/post-ra-machine-sink-spill.mir
===================================================================
--- /dev/null
+++ test/CodeGen/AArch64/post-ra-machine-sink-spill.mir
@@ -0,0 +1,188 @@
+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=postra-machine-sink -verify-machineinstrs  -o - %s | FileCheck %s
+
+---
+# Sink the store to stack.0 to %bb.1.
+# CHECK-LABEL: name: sinkspill1
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: STRXui $x0, %stack.0
+# CHECK-LABEL: bb.1:
+# CHECK: liveins: $x0
+# CHECK: STRXui $x0, %stack.0
+name: sinkspill1
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0 , $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    Bcc 11, %bb.1, implicit $nzcv
+    B %bb.2
+
+  bb.1:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+
+  bb.2:
+    $x0 = COPY $xzr
+    RET   $x0
+...
+
+---
+# Sink the store to stack.0 to %bb.2.
+# CHECK-LABEL: name: sinkspill2
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: STRXui $x0, %stack.0
+# CHECK-LABEL: bb.2:
+# CHECK: liveins:{{.*}} $x0
+# CHECK: STRXui $x0, %stack.0
+name: sinkspill2
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0, $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    Bcc 11, %bb.2, implicit $nzcv
+
+  bb.1:
+    liveins: $x0
+    RET $x0
+
+  bb.2:
+    liveins: $x1
+    $x0 = ADDXrr $x1, killed $x1
+
+  bb.3:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+...
+
+---
+# Sink the store to stack.0 to %bb.3.
+# CHECK-LABEL: name: sinkspill3
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: STRXui $x0, %stack.0
+# CHECK-LABEL: bb.2:
+# CHECK: liveins:{{.*}} $x0
+# CHECK-LABEL: bb.3:
+# CHECK: liveins:{{.*}} $x0
+# CHECK: STRXui $x0, %stack.0
+name: sinkspill3
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0, $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    Bcc 11, %bb.2, implicit $nzcv
+
+  bb.1:
+    liveins: $x0
+    RET $x0
+
+  bb.2:
+    liveins: $x1
+    $x1 = ADDXrr $x1, killed $x1
+
+  bb.3:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+...
+
+---
+# Keep the store to stack.0 in bb.0.
+# CHECK-LABEL: name: donotsinkspill1
+# CHECK-LABEL: bb.0:
+# CHECK: STRXui $x0, %stack.0
+name: donotsinkspill1
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0 , $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    Bcc 11, %bb.1, implicit $nzcv
+    B %bb.2
+
+  bb.1:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+
+  bb.2:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET   $x0
+...
+
+---
+# Keep the store to stack.0 in bb.0 due to register dependency on x0.
+# CHECK-LABEL: name: donotsinkspill2
+# CHECK-LABEL: bb.0:
+# CHECK: STRXui $x0, %stack.0
+name: donotsinkspill2
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0 , $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    $x0 = ADDXrr $x1, killed $x1
+    Bcc 11, %bb.1, implicit $nzcv
+    B %bb.2
+
+  bb.1:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+
+  bb.2:
+    $x0 = COPY $xzr
+    RET $x0
+...
+
+---
+# Keep the store to stack.0 in bb.0 due to LDRXui in bb.0.
+# CHECK-LABEL: name: donotsinkspill3
+# CHECK-LABEL: bb.0:
+# CHECK: STRXui $x0, %stack.0
+name: donotsinkspill3
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 8,
+      stack-id: 0, callee-saved-register: '', callee-saved-restored: true,
+      di-variable: '', di-expression: '', di-location: '' }
+body: |
+  bb.0:
+    liveins: $x0 , $w1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    STRXui $x0, %stack.0, 0 :: (store 8 into %stack.0)
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    Bcc 11, %bb.1, implicit $nzcv
+    B %bb.2
+
+  bb.1:
+    $x0 = LDRXui %stack.0, 0 :: (load 8 from %stack.0)
+    RET $x0
+
+  bb.2:
+    $x0 = COPY $xzr
+    RET $x0
+...