[coro] Fix rematerializable instruction sinking to coro.suspend blocks

[aschwaighofer]

Summary

There is a constraint that coro.suspend instructions need to be in their
own blocks. The coro split pass initially creates IR that obeys this constraint
(which is later checked). Sinking rematerializable instructions into these
blocks breaks that constraint.

Instead rematerialize in the predecessor block to the suspend's single
predecessor block.

[nate_chandler]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1932

This should this be getSinglePredecessor(), right?

[aschwaighofer]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1932

Thanks, for catching that. Yes we need to rematerialized the value in the predecessor block to the the suspend if the suspend uses it.

[aschwaighofer]

getSingleSuccessor call should be getSinglePredecessor

[rjmccall]

Ah, and this is one of those things that probably doesn't affect the C++ lowering because it doesn't pass information directly to and from coro.suspend. LGTM.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

static void rewriteMaterializableInstructions(IRBuilder<> &IRB,
                                               const SpillInfo &Spills) {
   for (const auto &E : Spills) {
     Value *Def = E.first;
     BasicBlock *CurrentBlock = nullptr;
     Instruction *CurrentMaterialization = nullptr;
     for (Instruction *U : E.second) {
       // If we have not seen this block, materialize the value.
       if (CurrentBlock != U->getParent()) {
-        CurrentBlock = U->getParent();
+
+        bool IsInCoroSuspendBlock = isa<AnyCoroSuspendInst>(U);
+        CurrentBlock = IsInCoroSuspendBlock
+                           ? U->getParent()->getSingleSuccessor()
+                           : U->getParent();
         CurrentMaterialization = cast<Instruction>(Def)->clone();
         CurrentMaterialization->setName(Def->getName());
         CurrentMaterialization->insertBefore(
-            &*CurrentBlock->getFirstInsertionPt());
+            IsInCoroSuspendBlock ? CurrentBlock->getTerminator()
+                                 : &*CurrentBlock->getFirstInsertionPt());
       }
       if (auto *PN = dyn_cast<PHINode>(U)) {
         assert(PN->getNumIncomingValues() == 1 &&
                "unexpected number of incoming "
                "values in the PHINode");
         PN->replaceAllUsesWith(CurrentMaterialization);
         PN->eraseFromParent();
         continue;

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

 ; RUN: opt < %s -enable-coroutines -O2 -S | FileCheck --check-prefixes=CHECK %s
 ; RUN: opt < %s -enable-coroutines -passes='default<O2>' -S | FileCheck --check-prefixes=CHECK %s
-
+; RUN: opt < %s -enable-coroutines -O0 -S
 target datalayout = "p:64:64:64"
 
 %async.task = type { i64 }
 %async.actor = type { i64 }
 %async.fp = type <{ i32, i32 }>
 
 %async.ctxt = type { i8*, void (i8*, %async.task*, %async.actor*)* }
 
entry:
   %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*
+  %task.casted = bitcast i8* %arg0 to %async.task*
   %res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 2,
                                                   i8* %resume.func_ptr,
                                                   i8* %resume_proj_fun,
                                                   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), !dbg !9
+                                                  i8* %callee, i8* %callee_context, %async.task* %task.casted, %async.actor *%actor), !dbg !9
 
   %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*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
                                                   i8* %resume.func_ptr.1,
                                                   i8* %resume_proj_fun.2,
                                                   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)
+                                                  i8* %callee.2, i8* %callee_context, %async.task* %task.casted, %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

 declare i1 @llvm.coro.end(i8*, i1)
 declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, 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*)
+declare i8* @hide(i8*)
 
 !llvm.dbg.cu = !{!2}
 !llvm.module.flags = !{!0}
 
 !0 = !{i32 2, !"Debug Info Version", i32 3}
 ; CHECK: ![[SP1]] = distinct !DISubprogram(name: "my_async_function",
 ; CHECK-SAME:                              linkageName: "my_async_function",
 ; CHECK-SAME:                              scopeLine: 1