[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,112 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   return FrameTy;
 }
 
+// Dominance issue fixer for each predecessor satisfying predicate function
+static void
+remapPredecessorIfPredicate(BasicBlock *InitialBlock,
+                            BasicBlock *ReplacementBlock,
+                            std::function<bool(BasicBlock *)> Predicate) {
+
+  SmallVector<BasicBlock *> InitialBlockDominanceViolatingPredecessors;
+
+  auto Predecessors = predecessors(InitialBlock);
+  std::copy_if(Predecessors.begin(), Predecessors.end(),
+               std::back_inserter(InitialBlockDominanceViolatingPredecessors),
+               Predicate);
+
+  // Tailor the predecessor terminator - redirect it to ReplacementBlock.
+  for (auto *InitialBlockDominanceViolatingPredecessor :
+       InitialBlockDominanceViolatingPredecessors) {
+    auto *Terminator =
+        InitialBlockDominanceViolatingPredecessor->getTerminator();
+    if (auto *TerminatorIsInvoke = dyn_cast<InvokeInst>(Terminator)) {
+      if (InitialBlock == TerminatorIsInvoke->getUnwindDest()) {
+        TerminatorIsInvoke->setUnwindDest(ReplacementBlock);
+      } else {
+        TerminatorIsInvoke->setNormalDest(ReplacementBlock);
+      }
+    } else if (auto *TerminatorIsBranch = dyn_cast<BranchInst>(Terminator)) {
+      for (unsigned int SuccessorIdx = 0;
+           SuccessorIdx < TerminatorIsBranch->getNumSuccessors();
+           ++SuccessorIdx) {
+        if (InitialBlock == TerminatorIsBranch->getSuccessor(SuccessorIdx)) {
+          TerminatorIsBranch->setSuccessor(SuccessorIdx, ReplacementBlock);
+        }
+      }
+    } else if (auto *TerminatorIsCleanupReturn =
+                   dyn_cast<CleanupReturnInst>(Terminator)) {
+      TerminatorIsCleanupReturn->setUnwindDest(ReplacementBlock);
+    } else {
+      Terminator->print(dbgs());
+      report_fatal_error("Terminator is not implemented");
+    }
+  }
+}
+
+// Fixer for the "Instruction does not dominate all uses!" bug
+// The fix consists of mapping dominance violating paths (where CoroBegin does
+// not dominate cleanup nodes), duplicating them and setup 2 flows - the one
+// that insertSpills intended to create (using the spill) and another one,
+// preserving the logics of pre-splitting, which would be executed if unwinding
+// happened before CoroBegin
+static void enforceDominationByCoroBegin(const FrameDataInfo &FrameData,
+                                         const coro::Shape &Shape, Function *F,
+                                         DominatorTree &DT) {
+  SmallPtrSet<BasicBlock *, 3> SpillUserDominanceViolatingBlocksSet;
+
+  // 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)) {
+        SpillUserDominanceViolatingBlocksSet.insert(CurrentBlock);
+      }
+    }
+  }
+
+  if (SpillUserDominanceViolatingBlocksSet.empty()) {
+    return;
+  }
+
+  // Run is in reversed post order, to enforce visiting predecessors before
+  // successors
+  for (BasicBlock *CurrentBlock : ReversePostOrderTraversal<Function *>(F)) {
+    if (!SpillUserDominanceViolatingBlocksSet.contains(CurrentBlock)) {
+      continue;
+    }
+    SpillUserDominanceViolatingBlocksSet.erase(CurrentBlock);
+
+    // The duplicate will become the unspilled alternative
+    ValueToValueMapTy VMap;
+    auto *UnspilledAlternativeBlock =
+        CloneBasicBlock(CurrentBlock, VMap, ".unspilled_alternative", F);
+
+    // Remap node instructions
+    for (Instruction &UnspilledAlternativeBlockInstruction :
+         *UnspilledAlternativeBlock) {
+      RemapInstruction(&UnspilledAlternativeBlockInstruction, VMap,
+                       RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+    }
+
+    // Clone only if predecessor breaks dominance against CoroBegin
+    remapPredecessorIfPredicate(CurrentBlock, UnspilledAlternativeBlock,
+                                [&DT, &Shape](BasicBlock *PredecessorNode) {
+                                  return !DT.dominates(Shape.CoroBegin,
+                                                       PredecessorNode);
+                                });
+
+    // We changed dominance tree, so recalculate it
+    DT.recalculate(*F);
+
+    if (SpillUserDominanceViolatingBlocksSet.empty()) {
+      return;
+    }
+  }
+
+  assert(SpillUserDominanceViolatingBlocksSet.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 +1160,8 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
     return GEP;
   };
 
+  enforceDominationByCoroBegin(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,18 @@ 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
+    // preserve the invariant of "Unwind edges out of a funclet pad must have
+    // the same unwind dest"
+    BasicBlock *UnwindDest = nullptr;
+    if (auto *Invoke =
+            dyn_cast<InvokeInst>(FromPad->getParent()->getTerminator())) {
+      UnwindDest = Invoke->getUnwindDest();
+    }
+
+    auto *CleanupRet = Builder.CreateCleanupRet(FromPad, UnwindDest);
     End->getParent()->splitBasicBlock(End);
     CleanupRet->getParent()->getTerminator()->eraseFromParent();
   }

llvm/test/Transforms/Coroutines/pr148035_inst_does_not_dominate.ll

@@ -0,0 +1,44 @@
+; In coro-split, this coroutine code reduced IR, produced using clang with async-exceptions
+; crashed before fix because of the validation mismatch of Instruction does not dominate all uses!
+; RUN: opt < %s -passes='coro-split' -S
+
+; Function Attrs: presplitcoroutine
+define i8 @"?resuming_on_new_thread@@YA?AUtask@@Vunique_ptr@@@Z"(ptr %0) #0 personality ptr null {
+  invoke void @llvm.seh.scope.begin()
+          to label %2 unwind label %14
+
+2:                                                ; preds = %1
+  %3 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %4 = load volatile ptr, ptr null, align 8
+  %5 = call ptr @llvm.coro.begin(token %3, ptr %4)
+  %6 = call token @llvm.coro.save(ptr null)
+  %7 = call i8 @llvm.coro.suspend(token none, i1 false)
+  invoke void @llvm.seh.try.begin()
+          to label %common.ret unwind label %8
+
+common.ret:                                       ; preds = %12, %10, %2
+  ret i8 0
+
+8:                                                ; preds = %2
+  %9 = catchswitch within none [label %10] unwind label %12
+
+10:                                               ; preds = %8
+  %11 = catchpad within %9 [ptr null, i32 0, ptr null]
+  br label %common.ret
+
+12:                                               ; preds = %8
+  %13 = cleanuppad within none []
+  invoke void @llvm.seh.scope.end()
+          to label %common.ret unwind label %14
+
+14:                                               ; preds = %12, %1
+  %15 = cleanuppad within none []
+  store i32 0, ptr %0, align 4
+  cleanupret from %15 unwind to caller
+}
+
+attributes #0 = { presplitcoroutine }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}

