[llvm] d861cc0 - [coro] Async coroutines: Make sure we can handle control flow in suspend point dispatch function

[Arnold Schwaighofer via llvm-commits]

Author: Arnold Schwaighofer
Date: 2020-11-16T11:59:02-08:00
New Revision: d861cc0e4311ec193cea9aaa5cfdc2b55563222c

URL: https://github.com/llvm/llvm-project/commit/d861cc0e4311ec193cea9aaa5cfdc2b55563222c
DIFF: https://github.com/llvm/llvm-project/commit/d861cc0e4311ec193cea9aaa5cfdc2b55563222c.diff

LOG: [coro] Async coroutines: Make sure we can handle control flow in suspend point dispatch function

Create a valid basic block with a terminator before we call
InlineFunction.

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

Added: 
    

Modified: 
    llvm/lib/Transforms/Coroutines/CoroSplit.cpp
    llvm/test/Transforms/Coroutines/coro-async.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index ddc7bc7752d9..34ee70a6d380 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1458,11 +1458,11 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
     TailCall->setDebugLoc(DbgLoc);
     TailCall->setTailCall();
     TailCall->setCallingConv(Fn->getCallingConv());
+    Builder.CreateRetVoid();
     InlineFunctionInfo FnInfo;
     auto InlineRes = InlineFunction(*TailCall, FnInfo);
     assert(InlineRes.isSuccess() && "Expected inlining to succeed");
     (void)InlineRes;
-    Builder.CreateRetVoid();
 
     // Replace the lvm.coro.async.resume intrisic call.
     replaceAsyncResumeFunction(Suspend, Continuation);

diff  --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll
index 0faff4aed2e0..f0036440c803 100644
--- a/llvm/test/Transforms/Coroutines/coro-async.ll
+++ b/llvm/test/Transforms/Coroutines/coro-async.ll
@@ -260,6 +260,60 @@ define swiftcc void @top_level_caller(i8* %ctxt, i8* %task, i8* %actor) {
 ; CHECK: tail call swiftcc void @asyncSuspend
 ; CHECK: ret void
 
+ at dont_crash_on_cf_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
+}>
+
+
+define swiftcc void @dont_crash_on_cf_dispatch(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
+  %isNull = icmp eq %async.task* %task, null
+  br i1 %isNull, label %is_null, label %is_not_null
+
+is_null:
+  ret void
+
+is_not_null:
+  %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
+}
+
+define swiftcc void @dont_crash_on_cf(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)  {
+entry:
+  %id = call token @llvm.coro.id.async(i32 128, i32 16, i32 0,
+          i8* bitcast (<{i32, i32}>* @dont_crash_on_cf_fp to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  %arg0 = bitcast %async.task* %task to i8*
+  %arg1 = bitcast <{ i32, i32}>* @my_other_async_function_fp to i8*
+  %callee_context = call i8* @llvm.coro.async.context.alloc(i8* %arg0, i8* %arg1)
+	%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*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
+                                                  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*
+  tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor)
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
 declare i8* @llvm.coro.prepare.async(i8*)
 declare token @llvm.coro.id.async(i32, i32, i32, i8*)
 declare i8* @llvm.coro.begin(token, i8*)