[WIP] [Coroutines] fix coroutines + std::unique_ptr with async exceptions validation errors

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

@@ -17,6 +17,7 @@
 
 #include "CoroInternal.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Analysis/StackLifetime.h"
 #include "llvm/IR/DIBuilder.h"
@@ -35,6 +36,7 @@
 #include "llvm/Transforms/Coroutines/SpillUtils.h"
 #include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <algorithm>
@@ -973,6 +975,178 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   return FrameTy;
 }
 
+// This function assumes that it is being called on basic block in reversed
+// post-order, meaning predecessors are visited before successors failing to do
+// so will cause VMap to be non-valid and will cause instructions to fail
+// mapping to their corresponding clones
+static void finalizeBasicBlockCloneAndTrackSuccessors(
+    BasicBlock *InitialBlock, BasicBlock *ClonedBlock, ValueToValueMapTy &VMap,
+    SmallSet<BasicBlock *, 20> &SuccessorBlocksSet) {
+  // This code will examine the basic block, fix issues caused by clones
+  //  for example - tailor cleanupret to the corresponding cleanuppad
+  //  it will use VMap to do so
+  // in addition, it will add the node successors to SuccessorBlocksSet
+
+        // Remap the instructions, VMap here aggregates instructions across
+        // multiple BasicBlocks, and we assume that traversal is reversed post-order,
+        // therefore successor blocks (for example instructions having funclet
+        // tags) will be mapped correctly to the new cloned cleanuppad
+        for (Instruction &ClonedBlockInstruction : *ClonedBlock) {
+          RemapInstruction(&ClonedBlockInstruction, VMap,
+RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+        }
+
+        const auto &InitialBlockTerminator = InitialBlock->getTerminator();
+
+        // If it's cleanupret, find the correspondant cleanuppad (use the VMap to
+        // find it).
+        if (auto *InitialBlockTerminatorCleanupReturn =
+dyn_cast<CleanupReturnInst>(InitialBlockTerminator)) {
+    // if none found do nothing
+    if (VMap.find(InitialBlockTerminatorCleanupReturn->getCleanupPad()) ==
+        VMap.end()) {
+      return;
+    }
+
+    auto *ClonedBlockTerminatorCleanupReturn =
+cast<CleanupReturnInst>(ClonedBlock->getTerminator());
+
+          // Assuming reversed post-order traversal
+    llvm::Value *ClonedBlockCleanupPadValue =
+        VMap[InitialBlockTerminatorCleanupReturn->getCleanupPad()];
+    auto *ClonedBlockCleanupPad =
+cast<CleanupPadInst>(ClonedBlockCleanupPadValue);
+            ClonedBlockTerminatorCleanupReturn->setCleanupPad(ClonedBlockCleanupPad);
+
+        // If it's a branch/invoke, keep track of its successors, we want to
+          // calculate dominance between CoroBegin and them also
+        } else if (auto *InitialBlockTerminatorBranch =
+dyn_cast<BranchInst>(InitialBlockTerminator)) {
+          for (unsigned int successorIdx = 0;
+successorIdx < InitialBlockTerminatorBranch->getNumSuccessors();
+++successorIdx) {
+            SuccessorBlocksSet.insert(
+          InitialBlockTerminatorBranch->getSuccessor(successorIdx));
+          }
+        } else if (auto *InitialBlockTerminatorInvoke =
+dyn_cast<InvokeInst>(InitialBlockTerminator)) {
+    SuccessorBlocksSet.insert(InitialBlockTerminatorInvoke->getUnwindDest());
+  } else if (isa<ReturnInst>(InitialBlockTerminator)) {
+    // No action needed
+        } else {
+          InitialBlockTerminator->print(dbgs());
+    report_fatal_error("Terminator is not implemented in "
+                       "finalizeBasicBlockCloneAndTrackSuccessors");
+  }
+}
+
+// Dominance issue fixer for each predecessor satisfying predicate function
+void splitIfBasicBlockPredecessors(
+    BasicBlock *InitialBlock, BasicBlock *ReplacementBlock,
+    std::function<bool(BasicBlock *)> Predicate) {
+
+  SmallVector<BasicBlock *> InitialBlockPredecessors;
+
+  auto Predecessors = predecessors(InitialBlock);
+  std::copy_if(Predecessors.begin(), Predecessors.end(),
+               std::back_inserter(InitialBlockPredecessors), Predicate);
+
+        // Fixups on the predecessors terminator - point them to ReplacementBlock.
+  for (auto *InitialBlockPredecessor : InitialBlockPredecessors) {
+    auto *InitialBlockPredecessorTerminator =
+        InitialBlockPredecessor->getTerminator();
+    if (auto *InitialBlockPredecessorTerminatorInvoke =
+dyn_cast<InvokeInst>(InitialBlockPredecessorTerminator)) {
+      if (InitialBlock ==
+          InitialBlockPredecessorTerminatorInvoke->getUnwindDest()) {
+        InitialBlockPredecessorTerminatorInvoke->setUnwindDest(
+            ReplacementBlock);
+      } else {
+        InitialBlockPredecessorTerminatorInvoke->setNormalDest(
+            ReplacementBlock);
+      }
+    } else if (auto *InitialBlockPredecessorTerminatorBranch =
+dyn_cast<BranchInst>(InitialBlockPredecessorTerminator)) {
+            for (unsigned int successorIdx = 0;
+                successorIdx <
+           InitialBlockPredecessorTerminatorBranch->getNumSuccessors();
+                ++successorIdx) {
+              if (InitialBlock ==
+InitialBlockPredecessorTerminatorBranch->getSuccessor(
+successorIdx)) {
+                InitialBlockPredecessorTerminatorBranch->setSuccessor(
+successorIdx, ReplacementBlock);
+        }
+      }
+    } else if (auto *InitialBlockPredecessorTerminatorCleanupReturn =
+                   dyn_cast<CleanupReturnInst>(
+                       InitialBlockPredecessorTerminator)) {
+      InitialBlockPredecessorTerminatorCleanupReturn->setUnwindDest(
+          ReplacementBlock);
+          } else {
+            InitialBlockPredecessorTerminator->print(dbgs());
+      report_fatal_error(
+          "Terminator is not implemented in splitIfBasicBlockPredecessors");
+    }
+  }
+}
+
+// Fixer for the "Instruction does not dominate all uses!" bug
+// The fix consists of mapping problematic paths (where CoroBegin does not
+// dominate cleanup nodes) and duplicates them to 2 flows - the one that
+// insertSpills intended to create (using the spill) and another one, preserving
+// the logics of pre-splitting, which would be triggered if unwinding happened
+// before CoroBegin
+static void
+splitBasicBlocksNotDominatedByCoroBegin(const FrameDataInfo &FrameData,
+                                        const coro::Shape &Shape, Function *F,
+                                        DominatorTree &DT) {
+  ValueToValueMapTy VMap;
+  SmallSet<BasicBlock *, 20> SpillUserBlocksSet;
+
+  // Prepare the node set, logics will be run only on those nodes
+  for (const auto &E : FrameData.Spills) {
+    for (auto *U : E.second) {
+      auto CurrentBlock = U->getParent();
+      if (!DT.dominates(Shape.CoroBegin, CurrentBlock)) {
+        SpillUserBlocksSet.insert(CurrentBlock);
+      }
+          }
+        }
+
+        // Run is in reversed post order, to enforce visiting predecessors before
+  // successors
+  for (BasicBlock *CurrentBlock : ReversePostOrderTraversal<Function *>(F)) {
+    if (!SpillUserBlocksSet.contains(CurrentBlock)) {
+        continue;
+      }
+    SpillUserBlocksSet.erase(CurrentBlock);
+
+    // Preserve the current node. the duplicate will become the unspilled
+    // alternative
+    auto UnspilledAlternativeBlock =
+        CloneBasicBlock(CurrentBlock, VMap, ".unspilled_alternative", F);
+
+    // Remap node instructions, keep track of successors to visit them in next
+    // iterations
+    finalizeBasicBlockCloneAndTrackSuccessors(
+        CurrentBlock, UnspilledAlternativeBlock, VMap, SpillUserBlocksSet);
+
+    // Split only if predecessor breaks dominance against CoroBegin
+    splitIfBasicBlockPredecessors(CurrentBlock, UnspilledAlternativeBlock,
+                                  [&DT, &Shape](BasicBlock *PredecessorNode) {
+                                    return !DT.dominates(Shape.CoroBegin,
+                                                         PredecessorNode);
+                                  });
+
+    // We changed dominance tree, so recalculate
+    DT.recalculate(*F);
+  }
+
+  assert(SpillUserBlocksSet.empty() &&
+         "Graph is corrupted by SpillUserBlocksSet");
+}
+
 // 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
@@ -1052,6 +1226,8 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
     return GEP;
   };
 
+  splitBasicBlocksNotDominatedByCoroBegin(FrameData, Shape, F, DT);
+
   for (auto const &E : FrameData.Spills) {
     Value *Def = E.first;
     auto SpillAlignment = Align(FrameData.getAlign(Def));

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -376,7 +376,16 @@ static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
   // If coro.end has an associated bundle, add cleanupret instruction.
   if (auto Bundle = End->getOperandBundle(LLVMContext::OB_funclet)) {
     auto *FromPad = cast<CleanupPadInst>(Bundle->Inputs[0]);
-    auto *CleanupRet = Builder.CreateCleanupRet(FromPad, nullptr);
+
+    // If the terminator is an invoke, 
+    // set the cleanupret unwind destination the same as the other edges, to
+    // avoid validation errors
+    BasicBlock *UBB = nullptr;
+    if (auto II = dyn_cast<InvokeInst>(FromPad->getParent()->getTerminator())) {
+      UBB = II->getUnwindDest();
+    }
+
+    auto *CleanupRet = Builder.CreateCleanupRet(FromPad, UBB);
     End->getParent()->splitBasicBlock(End);
     CleanupRet->getParent()->getTerminator()->eraseFromParent();
   }