[Coroutine] Split PHI Nodes in `cleanuppad` blocks in a way that obeys EH pad rules

[dpaoliello]

Summary

Issue Details:
In order to support coroutine splitting, any multi-value PHI node in a coroutine is split into multiple blocks with single-value PHI Nodes, which then allows a subsequent transform to generate reload instructions as required (i.e., to reload the value if required if the coroutine has been resumed). This causes issues with EH pads (catchswitch and catchpad) as all pads within a catchswitch must have the same unwind destination, but the coroutine splitting logic may modify them to each have a unique unwind destination if there is a PHI node in the unwind cleanuppad that is set from values in the catchswitch and cleanuppad blocks.

Fix Details:
During splitting, if such a PHI node is detected, then create a "dispatcher" cleanuppad as well as the blocks with single-value PHI Nodes: thus the "dispatcher" is the unwind destination and it will detect which predecessor called it and then branch to the appropriate single-value PHI node block, which will then branch back to the original cleanuppad block.

[dpaoliello]

Ran clang-format

[dpaoliello]

Added all context lines

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1211

Should this be

if (CatchSwitch->getUnwindDest() == BB)
  return true;

?
Or is it impossible to have multiple predecessors with CatchSwitch all pointing to BB somehow?

[rnk]

Seems reasonable, if the new multi-value phi works with later passes.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1145

Isn't this still a multi-valued phi node? Is it OK because it has constant arguments?

[GorNishanov]

LGTM.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1145

The use of %2 is in the same basic block as the definition for the phi, so it cannot cross a suspend point and therefore we will not try to spill and reload it.

[dpaoliello]

Addressed feedback

[dpaoliello]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1145

The subsequent phase that is transforming the single-valued phi nodes is trying to handle a value being written before a coro suspend and then read afterwards (hence the introduction of reload instructions). It is true that said subsequent transform will not handle this multi-valued phi node, *but* since the prior single-valued phi nodes are introduced by splitting an existing edge we are guaranteed to not have a coro suspend between the single-valued phi nodes and this multi-valued phi node, thus reading these values is safe (and does not require a reload).

1211

That's a fair point: I think having this in its own function creates some ambiguity, since what I'm specifically looking for is a cleanuppad with a catchswitch predecessor (therefore the cleanuppad is an unwind destination). I'm going this check inline to make it clearer.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1249

Could we add an assertion if BB is guaranteed to be the unwind dest?

[dpaoliello]

Add requested assert

[dpaoliello]

I believe that this commit has the required approvals: can someone please merge it?

[dpaoliello]

@lxfind Do you have commit permissions? If so, could you please merge this in?

[lxfind]

In D88059#2295409, @dpaoliello wrote:

@lxfind Do you have commit permissions? If so, could you please merge this in?

Sure! Could you please rebase? Somehow I couldn't apply this patch locally.

[dpaoliello]

Rebase to master

[dpaoliello]

In D88059#2295449, @lxfind wrote:

In D88059#2295409, @dpaoliello wrote:

@lxfind Do you have commit permissions? If so, could you please merge this in?

Sure! Could you please rebase? Somehow I couldn't apply this patch locally.

Done - thanks!

[lxfind]

In D88059#2295409, @dpaoliello wrote:

@lxfind Do you have commit permissions? If so, could you please merge this in?

hmm I still cannot patch this from master. Can you double check that you don't have a local commit in the tree?

[dpaoliello]

Regenerate diff

[dpaoliello]

In D88059#2295488, @lxfind wrote:

In D88059#2295409, @dpaoliello wrote:

@lxfind Do you have commit permissions? If so, could you please merge this in?

hmm I still cannot patch this from master. Can you double check that you don't have a local commit in the tree?

Looks like it's an issue with using PowerShell on Windows to redirect the output of git diff to a file.
I regenerated with git built-in --output option, and it seems to be working - please try again with the latest patch.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

113

CS gets unused in release build, leading to a warning. You will need (void) it.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

struct SuspendCrossingInfo {
 
   SuspendCrossingInfo(Function &F, coro::Shape &Shape);
 
   bool hasPathCrossingSuspendPoint(BasicBlock *DefBB, BasicBlock *UseBB) const {
     size_t const DefIndex = Mapping.blockToIndex(DefBB);
     size_t const UseIndex = Mapping.blockToIndex(UseBB);
 
     bool const Result = Block[UseIndex].Kills[DefIndex];
     LLVM_DEBUG(dbgs() << UseBB->getName() << " => " << DefBB->getName()
                       << " answer is " << Result << "\n");
     return Result;
   }
 
   bool isDefinitionAcrossSuspend(BasicBlock *DefBB, User *U) const {
     auto *I = cast<Instruction>(U);
 
     // We rewrote PHINodes, so that only the ones with exactly one incoming
else if (auto *CR = dyn_cast<CleanupReturnInst>(TI))
     CR->setUnwindDest(Succ);
   else
     llvm_unreachable("unexpected terminator instruction");
 }
 
