[Coroutines] Part 7: Split coroutine into subfunctions

[GorNishanov]

Summary

This patch adds simple coroutine splitting logic to CoroSplit pass.

Documentation and overview is here: http://llvm.org/docs/Coroutines.html.

Upstreaming sequence (rough plan)
1.Add documentation. (https://reviews.llvm.org/D22603)
2.Add coroutine intrinsics. (https://reviews.llvm.org/D22659)
...

7. Split coroutine into subfunctions <= we are here
8. Coroutine Frame Building algorithm
9. Handle coroutine with unwinds

10+. The rest of the logic

[GorNishanov]

Remove PermissiveArrayRef and inline replaceAndRemove into its callers. Saves us about 15 lines.

[majnemer]

lib/Transforms/Coroutines/CoroCleanup.cpp

30–37

I'd just make simplifyFunctionCFG a non-internal function and call it instead of making a new FPM.

44

auto *

44

Should this be a CallSite instead?

lib/Transforms/Coroutines/CoroEarly.cpp

52

Pointers lean right.

lib/Transforms/Coroutines/CoroFrame.cpp

29

auto *

30

You could use &BB->front().

46

Please have a period at the end of this comment.

63–66

auto *

68

We aren't huge users of std::numeric_limits, UINT32_MAX is a little more terse.

101

auto *

109

auto *

lib/Transforms/Coroutines/Coroutines.cpp

196

auto * although maybe we should use a CallSite?

[GorNishanov]

Feedback implemented.

[GorNishanov]

lib/Transforms/Coroutines/CoroCleanup.cpp

30–37

We may end up needed more cleanup passes there. I'll keep at is as is for now.

44

The only throwing coroutine intrinsics are coro.resume and coro.destroy are lowered in coro.early.

lib/Transforms/Coroutines/CoroFrame.cpp

63–66

Missed those. Will upload in a sec.

lib/Transforms/Coroutines/Coroutines.cpp

196

Only coro.resume and coro.destroy are throwing and they are replaced by CoroEarly Pass.

[GorNishanov]

All feedback implemented or rejected :)

[majnemer]

lib/Transforms/Coroutines/CoroFrame.cpp

23

Do we still need this include?

105

auto *

lib/Transforms/Coroutines/CoroSplit.cpp

172

A comment should be a sentence: start with a cap letter and end with a period.

lib/Transforms/Coroutines/Coroutines.cpp

204–206

This is not true, the basic blocks can be in a jumbled order. If you want to do this sort of check, you need to use a depth first iterator like depth_first and turn this assert into a report_fatal_error.

Alternatively, you can save the coro.frame you come across for later processing instead of a dfs walk.

209

Ditto.

226–227

I think this should be a report_fatal_error because nothing prohibits it upfront.

[GorNishanov]

reworked coro::Shape::buildFrom(Function& F) to address the feedback
+ little nits fixed

[majnemer]

lib/Transforms/Coroutines/CoroFrame.cpp

27–45

These functions seem dead.

lib/Transforms/Coroutines/CoroSplit.cpp

148

Caps + period.

172

Hmm, do you have a comment somewhere explaining why? Is it expected that there are no meaningful users of the args?

188

Period.

190–191

You could use changeToUnreachable.

201

Ditto.

287

Args.empty().

297

Period.

329

Caps + period.

lib/Transforms/Coroutines/Coroutines.cpp

215–217

Do we have a noduplicate for this case?

243–244

Is this the right message?

263

Maybe a comment here?

264

I'd sink this into the if.

264

auto *.

273–277

I'd just call changeToUnreachable here.

[GorNishanov]

• Mark final suspend and fallthrough coro end as no duplicate in coro early
• Remove unused functions from CoroFrame (but they are coming back in the next patch with vengeance)
• Use changeToUnreachable to simplify code in CoroSplit and Coroutines.cpp
• Bring more comments in compliance with coding style

[GorNishanov]

lib/Transforms/Coroutines/CoroSplit.cpp

172

Added a comment: The buildCoroutineFrame algorithm already rewritten access to the args that occurs after suspend points with loads and stores to/from the coroutine frame.

lib/Transforms/Coroutines/Coroutines.cpp

215–217

We do now. I added code to CoroEarly to mark them no duplicate

[GorNishanov]

add /*UseLLVMTrap=*/ to a bool parameter

[majnemer]

LGTM with nits :)

lib/Transforms/Coroutines/CoroFrame.cpp

94

Period.

101

Ditto.

lib/Transforms/Coroutines/CoroInternal.h

90

This comment is not very useful, I'd remove it.

lib/Transforms/Coroutines/Coroutines.cpp

185–186

auto *

[GorNishanov]

Nits addressed. Preparing to land

lib/Transforms/Coroutines/CMakeLists.txt

 add_llvm_library(LLVMCoroutines
   Coroutines.cpp
   CoroCleanup.cpp
   CoroEarly.cpp
   CoroElide.cpp
-  CoroSplit.cpp
+  CoroFrame.cpp
+  CoroSplit.cpp  
   )
 
 add_dependencies(LLVMCoroutines intrinsics_gen)

lib/Transforms/Coroutines/CoroCleanup.cpp

 //===- CoroCleanup.cpp - Coroutine Cleanup Pass ---------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass lowers all remaining coroutine intrinsics.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Pass.h"
