Skip to content

Commit 92f1f1e

Browse files
ChuanqiXu9ChuanqiXu9
committed
[Coroutines] Use to collect lifetime marker of in CoroFrame Differential Revision: https://reviews.llvm.org/D85279 
1 parent 0215ae9 commit 92f1f1e

File tree

3 files changed

+185
-24
lines changed

3 files changed

+185
-24
lines changed

‎llvm/lib/Transforms/Coroutines/CoroFrame.cpp

Lines changed: 44 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1567,50 +1567,70 @@ static void sinkLifetimeStartMarkers(Function &F, coro::Shape &Shape,
15671567
}
15681568

15691569
for (Instruction &I : instructions(F)) {
1570-
if (!isa<AllocaInst>(&I))
1570+
AllocaInst* AI = dyn_cast<AllocaInst>(&I);
1571+
if (!AI)
15711572
continue;
15721573

15731574
for (BasicBlock *DomBB : DomSet) {
15741575
bool Valid = true;
1575-
SmallVector<Instruction *, 1> BCInsts;
1576+
SmallVector<Instruction *, 1> Lifetimes;
15761577

1577-
auto isUsedByLifetimeStart = [&](Instruction *I) {
1578-
if (isa<BitCastInst>(I) && I->hasOneUse())
1579-
if (auto *IT = dyn_cast<IntrinsicInst>(I->user_back()))
1580-
return IT->getIntrinsicID() == Intrinsic::lifetime_start;
1578+
auto isLifetimeStart = [](Instruction* I) {
1579+
if (auto* II = dyn_cast<IntrinsicInst>(I))
1580+
return II->getIntrinsicID() == Intrinsic::lifetime_start;
15811581
return false;
15821582
};
15831583

1584-
for (User *U : I.users()) {
1584+
auto collectLifetimeStart = [&](Instruction *U, AllocaInst *AI) {
1585+
if (isLifetimeStart(U)) {
1586+
Lifetimes.push_back(U);
1587+
return true;
1588+
}
1589+
if (!U->hasOneUse() || U->stripPointerCasts() != AI)
1590+
return false;
1591+
if (isLifetimeStart(U->user_back())) {
1592+
Lifetimes.push_back(U->user_back());
1593+
return true;
1594+
}
1595+
return false;
1596+
};
1597+
1598+
for (User *U : AI->users()) {
15851599
Instruction *UI = cast<Instruction>(U);
15861600
// For all users except lifetime.start markers, if they are all
15871601
// dominated by one of the basic blocks and do not cross
15881602
// suspend points as well, then there is no need to spill the
15891603
// instruction.
15901604
if (!DT.dominates(DomBB, UI->getParent()) ||
1591-
Checker.isDefinitionAcrossSuspend(DomBB, U)) {
1592-
// Skip bitcast used by lifetime.start markers.
1593-
if (isUsedByLifetimeStart(UI)) {
1594-
BCInsts.push_back(UI);
1605+
Checker.isDefinitionAcrossSuspend(DomBB, UI)) {
1606+
// Skip lifetime.start, GEP and bitcast used by lifetime.start
1607+
// markers.
1608+
if (collectLifetimeStart(UI, AI))
15951609
continue;
1596-
}
15971610
Valid = false;
15981611
break;
15991612
}
16001613
}
16011614
// Sink lifetime.start markers to dominate block when they are
16021615
// only used outside the region.
1603-
if (Valid && BCInsts.size() != 0) {
1604-
auto *NewBitcast = BCInsts[0]->clone();
1605-
auto *NewLifetime = cast<Instruction>(BCInsts[0]->user_back())->clone();
1606-
NewLifetime->replaceUsesOfWith(BCInsts[0], NewBitcast);
1607-
NewBitcast->insertBefore(DomBB->getTerminator());
1616+
if (Valid && Lifetimes.size() != 0) {
1617+
// May be AI itself, when the type of AI is i8*
1618+
auto *NewBitCast = [&](AllocaInst *AI) -> Value* {
1619+
if (isa<AllocaInst>(Lifetimes[0]->getOperand(1)))
1620+
return AI;
1621+
auto *Int8PtrTy = Type::getInt8PtrTy(F.getContext());
1622+
return CastInst::Create(Instruction::BitCast, AI, Int8PtrTy, "",
1623+
DomBB->getTerminator());
1624+
}(AI);
1625+
1626+
auto *NewLifetime = Lifetimes[0]->clone();
1627+
NewLifetime->replaceUsesOfWith(NewLifetime->getOperand(1), NewBitCast);
16081628
NewLifetime->insertBefore(DomBB->getTerminator());
16091629

16101630
// All the outsided lifetime.start markers are no longer necessary.
1611-
for (Instruction *S : BCInsts) {
1612-
S->user_back()->eraseFromParent();
1613-
}
1631+
for (Instruction *S : Lifetimes)
1632+
S->eraseFromParent();
1633+
16141634
break;
16151635
}
16161636
}
@@ -1676,14 +1696,14 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
16761696
if (!II || II->getIntrinsicID() != Intrinsic::lifetime_start)
16771697
continue;
16781698

1679-
if (auto *OpInst = dyn_cast<BitCastInst>(I.getOperand(1)))
1680-
if (auto *AI = dyn_cast<AllocaInst>(OpInst->getOperand(0))) {
1699+
if (auto *OpInst = dyn_cast<Instruction>(II->getOperand(1))) {
1700+
if (auto *AI = dyn_cast<AllocaInst>(OpInst->stripPointerCasts())) {
16811701

16821702
if (LifetimeMap.find(AI) == LifetimeMap.end())
16831703
LifetimeMap[AI] = std::make_unique<LifetimeStart>();
1684-
1685-
LifetimeMap[AI]->insert(OpInst);
1704+
LifetimeMap[AI]->insert(isa<AllocaInst>(OpInst) ? II : OpInst);
16861705
}
1706+
}
16871707
}
16881708

16891709
// Collect the spills for arguments and other not-materializable values.

‎llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
; Corresponding to coro-split-sink-lifetime-01.ll. This file tests that whether the CoroFrame
2+
; pass knows the operand of lifetime.start intrinsic may be GEP as well.
3+
; RUN: opt < %s -coro-split -S | FileCheck %s
4+
; RUN: opt < %s -passes=coro-split -S | FileCheck %s
5+
6+
%"struct.std::coroutine_handle" = type { i8* }
7+
%"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
8+
%"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
9+
10+
declare i8* @malloc(i64)
11+
declare void @print(i32)
12+
13+
%i8.array = type { [100 x i8] }
14+
declare void @consume.i8.array(%i8.array*)
15+
16+
define void @a.gep() "coroutine.presplit"="1" {
17+
entry:
18+
%ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
19+
%testval = alloca %i8.array
20+
%cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0
21+
; lifetime of %testval starts here, but not used until await.ready.
22+
call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast)
23+
%id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
24+
%alloc = call i8* @malloc(i64 16) #3
25+
%vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
26+
27+
%save = call token @llvm.coro.save(i8* null)
28+
%Result.i19 = getelementptr inbounds %"struct.lean_future<int>::Awaiter", %"struct.lean_future<int>::Awaiter"* %ref.tmp7, i64 0, i32 0
29+
%suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
30+
switch i8 %suspend, label %exit [
31+
i8 0, label %await.ready
32+
i8 1, label %exit
33+
]
34+
await.ready:
35+
%StrayCoroSave = call token @llvm.coro.save(i8* null)
36+
%val = load i32, i32* %Result.i19
37+
call void @consume.i8.array(%i8.array* %testval)
38+
call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast)
39+
call void @print(i32 %val)
40+
br label %exit
41+
exit:
42+
call i1 @llvm.coro.end(i8* null, i1 false)
43+
ret void
44+
}
45+
; CHECK-LABEL: @a.gep.resume(
46+
; CHECK: %testval = alloca %i8.array
47+
; CHECK-NEXT: %0 = bitcast %i8.array* %testval to i8*
48+
; CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 100, i8* %0)
49+
; CHECK-NEXT: getelementptr inbounds %a.gep.Frame
50+
; CHECK-NEXT: getelementptr inbounds %"struct.lean_future<int>::Awaiter"
51+
; CHECK-NEXT: getelementptr inbounds %i8.array, %i8.array* %testval
52+
; CHECK-NEXT: %val = load i32, i32* %Result
53+
; CHECK-NEXT: call void @consume.i8.array(%i8.array* %testval)
54+
; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast1)
55+
; CHECK-NEXT: call void @print(i32 %val)
56+
; CHECK-NEXT: ret void
57+
58+
declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
59+
declare i1 @llvm.coro.alloc(token) #3
60+
declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
61+
declare i64 @llvm.coro.size.i64() #5
62+
declare i8* @llvm.coro.begin(token, i8* writeonly) #3
63+
declare void @"\01?puts@@YAXZZ"(...)
64+
declare token @llvm.coro.save(i8*) #3
65+
declare i8* @llvm.coro.frame() #5
66+
declare i8 @llvm.coro.suspend(token, i1) #3
67+
declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10
68+
declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
69+
declare i1 @llvm.coro.end(i8*, i1) #3
70+
declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
71+
declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4

‎llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-04.ll

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
; Tests that coro-split will optimize the lifetime.start maker of each local variable,
2+
; sink them to the places after the suspend block.
3+
; RUN: opt < %s -coro-split -S | FileCheck %s
4+
; RUN: opt < %s -passes=coro-split -S | FileCheck %s
5+
6+
%"struct.std::coroutine_handle" = type { i8* }
7+
%"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
8+
%"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
9+
10+
declare i8* @malloc(i64)
11+
declare void @print(i32)
12+
declare void @consume.i8(i8)
13+
14+
define void @a() "coroutine.presplit"="1" {
15+
entry:
16+
%ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
17+
%testval = alloca i8
18+
; lifetime of %testval starts here, but not used until await.ready.
19+
call void @llvm.lifetime.start.p0i8(i64 1, i8* %testval)
20+
%id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
21+
%alloc = call i8* @malloc(i64 16) #3
22+
%vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
23+
24+
%save = call token @llvm.coro.save(i8* null)
25+
%Result.i19 = getelementptr inbounds %"struct.lean_future<int>::Awaiter", %"struct.lean_future<int>::Awaiter"* %ref.tmp7, i64 0, i32 0
26+
%suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
27+
switch i8 %suspend, label %exit [
28+
i8 0, label %await.ready
29+
i8 1, label %exit
30+
]
31+
await.ready:
32+
%StrayCoroSave = call token @llvm.coro.save(i8* null)
33+
%val = load i32, i32* %Result.i19
34+
%test = load i8, i8* %testval
35+
call void @consume.i8(i8 %test)
36+
call void @llvm.lifetime.end.p0i8(i64 1, i8* %testval)
37+
call void @print(i32 %val)
38+
br label %exit
39+
exit:
40+
call i1 @llvm.coro.end(i8* null, i1 false)
41+
ret void
42+
}
43+
44+
; CHECK-LABEL: @a.resume(
45+
; CHECK: %testval = alloca i8, align 1
46+
; CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 1, i8* %testval)
47+
; CHECK-NEXT: getelementptr inbounds %a.Frame
48+
; CHECK-NEXT: getelementptr inbounds %"struct.lean_future<int>::Awaiter"
49+
; CHECK-NEXT: %val = load i32, i32* %Result
50+
; CHECK-NEXT: %test = load i8, i8* %testval
51+
; CHECK-NEXT: call void @consume.i8(i8 %test)
52+
; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 1, i8* %testval)
53+
; CHECK-NEXT: call void @print(i32 %val)
54+
; CHECK-NEXT: ret void
55+
56+
57+
declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
58+
declare i1 @llvm.coro.alloc(token) #3
59+
declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
60+
declare i64 @llvm.coro.size.i64() #5
61+
declare i8* @llvm.coro.begin(token, i8* writeonly) #3
62+
declare void @"\01?puts@@YAXZZ"(...)
63+
declare token @llvm.coro.save(i8*) #3
64+
declare i8* @llvm.coro.frame() #5
65+
declare i8 @llvm.coro.suspend(token, i1) #3
66+
declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10
67+
declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
68+
declare i1 @llvm.coro.end(i8*, i1) #3
69+
declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
70+
declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4

0 commit comments

Comments
 (0)