[coroutines] Handle spills before catchswitch

[GorNishanov]

Summary

If we need to spill the result of the PHI instruction, we insert the spill after
all of the PHIs and EHPads, however, in a catchswitch block there is no
room to insert the spill. Make room by splitting away catchswitch into a separate
block.

Before the fix:

catch.dispatch:
  %val = phi i32 [ 1, %if.then ], [ 2, %if.else ]
  %switch = catchswitch within none [label %catch] unwind label %cleanuppad

After:

catch.dispatch:
  %val = phi i32 [ 1, %if.then ], [ 2, %if.else ]
  %tok = cleanuppad within none []
  ; spill goes here
 cleanupret from %tok unwind label %catch.dispatch.switch
catch.dispatch.switch:
  %switch = catchswitch within none [label %catch] unwind label %cleanuppad

[majnemer]

lib/Transforms/Coroutines/CoroFrame.cpp

360–368

What if the catchswitch had a phi on it? You are adding a predecessor which would mean it needs new entries.

[GorNishanov]

lib/Transforms/Coroutines/CoroFrame.cpp

360–368

Not sure I understand.

PHIs are staying in place. The purpose of this change is to break the catchswitch away from its PHIs, so that we have room to insert spills. The catchswitch goes into a new block and refers to the original PHIs as before.

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroFrame.cpp

360–368

Ah, I understand now.

[GorNishanov]

lib/Transforms/Coroutines/CoroFrame.cpp

360–368

Aah. I understand now. Description was incorrectly stating that we are splitting the PHIs away.
Fixed the description to clarify that we splitting away the catchswitch keeping the PHIs in place.

[GorNishanov]

Thank you very much for the review!

Committed r303232

lib/Transforms/Coroutines/CoroFrame.cpp

for (auto const &S : Spills) {
 
     Types.push_back(Ty);
   }
   FrameTy->setBody(Types);
 
   return FrameTy;
 }
 
+// We need to make room to insert a spill after initial PHIs, but before
+// catchswitch instruction. Placing it before violates the requirement that
+// catchswitch, like all other EHPads must be the first nonPHI in a block.
+//
+// Split away catchswitch into a separate block and insert in its place:
+//
+//   cleanuppad <InsertPt> cleanupret.
+//
+// cleanupret instruction will act as an insert point for the spill.
+static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
+  BasicBlock *CurrentBlock = CatchSwitch->getParent();
+  BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
+  CurrentBlock->getTerminator()->eraseFromParent();
+
+  auto *CleanupPad =
+      CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
+  auto *CleanupRet =
+      CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
+  return CleanupRet;
+}
+
 // Replace all alloca and SSA values that are accessed across suspend points
 // with GetElementPointer from coroutine frame + loads and stores. Create an
 // AllocaSpillBB that will become the new entry block for the resume parts of
 // the coroutine:
 //
 //    %hdl = coro.begin(...)
 //    whatever
 //
if (CurrentValue != E.def()) {
           InsertPt = FramePtr->getNextNode();
         } else if (auto *II = dyn_cast<InvokeInst>(CurrentValue)) {
           // If we are spilling the result of the invoke instruction, split the
           // normal edge and insert the spill in the new block.
           auto NewBB = SplitEdge(II->getParent(), II->getNormalDest());
           InsertPt = NewBB->getTerminator();
         } else if (dyn_cast<PHINode>(CurrentValue)) {
           // Skip the PHINodes and EH pads instructions.
-          InsertPt =
-              &*cast<Instruction>(E.def())->getParent()->getFirstInsertionPt();
+          BasicBlock *DefBlock = cast<Instruction>(E.def())->getParent();
+          if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
+            InsertPt = splitBeforeCatchSwitch(CSI);
+          else
+            InsertPt = &*DefBlock->getFirstInsertionPt();
         } else {
           // For all other values, the spill is placed immediately after
           // the definition.
           assert(!isa<TerminatorInst>(E.def()) && "unexpected terminator");
           InsertPt = cast<Instruction>(E.def())->getNextNode();
         }
 
         Builder.SetInsertPoint(InsertPt);

test/Transforms/Coroutines/coro-catchswitch.ll

+; Verifies that we can insert the spill for a PHI preceding the catchswitch
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+target datalayout = "e-m:x-p:32:32-i64:64-f80:32-n8:16:32-a:0:32-S32"
+target triple = "i686-pc-windows-msvc"
+
+; CHECK-LABEL: define void @f(
+define void @f(i1 %cond) "coroutine.presplit"="1" personality i32 0 {
+entry:
+  %id = call token @llvm.coro.id(i32 8, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  br i1 %cond, label %if.else, label %if.then
+
+if.then:
+  invoke void @may_throw1()
+          to label %coro.ret unwind label %catch.dispatch
+
+if.else:
+  invoke void @may_throw2()
+          to label %coro.ret unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %if.else, %if.then
+  %val = phi i32 [ 1, %if.then ], [ 2, %if.else ]
+  %switch = catchswitch within none [label %catch] unwind label %cleanuppad
+
+; Verifies that we split out the PHI into a separate block
+; added a cleanuppad spill cleanupret unwinding into the catchswitch.
+
+; CHECK: catch.dispatch:
+; CHECK:  %val = phi i32 [ 2, %if.else ], [ 1, %if.then ]
+; CHECK:  %[[Pad:.+]] = cleanuppad within none []
+; CHECK:  %val.spill.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK:  store i32 %val, i32* %val.spill.addr
+; CHECK:  cleanupret from %[[Pad]] unwind label %[[Switch:.+]]
+
+; CHECK: [[Switch]]:
+; CHECK: %switch = catchswitch within none [label %catch] unwind to caller
+
+catch:                                            ; preds = %catch.dispatch
+  %pad = catchpad within %switch [i8* null, i32 64, i8* null]
+  catchret from %pad to label %suspend
+
+suspend:
+  %sp = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp, label %coro.ret [
+    i8 0, label %resume
+    i8 1, label %coro.ret
+  ]
+
+resume:                                   ; preds = %await2.suspend
+  call void @print(i32 %val)
+  br label %coro.ret
+
+coro.ret:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+    ret void
+
+cleanuppad:
+  %cpad = cleanuppad within none []
+  cleanupret from %cpad unwind to caller
+}
+
+; Function Attrs: argmemonly nounwind readonly
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
+
+; Function Attrs: nounwind
+declare i1 @llvm.coro.alloc(token) #2
+
+; Function Attrs: nobuiltin
+declare i32 @llvm.coro.size.i32() #4
+declare i8* @llvm.coro.begin(token, i8* writeonly) #2
+declare token @llvm.coro.save(i8*)
+declare i8 @llvm.coro.suspend(token, i1)
+
+declare void @may_throw1()
+declare void @may_throw2()
+declare void @print(i32)
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)
+
+declare i1 @llvm.coro.end(i8*, i1) #2
+
+; Function Attrs: nobuiltin nounwind
+
+; Function Attrs: argmemonly nounwind readonly
+declare i8* @llvm.coro.free(token, i8* nocapture readonly) #1
+