 // Replaces all uses of OldPred with the NewPred block in all PHINodes in a
 // block.
 static void updatePhiNodes(BasicBlock *DestBB, BasicBlock *OldPred,
-                           BasicBlock *NewPred,
-                           PHINode *LandingPadReplacement) {
+                           BasicBlock *NewPred, PHINode *Until = nullptr) {
   unsigned BBIdx = 0;
   for (BasicBlock::iterator I = DestBB->begin(); isa<PHINode>(I); ++I) {
     PHINode *PN = cast<PHINode>(I);
 
     // We manually update the LandingPadReplacement PHINode and it is the last
     // PHI Node. So, if we find it, we are done.
-    if (LandingPadReplacement == PN)
+    if (Until == PN)
       break;
 
     // Reuse the previous value of BBIdx if it lines up.  In cases where we
     // have multiple phi nodes with *lots* of predecessors, this is a speed
     // win because we don't have to scan the PHI looking for TIBB.  This
     // happens because the BB list of PHI nodes are usually in the same
     // order.
     if (PN->getIncomingBlock(BBIdx) != OldPred)
static BasicBlock *ehAwareSplitEdge(BasicBlock *BB, BasicBlock *Succ,
   else
     llvm_unreachable("handling for other EHPads not implemented yet");
 
   auto *NewCleanupPad = CleanupPadInst::Create(ParentPad, {}, "", NewBB);
   CleanupReturnInst::Create(NewCleanupPad, Succ, NewBB);
   return NewBB;
 }
 
+// Moves the values in the PHIs in SuccBB that correspong to PredBB into a new
+// PHI in InsertedBB.
+static void movePHIValuesToInsertedBlock(BasicBlock *SuccBB,
+                                         BasicBlock *InsertedBB,
+                                         BasicBlock *PredBB,
+                                         PHINode *UntilPHI = nullptr) {
+  auto *PN = cast<PHINode>(&SuccBB->front());
+  do {
+    int Index = PN->getBasicBlockIndex(InsertedBB);
+    Value *V = PN->getIncomingValue(Index);
+    PHINode *InputV = PHINode::Create(
+        V->getType(), 1, V->getName() + Twine(".") + SuccBB->getName(),
+        &InsertedBB->front());
+    InputV->addIncoming(V, PredBB);
+    PN->setIncomingValue(Index, InputV);
+    PN = dyn_cast<PHINode>(PN->getNextNode());
+  } while (PN != UntilPHI);
+}
+
+// Rewrites the PHI Nodes in a cleanuppad.
+static void rewritePHIsForCleanupPad(BasicBlock *CleanupPadBB,
+                                     CleanupPadInst *CleanupPad) {
+  // For every incoming edge to a CleanupPad we will create a new block holding
+  // all incoming values in single-value PHI nodes. We will then create another
+  // block to act as a dispather (as all unwind edges for related EH blocks
+  // must be the same).
+  //
+  // cleanuppad:
+  //    %2 = phi i32[%0, %catchswitch], [%1, %catch.1]
+  //    %3 = cleanuppad within none []
+  //
+  // It will create:
+  //
+  // cleanuppad.corodispatch
+  //    %2 = phi i8[0, %catchswitch], [1, %catch.1]
+  //    %3 = cleanuppad within none []
+  //    switch i8 % 2, label %unreachable
+  //            [i8 0, label %cleanuppad.from.catchswitch
+  //             i8 1, label %cleanuppad.from.catch.1]
+  // cleanuppad.from.catchswitch:
+  //    %4 = phi i32 [%0, %catchswitch]
+  //    br %label cleanuppad
+  // cleanuppad.from.catch.1:
+  //    %6 = phi i32 [%1, %catch.1]
+  //    br %label cleanuppad
+  // cleanuppad:
+  //    %8 = phi i32 [%4, %cleanuppad.from.catchswitch],
+  //                 [%6, %cleanuppad.from.catch.1]
+
+  // Unreachable BB, in case switching on an invalid value in the dispatcher.
+  auto *UnreachBB = BasicBlock::Create(
+      CleanupPadBB->getContext(), "unreachable", CleanupPadBB->getParent());
+  IRBuilder<> Builder(UnreachBB);
+  Builder.CreateUnreachable();
+
+  // Create a new cleanuppad which will be the dispatcher.
+  auto *NewCleanupPadBB =
+      BasicBlock::Create(CleanupPadBB->getContext(),
+                         CleanupPadBB->getName() + Twine(".corodispatch"),
+                         CleanupPadBB->getParent(), CleanupPadBB);
+  Builder.SetInsertPoint(NewCleanupPadBB);
+  auto *SwitchType = Builder.getInt8Ty();
+  auto *SetDispatchValuePN =
+      Builder.CreatePHI(SwitchType, pred_size(CleanupPadBB));
+  CleanupPad->removeFromParent();
+  CleanupPad->insertAfter(SetDispatchValuePN);
+  auto *SwitchOnDispatch = Builder.CreateSwitch(SetDispatchValuePN, UnreachBB,
+                                                pred_size(CleanupPadBB));
+
+  int SwitchIndex = 0;
+  SmallVector<BasicBlock *, 8> Preds(pred_begin(CleanupPadBB),
+                                     pred_end(CleanupPadBB));
+  for (BasicBlock *Pred : Preds) {
+    // Create a new cleanuppad and move the PHI values to there.
+    auto *CaseBB = BasicBlock::Create(CleanupPadBB->getContext(),
+                                      CleanupPadBB->getName() +
+                                          Twine(".from.") + Pred->getName(),
+                                      CleanupPadBB->getParent(), CleanupPadBB);
+    updatePhiNodes(CleanupPadBB, Pred, CaseBB);
+    CaseBB->setName(CleanupPadBB->getName() + Twine(".from.") +
+                    Pred->getName());
+    Builder.SetInsertPoint(CaseBB);
+    Builder.CreateBr(CleanupPadBB);
+    movePHIValuesToInsertedBlock(CleanupPadBB, CaseBB, NewCleanupPadBB);
+
+    // Update this Pred to the new unwind point.
+    setUnwindEdgeTo(Pred->getTerminator(), NewCleanupPadBB);
+
+    // Setup the switch in the dispatcher.
+    auto *SwitchConstant = ConstantInt::get(SwitchType, SwitchIndex);
+    SetDispatchValuePN->addIncoming(SwitchConstant, Pred);
+    SwitchOnDispatch->addCase(SwitchConstant, CaseBB);
+    SwitchIndex++;
+  }
+}
+
+static bool isCatchSwitchUnwindDest(BasicBlock *BB) {
+  SmallVector<BasicBlock *, 8> Preds(pred_begin(BB), pred_end(BB));
+  for (BasicBlock *Pred : Preds) {
+    if (auto *CatchSwitch = dyn_cast<CatchSwitchInst>(Pred->getTerminator())) {
+      return CatchSwitch->getUnwindDest() == BB;
+    }
+  }
+
+  return false;
+}
+
 static void rewritePHIs(BasicBlock &BB) {
   // For every incoming edge we will create a block holding all
   // incoming values in a single PHI nodes.
   //
   // loop:
   //    %n.val = phi i32[%n, %entry], [%inc, %loop]
   //
   // It will create:
   //
   // loop.from.entry:
   //    %n.loop.pre = phi i32 [%n, %entry]
   //    br %label loop
   // loop.from.loop:
   //    %inc.loop.pre = phi i32 [%inc, %loop]
   //    br %label loop
   //
   // After this rewrite, further analysis will ignore any phi nodes with more
   // than one incoming edge.
 
   // TODO: Simplify PHINodes in the basic block to remove duplicate
   // predecessors.
 
+  // Special case for CleanupPad: all EH blocks must have the same unwind edge
+  // so we need to create an additional "dispatcher" block.
+  if (auto *CleanupPad = dyn_cast_or_null<CleanupPadInst>(BB.getFirstNonPHI()))
+    if (isCatchSwitchUnwindDest(&BB)) {
+      rewritePHIsForCleanupPad(&BB, CleanupPad);
+      return;
+    }
+
   LandingPadInst *LandingPad = nullptr;
   PHINode *ReplPHI = nullptr;
   if ((LandingPad = dyn_cast_or_null<LandingPadInst>(BB.getFirstNonPHI()))) {
     // ehAwareSplitEdge will clone the LandingPad in all the edge blocks.
     // We replace the original landing pad with a PHINode that will collect the
     // results from all of them.
     ReplPHI = PHINode::Create(LandingPad->getType(), 1, "", LandingPad);
     ReplPHI->takeName(LandingPad);
     LandingPad->replaceAllUsesWith(ReplPHI);
     // We will erase the original landing pad at the end of this function after
     // ehAwareSplitEdge cloned it in the transition blocks.
   }
 
   SmallVector<BasicBlock *, 8> Preds(pred_begin(&BB), pred_end(&BB));
   for (BasicBlock *Pred : Preds) {
     auto *IncomingBB = ehAwareSplitEdge(Pred, &BB, LandingPad, ReplPHI);
     IncomingBB->setName(BB.getName() + Twine(".from.") + Pred->getName());
-    auto *PN = cast<PHINode>(&BB.front());
-    do {
-      int Index = PN->getBasicBlockIndex(IncomingBB);
-      Value *V = PN->getIncomingValue(Index);
-      PHINode *InputV = PHINode::Create(
-          V->getType(), 1, V->getName() + Twine(".") + BB.getName(),
-          &IncomingBB->front());
-      InputV->addIncoming(V, Pred);
-      PN->setIncomingValue(Index, InputV);
-      PN = dyn_cast<PHINode>(PN->getNextNode());
-    } while (PN != ReplPHI); // ReplPHI is either null or the PHI that replaced
-                             // the landing pad.
+
+    // Stop the moving of values at ReplPHI, as this is either null or the PHI
+    // that replaced the landing pad.
+    movePHIValuesToInsertedBlock(&BB, IncomingBB, Pred, ReplPHI);
   }
 
   if (LandingPad) {
     // Calls to ehAwareSplitEdge function cloned the original lading pad.
     // No longer need it.
     LandingPad->eraseFromParent();
   }
 }

llvm/test/Transforms/Coroutines/coro-catchswitch-cleanuppad.ll

+; Tests the PHI nodes in cleanuppads for catchswitch instructions are correctly
+; split up.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+declare i32 @__CxxFrameHandler3(...)
+define i8* @f2(i1 %val) "coroutine.presplit"="1" personality i32 (...)* @__CxxFrameHandler3 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %valueA = call i32 @f();
+  %valueB = call i32 @f();
+  %need.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.alloc, label %dyn.alloc, label %dowork.0
+
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  br label %dowork.0
+
+dowork.0:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
+  invoke void @print(i32 0)
+    to label %checksuspend unwind label %catch.dispatch.1
+
+checksuspend:
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %dowork.1
+                                i8 1, label %cleanup]
+
+dowork.1:
+  invoke void @print(i32 0)
+    to label %checksuspend unwind label %catch.dispatch.1
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+
+catch.dispatch.1:
+  %cs1 = catchswitch within none [label %handler1] unwind to caller
+handler1:
+  %h1 = catchpad within %cs1 [i8* null, i32 64, i8* null]
+  invoke void @print(i32 2) [ "funclet"(token %h1) ]
+          to label %catchret1 unwind label %catch.dispatch.2
+catchret1:
+  catchret from %h1 to label %cleanup
+
+catch.dispatch.2:
+  %cs2 = catchswitch within %h1 [label %handler2] unwind label %cleanup2
+handler2:
+  %h2 = catchpad within %cs2 [i8* null, i32 64, i8* null]
+  invoke void @print(i32 3) [ "funclet"(token %h2) ]
+          to label %cleanup unwind label %cleanup2
+cleanup2:
+  %cleanupval2 = phi i32 [%valueA, %catch.dispatch.2], [%valueB, %handler2]
+  cleanuppad within %h1 []
+  call void @print(i32 %cleanupval2)
+  br label %cleanup
+
+; Verifiers that a "dispatcher" cleanuppad is created.
+
+; catchswitch and all associated catchpads are required to have the same unwind
+; edge, but coro requires that PHI nodes are split up so that reload
+; instructions can be generated, therefore we create a new "dispatcher"
+; cleanuppad which forwards to individual blocks that contain the reload
+; instructions per catchswitch/catchpad and then all branch back to the
+; original cleanuppad block.
+
+; CHECK: catch.dispatch.2:
+; CHECK:   %cs2 = catchswitch within %h1 [label %handler2] unwind label %cleanup2.corodispatch
+
+; CHECK: handler2:
+; CHECK:   invoke void @print(i32 3)
+; CHECK:           to label %cleanup unwind label %cleanup2.corodispatch
+
+; CHECK: cleanup2.corodispatch:
+; CHECK:   %1 = phi i8 [ 0, %handler2 ], [ 1, %catch.dispatch.2 ]
+; CHECK:   %2 = cleanuppad within %h1 []
+; CHECK:   %switch = icmp ult i8 %1, 1
+; CHECK:   br i1 %switch, label %cleanup2.from.handler2, label %cleanup2.from.catch.dispatch.2
+
+; CHECK: cleanup2.from.handler2:
+; CHECK:   %valueB.reload = load i32, i32* %valueB.spill.addr, align 4
+; CHECK:   br label %cleanup2
+
+; CHECK: cleanup2.from.catch.dispatch.2:
+; CHECK:   %valueA.reload = load i32, i32* %valueA.spill.addr, align 4
+; CHECK:   br label %cleanup2
+
+; CHECK: cleanup2:
+; CHECK:   %cleanupval2 = phi i32 [ %valueA.reload, %cleanup2.from.catch.dispatch.2 ], [ %valueB.reload, %cleanup2.from.handler2 ]
+; CHECK:   call void @print(i32 %cleanupval2)
+; CHECK:   br label %cleanup
+}
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1) 
+
+declare noalias i8* @malloc(i32)
+declare void @print(i32)
+declare void @free(i8*)
+
+declare i32 @f()
+