Add @llvm.coro.async.size.replace intrinsic.

The new intrinsic replaces the size in one specified AsyncFunctionPointer with
the size in another.  This ability is necessary for functions which merely
forward to async functions such as those defined for partial applications.

Reviewed By: aschwaighofer

Differential Revision: https://reviews.llvm.org/D97229

3 files changed, 34 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
index 298149f..5b09cdb 100644
--- a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
@@ -80,6 +80,23 @@ bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
       case Intrinsic::coro_subfn_addr:
         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
         break;
+      case Intrinsic::coro_async_size_replace:
+        auto *Target = cast<ConstantStruct>(
+            cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
+                ->getInitializer());
+        auto *Source = cast<ConstantStruct>(
+            cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
+                ->getInitializer());
+        auto *TargetSize = Target->getOperand(1);
+        auto *SourceSize = Source->getOperand(1);
+        if (TargetSize->isElementWiseEqual(SourceSize)) {
+          break;
+        }
+        auto *TargetRelativeFunOffset = Target->getOperand(0);
+        auto *NewFuncPtrStruct = ConstantStruct::get(
+            Target->getType(), TargetRelativeFunOffset, SourceSize);
+        Target->replaceAllUsesWith(NewFuncPtrStruct);
+        break;
       }
       II->eraseFromParent();
       Changed = true;
@@ -95,10 +112,10 @@ bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
 }
 
 static bool declaresCoroCleanupIntrinsics(const Module &M) {
-  return coro::declaresIntrinsics(M, {"llvm.coro.alloc", "llvm.coro.begin",
-                                      "llvm.coro.subfn.addr", "llvm.coro.free",
-                                      "llvm.coro.id", "llvm.coro.id.retcon",
-                                      "llvm.coro.id.retcon.once"});
+  return coro::declaresIntrinsics(
+      M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
+          "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
+          "llvm.coro.id.retcon.once", "llvm.coro.async.size.replace"});
 }
 
 PreservedAnalyses CoroCleanupPass::run(Function &F,
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index 9fa2fd1..7f6e95f 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -376,6 +376,18 @@ public:
   }
 };
 
+/// This represents the llvm.coro.async.size.replace instruction.
+class LLVM_LIBRARY_VISIBILITY CoroAsyncSizeReplace : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_async_size_replace;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 /// This represents the llvm.coro.frame instruction.
 class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst {
 public:
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 3b2b3c6..e35875b 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -126,6 +126,7 @@ static bool isCoroutineIntrinsicName(StringRef Name) {
       "llvm.coro.alloc",
       "llvm.coro.async.context.alloc",
       "llvm.coro.async.context.dealloc",
+      "llvm.coro.async.size.replace",
       "llvm.coro.async.store_resume",
       "llvm.coro.begin",
       "llvm.coro.destroy",