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

diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index 99e18d0..582fe61 100644 (file)
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -1216,6 +1216,7 @@ def int_coro_async_context_dealloc : Intrinsic<[],
 def int_coro_async_resume : Intrinsic<[llvm_ptr_ty],
     [],
     []>;
+def int_coro_async_size_replace : Intrinsic<[], [llvm_ptr_ty, llvm_ptr_ty], []>;
 def int_coro_suspend_async
     : Intrinsic<[llvm_any_ty], [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
 def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],

diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
index 298149f..5b09cdb 100644 (file)
--- 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 (file)
--- 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 (file)
--- 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",

diff --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll
index aa83405..fb00a47 100644 (file)
--- a/llvm/test/Transforms/Coroutines/coro-async.ll
+++ b/llvm/test/Transforms/Coroutines/coro-async.ll
@@ -27,6 +27,15 @@ declare void @my_other_async_function(i8* %async.ctxt)
      to i32),
      i32 128    ; Initial async context size without space for frame
 }>
+@my_async_function_pa_fp = constant <{ i32, i32 }>
+  <{ i32 trunc (
+       i64 sub (
+         i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @my_async_function_pa to i64),
+         i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_pa_fp, i32 0, i32 1) to i64)
+       )
+     to i32),
+     i32 8
+}>
 
 ; Function that implements the dispatch to the callee function.
 define swiftcc void @my_async_function.my_other_async_function_fp.apply(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
@@ -99,8 +108,15 @@ entry:
   unreachable
 }
 
+define void @my_async_function_pa(i8* %ctxt, %async.task* %task, %async.actor* %actor) {
+  call void @llvm.coro.async.size.replace(i8* bitcast (<{i32, i32}>* @my_async_function_pa_fp to i8*), i8* bitcast (<{i32, i32}>* @my_async_function_fp to i8*))
+  call swiftcc void @my_async_function(i8* %ctxt, %async.task* %task, %async.actor* %actor)
+  ret void
+}
+
 ; Make sure we update the async function pointer
 ; CHECK: @my_async_function_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 }
+; CHECK: @my_async_function_pa_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 }
 ; CHECK: @my_async_function2_fp = constant <{ i32, i32 }> <{ {{.*}}, i32 176 }
 
 ; CHECK-LABEL: define swiftcc void @my_async_function(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
@@ -511,3 +527,4 @@ declare void @llvm.coro.async.context.dealloc(i8*)
 declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)
 declare swiftcc void @asyncSuspend(i8*, %async.task*, %async.actor*)
 declare i8* @llvm.coro.async.resume()
+declare void @llvm.coro.async.size.replace(i8*, i8*)