llvm/test/Transforms/Coroutines/pr148035_inst_does_not_dominate_2.ll

@@ -0,0 +1,44 @@
+; In coro-split, this coroutine code reduced IR, produced using clang with async-exceptions
+; crashed before fix because of the validation mismatch of Instruction does not dominate all uses!
+; RUN: opt < %s -passes='coro-split' -S
+
+; Function Attrs: presplitcoroutine
+define i8 @"?resuming_on_new_thread@@YA?AUtask@@Vunique_ptr@@@Z"(ptr %0) #0 personality ptr null {
+  invoke void @llvm.seh.scope.begin()
+          to label %2 unwind label %14
+
+2:                                                ; preds = %1
+  %3 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %4 = load volatile ptr, ptr null, align 8
+  %5 = call ptr @llvm.coro.begin(token %3, ptr %4)
+  %6 = call token @llvm.coro.save(ptr null)
+  %7 = call i8 @llvm.coro.suspend(token none, i1 false)
+  invoke void @llvm.seh.try.begin()
+          to label %common.ret unwind label %8
+
+common.ret:                                       ; preds = %12, %10, %2
+  ret i8 0
+
+8:                                                ; preds = %2
+  %9 = catchswitch within none [label %10] unwind label %12
+
+10:                                               ; preds = %8
+  %11 = catchpad within %9 [ptr null, i32 0, ptr null]
+  br label %common.ret
+
+12:                                               ; preds = %8
+  %13 = cleanuppad within none []
+  invoke void @llvm.seh.scope.end()
+          to label %common.ret unwind label %14
+
+14:                                               ; preds = %12, %1
+  %15 = cleanuppad within none []
+  store i32 0, ptr %0, align 4
+  br label %common.ret
+}
+
+attributes #0 = { presplitcoroutine }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}

llvm/test/Transforms/Coroutines/pr148035_not_implemented_terminator_for_pred.ll

