Index: lib/Transforms/Coroutines/CoroCleanup.cpp =================================================================== --- lib/Transforms/Coroutines/CoroCleanup.cpp +++ lib/Transforms/Coroutines/CoroCleanup.cpp @@ -49,7 +49,7 @@ II->replaceAllUsesWith(II->getArgOperand(1)); break; case Intrinsic::coro_free: - II->replaceAllUsesWith(II->getArgOperand(0)); + II->replaceAllUsesWith(II->getArgOperand(1)); break; case Intrinsic::coro_alloc: II->replaceAllUsesWith(ConstantInt::getTrue(Context)); Index: lib/Transforms/Coroutines/CoroEarly.cpp =================================================================== --- lib/Transforms/Coroutines/CoroEarly.cpp +++ lib/Transforms/Coroutines/CoroEarly.cpp @@ -115,12 +115,17 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) { bool Changed = false; + CoroIdInst *CoroId = nullptr; + SmallVector CoroFrees; for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) { Instruction &I = *IB++; if (auto CS = CallSite(&I)) { switch (CS.getIntrinsicID()) { default: continue; + case Intrinsic::coro_free: + CoroFrees.push_back(cast(&I)); + break; case Intrinsic::coro_suspend: // Make sure that final suspend point is not duplicated as CoroSplit // pass expects that there is at most one final suspend point. @@ -141,6 +146,7 @@ F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT); setCannotDuplicate(CII); CII->setCoroutineSelf(); + CoroId = cast(&I); } } break; @@ -160,6 +166,12 @@ Changed = true; } } + // Make sure that all CoroFree reference the coro.id intrinsic. + // Token type is not exposed through coroutine C/C++ builtins to plain C, so + // we allow specifying none and fixing it up here. + if (CoroId) + for (CoroFreeInst *CF : CoroFrees) + CF->setArgOperand(0, CoroId); return Changed; } @@ -178,9 +190,10 @@ // This pass has work to do only if we find intrinsics we are going to lower // in the module. bool doInitialization(Module &M) override { - if (coro::declaresIntrinsics(M, {"llvm.coro.begin", "llvm.coro.end", - "llvm.coro.resume", "llvm.coro.destroy", - "llvm.coro.done", "llvm.coro.suspend"})) + if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy", + "llvm.coro.done", "llvm.coro.end", + "llvm.coro.free", "llvm.coro.promise", + "llvm.coro.resume", "llvm.coro.suspend"})) L = llvm::make_unique(M); return false; } Index: lib/Transforms/Coroutines/CoroFrame.cpp =================================================================== --- lib/Transforms/Coroutines/CoroFrame.cpp +++ lib/Transforms/Coroutines/CoroFrame.cpp @@ -438,6 +438,17 @@ CurrentReload = CreateReload(&*CurrentBlock->getFirstInsertionPt()); } + // If we have a single edge PHINode, remove it and replace it with a reload + // from the coroutine frame. (We already took care of multi edge PHINodes + // by rewriting them in the rewritePHIs function). + if (auto *PN = dyn_cast(E.user())) { + assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming " + "values in the PHINode"); + PN->replaceAllUsesWith(CurrentReload); + PN->eraseFromParent(); + continue; + } + // Replace all uses of CurrentValue in the current instruction with reload. E.user()->replaceUsesOfWith(CurrentValue, CurrentReload); } Index: test/Transforms/Coroutines/phi-coro-end.ll =================================================================== --- /dev/null +++ test/Transforms/Coroutines/phi-coro-end.ll @@ -0,0 +1,48 @@ +; Verify that we correctly handle suspend when the coro.end block contains phi +; RUN: opt < %s -O2 -enable-coroutines -S | FileCheck %s + +define i8* @f(i32 %n) { +entry: + %id = call token @llvm.coro.id(i32 0, 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) + %0 = call i8 @llvm.coro.suspend(token none, i1 false) + switch i8 %0, label %suspend [i8 0, label %cleanup i8 1, label %cleanup] + +cleanup: + %mem = call i8* @llvm.coro.free(token %id, i8* %hdl) + call void @free(i8* %mem) + br label %suspend + +suspend: + %r = phi i32 [%n, %entry], [1, %cleanup] + call void @llvm.coro.end(i8* %hdl, i1 false) + call void @print(i32 %r) + ret i8* %hdl +} + +; CHECK-LABEL: @main +define i32 @main() { +entry: + %hdl = call i8* @f(i32 4) + call void @llvm.coro.resume(i8* %hdl) + ret i32 0 +;CHECK: call void @print(i32 4) +;CHECK: ret i32 0 +} + +declare i8* @llvm.coro.alloc() +declare i32 @llvm.coro.size.i32() +declare i8* @llvm.coro.free(token, i8*) +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 i8* @llvm.coro.begin(token, i8*) +declare void @llvm.coro.end(i8*, i1) + +declare noalias i8* @malloc(i32) +declare void @print(i32) +declare void @free(i8*)