[coro] Async coroutines: Allow more than 3 arguments in the dispatch function

[aschwaighofer]

Summary

We need to be able to call function pointers. Inline the dispatch
function.

rdar://70097093

[rjmccall]

I don't quite understand the need to inline here.

[aschwaighofer]

Semantically, it is not needed. But without inlining here this helper dispatch function thunk will remain at O0 because inlining is not run again.

[aschwaighofer]

I am worried about introducing unnecessary thunks that the debugger has to step through.

[aschwaighofer]

Also inline the context project function used by the resume function.

[rjmccall]

What is this function being inlined? We expect the frontend to emit a thunk that glues things together?

[aschwaighofer]

Yes, frontend is expected to emit a thunk that models the tail call to the function at the suspend point. This prevents code moving in between the tail call and return before coroutine splitting and allows for specifying how to perform the call (e.g pointer authentication).

define internal hidden void @__suspend_dispatch_3.1(i8* %0, %task* %1, %executor* %2, %context* %3)  {
entry:
  %4 = bitcast i8* %0 to void (%task*, %executor*, %context*)*
  tail call swiftcc void %4(%task* %1, %executor* %2, %context* %3)
  ret void
}

define linkonce_odr hidden i8* @__resume_project_context(i8* %0) {
entry:
  %1 = bitcast i8* %0 to i8**
  %2 = load i8*, i8** %1, align 8
  ret i8* %2
}

%resume_func = call i8* @llvm.coro.async.resume()
%24 = call { i8*, i8*, i8* } (i8*, i8*, ...) @llvm.coro.suspend.async(
  i8* %resume_func, 
  i8* bitcast (i8* (i8*)* @__resume_project_context to i8*), 
  i8* bitcast (void (i8*, %task*, %executor*, %context*)* @__suspend_dispatch_3.1 to i8*), 
  i8* %fun_ptr, %task* %0, %executor* %1, %context* %22)

[aschwaighofer]

Transfer the debug location from the suspend point to the functions called at and after the suspend point (context projection).

[rjmccall]

I see. This LGTM as an immediate fix, then.

I wonder if we shouldn't have a different IR design, though. Maybe we should lower directly to IR that uses a pattern more like the SIL get_unsafe_continuation pattern? i.e.

%0 = call token @llvm.coro.async.continuation()  // returns the continuation function pointer; must be used by exactly one suspend
... // somewhere in here we store the continuation into the child context
call void @llvm.coro.async.suspend(token %0)

I can't remember if we discussed and discarded this.

[rjmccall]

Oh, unfortunately llvm.coro.async.continuation wouldn't be able to just return the continuation function pointer; there would have to be a llvm.coro.async.get_continuation_function intrinsic. But we do something similar with the alloc intrinsics.

[aschwaighofer]

Use the function argument index instead of the function argument in
coro.id.async. This solves any spurious use issues.

Coerce the arguments of the tail call function at a suspend point. The LLVM
optimizer seems to drop casts leading to a vararg intrinsic.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

679

InlineRes will be unused in a release build. Add (void)?

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

case coro::ABI::Async: {
     CallerContext->setCallingConv(ProjectionFunc->getCallingConv());
     // The frame is located after the async_context header.
     auto &Context = Builder.getContext();
     auto *FramePtrAddr = Builder.CreateConstInBoundsGEP1_32(
         Type::getInt8Ty(Context), CallerContext,
         Shape.AsyncLowering.FrameOffset, "async.ctx.frameptr");
     return Builder.CreateBitCast(FramePtrAddr, FramePtrTy);
   }
   // In continuation-lowering, the argument is the opaque storage.
   case coro::ABI::Retcon:
   case coro::ABI::RetconOnce: {
     Argument *NewStorage = &*NewF->arg_begin();
     auto FramePtrTy = Shape.FrameTy->getPointerTo();
 
     // If the storage is inline, just bitcast to the storage to the frame type.
     if (Shape.RetconLowering.IsFrameInlineInStorage)
       return Builder.CreateBitCast(NewStorage, FramePtrTy);
for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
 
     // Place it before the first suspend.
     auto *ReturnBB =
         BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);
     Branch->setSuccessor(0, ReturnBB);
 
     IRBuilder<> Builder(ReturnBB);
 