+#include "llvm/Transforms/Scalar.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-cleanup"
 
+namespace {
+// Created on demand if CoroCleanup pass has work to do.
+struct Lowerer : coro::LowererBase {
+  Lowerer(Module &M) : LowererBase(M) {}
+  bool lowerRemainingCoroIntrinsics(Function &F);
+};
+}
+
+static void simplifyCFG(Function &F) {
+  llvm::legacy::FunctionPassManager FPM(F.getParent());
+  FPM.add(createCFGSimplificationPass());
+
+  FPM.doInitialization();
+  FPM.run(F);
+  FPM.doFinalization();
+}
+
+bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
+  bool Changed = false;
+
+  for (auto IB = inst_begin(F), E = inst_end(F); IB != E;) {
+    Instruction &I = *IB++;
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I)) {
+      switch (II->getIntrinsicID()) {
+      default:
+        continue;
+      case Intrinsic::coro_begin:
+        II->replaceAllUsesWith(II->getArgOperand(1));
+        break;
+      case Intrinsic::coro_free:
+        II->replaceAllUsesWith(II->getArgOperand(0));
+        break;
+      case Intrinsic::coro_alloc:
+        II->replaceAllUsesWith(ConstantInt::getTrue(Context));
+        break;
+      case Intrinsic::coro_id:
+        II->replaceAllUsesWith(ConstantTokenNone::get(Context));
+        break;
+      }
+      II->eraseFromParent();
+      Changed = true;
+    }
+  }
+
+  if (Changed) {
+    // After replacement were made we can cleanup the function body a little.
+    simplifyCFG(F);
+  }
+  return Changed;
+}
+
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 struct CoroCleanup : FunctionPass {
   static char ID; // Pass identification, replacement for typeid
 
   CoroCleanup() : FunctionPass(ID) {}
 
-  bool runOnFunction(Function &F) override { return false; }
+  std::unique_ptr<Lowerer> L;
+
+  // 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.alloc", "llvm.coro.begin",
+                                     "llvm.coro.free", "llvm.coro.id"}))
+      L = llvm::make_unique<Lowerer>(M);
+    return false;
+  }
+
+  bool runOnFunction(Function &F) override {
+    if (L)
+      return L->lowerRemainingCoroIntrinsics(F);
+    return false;
+  }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    if (!L)
       AU.setPreservesAll();
   }
 };
-
 }
 
 char CoroCleanup::ID = 0;
 INITIALIZE_PASS(CoroCleanup, "coro-cleanup",
                 "Lower all coroutine related intrinsics", false, false)
 
 Pass *llvm::createCoroCleanupPass() { return new CoroCleanup(); }

lib/Transforms/Coroutines/CoroEarly.cpp

 void Lowerer::lowerResumeOrDestroy(CallSite CS,
                                    CoroSubFnInst::ResumeKind Index) {
   Value *ResumeAddr =
       makeSubFnCall(CS.getArgOperand(0), Index, CS.getInstruction());
   CS.setCalledFunction(ResumeAddr);
   CS.setCallingConv(CallingConv::Fast);
 }
 
+// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
+// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
+// NoDuplicate attribute will be removed from coro.begin otherwise, it will
+// interfere with inlining.
+static void setCannotDuplicate(CoroIdInst *CoroId) {
+  for (User* U : CoroId->users())
+    if (auto *CB = dyn_cast<CoroBeginInst>(U))
+      CB->setCannotDuplicate();
+}
+
 bool Lowerer::lowerEarlyIntrinsics(Function &F) {
   bool Changed = false;
   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_id:
         // Mark a function that comes out of the frontend that has a coro.begin
         // with a coroutine attribute.
         if (auto *CII = cast<CoroIdInst>(&I)) {
-          if (CII->getInfo().isPreSplit())
+          if (CII->getInfo().isPreSplit()) {
             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
+            setCannotDuplicate(CII);
+          }
         }
         break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
         break;

lib/Transforms/Coroutines/CoroFrame.cpp

+//===- CoroFrame.cpp - Builds and manipulates coroutine frame -------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+// This file contains classes used to discover if for a particular value
+// there from sue to definition that crosses a suspend block.
+//
+// Using the information discovered we form a Coroutine Frame structure to
+// contain those values. All uses of those values are replaced with appropriate
+// GEP + load from the coroutine frame. At the point of the definition we spill
+// the value into the coroutine frame.
+//
+// TODO: pack values tightly using liveness info.
+//===----------------------------------------------------------------------===//
+
+#include "CoroInternal.h"
+#include "llvm/IR/IRBuilder.h"
+#include <limits>
+
+using namespace llvm;
+
+// Splits the block at a particular instruction unless it is the first
+// instruction in the block with a single predecessor.
+static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
+  auto BB = I->getParent();
+  if (&*BB->begin() == I) {
+    if (BB->getSinglePredecessor()) {
+      BB->setName(Name);
+      return BB;
+    }
+  }
+  return BB->splitBasicBlock(I, Name);
+}
+
+// Split above and below a particular instruction so that it
+// will be all alone by itself in a block.
+static void splitAround(Instruction *I, const Twine &Name) {
+  splitBlockIfNotFirst(I, Name);
+  splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
+}
+
+// TODO: Implement in future patches
+struct SpillInfo {};
+
+// Build a struct that will keep state for an active coroutine.
+//   struct f.frame {
+//     ResumeFnTy ResumeFnAddr;
+//     ResumeFnTy DestroyFnAddr;
+//     int ResumeIndex;
+//     ... promise (if present) ...
+//     ... spills ...
+//   };
+static StructType *buildFrameType(Function &F, coro::Shape &Shape,
+                                  SpillInfo &Spills) {
+  LLVMContext &C = F.getContext();
+  SmallString<32> Name(F.getName());
+  Name.append(".Frame");
+  StructType *FrameTy = StructType::create(C, Name);
+  auto FramePtrTy = FrameTy->getPointerTo();
+  auto FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
+                                /*IsVarArgs=*/false);
+  auto FnPtrTy = FnTy->getPointerTo();
+
+  if (Shape.CoroSuspends.size() > std::numeric_limits<uint32_t>::max())
+    report_fatal_error("Cannot handle coroutine with this many suspend points");
+
+  SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, Type::getInt32Ty(C)};
+  // TODO: Populate from Spills.
+  FrameTy->setBody(Types);
+
+  return FrameTy;
+}
+
+// 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
+//
+// becomes:
+//
+//    %hdl = coro.begin(...)
+//    %FramePtr = bitcast i8* hdl to %f.frame*
+//    br label %AllocaSpillBB
+//
+//  AllocaSpillBB:
+//    ; geps corresponding to allocas that were moved to coroutine frame
+//    br label PostSpill
+//
+//  PostSpill:
+//    whatever
+//
+//
+static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
+  auto CB = Shape.CoroBegin;
+  IRBuilder<> Builder(CB->getNextNode());
+  PointerType *FramePtrTy = Shape.FrameTy->getPointerTo();
+  Instruction *FramePtr =
+      cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
+
+  // TODO: Insert Spills.
+
+  auto FramePtrBB = FramePtr->getParent();
+  Shape.AllocaSpillBlock =
+      FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
+  Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),
+                                          "PostSpill");
+  // TODO: Insert geps for alloca moved to coroutine frame
+
+  return FramePtr;
+}
+
+void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
+  SpillInfo Spills;
+  // TODO: Compute Spills (incoming in later patches)
+
+  Shape.FrameTy = buildFrameType(F, Shape, Spills);
+  Shape.FramePtr = insertSpills(Spills, Shape);
+}

lib/Transforms/Coroutines/CoroInstr.h

public:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_begin;
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
+/// This represents the llvm.coro.save instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSaveInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_save;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.suspend instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSuspendInst : public IntrinsicInst {
+  enum { SaveArg, FinalArg };
+
+public:
+  CoroSaveInst *getCoroSave() const {
+    Value *Arg = getArgOperand(SaveArg);
+    if (auto *SI = dyn_cast<CoroSaveInst>(Arg))
+      return SI;
+    assert(isa<ConstantTokenNone>(Arg));
+    return nullptr;
+  }
+  bool isFinal() const {
+    return cast<Constant>(getArgOperand(FinalArg))->isOneValue();
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_suspend;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.size instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSizeInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_size;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.end instruction.
+class LLVM_LIBRARY_VISIBILITY CoroEndInst : public IntrinsicInst {
+  enum { FrameArg, UnwindArg };
+
+public:
+  bool isFallthrough() const { return !isUnwind(); }
+  bool isUnwind() const {
+    return cast<Constant>(getArgOperand(UnwindArg))->isOneValue();
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_end;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 } // End namespace llvm.

lib/Transforms/Coroutines/CoroInternal.h

 #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
 #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
 
 #include "CoroInstr.h"
 #include "llvm/Transforms/Coroutines.h"
 
 namespace llvm {
 
+class CallGraph;
+class CallGraphSCC;
 class PassRegistry;
 
 void initializeCoroEarlyPass(PassRegistry &);
 void initializeCoroSplitPass(PassRegistry &);
 void initializeCoroElidePass(PassRegistry &);
 void initializeCoroCleanupPass(PassRegistry &);
 
 // CoroEarly pass marks every function that has coro.begin with a string

 
 #define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger"
 
 namespace coro {
 
 bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
 void replaceAllCoroAllocs(CoroBeginInst *CB, bool Replacement);
 void replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement);
+void updateCallGraph(Function &Caller, ArrayRef<Function *> Funcs,
+                     CallGraph &CG, CallGraphSCC &SCC);
 
 // Keeps data and helper functions for lowering coroutine intrinsics.
 struct LowererBase {
   Module &TheModule;
   LLVMContext &Context;
+  PointerType *const Int8Ptr;
   FunctionType *const ResumeFnType;
+  ConstantPointerNull *const NullPtr;
 
   LowererBase(Module &M);
   Value *makeSubFnCall(Value *Arg, int Index, Instruction *InsertPt);
 };
 
+// Holds structural Coroutine Intrinsics for a particular function and other
+// values used during CoroSplit pass.
+struct LLVM_LIBRARY_VISIBILITY Shape {
+  CoroBeginInst *CoroBegin;
+  SmallVector<CoroEndInst *, 4> CoroEnds;
+  SmallVector<CoroSizeInst *, 2> CoroSizes;
+  SmallVector<CoroSuspendInst *, 4> CoroSuspends;
+
+  // Field Indexes for known coroutine frame fields.
+  enum {
+    ResumeField = 0,
+    DestroyField = 1,
+    IndexField = 2,
+  };
+
+  StructType *FrameTy;
+  Instruction *FramePtr;
+  BasicBlock* AllocaSpillBlock;
+  SwitchInst* ResumeSwitch;
+  bool HasFinalSuspend;
+
+  Shape() = default;
+  explicit Shape(Function &F) { buildFrom(F); }
+  void buildFrom(Function &F);
+};
+
+// Defined in CoroFrame.cpp
+void buildCoroutineFrame(Function& F, Shape& Shape);
+
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

lib/Transforms/Coroutines/CoroSplit.cpp

 //===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass builds the coroutine frame and outlines resume and destroy parts
 // of the coroutine into separate functions.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
 // We present a coroutine to an LLVM as an ordinary function with suspension
 // points marked up with intrinsics. We let the optimizer party on the coroutine
 // as a single function for as long as possible. Shortly before the coroutine is
 // eligible to be inlined into its callers, we split up the coroutine into parts
 // corresponding to an initial, resume and destroy invocations of the coroutine,
 // add them to the current SCC and restart the IPO pipeline to optimize the
 // coroutine subfunctions we extracted before proceeding to the caller of the
 // coroutine.
 
+template <typename T> struct PermissiveArrayRef : ArrayRef<T> {
+  template <typename U>
+  PermissiveArrayRef(
+      const ArrayRef<U *> &A,
+      typename std::enable_if<std::is_convertible<U *, T>::value>::type * =
+          nullptr)
+      : ArrayRef((T const *)A.data(), A.size()) {}
+
+  template <typename U, typename V>
+  PermissiveArrayRef(
+      const SmallVectorTemplateCommon<U *, V> &A,
+      typename std::enable_if<std::is_convertible<U *, T>::value>::type * =
+          nullptr)
+      : ArrayRef((T const *)A.data(), A.size()) {}
+};
+
+// Replace all instructions provided with the new value. If optional VMap was
+// provided, replace mapped instructions with the new value.
+static void replaceAndRemove(PermissiveArrayRef<Instruction *> Instrs,
+                             Value *NewValue,
+                             ValueToValueMapTy *VMap = nullptr) {
+  for (Instruction *I : Instrs) {
+    if (VMap)
+      I = cast<Instruction>((*VMap)[I]);
+    I->replaceAllUsesWith(NewValue);
+    I->eraseFromParent();
+  }
+}
+
+// Create an entry block for a resume function with a switch that will jump to
+// suspend points.
+static BasicBlock *createResumeEntryBlock(Function &F, coro::Shape &Shape) {
+  LLVMContext &C = F.getContext();
+
+  // resume.entry:
+  //  %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0,
+  //  i32 2
+  //  % index = load i32, i32* %index.addr
+  //  switch i32 %index, label %unreachable [
+  //    i32 0, label %resume.0
+  //    i32 1, label %resume.1
+  //    ...
+  //  ]
+
+  auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);
+  auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);
+
+  IRBuilder<> Builder(NewEntry);
+  auto *FramePtr = Shape.FramePtr;
+  auto *FrameTy = Shape.FrameTy;
+  auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
+      FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
+  auto *Index = Builder.CreateLoad(GepIndex, "index");
+  auto *Switch =
+      Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());
+  Shape.ResumeSwitch = Switch;
+
+  uint32_t SuspendIndex = 0;
+  for (auto S : Shape.CoroSuspends) {
+    ConstantInt *IndexVal = Builder.getInt32(SuspendIndex);
+
+    // Replace CoroSave with a store to Index:
+    //    %index.addr = getelementptr %f.frame... (index field number)
+    //    store i32 0, i32* %index.addr1
+    auto *Save = S->getCoroSave();
+    Builder.SetInsertPoint(Save);
+    auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
+        FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
+    Builder.CreateStore(IndexVal, GepIndex);
+    Save->replaceAllUsesWith(ConstantTokenNone::get(C));
+    Save->eraseFromParent();
+
+    // Split block before and after coro.suspend and add a jump from an entry
+    // switch:
+    //
+    //  whateverBB:
+    //    whatever
+    //    %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+    //    switch i8 %0, label %suspend[i8 0, label %resume
+    //                                 i8 1, label %cleanup]
+    // becomes:
+    //
+    //  whateverBB:
+    //     whatever
+    //     br label %resume.0.landing
+    //
+    //  resume.0: ; <--- jump from the switch in the resume.entry
+    //     %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)
+    //     br label %resume.0.landing
+    //
+    //  resume.0.landing:
+    //     %1 = phi i8[-1, %whateverBB], [%0, %resume.0]
+    //     switch i8 % 1, label %suspend [i8 0, label %resume
+    //                                    i8 1, label %cleanup]
+
+    auto *SuspendBB = S->getParent();
+    auto *ResumeBB =
+        SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));
+    auto *LandingBB = ResumeBB->splitBasicBlock(
+        S->getNextNode(), ResumeBB->getName() + Twine(".landing"));
+    Switch->addCase(IndexVal, ResumeBB);
+
+    cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);
+    auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "", &LandingBB->front());
+    S->replaceAllUsesWith(PN);
+    PN->addIncoming(Builder.getInt8(-1), SuspendBB);
+    PN->addIncoming(S, ResumeBB);
+
+    ++SuspendIndex;
+  }
+
+  Builder.SetInsertPoint(UnreachBB);
+  Builder.CreateUnreachable();
+
+  return NewEntry;
+}
+
+// In Resumers, we replace fallthrough coro.end with ret void and delete the
+// rest of the block.
+static void replaceFallthroughCoroEnd(IntrinsicInst *End,
+                                      ValueToValueMapTy &VMap) {
+  auto *NewE = cast<IntrinsicInst>(VMap[End]);
+  ReturnInst::Create(NewE->getContext(), nullptr, NewE);
+
+  // remove the rest of the block, by splitting it into an unreachable block
+  auto *BB = NewE->getParent();
+  BB->splitBasicBlock(NewE);
+  BB->getTerminator()->eraseFromParent();
+}
+
+// Create a resume clone by cloning the body of the original function, setting
+// new entry block and replacing coro.suspend an appropriate value to force
+// resume or cleanup pass for every suspend point.
+static Function *createClone(Function &F, Twine Suffix, coro::Shape &Shape,
+                             BasicBlock *ResumeEntry, int8_t FnIndex) {
+
+  Module *M = F.getParent();
+  auto *FrameTy = Shape.FrameTy;
+  auto *FnPtrTy = cast<PointerType>(FrameTy->getElementType(0));
+  auto *FnTy = cast<FunctionType>(FnPtrTy->getElementType());
+
+  Function *NewF =
+      Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,
+                       F.getName() + Suffix, M);
+  NewF->addAttribute(1, Attribute::NonNull);
+  NewF->addAttribute(1, Attribute::NoAlias);
+
+  ValueToValueMapTy VMap;
+  // replace all args with undefs
+  for (Argument &A : F.getArgumentList())
+    VMap[&A] = UndefValue::get(A.getType());
+
+  SmallVector<ReturnInst *, 4> Returns;
+
+  CloneFunctionInto(NewF, &F, VMap, /*ModuleLevelChanges=*/true, Returns);
+
+  // If we have debug info, update it. ModuleLevelChanges = true above, does
+  // the heavy lifting, we just need to repoint subprogram at the same
+  // DICompileUnit as the original function F.
+  if (DISubprogram *SP = F.getSubprogram())
+    NewF->getSubprogram()->replaceUnit(SP->getUnit());
+
+  LLVMContext &C = NewF->getContext();
+
+  // Remove old returns
+  for (ReturnInst *Return : Returns) {
+    new UnreachableInst(C, Return);
+    Return->eraseFromParent();
+  }
+
+  // Remove old return attributes.
+  NewF->removeAttributes(
+      AttributeSet::ReturnIndex,
+      AttributeSet::get(
+          NewF->getContext(), AttributeSet::ReturnIndex,
+          AttributeFuncs::typeIncompatible(NewF->getReturnType())));
+
+  // Make AllocaSpillBlock the new entry block
+  auto *SwitchBB = cast<BasicBlock>(VMap[ResumeEntry]);
+  auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
+  Entry->moveBefore(&NewF->getEntryBlock());
+  Entry->getTerminator()->eraseFromParent();
+  BranchInst::Create(SwitchBB, Entry);
+  Entry->setName("entry" + Suffix);
+
+  // Clear all predecessors of the new entry block.
+  auto *Switch = cast<SwitchInst>(VMap[Shape.ResumeSwitch]);
+  Entry->replaceAllUsesWith(Switch->getDefaultDest());
+
+  IRBuilder<> Builder(&NewF->getEntryBlock().front());
+
+  // Remap frame pointer.
+  Argument *NewFramePtr = &NewF->getArgumentList().front();
+  Value *OldFramePtr = cast<Value>(VMap[Shape.FramePtr]);
+  NewFramePtr->takeName(OldFramePtr);
+  OldFramePtr->replaceAllUsesWith(NewFramePtr);
+
+  // Remap vFrame pointer.
+  auto *NewVFrame = Builder.CreateBitCast(
+      NewFramePtr, Type::getInt8PtrTy(Builder.getContext()), "vFrame");
+  Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
+  OldVFrame->replaceAllUsesWith(NewVFrame);
+
+  // Replace coro suspend with the appropriate resume index.
+  auto *NewValue = Builder.getInt8(FnIndex);
+  replaceAndRemove(Shape.CoroSuspends, NewValue, &VMap);
+
+  // Remove coro.end intrinsics.
+  replaceFallthroughCoroEnd(Shape.CoroEnds.front(), VMap);
+  // FIXME: coming in upcoming patches:
+  // replaceUnwindCoroEnds(Shape.CoroEnds, VMap);
+
+  // Store the address of this clone in the coroutine frame.
+  Builder.SetInsertPoint(Shape.FramePtr->getNextNode());
+  auto *G = Builder.CreateConstInBoundsGEP2_32(Shape.FrameTy, Shape.FramePtr, 0,
+                                               FnIndex, "fn.addr");
+  Builder.CreateStore(NewF, G);
+  NewF->setCallingConv(CallingConv::Fast);
+
+  return NewF;
+}
+
+static void removeCoroEnds(coro::Shape &Shape) {
+  for (CoroEndInst *CE : Shape.CoroEnds)
+    CE->eraseFromParent();
+}
+
+static void replaceFrameSize(coro::Shape &Shape) {
+  if (Shape.CoroSizes.empty())
+    return;
+
+  // In the same function all coro.sizes should have the same result type.
+  auto *SizeIntrin = Shape.CoroSizes.back();
+  Module *M = SizeIntrin->getModule();
+  const DataLayout &DL = M->getDataLayout();
+  auto Size = DL.getTypeAllocSize(Shape.FrameTy);
+  auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
+
+  replaceAndRemove(Shape.CoroSizes, SizeConstant);
+}
+
+// Create a global constant array containing pointers to functions provided and
+// set Info parameter of CoroBegin to point at this constant. Example:
+//
+//   @f.resumers = internal constant [2 x void(%f.frame*)*]
+//                    [void(%f.frame*)* @f.resume, void(%f.frame*)* @f.destroy]
+//   define void @f() {
+//     ...
+//     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
+//                    i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to i8*))
+//
+// Assumes that all the functions have the same signature.
+static void setCoroInfo(Function &F, CoroBeginInst *CoroBegin,
+                        std::initializer_list<Function *> Fns) {
+
+  SmallVector<Constant *, 4> Args(Fns.begin(), Fns.end());
+  assert(Args.size() > 0);
+  Function *Part = *Fns.begin();
+  Module *M = Part->getParent();
+  auto *ArrTy = ArrayType::get(Part->getType(), Args.size());
+
+  auto *ConstVal = ConstantArray::get(ArrTy, Args);
+  auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
+                                GlobalVariable::PrivateLinkage, ConstVal,
+                                F.getName() + Twine(".resumers"));
+
+  // Update coro.begin instruction to refer to this constant
+  LLVMContext &C = F.getContext();
+  auto *BC = ConstantExpr::getPointerCast(GV, Type::getInt8PtrTy(C));
+  CoroBegin->getId()->setInfo(BC);
+}
+
+static void postSplitCleanup(Function &F) {
+  removeUnreachableBlocks(F);
+  llvm::legacy::FunctionPassManager FPM(F.getParent());
+
+  FPM.add(createVerifierPass());
+  FPM.add(createSCCPPass());
+  FPM.add(createCFGSimplificationPass());
+  FPM.add(createEarlyCSEPass());
+  FPM.add(createCFGSimplificationPass());
+
+  FPM.doInitialization();
+  FPM.run(F);
+  FPM.doFinalization();
+}
+
+static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
+  coro::Shape Shape(F);
+  buildCoroutineFrame(F, Shape);
+
+  auto *ResumeEntry = createResumeEntryBlock(F, Shape);
+  auto *ResumeClone = createClone(F, ".resume", Shape, ResumeEntry, 0);
+  auto *DestroyClone = createClone(F, ".destroy", Shape, ResumeEntry, 1);
+
+  // we no longer need coro.end in F
+  removeCoroEnds(Shape);
+
+  postSplitCleanup(F);
+  postSplitCleanup(*ResumeClone);
+  postSplitCleanup(*DestroyClone);
+
+  replaceFrameSize(Shape);
+
+  setCoroInfo(F, Shape.CoroBegin, {ResumeClone, DestroyClone});
+  coro::updateCallGraph(F, {ResumeClone, DestroyClone}, CG, SCC);
+}
+
 // When we see the coroutine the first time, we insert an indirect call to a
 // devirt trigger function and mark the coroutine that it is now ready for
 // split.
 static void prepareForSplit(Function &F, CallGraph &CG) {
   Module &M = *F.getParent();
 #ifndef NDEBUG
   Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
   assert(DevirtFn && "coro.devirt.trigger function not found");

 // found, we will create one and add it to the current SCC.
 static void createDevirtTriggerFunc(CallGraph &CG, CallGraphSCC &SCC) {
   Module &M = CG.getModule();
   if (M.getFunction(CORO_DEVIRT_TRIGGER_FN))
     return;
 
   LLVMContext &C = M.getContext();
   auto *FnTy = FunctionType::get(Type::getVoidTy(C), Type::getInt8PtrTy(C),
                                  /*IsVarArgs=*/false);
   Function *DevirtFn =
       Function::Create(FnTy, GlobalValue::LinkageTypes::PrivateLinkage,
                        CORO_DEVIRT_TRIGGER_FN, &M);
   DevirtFn->addFnAttr(Attribute::AlwaysInline);
   auto *Entry = BasicBlock::Create(C, "entry", DevirtFn);
   ReturnInst::Create(C, Entry);
 
   auto *Node = CG.getOrInsertFunction(DevirtFn);
for (Function *F : Coroutines) {
       StringRef Value = Attr.getValueAsString();
       DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName()
                    << "' state: " << Value << "\n");
       if (Value == UNPREPARED_FOR_SPLIT) {
         prepareForSplit(*F, CG);
         continue;
       }
       F->removeFnAttr(CORO_PRESPLIT_ATTR);
+      splitCoroutine(*F, CG, SCC);
     }
     return true;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesAll();
     CallGraphSCCPass::getAnalysisUsage(AU);
   }
 };
 }
 
 char CoroSplit::ID = 0;
 INITIALIZE_PASS(
     CoroSplit, "coro-split",
     "Split coroutine into a set of functions driving its state machine", false,
     false)
 
 Pass *llvm::createCoroSplitPass() { return new CoroSplit(); }

