diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp index 6b68cf5bc2c20..c44eaddd7ee55 100644 --- a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp @@ -8,6 +8,7 @@ #include "llvm/Transforms/Coroutines/CoroCleanup.h" #include "CoroInternal.h" +#include "llvm/Analysis/PtrUseVisitor.h" #include "llvm/IR/DIBuilder.h" #include "llvm/IR/Function.h" #include "llvm/IR/IRBuilder.h" @@ -15,6 +16,7 @@ #include "llvm/IR/Module.h" #include "llvm/IR/PassManager.h" #include "llvm/Transforms/Scalar/SimplifyCFG.h" +#include "llvm/Transforms/Utils/Local.h" using namespace llvm; @@ -30,9 +32,28 @@ struct Lowerer : coro::LowererBase { bool lower(Function &F); private: - void elideCoroNoop(IntrinsicInst *II); void lowerCoroNoop(IntrinsicInst *II); }; + +// Recursively walk and eliminate resume/destroy call on noop coro +class NoopCoroElider : public PtrUseVisitor { + using Base = PtrUseVisitor; + + IRBuilder<> Builder; + +public: + NoopCoroElider(const DataLayout &DL, LLVMContext &C) : Base(DL), Builder(C) {} + + void run(IntrinsicInst *II); + + void visitLoadInst(LoadInst &I) { enqueueUsers(I); } + void visitCallBase(CallBase &CB); + void visitIntrinsicInst(IntrinsicInst &II); + +private: + bool tryEraseCallInvoke(Instruction *I); + void eraseFromWorklist(Instruction *I); +}; } static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) { @@ -73,6 +94,7 @@ bool Lowerer::lower(Function &F) { bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage(); bool Changed = false; + NoopCoroElider NCE(F.getDataLayout(), F.getContext()); SmallPtrSet DeadInsts{}; for (Instruction &I : instructions(F)) { if (auto *II = dyn_cast(&I)) { @@ -100,7 +122,7 @@ bool Lowerer::lower(Function &F) { II->replaceAllUsesWith(ConstantTokenNone::get(Context)); break; case Intrinsic::coro_noop: - elideCoroNoop(II); + NCE.run(II); if (!II->user_empty()) lowerCoroNoop(II); break; @@ -142,28 +164,6 @@ bool Lowerer::lower(Function &F) { return Changed; } -void Lowerer::elideCoroNoop(IntrinsicInst *II) { - for (User *U : make_early_inc_range(II->users())) { - auto *Fn = dyn_cast(U); - if (Fn == nullptr) - continue; - - auto *User = Fn->getUniqueUndroppableUser(); - if (auto *Call = dyn_cast(User)) { - Call->eraseFromParent(); - Fn->eraseFromParent(); - continue; - } - - if (auto *I = dyn_cast(User)) { - Builder.SetInsertPoint(I); - Builder.CreateBr(I->getNormalDest()); - I->eraseFromParent(); - Fn->eraseFromParent(); - } - } -} - void Lowerer::lowerCoroNoop(IntrinsicInst *II) { if (!NoopCoro) { LLVMContext &C = Builder.getContext(); @@ -200,6 +200,60 @@ void Lowerer::lowerCoroNoop(IntrinsicInst *II) { II->replaceAllUsesWith(NoopCoroVoidPtr); } +void NoopCoroElider::run(IntrinsicInst *II) { + visitPtr(*II); + + Worklist.clear(); + VisitedUses.clear(); +} + +void NoopCoroElider::visitCallBase(CallBase &CB) { + auto *V = U->get(); + bool ResumeOrDestroy = V == CB.getCalledOperand(); + if (ResumeOrDestroy) { + [[maybe_unused]] bool Success = tryEraseCallInvoke(&CB); + assert(Success && "Unexpected CallBase"); + + auto AboutToDeleteCallback = [this](Value *V) { + eraseFromWorklist(cast(V)); + }; + RecursivelyDeleteTriviallyDeadInstructions(V, nullptr, nullptr, + AboutToDeleteCallback); + } +} + +void NoopCoroElider::visitIntrinsicInst(IntrinsicInst &II) { + if (auto *SubFn = dyn_cast(&II)) { + auto *User = SubFn->getUniqueUndroppableUser(); + if (!tryEraseCallInvoke(cast(User))) + return; + SubFn->eraseFromParent(); + } +} + +bool NoopCoroElider::tryEraseCallInvoke(Instruction *I) { + if (auto *Call = dyn_cast(I)) { + eraseFromWorklist(Call); + Call->eraseFromParent(); + return true; + } + + if (auto *II = dyn_cast(I)) { + Builder.SetInsertPoint(II); + Builder.CreateBr(II->getNormalDest()); + eraseFromWorklist(II); + II->eraseFromParent(); + return true; + } + return false; +} + +void NoopCoroElider::eraseFromWorklist(Instruction *I) { + erase_if(Worklist, [I](UseToVisit &U) { + return I == U.UseAndIsOffsetKnown.getPointer()->getUser(); + }); +} + static bool declaresCoroCleanupIntrinsics(const Module &M) { return coro::declaresIntrinsics( M, diff --git a/llvm/test/Transforms/Coroutines/coro-cleanup-noop-elide.ll b/llvm/test/Transforms/Coroutines/coro-cleanup-noop-elide.ll new file mode 100644 index 0000000000000..6d9dd654b914a --- /dev/null +++ b/llvm/test/Transforms/Coroutines/coro-cleanup-noop-elide.ll @@ -0,0 +1,51 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6 +; RUN: opt < %s -S -passes='coro-cleanup' | FileCheck %s + +; Tests that resume or destroy a no-op coroutine can be erased; Finally, erase coro.noop if it has no users. +define void @erase() personality i32 0 { +; CHECK-LABEL: define void @erase() personality i32 0 { +; CHECK-NEXT: [[DONE:.*:]] +; CHECK-NEXT: ret void +; + %frame = call noundef ptr @llvm.coro.noop() + %resume = call ptr @llvm.coro.subfn.addr(ptr %frame, i8 0) + call fastcc void %resume(ptr %frame) + %destroy = call ptr @llvm.coro.subfn.addr(ptr %frame, i8 1) + invoke fastcc void %destroy(ptr %frame) + to label %done unwind label %unwind + +done: + ret void + +unwind: + %pad = landingpad { ptr, i32 } + catch ptr null + call void @terminate() + unreachable +} + +; Tests the load-and-call pattern despite mismatched calling conventions. Prevent instcombine from breaking code. +define void @load() personality i32 0 { +; CHECK-LABEL: define void @load() personality i32 0 { +; CHECK-NEXT: [[DONE:.*:]] +; CHECK-NEXT: ret void +; + %frame = call noundef ptr @llvm.coro.noop() + %resume = load ptr, ptr %frame, align 8 + call void %resume(ptr %frame) + %destroy.addr = getelementptr inbounds nuw i8, ptr %frame, i64 8 + %destroy = load ptr, ptr %destroy.addr, align 8 + invoke void %destroy(ptr %frame) + to label %done unwind label %unwind + +done: + ret void + +unwind: + %pad = landingpad { ptr, i32 } + catch ptr null + call void @terminate() + unreachable +} + +declare void @terminate() noreturn diff --git a/llvm/test/Transforms/Coroutines/coro-cleanup-noop-erase.ll b/llvm/test/Transforms/Coroutines/coro-cleanup-noop-erase.ll deleted file mode 100644 index 7fd9dc900ddb2..0000000000000 --- a/llvm/test/Transforms/Coroutines/coro-cleanup-noop-erase.ll +++ /dev/null @@ -1,24 +0,0 @@ -; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6 -; Tests that resume or destroy a no-op coroutine can be erased; Finally, erase coro.noop if it has no users. -; RUN: opt < %s -S -passes='coro-cleanup' | FileCheck %s - -define void @fn() personality i32 0 { -; CHECK-LABEL: define void @fn() personality i32 0 { -; CHECK-NEXT: [[DONE:.*:]] -; CHECK-NEXT: ret void -; - %frame = call noundef ptr @llvm.coro.noop() - %resume = call ptr @llvm.coro.subfn.addr(ptr %frame, i8 0) - call fastcc void %resume(ptr %frame) - %destroy = call ptr @llvm.coro.subfn.addr(ptr %frame, i8 1) - invoke fastcc void %destroy(ptr %frame) - to label %done unwind label %unwind - -done: - ret void - -unwind: - %pad = landingpad { ptr, i32 } - catch ptr null - unreachable -}