-    // Insert the call to the tail call function.
+    // Insert the call to the tail call function and inline it.
     auto *Fun = Suspend->getMustTailCallFunction();
     SmallVector<Value *, 8> Args(Suspend->operand_values());
     auto *TailCall = Builder.CreateCall(
         cast<FunctionType>(Fun->getType()->getPointerElementType()), Fun,
         ArrayRef<Value *>(Args).drop_front(3).drop_back(1));
-    TailCall->setTailCallKind(CallInst::TCK_MustTail);
+    TailCall->setTailCall();
     TailCall->setCallingConv(Fun->getCallingConv());
+    InlineFunctionInfo FnInfo;
+    auto InlineRes = InlineFunction(*TailCall, FnInfo);
+    assert(InlineRes.isSuccess() && "Expected inlining to succeed");
     Builder.CreateRetVoid();
 
     // Replace the lvm.coro.async.resume intrisic call.
     replaceAsyncResumeFunction(Suspend, Continuation);
   }
 
   assert(Clones.size() == Shape.CoroSuspends.size());
   for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {

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

<{ i32 trunc ( ; Relative pointer to async function
          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
 }>
 
 ; Function that implements the dispatch to the callee function.
-define swiftcc void @my_async_function.my_other_async_function_fp.apply(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
-  musttail call swiftcc void @asyncSuspend(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)
+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
 }
 
 declare void @some_user(i64)
 declare void @some_may_write(i64*)
 
 define i8* @resume_context_projection(i8* %ctxt) {
 entry:
entry:
   %return_to_caller.addr = bitcast void(i8*, %async.task*, %async.actor*)** %callee_context.return_to_caller.addr to i8**
   %resume.func_ptr = call i8* @llvm.coro.async.resume()
   store i8* %resume.func_ptr, i8** %return_to_caller.addr
 
   ; store caller context into callee context
   %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr
   %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
   %res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
                                                   i8* %resume.func_ptr,
                                                   i8* %resume_proj_fun,
-                                                  void (i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
-                                                  i8* %callee_context, %async.task* %task, %async.actor *%actor)
+                                                  void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
+                                                  i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
 
   call void @llvm.coro.async.context.dealloc(i8* %callee_context)
   %continuation_task_arg = extractvalue {i8*, i8*, i8*} %res, 1
   %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
   %val = load i64, i64* %proj.1
   call void @some_user(i64 %val)
   %val.2 = load i64, i64* %proj.2
   call void @some_user(i64 %val.2)

 ; CHECK:   [[CALLEE_CTXT_SPILL:%.*]] = getelementptr inbounds i8, i8* %async.ctxt, i64 160
 ; CHECK:   [[CAST2:%.*]] = bitcast i8* [[CALLEE_CTXT_SPILL]] to i8**
 ; CHECK:   store i8* [[CALLEE_CTXT]], i8** [[CAST2]]
 ; CHECK:   [[TYPED_RETURN_TO_CALLER_ADDR:%.*]] = getelementptr inbounds i8, i8* [[CALLEE_CTXT]], i64 8
 ; CHECK:   [[RETURN_TO_CALLER_ADDR:%.*]] = bitcast i8* [[TYPED_RETURN_TO_CALLER_ADDR]] to i8**
 ; CHECK:   store i8* bitcast (void (i8*, i8*, i8*)* @my_async_function.resume.0 to i8*), i8** [[RETURN_TO_CALLER_ADDR]]
 ; CHECK:   [[CALLER_CONTEXT_ADDR:%.*]] = bitcast i8* [[CALLEE_CTXT]] to i8**
 ; CHECK:   store i8* %async.ctxt, i8** [[CALLER_CONTEXT_ADDR]]
-; CHECK:   musttail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor)
+; CHECK:   tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor)
 ; CHECK:   ret void
 ; CHECK: }
 
 ; CHECK-LABEL: define internal swiftcc void @my_async_function.resume.0(i8* nocapture readonly %0, i8* %1, i8* nocapture readnone %2) {
 ; CHECK: entryresume.0:
 ; CHECK:   [[CALLER_CONTEXT_ADDR:%.*]] = bitcast i8* %0 to i8**
 ; CHECK:   [[CALLER_CONTEXT:%.*]] = load i8*, i8** [[CALLER_CONTEXT_ADDR]]
 ; CHECK:   [[FRAME_PTR:%.*]] = getelementptr inbounds i8, i8* [[CALLER_CONTEXT]], i64 128
entry:
   %callee_context.0 = bitcast i8* %callee_context to %async.ctxt*
   %callee_context.return_to_caller.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 1
   %return_to_caller.addr = bitcast void(i8*, %async.task*, %async.actor*)** %callee_context.return_to_caller.addr to i8**
   %resume.func_ptr = call i8* @llvm.coro.async.resume()
   store i8* %resume.func_ptr, i8** %return_to_caller.addr
   %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr
   %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
   %res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
                                                   i8* %resume.func_ptr,
                                                   i8* %resume_proj_fun,
-                                                  void (i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
-                                                  i8* %callee_context, %async.task* %task, %async.actor *%actor)
+                                                  void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
+                                                  i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
 
   %continuation_task_arg = extractvalue {i8*, i8*, i8*} %res, 0
   %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
 
   %callee_context.0.1 = bitcast i8* %callee_context to %async.ctxt*
   %callee_context.return_to_caller.addr.1 = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0.1, i32 0, i32 1
   %return_to_caller.addr.1 = bitcast void(i8*, %async.task*, %async.actor*)** %callee_context.return_to_caller.addr.1 to i8**
   %resume.func_ptr.1 = call i8* @llvm.coro.async.resume()
   store i8* %resume.func_ptr.1, i8** %return_to_caller.addr.1
   %callee_context.caller_context.addr.1 = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0.1, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr.1
   %resume_proj_fun.2 = bitcast i8*(i8*)* @resume_context_projection to i8*
+  %callee.2 = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
   %res.2 = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
                                                   i8* %resume.func_ptr.1,
                                                   i8* %resume_proj_fun.2,
-                                                  void (i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
-                                                  i8* %callee_context, %async.task* %task, %async.actor *%actor)
+                                                  void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
+                                                  i8* %callee.2, i8* %callee_context, %async.task* %task, %async.actor *%actor)
 
   call void @llvm.coro.async.context.dealloc(i8* %callee_context)
   %continuation_actor_arg = extractvalue {i8*, i8*, i8*} %res.2, 1
   %actor.2 =  bitcast i8* %continuation_actor_arg to %async.actor*
 
   tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor.2)
   call i1 @llvm.coro.end(i8* %hdl, i1 0)
   unreachable
 }
 
 ; CHECK-LABEL: define swiftcc void @my_async_function2(%async.task* %task, %async.actor* %actor, i8* %async.ctxt) {
 ; CHECK: store i8* %async.ctxt,
 ; CHECK: store %async.actor* %actor,
 ; CHECK: store %async.task* %task,
 ; CHECK: [[CALLEE_CTXT:%.*]] =  tail call i8* @llvm.coro.async.context.alloc(
 ; CHECK: store i8* [[CALLEE_CTXT]],
 ; CHECK: store i8* bitcast (void (i8*, i8*, i8*)* @my_async_function2.resume.0 to i8*),
 ; CHECK: store i8* %async.ctxt,
-; CHECK: musttail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor)
+; CHECK: tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor)
 ; CHECK: ret void
 
 ; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.0(i8* %0, i8* nocapture readnone %1, i8* nocapture readonly %2) {
 ; CHECK: [[CALLEE_CTXT_ADDR:%.*]] = bitcast i8* %2 to i8**
 ; CHECK: [[CALLEE_CTXT:%.*]] = load i8*, i8** [[CALLEE_CTXT_ADDR]]
 ; CHECK: [[CALLEE_CTXT_SPILL_ADDR:%.*]] = getelementptr inbounds i8, i8* [[CALLEE_CTXT]], i64 152
 ; CHECK: [[CALLEE_CTXT_SPILL_ADDR2:%.*]] = bitcast i8* [[CALLEE_CTXT_SPILL_ADDR]] to i8**
 ; CHECK: store i8* bitcast (void (i8*, i8*, i8*)* @my_async_function2.resume.1 to i8*),
 ; CHECK: [[CALLLE_CTXT_RELOAD:%.*]] = load i8*, i8** [[CALLEE_CTXT_SPILL_ADDR2]]
-; CHECK: musttail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT_RELOAD]]
+; CHECK: tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT_RELOAD]]
 ; CHECK: ret void
 
 ; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.1(i8* nocapture readnone %0, i8* %1, i8* nocapture readonly %2) {
 ; CHECK: bitcast i8* %2 to i8**
 ; CHECK: [[ACTOR_ARG:%.*]] = bitcast i8* %1
 ; CHECK: tail call swiftcc void @asyncReturn({{.*}}[[ACTOR_ARG]])
 ; CHECK: ret void