[CoroCleanup] Noop coroutine elision for load-and-call pattern

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

@@ -8,13 +8,15 @@
 
 #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"
 #include "llvm/IR/InstIterator.h"
 #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,27 @@ 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<NoopCoroElider> {
+  using Base = PtrUseVisitor<NoopCoroElider>;
+
+  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);
+};
 }
 
 static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
@@ -73,6 +93,7 @@ bool Lowerer::lower(Function &F) {
   bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage();
   bool Changed = false;
 
+  NoopCoroElider NCE(F.getDataLayout(), F.getContext());
   SmallPtrSet<Instruction *, 8> DeadInsts{};
   for (Instruction &I : instructions(F)) {
     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
@@ -100,7 +121,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 +163,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<CoroSubFnInst>(U);
-    if (Fn == nullptr)
-      continue;
-
-    auto *User = Fn->getUniqueUndroppableUser();
-    if (auto *Call = dyn_cast<CallInst>(User)) {
-      Call->eraseFromParent();
-      Fn->eraseFromParent();
-      continue;
-    }
-
-    if (auto *I = dyn_cast<InvokeInst>(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 +199,47 @@ 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");
+    RecursivelyDeleteTriviallyDeadInstructions(V);
+  }
+}
+
+void NoopCoroElider::visitIntrinsicInst(IntrinsicInst &II) {
+  if (auto *SubFn = dyn_cast<CoroSubFnInst>(&II)) {
+    auto *User = SubFn->getUniqueUndroppableUser();
+    if (!tryEraseCallInvoke(cast<Instruction>(User)))
+      return;
+    SubFn->eraseFromParent();
+  }
+}
+
+bool NoopCoroElider::tryEraseCallInvoke(Instruction *I) {
+  if (auto *Call = dyn_cast<CallInst>(I)) {
+    Call->eraseFromParent();
+    return true;
+  }
+
+  if (auto *II = dyn_cast<InvokeInst>(I)) {
+    Builder.SetInsertPoint(II);
+    Builder.CreateBr(II->getNormalDest());
+    II->eraseFromParent();
+    return true;
+  }
+  return false;
+}
+
 static bool declaresCoroCleanupIntrinsics(const Module &M) {
   return coro::declaresIntrinsics(
       M,

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