[CodeExtractor] Do not extract unsafe lifetime markers

Lifetime markers which reference inputs to the extraction region are not
safe to extract. Example ('rhs' will be extracted):

```
               entry:
              +------------+
              | x = alloca |
              | y = alloca |
              +------------+
             /              \
   lhs:                      rhs:
  +-------------------+     +-------------------+
  | lifetime_start(x) |     | lifetime_start(x) |
  | use(x)            |     | lifetime_start(y) |
  | lifetime_end(x)   |     | use(x, y)         |
  | lifetime_start(y) |     | lifetime_end(y)   |
  | use(y)            |     | lifetime_end(x)   |
  | lifetime_end(y)   |     +-------------------+
  +-------------------+
```

Prior to extraction, the stack coloring pass sees that the slots for 'x'
and 'y' are in-use at the same time. After extraction, the coloring pass
infers that 'x' and 'y' are *not* in-use concurrently, because markers
from 'rhs' are no longer available to help decide otherwise.

This leads to a miscompile, because the stack slots actually are in-use
concurrently in the extracted function.

Fix this by moving lifetime start/end markers for memory regions defined
in the calling function around the call to the extracted function.

Fixes llvm.org/PR39671 (rdar://45939472).

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

llvm-svn: 350420

3 files changed, 71 insertions, 2 deletions

diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll
index 6bb38d44f46..04789eaad29 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll
@@ -6,10 +6,14 @@
 
 @g = external local_unnamed_addr global i32, align 4
 
+; CHECK-LABEL: define{{.*}}@caller(
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* %tmp.i)
+; CHECK-NEXT: call void @callee_unknown_use1.{{.*}}(i8* %tmp.i
+; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 -1, i8* %tmp.i)
+
 define i32 @callee_unknown_use1(i32 %arg) local_unnamed_addr #0 {
 ; CHECK-LABEL:define{{.*}}@callee_unknown_use1.{{[0-9]}}
 ; CHECK-NOT: alloca
-; CHECK: call void @llvm.lifetime
 bb:
   %tmp = alloca  i8, align 4
   %tmp2 = load i32, i32* @g, align 4, !tbaa !2
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca5.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca5.ll
index 9c53496e1ce..0bde58fbccd 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca5.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineAlloca5.ll
@@ -9,7 +9,6 @@
 define i32 @callee_unknown_use2(i32 %arg) local_unnamed_addr #0 {
 ; CHECK-LABEL:define{{.*}}@callee_unknown_use2.{{[0-9]}}
 ; CHECK-NOT: alloca
-; CHECK: call void @llvm.lifetime
 bb:
   %tmp = alloca i32, align 4
   %tmp1 = bitcast i32* %tmp to i8*
diff --git a/llvm/test/Transforms/HotColdSplit/lifetime-markers-on-inputs.ll b/llvm/test/Transforms/HotColdSplit/lifetime-markers-on-inputs.ll
new file mode 100644
index 00000000000..c6482f823b4
--- /dev/null
+++ b/llvm/test/Transforms/HotColdSplit/lifetime-markers-on-inputs.ll
@@ -0,0 +1,66 @@
+; RUN: opt -S -hotcoldsplit < %s 2>&1 | FileCheck %s
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
+
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
+
+declare void @use(i8*)
+
+declare void @cold_use2(i8*, i8*) cold
+
+; CHECK-LABEL: define {{.*}}@foo(
+define void @foo() {
+entry:
+  %local1 = alloca i256
+  %local2 = alloca i256
+  %local1_cast = bitcast i256* %local1 to i8*
+  %local2_cast = bitcast i256* %local2 to i8*
+  br i1 undef, label %normalPath, label %outlinedPath
+
+normalPath:
+  ; These two uses of stack slots are non-overlapping. Based on this alone,
+  ; the stack slots could be merged.
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local1_cast)
+  call void @use(i8* %local1_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local1_cast)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @use(i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  ret void
+
+; CHECK-LABEL: codeRepl:
+; CHECK: [[local1_cast:%.*]] = bitcast i256* %local1 to i8*
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local1_cast]])
+; CHECK: [[local2_cast:%.*]] = bitcast i256* %local2 to i8*
+; CHECK: call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[local2_cast]])
+; CHECK: call i1 @foo.cold.1(i8* %local1_cast, i8* %local2_cast)
+; CHECK: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local2_cast]])
+; CHECK: call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[local1_cast]])
+; CHECK: br i1
+
+outlinedPath:
+  ; These two uses of stack slots are overlapping. This should prevent
+  ; merging of stack slots. CodeExtractor must replicate the effects of
+  ; these markers in the caller to inhibit stack coloring.
+  %gep1 = getelementptr inbounds i8, i8* %local1_cast, i64 1
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %gep1)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @cold_use2(i8* %local1_cast, i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %gep1)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  br i1 undef, label %outlinedPath2, label %outlinedPathExit
+
+outlinedPath2:
+  ; These extra lifetime markers are used to test that we emit only one
+  ; pair of guard markers in the caller per memory object.
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %local2_cast)
+  call void @use(i8* %local2_cast)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %local2_cast)
+  ret void
+
+outlinedPathExit:
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@foo.cold.1(
+; CHECK-NOT: @llvm.lifetime