Index: lib/Transforms/Coroutines/CoroElide.cpp =================================================================== --- lib/Transforms/Coroutines/CoroElide.cpp +++ lib/Transforms/Coroutines/CoroElide.cpp @@ -14,6 +14,7 @@ #include "CoroInternal.h" #include "llvm/Analysis/AliasAnalysis.h" #include "llvm/Analysis/InstructionSimplify.h" +#include "llvm/IR/Dominators.h" #include "llvm/IR/InstIterator.h" #include "llvm/Pass.h" #include "llvm/Support/ErrorHandling.h" @@ -35,8 +36,8 @@ Lowerer(Module &M) : LowererBase(M) {} void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA); - bool shouldElide() const; - bool processCoroId(CoroIdInst *, AAResults &AA); + bool shouldElide(Function *F, DominatorTree &DT) const; + bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT); }; } // end anonymous namespace @@ -142,33 +143,59 @@ removeTailCallAttribute(Frame, AA); } -bool Lowerer::shouldElide() const { +bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const { // If no CoroAllocs, we cannot suppress allocation, so elision is not // possible. if (CoroAllocs.empty()) return false; // Check that for every coro.begin there is a coro.destroy directly - // referencing the SSA value of that coro.begin. If the value escaped, then - // coro.destroy would have been referencing a memory location storing that - // value and not the virtual register. + // referencing the SSA value of that coro.begin for every reachable exiting + // path out of the function. If the value escaped, then coro.destroy would + // have been referencing a memory location storing that value and not the + // virtual register. + + SmallPtrSet TerminatorInsts; + for (BasicBlock &B : *F) { + auto *TI = B.getTerminator(); + if (isa(TI) || isa(TI)) + TerminatorInsts.insert(TI); + } SmallPtrSet ReferencedCoroBegins; - + SmallPtrSet DominatedTerminatorInsts; for (CoroSubFnInst *DA : DestroyAddr) { - if (auto *CB = dyn_cast(DA->getFrame())) + if (auto *CB = dyn_cast(DA->getFrame())) { ReferencedCoroBegins.insert(CB); - else + + // Update the set of terminator instructions that are dominated by one of + // these destroy instructions. + for (Instruction *I : TerminatorInsts) { + // No need to check if the function has already determined that all + // terminator instructions are dominated. + if (DominatedTerminatorInsts.size() == TerminatorInsts.size()) + break; + // Skip instructions that we know are dominated. + if (DominatedTerminatorInsts.count(I) != 0) + continue; + + if (DT.dominates(DA, I)) + DominatedTerminatorInsts.insert(I); + } + } else return false; } - // If size of the set is the same as total number of CoroBegins, means we - // found a coro.free or coro.destroy mentioning a coro.begin and we can - // perform heap elision. - return ReferencedCoroBegins.size() == CoroBegins.size(); + // If size of the set is the same as total number of CoroBegins, that means + // that we found a coro.free or coro.destroy that directly refences the + // coro.begin. If that set of coro.free or coro.destroy dominates each + // terminator instruction in the function, we can perform heap elision. + return DominatedTerminatorInsts.size() == TerminatorInsts.size() && + ReferencedCoroBegins.size() == CoroBegins.size(); } -bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA) { +bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA, + DominatorTree &DT) { CoroBegins.clear(); CoroAllocs.clear(); CoroFrees.clear(); @@ -214,7 +241,7 @@ replaceWithConstant(ResumeAddrConstant, ResumeAddr); - bool ShouldElide = shouldElide(); + bool ShouldElide = shouldElide(CoroId->getFunction(), DT); auto *DestroyAddrConstant = ConstantExpr::getExtractValue( Resumers, @@ -294,14 +321,16 @@ return Changed; AAResults &AA = getAnalysis().getAAResults(); + DominatorTree &DT = getAnalysis().getDomTree(); for (auto *CII : L->CoroIds) - Changed |= L->processCoroId(CII, AA); + Changed |= L->processCoroId(CII, AA, DT); return Changed; } void getAnalysisUsage(AnalysisUsage &AU) const override { AU.addRequired(); + AU.addRequired(); } StringRef getPassName() const override { return "Coroutine Elision"; } }; Index: test/Transforms/Coroutines/coro-heap-elide.ll =================================================================== --- test/Transforms/Coroutines/coro-heap-elide.ll +++ test/Transforms/Coroutines/coro-heap-elide.ll @@ -81,6 +81,40 @@ ret void } +; CHECK-LABEL: @callResume_PR34897_no_elision( +define void @callResume_PR34897_no_elision(i1 %cond) { +; CHECK-LABEL: entry: +entry: +; CHECK: call i8* @CustomAlloc( + %hdl = call i8* @f() +; CHECK: tail call void @bar( + tail call void @bar(i8* %hdl) +; CHECK: tail call void @bar( + tail call void @bar(i8* null) + br i1 %cond, label %if.then, label %if.else + +; CHECK-LABEL: if.then: +if.then: +; CHECK: call fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8* + %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0) + %1 = bitcast i8* %0 to void (i8*)* + call fastcc void %1(i8* %hdl) +; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.destroy to void (i8*)*)(i8* + %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1) + %3 = bitcast i8* %2 to void (i8*)* + call fastcc void %3(i8* %hdl) + br label %return + +if.else: + br label %return + +; CHECK-LABEL: return: +return: +; CHECK: ret void + ret void +} + + ; a coroutine start function (cannot elide heap alloc, due to second argument to ; coro.begin not pointint to coro.alloc) define i8* @f_no_elision() personality i8* null {