lib/Transforms/Coroutines/Coroutines.cpp

 //===-- Coroutines.cpp ----------------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This file implements the common infrastructure for Coroutine Passes.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 
 using namespace llvm;
 
Builder.addExtension(PassManagerBuilder::EP_ScalarOptimizerLate,
                        addCoroutineScalarOptimizerPasses);
   Builder.addExtension(PassManagerBuilder::EP_OptimizerLast,
                        addCoroutineOptimizerLastPasses);
 }
 
 // Construct the lowerer base class and initialize its members.
 coro::LowererBase::LowererBase(Module &M)
     : TheModule(M), Context(M.getContext()),
-      ResumeFnType(FunctionType::get(Type::getVoidTy(Context),
-                                     Type::getInt8PtrTy(Context),
-                                     /*isVarArg=*/false)) {}
+      Int8Ptr(Type::getInt8PtrTy(Context)),
+      ResumeFnType(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
+                                     /*isVarArg=*/false)),
+      NullPtr(ConstantPointerNull::get(Int8Ptr)) {}
 
 // Creates a sequence of instructions to obtain a resume function address using
 // llvm.coro.subfn.addr. It generates the following sequence:
 //
 //    call i8* @llvm.coro.subfn.addr(i8* %Arg, i8 %index)
 //    bitcast i8* %2 to void(i8*)*
 
 Value *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
