[CodeExtractor] Use 'normal destination' BB as insert point to store invoke results.

Currently CodeExtractor tries to use the next node after an invoke to
place the store for the result of the invoke, if it is an out parameter
of the region. This fails, as the invoke terminates the current BB.
In that case, we can place the store in the 'normal destination' BB, as
the result will only be available in that case.


Reviewers: davidxl, davide, efriedma

Reviewed By: davidxl

Differential Revision: https://reviews.llvm.org/D51037

llvm-svn: 340331

2 files changed, 56 insertions, 1 deletions

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index cb349e34606..eeb5d2bf5eb 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -925,8 +925,16 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
     auto *OutI = dyn_cast<Instruction>(outputs[i]);
     if (!OutI)
       continue;
+
     // Find proper insertion point.
-    Instruction *InsertPt = OutI->getNextNode();
+    Instruction *InsertPt;
+    // In case OutI is an invoke, we insert the store at the beginning in the
+    // 'normal destination' BB. Otherwise we insert the store right after OutI.
+    if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
+      InsertPt = InvokeI->getNormalDest()->getFirstNonPHI();
+    else
+      InsertPt = OutI->getNextNode();
+
     // Let's assume that there is no other guy interleave non-PHI in PHIs.
     if (isa<PHINode>(InsertPt))
       InsertPt = InsertPt->getParent()->getFirstNonPHI();
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll
new file mode 100644
index 00000000000..bba7ad05fac
--- /dev/null
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll
@@ -0,0 +1,47 @@
+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+
+; Function Attrs: nounwind uwtable
+define dso_local i8* @bar(i32 %arg) local_unnamed_addr #0 personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb1, label %bb5
+
+bb1:                                              ; preds = %bb
+  %call26 = invoke i8* @invoke_callee() #2
+          to label %cont unwind label %lpad
+lpad:                                            ; preds = %if.end
+  %0 = landingpad { i8*, i32 }
+         cleanup
+  resume { i8*, i32 } undef
+
+cont:
+    br label %bb5
+
+bb5:                                              ; preds = %bb4, %bb1, %bb
+  %retval = phi i8* [ %call26, %cont ], [ undef, %bb]
+  ret i8* %retval
+}
+
+; CHECK-LABEL: @dummy_caller
+; CHECK-LABEL: bb:
+; CHECK-NEXT:  [[CALL26LOC:%.*]] = alloca i8*
+; CHECK-LABEL: codeRepl.i:
+; CHECK-NEXT:   call void @bar.1_bb1(i8** [[CALL26LOC]])
+define i8* @dummy_caller(i32 %arg) {
+bb:
+  %tmp = tail call i8* @bar(i32 %arg)
+  ret i8* %tmp
+}
+
+; CHECK-LABEL: define internal void @bar.1_bb1
+; CHECK-LABEL: bb1:
+; CHECK-NEXT:    %call26 = invoke i8* @invoke_callee()
+; CHECK-NEXT:            to label %cont unwind label %lpad
+; CHECK-LABEL: cont:
+; CHECK-NEXT:    store i8* %call26, i8** %call26.out
+; CHECK-NEXT:    br label %bb5.exitStub
+
+; Function Attrs: nobuiltin
+declare dso_local noalias nonnull i8* @invoke_callee() local_unnamed_addr #1
+
+declare dso_local i32 @__gxx_personality_v0(...)