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

[nate_chandler]

Summary

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.

[aschwaighofer]

looks good to me

llvm/include/llvm/IR/Intrinsics.td

def int_coro_async_context_alloc : Intrinsic<[llvm_ptr_ty],
     [llvm_ptr_ty, llvm_ptr_ty],
     []>;
 def int_coro_async_context_dealloc : Intrinsic<[],
     [llvm_ptr_ty],
     []>;
 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],
                                        [IntrNoMem]>;
 def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                [WriteOnly<ArgIndex<1>>]>;
 
 def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
       case Intrinsic::coro_id_async:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
       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;
     }
   }
 
   if (Changed) {
     // After replacement were made we can cleanup the function body a little.
     simplifyCFG(F);
   }
 
   return Changed;
 }
 
 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,
                                        FunctionAnalysisManager &AM) {
   auto &M = *F.getParent();
   if (!declaresCoroCleanupIntrinsics(M) ||
       !Lowerer(M).lowerRemainingCoroIntrinsics(F))
     return PreservedAnalyses::all();

llvm/lib/Transforms/Coroutines/CoroInstr.h

public:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_async_resume;
   }
   static bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
+/// 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:
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_frame;
   }
   static bool classof(const Value *V) {

llvm/lib/Transforms/Coroutines/Coroutines.cpp

 
 #ifndef NDEBUG
 static bool isCoroutineIntrinsicName(StringRef Name) {
   // NOTE: Must be sorted!
   static const char *const CoroIntrinsics[] = {
       "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",
       "llvm.coro.done",
       "llvm.coro.end",
       "llvm.coro.end.async",
       "llvm.coro.frame",
       "llvm.coro.free",

llvm/test/Transforms/Coroutines/coro-async.ll

@my_async_function_fp = constant <{ i32, i32 }>
   <{ i32 trunc ( ; Relative pointer to async function
        i64 sub (
          i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @my_async_function to i64),
          i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_fp, i32 0, i32 1) to i64)
        )
      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) {
   %callee = bitcast i8* %fnPtr to void(i8*, %async.task*, %async.actor*)*
   tail call swiftcc void %callee(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)
   ret void
 }
 
entry:
   %val.2 = load i64, i64* %proj.2
   call void @some_user(i64 %val.2)
 
   tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor)
   call i1 (i8*, i1, ...) @llvm.coro.end.async(i8* %hdl, i1 0)
   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) {
 ; CHECK: entry:
 ; CHECK:   [[FRAMEPTR:%.*]] = getelementptr inbounds i8, i8* %async.ctxt, i64 128
 ; CHECK:   [[ACTOR_SPILL_ADDR:%.*]] = getelementptr inbounds i8, i8* %async.ctxt, i64 152
 ; CHECK:   [[CAST1:%.*]] = bitcast i8* [[ACTOR_SPILL_ADDR]] to %async.actor**
 ; CHECK:   store %async.actor* %actor, %async.actor** [[CAST1]]

 declare i1 @llvm.coro.end.async(i8*, i1, ...)
 declare i1 @llvm.coro.end(i8*, i1)
 declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
 declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
 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*)