bool coro::declaresIntrinsics(Module &M,
   for (StringRef Name : List) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
     if (M.getNamedValue(Name))
       return true;
   }
 
   return false;
 }
+
+// FIXME: This code is stolen from CallGraph::addToCallGraph(Function *F), which
+// happens to be private. It is better for this functionality exposed by the
+// CallGraph.
+static void buildCGN(CallGraph &CG, CallGraphNode *Node) {
+  Function *F = Node->getFunction();
+
+  // Look for calls by this function.
+  for (Instruction &I : instructions(F))
+    if (CallSite CS = CallSite(cast<Value>(&I))) {
+      const Function *Callee = CS.getCalledFunction();
+      if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
+        // Indirect calls of intrinsics are not allowed so no need to check.
+        // We can be more precise here by using TargetArg returned by
+        // Intrinsic::isLeaf.
+        Node->addCalledFunction(CS, CG.getCallsExternalNode());
+      else if (!Callee->isIntrinsic())
+        Node->addCalledFunction(CS, CG.getOrInsertFunction(Callee));
+    }
+}
+
+// Rebuild CGN after we extracted parts of the code from ParentFunc into
+// NewFuncs. Builds CGNs for the NewFuncs and adds them to the current SCC.
+void coro::updateCallGraph(Function &ParentFunc, ArrayRef<Function *> NewFuncs,
+                           CallGraph &CG, CallGraphSCC &SCC) {
+  // Rebuild CGN from scratch for the ParentFunc
+  auto *ParentNode = CG[&ParentFunc];
+  ParentNode->removeAllCalledFunctions();
+  buildCGN(CG, ParentNode);
+
+  SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
+
+  for (Function *F : NewFuncs) {
+    CallGraphNode *Callee = CG.getOrInsertFunction(F);
+    Nodes.push_back(Callee);
+    buildCGN(CG, Callee);
+  }
+
+  SCC.initialize(Nodes);
+}
+
+static void clear(coro::Shape &Shape) {
+  Shape.CoroBegin = nullptr;
+  Shape.CoroEnds.clear();
+  Shape.CoroSizes.clear();
+  Shape.CoroSuspends.clear();
+
+  Shape.FrameTy = nullptr;
+  Shape.FramePtr = nullptr;
+  Shape.AllocaSpillBlock = nullptr;
+  Shape.ResumeSwitch = nullptr;
+  Shape.HasFinalSuspend = false;
+}
+
+static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
+                                    CoroSuspendInst *SuspendInst) {
+  Module *M = SuspendInst->getModule();
+  auto Fn = Intrinsic::getDeclaration(M, Intrinsic::coro_save);
+  auto SaveInst =
+      cast<CoroSaveInst>(CallInst::Create(Fn, CoroBegin, "", SuspendInst));
+  assert(!SuspendInst->getCoroSave());
+  SuspendInst->setArgOperand(0, SaveInst);
+  return SaveInst;
+}
+
+void coro::Shape::buildFrom(Function &F) {
+  clear(*this);
+  for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
+    if (auto II = dyn_cast<IntrinsicInst>(&*IB++)) {
+      switch (II->getIntrinsicID()) {
+      default:
+        continue;
+      case Intrinsic::coro_size:
+        CoroSizes.push_back(cast<CoroSizeInst>(II));
+        break;
+      case Intrinsic::coro_frame:
+        assert(CoroBegin && "coro.frame should not appear before coro.begin");
+        II->replaceAllUsesWith(CoroBegin);
+        II->eraseFromParent();
+        break;
+      case Intrinsic::coro_suspend:
+        assert(CoroBegin && "coro.suspend should not appear before coro.begin");
+        CoroSuspends.push_back(cast<CoroSuspendInst>(II));
+        if (!CoroSuspends.back()->getCoroSave()) {
+          createCoroSave(CoroBegin, CoroSuspends.back());
+        }
+        if (CoroSuspends.back()->isFinal()) {
+          HasFinalSuspend = true;
+          if (CoroSuspends.size() > 1) {
+            assert(!CoroSuspends.front()->isFinal() &&
+              "Only one suspend point can be marked as final");
+            std::swap(CoroSuspends.front(), CoroSuspends.back());
+          }
+        }
+        break;
+      case Intrinsic::coro_begin: {
+        auto CB = cast<CoroBeginInst>(II);
+        if (CB->getId()->getInfo().isPreSplit()) {
+          assert(!CoroBegin &&
+            "coroutine should have exactly one defining @llvm.coro.begin");
+          CB->addAttribute(AttributeSet::ReturnIndex, Attribute::NonNull);
+          CB->addAttribute(AttributeSet::ReturnIndex, Attribute::NoAlias);
+          CB->removeAttribute(AttributeSet::FunctionIndex,
+                              Attribute::NoDuplicate);
+          CoroBegin = CB;
+        }
+        break;
+      }
+      case Intrinsic::coro_end:
+        CoroEnds.push_back(cast<CoroEndInst>(II));
+        if (CoroEnds.back()->isFallthrough()) {
+          if (CoroEnds.size() > 1) {
+            assert(!CoroEnds.front()->isFallthrough() &&
+              "Only one suspend point can be marked as final");
+            std::swap(CoroEnds.front(), CoroEnds.back());
+          }
+        }
+        break;
+      }
+    }
+  }
+}
 No newline at end of file

