diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -75,7 +75,7 @@ namespace { -/// A little helper class for building +/// A little helper class for building class CoroCloner { public: enum class Kind { @@ -563,7 +563,7 @@ // In the original function, the AllocaSpillBlock is a block immediately // following the allocation of the frame object which defines GEPs for // all the allocas that have been moved into the frame, and it ends by - // branching to the original beginning of the coroutine. Make this + // branching to the original beginning of the coroutine. Make this // the entry block of the cloned function. auto *Entry = cast(VMap[Shape.AllocaSpillBlock]); auto *OldEntry = &NewF->getEntryBlock(); @@ -1239,6 +1239,106 @@ S.resize(N); } +/// For every local variable that has lifetime intrinsics markers, we sink +/// their lifetime.start marker to the places where the variable is being +/// used for the first time. Doing so minimizes the lifetime of each variable, +/// hence minimizing the amount of data we end up putting on the frame. +static void sinkLifetimeStartMarkers(Function &F) { + DominatorTree Dom(F); + for (Instruction &I : instructions(F)) { + // We look for this particular pattern: + // %tmpX = alloca %.., align ... + // %0 = bitcast %...* %tmpX to i8* + // call void @llvm.lifetime.start.p0i8(i64 ..., i8* nonnull %0) #2 + if (!isa(&I)) + continue; + BitCastInst *CastInst = nullptr; + // There can be multiple lifetime start markers for the same variable. + SmallPtrSet LifetimeStartInsts; + // SinkBarriers stores all instructions that use this local variable. + // When sinking the lifetime start intrinsics, we can never sink past + // these barriers. + SmallPtrSet SinkBarriers; + bool Valid = true; + auto addSinkBarrier = [&](Instruction *I) { + // When adding a new barrier to SinkBarriers, we maintain the case + // that no instruction in SinkBarriers dominates another instruction. + bool FoundDom = false; + SmallPtrSet ToRemove; + for (auto *S : SinkBarriers) { + if (Dom.dominates(S, I)) { + FoundDom = true; + break; + } else if (Dom.dominates(I, S)) { + ToRemove.insert(S); + } + } + if (!FoundDom) { + SinkBarriers.insert(I); + for (auto *R : ToRemove) { + SinkBarriers.erase(R); + } + } + }; + for (User *U : I.users()) { + if (!isa(U)) + continue; + if (CastInst) { + // If we have multiple cast instructions for the alloca, don't + // deal with it beause it's too complex. + Valid = false; + break; + } + CastInst = cast(U); + for (User *CU : CastInst->users()) { + // If we see any user of CastInst that's not lifetime start/end + // intrinsics, give up because it's too complex. + if (auto *CUI = dyn_cast(CU)) { + if (CUI->getIntrinsicID() == Intrinsic::lifetime_start) + LifetimeStartInsts.insert(CUI); + else if (CUI->getIntrinsicID() == Intrinsic::lifetime_end) + addSinkBarrier(CUI); + else + Valid = false; + } else { + Valid = false; + } + } + } + if (!Valid || LifetimeStartInsts.empty()) + continue; + + for (User *U : I.users()) { + if (U == CastInst) + continue; + // Every user of the variable is also a sink barrier. + addSinkBarrier(cast(U)); + } + + // For each sink barrier, we insert a lifetime start marker right + // before it. + const auto *LifetimeStartInst = *LifetimeStartInsts.begin(); + for (auto *S : SinkBarriers) { + if (auto *IS = dyn_cast(S)) { + if (IS->getIntrinsicID() == Intrinsic::lifetime_end) { + // If we have a lifetime end marker in SinkBarriers, meaning it's + // not dominated by any other users, we can safely delete it. + IS->eraseFromParent(); + continue; + } + } + LifetimeStartInst->clone()->insertBefore(S); + } + // All the old markers are no longer necessary. + for (auto *S : LifetimeStartInsts) { + S->eraseFromParent(); + } + // Put the cast instruction always right after variable declaration + // to be safe. + CastInst->moveAfter(&I); + } +} + static void splitSwitchCoroutine(Function &F, coro::Shape &Shape, SmallVectorImpl &Clones) { assert(Shape.ABI == coro::ABI::Switch); @@ -1428,6 +1528,7 @@ return Shape; simplifySuspendPoints(Shape); + sinkLifetimeStartMarkers(F); buildCoroutineFrame(F, Shape); replaceFrameSize(Shape);