@@ -0,0 +1,50 @@
+; In coro-split, this coroutine code reduced IR, produced using clang with async-exceptions
+; crashed after first phase of fix because the terminator cleanupret was not implemented on predecessor fixer at the time
+; RUN: opt < %s -passes='coro-split' -S
+
+; Function Attrs: presplitcoroutine
+define i8 @"?resuming_on_new_thread@@YA?AUtask@@V?$unique_ptr@HU?$default_delete@H@std@@@std@@0@Z"(ptr %0) #0 personality ptr null {
+  invoke void @llvm.seh.scope.begin()
+          to label %2 unwind label %15
+
+2:                                                ; preds = %1
+  %3 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %4 = call ptr @llvm.coro.begin(token %3, ptr null)
+  %5 = call token @llvm.coro.save(ptr null)
+  %6 = call i8 @llvm.coro.suspend(token none, i1 false)
+  invoke void @llvm.seh.try.begin()
+          to label %7 unwind label %8
+
+7:                                                ; preds = %2
+  ret i8 0
+
+8:                                                ; preds = %2
+  %9 = catchswitch within none [label %10] unwind label %12
+
+10:                                               ; preds = %8
+  %11 = catchpad within %9 [ptr null, i32 0, ptr null]
+  ret i8 0
+
+12:                                               ; preds = %8
+  %13 = cleanuppad within none []
+  invoke void @llvm.seh.scope.end()
+          to label %14 unwind label %15
+
+14:                                               ; preds = %12
+  ret i8 0
+
+15:                                               ; preds = %12, %1
+  %16 = cleanuppad within none []
+  cleanupret from %16 unwind label %17
+
+17:                                               ; preds = %15
+  %18 = cleanuppad within none []
+  store i32 0, ptr %0, align 4
+  cleanupret from %18 unwind to caller
+}
+
+attributes #0 = { presplitcoroutine }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}

llvm/test/Transforms/Coroutines/pr148035_unwind_edges_not_having_same_dest.ll

@@ -0,0 +1,42 @@
+; In coro-split, this coroutine standard code reduced IR, produced using clang with async-exceptions
+; crashed before fix because of the validation mismatch of Unwind edges out of a funclet pad must have the same unwind dest
+; RUN: opt < %s -passes='coro-split' -S
+
+; Function Attrs: presplitcoroutine
+define i8 @"?resuming_on_new_thread@@YA?AUtask@@AEAVjthread@std@@@Z"() #0 personality ptr null {
+  %1 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %2 = call ptr @llvm.coro.begin(token %1, ptr null)
+  %3 = call token @llvm.coro.save(ptr null)
+  %4 = call i8 @llvm.coro.suspend(token none, i1 false)
+  invoke void @llvm.seh.try.begin()
+          to label %common.ret unwind label %5
+
+common.ret:                                       ; preds = %13, %7, %0
+  ret i8 0
+
+5:                                                ; preds = %0
+  %6 = catchswitch within none [label %7] unwind label %9
+
+7:                                                ; preds = %5
+  %8 = catchpad within %6 [ptr null, i32 0, ptr null]
+  br label %common.ret
+
+9:                                                ; preds = %5
+  %10 = cleanuppad within none []
+  invoke void @llvm.seh.scope.end() [ "funclet"(token %10) ]
+          to label %11 unwind label %13
+
+11:                                               ; preds = %9
+  %12 = call i1 @llvm.coro.end(ptr null, i1 true, token none) [ "funclet"(token %10) ]
+  cleanupret from %10 unwind label %13
+
+13:                                               ; preds = %11, %9
+  %14 = cleanuppad within none []
+  br label %common.ret
+}
+
+attributes #0 = { presplitcoroutine }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}

llvm/test/Transforms/Coroutines/pr148035_unwind_edges_not_having_same_dest_2.ll

@@ -0,0 +1,43 @@
+; In coro-split, this coroutine standard code reduced IR, produced using clang with async-exceptions
+; crashed before fix because of the validation mismatch of Unwind edges out of a funclet pad must have the same unwind dest
+; RUN: opt < %s -passes='coro-split' -S
+
+; Function Attrs: presplitcoroutine
+define i1 @"?resuming_on_new_thread@@YA?AUtask@@AEAVjthread@std@@@Z"() #0 personality ptr null {
+  %1 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %2 = call ptr @llvm.coro.begin(token %1, ptr null)
+  %3 = call token @llvm.coro.save(ptr null)
+  %4 = call i8 @llvm.coro.suspend(token none, i1 false)
+  invoke void @llvm.seh.try.begin()
+          to label %common.ret unwind label %5
+
+common.ret:                                       ; preds = %11, %13, %7, %0
+  %common.ret.op = phi i1 [ false, %11], [false, %13 ], [false, %7], [false, %0]
+  ret i1 %common.ret.op
+
+5:                                                ; preds = %0
+  %6 = catchswitch within none [label %7] unwind label %9
+
+7:                                                ; preds = %5
+  %8 = catchpad within %6 [ptr null, i32 0, ptr null]
+  br label %common.ret
+
+9:                                                ; preds = %5
+  %10 = cleanuppad within none []
+  invoke void @llvm.seh.scope.end() [ "funclet"(token %10) ]
+          to label %11 unwind label %13
+
+11:                                               ; preds = %9
+  %12 = call i1 @llvm.coro.end(ptr null, i1 true, token none) [ "funclet"(token %10) ]
+  br label %common.ret
+
+13:                                               ; preds = %9
+  %14 = cleanuppad within none []
+  br label %common.ret
+}
+
+attributes #0 = { presplitcoroutine }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 2, !"eh-asynch", i32 1}