test/Transforms/Coroutines/coro-split-00.ll

+; Tests that coro-split pass splits the coroutine into f, f.resume and f.destroy
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+define i8* @f() "coroutine.presplit"="1" {
+entry:
+  %id = call token @llvm.coro.id(i32 0, 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)
+  call void @print(i32 0)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume 
+                                i8 1, label %cleanup]
+resume:
+  call void @print(i32 1)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 0)  
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @f(
+; CHECK: call i8* @malloc
+; CHECK: call void @print(i32 0)
+; CHECK-NOT: call void @print(i32 1)
+; CHECK-NOT: call void @free(
+; CHECK: ret i8* %hdl
+
+; CHECK-LABEL: @f.resume(
+; CHECK-NOT: call i8* @malloc
+; CHECK-NOT: call void @print(i32 0)
+; CHECK: call void @print(i32 1)
+; CHECK-NOT: call void @print(i32 0)
+; CHECK: call void @free(
+; CHECK: ret void
+
+; CHECK-LABEL: @f.destroy(
+; CHECK-NOT: call i8* @malloc
+; CHECK-NOT: call void @print(
+; CHECK: call void @free(
+; CHECK: ret void
+
+declare i8* @llvm.coro.free(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*)
+declare i8* @llvm.coro.alloc(token)
+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*)

test/Transforms/Coroutines/coro-split-01.ll

+; Tests that a coroutine is split, inlined into the caller and devirtualized.
+; RUN: opt < %s -S -enable-coroutines -O2 | FileCheck %s
+
+define i8* @f() {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+  %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  br label %coro.begin
+coro.begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
+  call void @print(i32 0)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume 
+                                i8 1, label %cleanup]
+resume:
+  call void @print(i32 1)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 0)  
+  ret i8* %hdl
+}
+define i32 @main() {
+entry:
+  %hdl = call i8* @f()
+  call void @llvm.coro.resume(i8* %hdl)
+  ret i32 0
+; CHECK-LABEL: @main(
+; CHECK:      call i8* @malloc
+; CHECK-NOT:  call void @free
+; CHECK: call void @print(i32 0)
+; CHECK-NOT:  call void @free
+; CHECK: call void @print(i32 1)
+; CHECK:      call void @free
+; CHECK:      ret i32 0
+}
+
+declare i8* @llvm.coro.free(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*)
+declare i1 @llvm.coro.alloc(token)
+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*)

test/Transforms/Coroutines/restart-trigger.ll

 ; Verifies that restart trigger forces IPO pipelines restart and the same
 ; coroutine is looked at by CoroSplit pass twice.
 ; REQUIRES: asserts
 ; RUN: opt < %s -S -O0 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -O1 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
 
 ; CHECK:      CoroSplit: Processing coroutine 'f' state: 0
 ; CHECK-NEXT: CoroSplit: Processing coroutine 'f' state: 1
 
-declare token @llvm.coro.id(i32, i8*, i8*)
-declare i8* @llvm.coro.begin(token, i8*)
-
-; a coroutine start function
 define void @f() {
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
-  call i8* @llvm.coro.begin(token %id, i8* null)
-  ret void
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  call void @print(i32 0)
+  %s1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %s1, label %suspend [i8 0, label %resume 
+                                 i8 1, label %cleanup]
+resume:
+  call void @print(i32 1)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 0)
+  ret void  
 }
+
+declare token @llvm.coro.id(i32, i8*, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i8* @llvm.coro.free(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 void @llvm.coro.end(i8*, i1) 
+
+declare noalias i8* @malloc(i32)
+declare void @print(i32)
+declare void @free(i8*)