diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp --- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp @@ -2108,6 +2108,57 @@ } } +// Check whether the lifetime of an alloca might cross a coroutine suspension +// point. +static bool checkAllocaLifetimeMayCrossingSuspend(const AllocaInst &AI) { + llvm::SmallVector WorkList; + SmallPtrSet LifetimeEnds; + // Collect all lifetime.start and lifetime.end intrinsics for AI. + // All lifetime.start will be put in WorkList to traverse. + for (const auto &I : instructions(AI.getFunction())) { + if (!I.isLifetimeStartOrEnd()) + continue; + auto *II = cast(&I); + if (const auto *OpInst = dyn_cast(I.getOperand(1))) + if (const auto *AII = dyn_cast(OpInst->stripPointerCasts())) + if (AII == &AI) { + if (II->getIntrinsicID() == Intrinsic::lifetime_start) + WorkList.push_back(II); + else + LifetimeEnds.insert(&I); + } + } + // If there is no lifetime information, we assume conservatively that the + // lifetime of alloca might cross suspension points. + if (WorkList.empty()) + return true; + + // From each lifetime.start, if control flow ever reaches a suspension point + // before seen a lifetime.end, we can conclude the alloca's lifetime crosses + // coroutine suspensions. + llvm::SmallPtrSet Visited; + while (!WorkList.empty()) { + const auto *I = WorkList.back(); + WorkList.pop_back(); + if (!Visited.insert(I).second) + continue; + + while (!I->isTerminator()) { + if (LifetimeEnds.count(I)) + break; + if (isa(I)) + return true; + I = I->getNextNode(); + } + + if (I->isTerminator()) + for (unsigned Index = 0; Index < I->getNumSuccessors(); ++Index) + WorkList.push_back(&I->getSuccessor(Index)->front()); + } + + return false; +} + static void collectFrameAllocas(Function &F, coro::Shape &Shape, const SuspendCrossingInfo &Checker, SmallVectorImpl &Allocas) { @@ -2117,9 +2168,14 @@ continue; // The PromiseAlloca will be specially handled since it needs to be in a // fixed position in the frame. - if (AI == Shape.SwitchLowering.PromiseAlloca) { + if (AI == Shape.SwitchLowering.PromiseAlloca) continue; - } + + // If lifetime information guarantees the alloca will not live across + // suspensions, it's safe to put it on the stack. + if (!checkAllocaLifetimeMayCrossingSuspend(*AI)) + continue; + DominatorTree DT(F); AllocaUseVisitor Visitor{F.getParent()->getDataLayout(), DT, *Shape.CoroBegin, Checker}; diff --git a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll b/llvm/test/Transforms/Coroutines/coro-alloca-08.ll copy from llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll copy to llvm/test/Transforms/Coroutines/coro-alloca-08.ll --- a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll +++ b/llvm/test/Transforms/Coroutines/coro-alloca-08.ll @@ -1,5 +1,3 @@ -; Corresponding to coro-split-sink-lifetime-01.ll. This file tests that whether the CoroFrame -; pass knows the operand of lifetime.start intrinsic may be GEP as well. ; RUN: opt < %s -coro-split -S | FileCheck %s ; RUN: opt < %s -passes=coro-split -S | FileCheck %s @@ -8,24 +6,25 @@ %"struct.lean_future::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" } declare i8* @malloc(i64) -declare void @print(i32) %i8.array = type { [100 x i8] } -declare void @consume.i8.array(%i8.array* nocapture) +declare void @consume.i8.array(%i8.array*) -define void @a.gep() "coroutine.presplit"="1" { +; The lifetime of testval starts and ends before coro.suspend. Even though consume.i8.array +; might capture it, we can safely say it won't live across suspension. +define void @foo() "coroutine.presplit"="1" { entry: - %ref.tmp7 = alloca %"struct.lean_future::Awaiter", align 8 %testval = alloca %i8.array %cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0 - ; lifetime of %testval starts here, but not used until await.ready. - call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast) %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null) %alloc = call i8* @malloc(i64 16) #3 %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc) + call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast) + call void @consume.i8.array(%i8.array* %testval) + call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast) + %save = call token @llvm.coro.save(i8* null) - %Result.i19 = getelementptr inbounds %"struct.lean_future::Awaiter", %"struct.lean_future::Awaiter"* %ref.tmp7, i64 0, i32 0 %suspend = call i8 @llvm.coro.suspend(token %save, i1 false) switch i8 %suspend, label %exit [ i8 0, label %await.ready @@ -33,38 +32,51 @@ ] await.ready: %StrayCoroSave = call token @llvm.coro.save(i8* null) - %val = load i32, i32* %Result.i19 + br label %exit +exit: + call i1 @llvm.coro.end(i8* null, i1 false) + ret void +} + +; The lifetime of testval starts after coro.suspend. So it will never live across suspension +; points. +define void @bar() "coroutine.presplit"="1" { +entry: + %testval = alloca %i8.array + %cast = getelementptr inbounds %i8.array, %i8.array* %testval, i64 0, i32 0, i64 0 + %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null) + %alloc = call i8* @malloc(i64 16) #3 + %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc) + %save = call token @llvm.coro.save(i8* null) + %suspend = call i8 @llvm.coro.suspend(token %save, i1 false) + switch i8 %suspend, label %exit [ + i8 0, label %await.ready + i8 1, label %exit + ] +await.ready: + %StrayCoroSave = call token @llvm.coro.save(i8* null) + + call void @llvm.lifetime.start.p0i8(i64 100, i8* %cast) call void @consume.i8.array(%i8.array* %testval) call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast) - call void @print(i32 %val) + br label %exit exit: call i1 @llvm.coro.end(i8* null, i1 false) ret void } -; CHECK-LABEL: @a.gep.resume( -; CHECK: %testval = alloca %i8.array -; CHECK-NEXT: getelementptr inbounds %a.gep.Frame -; CHECK-NEXT: %0 = bitcast %i8.array* %testval to i8* -; CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 100, i8* %0) -; CHECK-NEXT: getelementptr inbounds %"struct.lean_future::Awaiter" -; CHECK-NEXT: getelementptr inbounds %i8.array, %i8.array* %testval -; CHECK-NEXT: %val = load i32, i32* %Result -; CHECK-NEXT: call void @consume.i8.array(%i8.array* %testval) -; CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast1) -; CHECK-NEXT: call void @print(i32 %val) -; CHECK-NEXT: ret void + +; Verify that for both foo and bar, testval isn't put on the frame. +; CHECK: %foo.Frame = type { void (%foo.Frame*)*, void (%foo.Frame*)*, i1 } +; CHECK: %bar.Frame = type { void (%bar.Frame*)*, void (%bar.Frame*)*, i1 } declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) declare i1 @llvm.coro.alloc(token) #3 -declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr declare i64 @llvm.coro.size.i64() #5 declare i8* @llvm.coro.begin(token, i8* writeonly) #3 -declare void @"\01?puts@@YAXZZ"(...) declare token @llvm.coro.save(i8*) #3 declare i8* @llvm.coro.frame() #5 declare i8 @llvm.coro.suspend(token, i1) #3 -declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10 declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2 declare i1 @llvm.coro.end(i8*, i1) #3 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4 diff --git a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll --- a/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll +++ b/llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll @@ -11,7 +11,7 @@ declare void @print(i32) %i8.array = type { [100 x i8] } -declare void @consume.i8.array(%i8.array* nocapture) +declare void @consume.i8.array(%i8.array*) define void @a.gep() "coroutine.presplit"="1" { entry: