From 9dd8c146746a0435115a461bd19d9aede965800c Mon Sep 17 00:00:00 2001
From: Manman Ren <manman.ren@gmail.com>
Date: Thu, 24 Mar 2016 23:21:29 +0000
Subject: [PATCH] CXX TLS: collect return blocks after SelectAllBasicBlocks.

It is incorrect to get the corresponding MBB for a ReturnInst before
SelectAllBasicBlocks since SelectAllBasicBlocks can change the
correspondence between a ReturnInst and the MBB it is in.

PR27062

llvm-svn: 264358
---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp | 22 +++++++++++++++-------
 llvm/test/CodeGen/X86/cxx_tlscc64.ll               | 16 ++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
index 8c0ff29..7dd02e4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
@@ -469,7 +469,6 @@ bool SelectionDAGISel::runOnMachineFunction(MachineFunction &mf) {
   MF->setHasInlineAsm(false);
 
   FuncInfo->SplitCSR = false;
-  SmallVector<MachineBasicBlock*, 4> Returns;
 
   // We split CSR if the target supports it for the given function
   // and the function has only return exits.
@@ -482,12 +481,8 @@ bool SelectionDAGISel::runOnMachineFunction(MachineFunction &mf) {
         continue;
 
       const TerminatorInst *Term = BB.getTerminator();
-      if (isa<UnreachableInst>(Term))
+      if (isa<UnreachableInst>(Term) || isa<ReturnInst>(Term))
         continue;
-      if (isa<ReturnInst>(Term)) {
-        Returns.push_back(FuncInfo->MBBMap[&BB]);
-        continue;
-      }
 
       // Bail out if the exit block is not Return nor Unreachable.
       FuncInfo->SplitCSR = false;
@@ -509,8 +504,21 @@ bool SelectionDAGISel::runOnMachineFunction(MachineFunction &mf) {
   RegInfo->EmitLiveInCopies(EntryMBB, TRI, *TII);
 
   // Insert copies in the entry block and the return blocks.
-  if (FuncInfo->SplitCSR)
+  if (FuncInfo->SplitCSR) {
+    SmallVector<MachineBasicBlock*, 4> Returns;
+    // Collect all the return blocks.
+    for (MachineBasicBlock &MBB : mf) {
+      if (!MBB.succ_empty())
+        continue;
+
+      MachineBasicBlock::iterator Term = MBB.getFirstTerminator();
+      if (Term != MBB.end() && Term->isReturn()) {
+        Returns.push_back(&MBB);
+        continue;
+      }
+    }
     TLI->insertCopiesSplitCSR(EntryMBB, Returns);
+  }
 
   DenseMap<unsigned, unsigned> LiveInMap;
   if (!FuncInfo->ArgDbgValues.empty())
diff --git a/llvm/test/CodeGen/X86/cxx_tlscc64.ll b/llvm/test/CodeGen/X86/cxx_tlscc64.ll
index c9be4a2..1915f22 100644
--- a/llvm/test/CodeGen/X86/cxx_tlscc64.ll
+++ b/llvm/test/CodeGen/X86/cxx_tlscc64.ll
@@ -151,5 +151,21 @@ entry:
   ret void
 }
 
+@ssp_var = internal thread_local global i8 0, align 1
+
+; CHECK-LABEL: test_ssp
+; CHECK-NOT: pushq %r11
+; CHECK-NOT: pushq %r10
+; CHECK-NOT: pushq %r9
+; CHECK-NOT: pushq %r8
+; CHECK-NOT: pushq %rsi
+; CHECK-NOT: pushq %rdx
+; CHECK-NOT: pushq %rcx
+; CHECK-NOT: pushq %rbx
+; CHECK: callq
+define cxx_fast_tlscc nonnull i8* @test_ssp() #2 {
+  ret i8* @ssp_var
+}
 attributes #0 = { nounwind "no-frame-pointer-elim"="true" }
 attributes #1 = { nounwind }
+attributes #2 = { nounwind sspreq }
-- 
2.7.4