[Coroutines] Part 9: Add cleanup subfunction.

[GorNishanov]

Summary

[Coroutines] Part 9: Add cleanup subfunction.

This patch completes coroutine heap allocation elision. Now, the heap elision example from docs\Coroutines.rst compiles and produces expected result (see test/Transform/Coroutines/ex3.ll)

Intrinsic Changes:

• coro.free gets a token parameter tying it to coro.id to allow reliably discovering all coro.frees associated with a particular coroutine.
• coro.id gets an extra parameter that points back to a coroutine function. This allows to check whether a coro.id describes the enclosing function or it belongs to a different function that was later inlined.

CoroSplit now creates three subfunctions:

1. f$resume - resume logic
2. f$destroy - cleanup logic, followed by a deallocation code
3. f$cleanup - just the cleanup code

CoroElide pass during devirtualization replaces coro.destroy with either f$destroy or f$cleanup depending whether heap elision is performed or not.

Other fixes, improvements:

• Fixed buglet in Shape::buildFrame that was not creating coro.save properly if coroutine has more than one suspend point.

• Switched to using variable width suspend index field (no longer limited to 32 bit index field can be as little as i1 or as large as i<whatever-size_t-is>)

[majnemer]

docs/Coroutines.rst

806–810

Hmm, coro.begin also takes the id. coro.free claims it reads the frame parameter, would anything break if we stopped doing that?

934–935

This is a little strange...

lib/Transforms/Coroutines/CoroElide.cpp

150–158

This looks like a duplicated version of Instruction::mayThrow

172–173

Terminators are not the only instruction which may throw. Calls may throw as well.

[GorNishanov]

docs/Coroutines.rst

806–810

I was thinking that if an allocation code gets unswitched we may get more than one coro.begin for the same coro.id and we would need to find all coro.begins and tied them back with a PHINode and use it as a replacement value for coro.free if heap elision is not performed.

Having it as a parameter allows LLVM to do it automatically for us.

934–935

It was one of the ways I thought of how we can detect whether a coro.id we are looking at describes the enclosing function or it describes a different function due to inline coroutine start function into its caller.

This used to be important in the design where I was stroring "elide-or-not" boolean in the coroutine frame. Since now we switched to separate f$destroy and f$cleanup, identifying whether coro.id corresponds to the current function or not is no longer important.

I can take that parameter out altogether.

lib/Transforms/Coroutines/CoroElide.cpp

150–158

Cool. I can inline this function into shouldElide as:

isa<ReturnInst> || mayThrow(T)

172–173

After we talked on IRC, I am going to redo the shouldElide into something simpler. Something like:

If every coro.destroy points at coro.begin (that takes care of the escape, since if there is no direct SSA to coro.begin), it did escape and we were not able to put coro.begin in the vreg
<gor> majnemer: More precisely, if for every coro.begin I found a coro.destroy directly pointing at it, we are good for elision

[GorNishanov]

Simplify shouldElide() to just check that for every coro.begin there is a coro.destroy referencing it directly.

[GorNishanov]

docs/Coroutines.rst

934–935

I'll keep if for now. We may run into another case where we will need it.

lib/Transforms/Coroutines/CoroElide.cpp

150–158

This code was simplified and no longer require these.

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroElide.cpp

147–203

This seems dead.

175

Pointers lean right.

[GorNishanov]

Nits addressed. Preparing to land.

docs/Coroutines.rst

 and 6 after which the coroutine will be destroyed.
 
 The LLVM IR for this coroutine looks like this:
 
 .. code-block:: none
 
   define i8* @f(i32 %n) {
   entry:
-    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+    %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 noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
     br label %loop
   loop:
     %n.val = phi i32 [ %n, %entry ], [ %inc, %loop ]
     %inc = add nsw i32 %n.val, 1
     call void @print(i32 %n.val)
     %0 = call i8 @llvm.coro.suspend(token none, i1 false)
     switch i8 %0, label %suspend [i8 0, label %loop
                                   i8 1, label %cleanup]
   cleanup:
-    %mem = call i8* @llvm.coro.free(i8* %hdl)
+    %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
     call void @free(i8* %mem)
     br label %suspend
   suspend:
     call void @llvm.coro.end(i8* %hdl, i1 false)
     ret i8* %hdl
   }
 
 The `entry` block establishes the coroutine frame. The `coro.size`_ intrinsic is

 After resume and destroy parts are outlined, function `f` will contain only the 
 code responsible for creation and initialization of the coroutine frame and 
 execution of the coroutine until a suspend point is reached:
 
 .. code-block:: none
 
   define i8* @f(i32 %n) {
   entry:
-    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
     %alloc = call noalias i8* @malloc(i32 24)
     %0 = call noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
     %frame = bitcast i8* %0 to %f.frame*
     %1 = getelementptr %f.frame, %f.frame* %frame, i32 0, i32 0
     store void (%f.frame*)* @f.resume, void (%f.frame*)** %1
     %2 = getelementptr %f.frame, %f.frame* %frame, i32 0, i32 1
     store void (%f.frame*)* @f.destroy, void (%f.frame*)** %2
    

 
 In the entry block, we will call `coro.alloc`_ intrinsic that will return `true`
 when dynamic allocation is required, and `false` if dynamic allocation is 
 elided.
 
 .. code-block:: none
 
   entry:
-    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+    %id = call token @llvm.coro.id(i32 0, i8* null, 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* @CustomAlloc(i32 %size)
     br label %coro.begin
   coro.begin:
     %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
     %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %phi)
 
 In the cleanup block, we will make freeing the coroutine frame conditional on
 `coro.free`_ intrinsic. If allocation is elided, `coro.free`_ returns `null`
 thus skipping the deallocation code:
 
 .. code-block:: llvm
 
   cleanup:
-    %mem = call i8* @llvm.coro.free(i8* %hdl)
+    %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
     %need.dyn.free = icmp ne i8* %mem, null
     br i1 %need.dyn.free, label %dyn.free, label %if.end
   dyn.free:
     call void @CustomFree(i8* %mem)
     br label %if.end
   if.end:
     ...
 

 store the current value produced by a coroutine.
 
 .. code-block:: text
 
   define i8* @f(i32 %n) {
   entry:
     %promise = alloca i32
     %pv = bitcast i32* %promise to i8*
-    %id = call token @llvm.coro.id(i32 0, i8* %pv, i8* null)
+    %id = call token @llvm.coro.id(i32 0, i8* %pv, 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 noalias i8* @llvm.coro.begin(token %id, i8* %phi)
     br label %loop
   loop:
     %n.val = phi i32 [ %n, %coro.begin ], [ %inc, %loop ]
     %inc = add nsw i32 %n.val, 1
     store i32 %n.val, i32* %promise
     %0 = call i8 @llvm.coro.suspend(token none, i1 false)
     switch i8 %0, label %suspend [i8 0, label %loop
                                   i8 1, label %cleanup]
   cleanup:
-    %mem = call i8* @llvm.coro.free(i8* %hdl)
+    %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
     call void @free(i8* %mem)
     br label %suspend
   suspend:
     call void @llvm.coro.end(i8* %hdl, i1 false)
     ret i8* %hdl
   }
 
 A coroutine consumer can rely on the `coro.promise`_ intrinsic to access the

 
 .. code-block:: text
 
   define i8* @f(i32 %n) {
   entry:
     %promise = alloca i32
     %pv = bitcast i32* %promise to i8*
     ; the second argument to coro.id points to the coroutine promise.
-    %id = call token @llvm.coro.id(i32 0, i8* %pv, i8* null)
+    %id = call token @llvm.coro.id(i32 0, i8* %pv, i8* null, i8* null)
     ...
     %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
     ...
     store i32 42, i32* %promise ; store something into the promise
     ...
     ret i8* %hdl
   }
 

 A frontend should emit exactly one `coro.begin` intrinsic per coroutine.
 
 .. _coro.free:
 
 'llvm.coro.free' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
-  declare i8* @llvm.coro.free(i8* <frame>)
+  declare i8* @llvm.coro.free(token %id, i8* <frame>)
 
 Overview:
 """""""""
 
 The '``llvm.coro.free``' intrinsic returns a pointer to a block of memory where 
 coroutine frame is stored or `null` if this instance of a coroutine did not use
 dynamically allocated memory for its coroutine frame.
 
 Arguments:
 """"""""""
 
-A pointer to the coroutine frame. This should be the same pointer that was 
-returned by prior `coro.begin` call.
+The first argument is a token returned by a call to '``llvm.coro.id``' 
+identifying the coroutine.
+
+The second argument is a pointer to the coroutine frame. This should be the same
+pointer that was returned by prior `coro.begin` call.
 
 Example (custom deallocation function):
 """""""""""""""""""""""""""""""""""""""
 
 .. code-block:: llvm
 
   cleanup:
-    %mem = call i8* @llvm.coro.free(i8* %frame)
+    %mem = call i8* @llvm.coro.free(token %id, i8* %frame)
     %mem_not_null = icmp ne i8* %mem, null
     br i1 %mem_not_null, label %if.then, label %if.end
   if.then:
     call void @CustomFree(i8* %mem)
     br label %if.end
   if.end:
     ret void
 
 Example (standard deallocation functions):
 """"""""""""""""""""""""""""""""""""""""""
 
 .. code-block:: llvm
 
   cleanup:
-    %mem = call i8* @llvm.coro.free(i8* %frame)
+    %mem = call i8* @llvm.coro.free(token %id, i8* %frame)
     call void @free(i8* %mem)
     ret void
 
 .. _coro.alloc:
 
 'llvm.coro.alloc' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::

 when possible.
 
 Example:
 """"""""
 
 .. code-block:: text
 
   entry:
-    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+    %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
     %dyn.alloc.required = call i1 @llvm.coro.alloc(token %id)
     br i1 %dyn.alloc.required, label %coro.alloc, label %coro.begin
 
   coro.alloc:
     %frame.size = call i32 @llvm.coro.size()
     %alloc = call i8* @MyAlloc(i32 %frame.size)
     br label %coro.begin
 

 coroutine frame.
 
 .. _coro.id:
 
 'llvm.coro.id' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
-  declare token @llvm.coro.id(i32 <align>, i8* <promise>, i8* <fnaddr>)
+  declare token @llvm.coro.id(i32 <align>, i8* <promise>, i8* <coroaddr>, 
+                                                          i8* <fnaddrs>)
 
 Overview:
 """""""""
 
 The '``llvm.coro.id``' intrinsic returns a token identifying a coroutine.
 
 Arguments:
 """"""""""
 
 The first argument provides information on the alignment of the memory returned 
 by the allocation function and given to `coro.begin` by the first argument. If 
 this argument is 0, the memory is assumed to be aligned to 2 * sizeof(i8*).
 This argument only accepts constants.
 
 The second argument, if not `null`, designates a particular alloca instruction
 to be a `coroutine promise`_.
 
-The third argument is `null` before coroutine is split, and later is replaced 
+The third argument is `null` coming out of the frontend. The CoroEarly pass sets
+this argument to point to the function this coro.id belongs to. 
+
+The fourth argument is `null` before coroutine is split, and later is replaced 
 to point to a private global constant array containing function pointers to 
 outlined resume and destroy parts of the coroutine.
 
 
 Semantics:
 """"""""""
 
 The purpose of this intrinsic is to tie together `coro.id`, `coro.alloc` and

 This pass also replaces `coro.resume` and `coro.destroy` intrinsics with direct 
 calls to resume and destroy functions for a particular coroutine where possible.
 
 CoroCleanup
 -----------
 This pass runs late to lower all coroutine related intrinsics not replaced by
 earlier passes.
 
-Upstreaming sequence (rough plan)
-=================================
-#. Add documentation.
-#. Add coroutine intrinsics.
-#. Add empty coroutine passes.
-#. Add coroutine devirtualization + tests.
-#. Add CGSCC restart trigger + tests.
-#. Add coroutine heap elision + tests.
-#. Add custom allocation heap elision + tests. <== we are here
-#. Add coroutine splitting logic + tests.
-#. Add simple coroutine frame builder + tests.
-#. Add the rest of the logic + tests. (Maybe split further as needed).
-
 Areas Requiring Attention
 =========================
 #. A coroutine frame is bigger than it could be. Adding stack packing and stack 
    coloring like optimization on the coroutine frame will result in tighter
    coroutine frames.
 
 #. Take advantage of the lifetime intrinsics for the data that goes into the
    coroutine frame. Leave lifetime intrinsics as is for the data that stays in

include/llvm/IR/Intrinsics.td

def int_experimental_gc_relocate : Intrinsic<[llvm_any_ty],
                                 [IntrReadMem]>;
 
 //===------------------------ Coroutine Intrinsics ---------------===//
 // These are documented in docs/Coroutines.rst
 
 // Coroutine Structure Intrinsics.
 
 def int_coro_id : Intrinsic<[llvm_token_ty], [llvm_i32_ty, llvm_ptr_ty, 
-                             llvm_ptr_ty], 
+                             llvm_ptr_ty, llvm_ptr_ty], 
                             [IntrArgMemOnly, IntrReadMem, 
                              ReadNone<1>, ReadOnly<2>, NoCapture<2>]>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                [WriteOnly<1>]>;
 
-def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
-                              [IntrArgMemOnly, ReadOnly<0>, NoCapture<0>]>;
+def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty], 
+                              [IntrReadMem, IntrArgMemOnly, ReadOnly<1>, 
+                               NoCapture<1>]>;
 def int_coro_end : Intrinsic<[], [llvm_ptr_ty, llvm_i1_ty], []>;
 
 def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
 
 def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>;
 def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>;
 

lib/IR/Verifier.cpp

void Verifier::visitIntrinsicCallSite(Intrinsic::ID ID, CallSite CS) {
   for (Value *V : CS.args()) 
     if (auto *MD = dyn_cast<MetadataAsValue>(V))
       visitMetadataAsValue(*MD, CS.getCaller());
 
   switch (ID) {
   default:
     break;
   case Intrinsic::coro_id: {
-    auto *InfoArg = CS.getArgOperand(2)->stripPointerCasts();
+    auto *InfoArg = CS.getArgOperand(3)->stripPointerCasts();
     if (isa<ConstantPointerNull>(InfoArg))
       break;
     auto *GV = dyn_cast<GlobalVariable>(InfoArg);
     Assert(GV && GV->isConstant() && GV->hasDefinitiveInitializer(),
       "info argument of llvm.coro.begin must refer to an initialized "
       "constant");
     Constant *Init = GV->getInitializer();
     Assert(isa<ConstantStruct>(Init) || isa<ConstantArray>(Init),

lib/Transforms/Coroutines/CoroEarly.cpp

if (auto CS = CallSite(&I)) {
         break;
       case Intrinsic::coro_id:
         // Mark a function that comes out of the frontend that has a coro.id
         // with a coroutine attribute.
         if (auto *CII = cast<CoroIdInst>(&I)) {
           if (CII->getInfo().isPreSplit()) {
             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
             setCannotDuplicate(CII);
+            CII->setCoroutineSelf();
           }
         }
         break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);

lib/Transforms/Coroutines/CoroElide.cpp

 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass replaces dynamic allocation of coroutine frame with alloca and
 // replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls
 // to coroutine sub-functions.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-elide"
 
 namespace {
 // Created on demand if CoroElide pass has work to do.
 struct Lowerer : coro::LowererBase {
   SmallVector<CoroIdInst *, 4> CoroIds;
   SmallVector<CoroBeginInst *, 1> CoroBegins;
   SmallVector<CoroAllocInst *, 1> CoroAllocs;
   SmallVector<CoroSubFnInst *, 4> ResumeAddr;
   SmallVector<CoroSubFnInst *, 4> DestroyAddr;
   SmallVector<CoroFreeInst *, 1> CoroFrees;
 
   Lowerer(Module &M) : LowererBase(M) {}
 
   void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
-  bool processCoroId(CoroIdInst *, AAResults &AA);
+  bool shouldElide(Function *F, DominatorTree &DT,
+                   PostDominatorTree &PDT) const;
+  bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT,
+                     PostDominatorTree &PDT);
 };
 } // end anonymous namespace
 
 // Go through the list of coro.subfn.addr intrinsics and replace them with the
 // provided constant.
 static void replaceWithConstant(Constant *Value,
                                 SmallVectorImpl<CoroSubFnInst *> &Users) {
   if (Users.empty())
void Lowerer::elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA) {
   //   mem = coro.alloc(id) ? malloc(coro.size()) : 0;
   //   coro.begin(id, mem)
   auto *False = ConstantInt::getFalse(C);
   for (auto *CA : CoroAllocs) {
     CA->replaceAllUsesWith(False);
     CA->eraseFromParent();
   }
 
-  // To suppress deallocation code, we replace all llvm.coro.free intrinsics
-  // associated with this coro.begin with null constant.
-  auto *NullPtr = ConstantPointerNull::get(Type::getInt8PtrTy(C));
-  for (auto *CF : CoroFrees) {
-    CF->replaceAllUsesWith(NullPtr);
-    CF->eraseFromParent();
-  }
-
   // FIXME: Design how to transmit alignment information for every alloca that
   // is spilled into the coroutine frame and recreate the alignment information
   // here. Possibly we will need to do a mini SROA here and break the coroutine
   // frame into individual AllocaInst recreating the original alignment.
   auto *Frame = new AllocaInst(FrameTy, "", InsertPt);
   auto *FrameVoidPtr =
       new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt);
 
   for (auto *CB : CoroBegins) {
     CB->replaceAllUsesWith(FrameVoidPtr);
     CB->eraseFromParent();
   }
 
   // Since now coroutine frame lives on the stack we need to make sure that
   // any tail call referencing it, must be made non-tail call.
   removeTailCallAttribute(Frame, AA);
 }
 
-bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA) {
+// This function is used to check whether a terminator causes control flow to
+// leaves the function in some way.
+static bool returnsOrUnwindsToCaller(TerminatorInst *T) {
+  if (isa<ReturnInst>(T) || isa<ResumeInst>(T))
+    return true;
+
+  if (auto *CR = dyn_cast<CleanupReturnInst>(T))
+    return CR->unwindsToCaller();
+  if (auto *CS = dyn_cast<CatchSwitchInst>(T))
+    return CS->unwindsToCaller();
+
+  return false;
+}
+
+// Check that on every code path leaving the function dominated by a coro.begin,
+// there is either a coro.free or coro.destroy postdominating a coro.begin.
+bool Lowerer::shouldElide(Function *F, DominatorTree &DT,
+                          PostDominatorTree &PDT) const {
+  // If no CoroAllocs, we cannot suppress allocation, so elision is not
+  // possible.
+  if (CoroAllocs.empty())
+    return false;
+
+  for (BasicBlock &BB : *F) {
+    auto *T = BB.getTerminator();
+    if (!returnsOrUnwindsToCaller(T))
+      continue;
+
+    for (CoroBeginInst *CB : CoroBegins) {
+      if (DT.dominates(CB, T)) {
+        bool PostDominated = false;
+        for (CoroFreeInst *CF : CoroFrees) {
+          if (PDT.dominates(&BB, CF->getParent())) {
+            PostDominated = true;
+            break;
+          }
+        }
+        if (PostDominated)
+          continue;
+
+        for (CoroSubFnInst *SF : DestroyAddr) {
+          if (PDT.dominates(&BB, SF->getParent())) {
+            PostDominated = true;
+            break;
+          }
+        }
+        if (!PostDominated)
+          return false;
+      }
+    }
+  }
+  return true;
+}
+
+bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
+                            DominatorTree &DT, PostDominatorTree &PDT) {
   CoroBegins.clear();
   CoroAllocs.clear();
+  CoroFrees.clear();
   ResumeAddr.clear();
   DestroyAddr.clear();
 
   // Collect all coro.begin and coro.allocs associated with this coro.id.
   for (User *U : CoroId->users()) {
     if (auto *CB = dyn_cast<CoroBeginInst>(U))
       CoroBegins.push_back(CB);
     else if (auto *CA = dyn_cast<CoroAllocInst>(U))
       CoroAllocs.push_back(CA);
+    else if (auto *CF = dyn_cast<CoroFreeInst>(U))
+      CoroFrees.push_back(CF);
   }
 
   // Collect all coro.subfn.addrs associated with coro.begin.
   // Note, we only devirtualize the calls if their coro.subfn.addr refers to
   // coro.begin directly. If we run into cases where this check is too
   // conservative, we can consider relaxing the check.
   for (CoroBeginInst *CB : CoroBegins) {
     for (User *U : CB->users())
bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
   ConstantArray *Resumers = CoroId->getInfo().Resumers;
   assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
                      "of coroutine subfunctions");
   auto *ResumeAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
 
   replaceWithConstant(ResumeAddrConstant, ResumeAddr);
 
-  if (DestroyAddr.empty())
-    return true;
+  bool ShouldElide = shouldElide(CoroId->getFunction(), DT, PDT);
 
-  auto *DestroyAddrConstant =
-      ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::DestroyIndex);
+  auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
+      Resumers,
+      ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
 
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
 
-  // If there is a coro.alloc that llvm.coro.id refers to, we have the ability
-  // to suppress dynamic allocation.
-  if (!CoroAllocs.empty()) {
-    // FIXME: The check above is overly lax. It only checks for whether we have
-    // an ability to elide heap allocations, not whether it is safe to do so.
-    // We need to do something like:
-    // If for every exit from the function where coro.begin is
-    // live, there is a coro.free or coro.destroy dominating that exit block,
-    // then it is safe to elide heap allocation, since the lifetime of coroutine
-    // is fully enclosed in its caller.
+  if (ShouldElide) {
     auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
     elideHeapAllocations(CoroId->getFunction(), FrameTy, AA);
+    coro::replaceCoroFree(CoroId, /*Elide=*/true);
   }
+
   return true;
 }
 
 // See if there are any coro.subfn.addr instructions referring to coro.devirt
 // trigger, if so, replace them with a direct call to devirt trigger function.
 static bool replaceDevirtTrigger(Function &F) {
   SmallVector<CoroSubFnInst *, 1> DevirtAddr;
   for (auto &I : instructions(F))
if (!L)
       return false;
 
     bool Changed = false;
 
     if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
       Changed = replaceDevirtTrigger(F);
 
     L->CoroIds.clear();
-    L->CoroFrees.clear();
 
-    // Collect all PostSplit coro.ids and all coro.free.
+    // Collect all PostSplit coro.ids.
     for (auto &I : instructions(F))
-      if (auto *CF = dyn_cast<CoroFreeInst>(&I))
-        L->CoroFrees.push_back(CF);
-      else if (auto *CII = dyn_cast<CoroIdInst>(&I))
+      if (auto *CII = dyn_cast<CoroIdInst>(&I))
         if (CII->getInfo().isPostSplit())
+          // If it is the coroutine itself, don't touch it.
+          if (CII->getCoroutine() != CII->getFunction())
             L->CoroIds.push_back(CII);
 
     // If we did not find any coro.id, there is nothing to do.
     if (L->CoroIds.empty())
       return Changed;
 
     AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+    DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+    PostDominatorTree &PDT =
+        getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
+
     for (auto *CII : L->CoroIds)
-      Changed |= L->processCoroId(CII, AA);
+      Changed |= L->processCoroId(CII, AA, DT, PDT);
 
     return Changed;
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AAResultsWrapperPass>();
-    AU.setPreservesCFG();
+    AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<PostDominatorTreeWrapperPass>();
   }
 };
 }
 
 char CoroElide::ID = 0;
 INITIALIZE_PASS_BEGIN(
     CoroElide, "coro-elide",
     "Coroutine frame allocation elision and indirect calls replacement", false,
     false)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
 INITIALIZE_PASS_END(
     CoroElide, "coro-elide",
     "Coroutine frame allocation elision and indirect calls replacement", false,
     false)
 
 Pass *llvm::createCoroElidePass() { return new CoroElide(); }

lib/Transforms/Coroutines/CoroFrame.cpp

 
 #include "CoroInternal.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/circular_raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
 
 // The "coro-suspend-crossing" flag is very noisy. There is another debug type,
 // "coro-frame", which results in leaner debug spew.
 #define DEBUG_TYPE "coro-suspend-crossing"
for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
     CoroSaveInst *const CoroSave = CSI->getCoroSave();
     BasicBlock *const CoroSaveBB = CoroSave->getParent();
     auto &B = getBlockData(CoroSaveBB);
     B.Suspend = true;
     B.Kills |= B.Consumes;
   }
 
   // Iterate propagating consumes and kills until they stop changing
-  int Iteration = 0; (void)Iteration;
-  
+  int Iteration = 0;
+  (void)Iteration;
   bool Changed;
   do {
     DEBUG(dbgs() << "iteration " << ++Iteration);
     DEBUG(dbgs() << "==============\n");
 
     Changed = false;
     for (size_t I = 0; I < N; ++I) {
       auto &B = Block[I];
static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   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() > UINT32_MAX)
-    report_fatal_error("Cannot handle coroutine with this many suspend points");
+  // Figure out how wide should be an integer type storing the suspend index.
+  unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
 
-  SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, Type::getInt32Ty(C)};
+  SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, Type::getIntNTy(C, IndexBits)};
   Value *CurrentDef = nullptr;
 
   // Create an entry for every spilled value.
   for (auto const &S : Spills) {
     if (CurrentDef == S.def())
       continue;
 
     CurrentDef = S.def();
 
     Type *Ty = nullptr;
     if (auto *AI = dyn_cast<AllocaInst>(CurrentDef))
       Ty = AI->getAllocatedType();
     else
       Ty = CurrentDef->getType();
 
     Types.push_back(Ty);
   }
   FrameTy->setBody(Types);
 
   return FrameTy;
 }
 
-// Returns the index of the last non-spill field in the coroutine frame.
-//  2 - if there is no coroutine promise specified or 3, if there is.
-static unsigned getLastNonSpillIndex(coro::Shape &Shape) {
-  // TODO: Add support for coroutine promise.
-  return 2;
-}
-
 // 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
 //
static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
   PointerType *FramePtrTy = Shape.FrameTy->getPointerTo();
   auto *FramePtr =
       cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
   Type *FrameTy = FramePtrTy->getElementType();
 
   Value *CurrentValue = nullptr;
   BasicBlock *CurrentBlock = nullptr;
   Value *CurrentReload = nullptr;
-  unsigned Index = getLastNonSpillIndex(Shape);
+  unsigned Index = coro::Shape::LastKnownField;
 
   // We need to keep track of any allocas that need "spilling"
   // since they will live in the coroutine frame now, all access to them
   // need to be changed, not just the access across suspend points
   // we remember allocas and their indices to be handled once we processed
   // all the spills.
   SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
 
for (Instruction &I : instructions(F)) {
     // pairs and should not be part of the Coroutine Frame
     if (isa<CoroSaveInst>(&I))
       continue;
     // CoroBeginInst returns a handle to a coroutine which is passed as a sole
     // parameter to .resume and .cleanup parts and should not go into coroutine
     // frame.
     if (isa<CoroBeginInst>(&I))
       continue;
+    // A token returned CoroIdInst is used to tie together structural intrinsics
+    // in a coroutine. It should not be saved to the coroutine frame.
+    if (isa<CoroIdInst>(&I))
+      continue;
 
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {
         // We cannot spill a token.
         if (I.getType()->isTokenTy())
           report_fatal_error(
               "token definition is separated from the use by a suspend point");
         assert(!materializable(I) &&

lib/Transforms/Coroutines/CoroInstr.h

 //===-- CoroInstr.h - Coroutine Intrinsics Instruction Wrappers -*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This file defines classes that make it really easy to deal with intrinsic
 // functions with the isa/dyncast family of functions.  In particular, this
 // allows you to do things like:
 //
 //     if (auto *SF = dyn_cast<CoroSubFnInst>(Inst))
-//        ... SF->getFrame() ... 
+//        ... SF->getFrame() ...
 //
 // All intrinsic function calls are instances of the call instruction, so these
 // are all subclasses of the CallInst class.  Note that none of these classes
 // has state or virtual methods, which is an important part of this gross/neat
 // hack working.
 //
 // The helpful comment above is borrowed from llvm/IntrinsicInst.h, we keep
 // coroutine intrinsic wrappers here since they are only used by the passes in

 class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst {
   enum { FrameArg, IndexArg };
 
 public:
   enum ResumeKind {
     RestartTrigger = -1,
     ResumeIndex,
     DestroyIndex,
+    CleanupIndex,
     IndexLast,
     IndexFirst = RestartTrigger
   };
 
   Value *getFrame() const { return getArgOperand(FrameArg); }
   ResumeKind getIndex() const {
     int64_t Index = getRawIndex()->getValue().getSExtValue();
     assert(Index >= IndexFirst && Index < IndexLast &&
public:
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
 /// This represents the llvm.coro.alloc instruction.
 class LLVM_LIBRARY_VISIBILITY CoroIdInst : public IntrinsicInst {
-  enum { AlignArg, PromiseArg, InfoArg };
+  enum { AlignArg, PromiseArg, CoroutineArg, InfoArg };
+
 public:
   // Info argument of coro.id is
   //   fresh out of the frontend: null ;
   //   outlined                 : {Init, Return, Susp1, Susp2, ...} ;
   //   postsplit                : [resume, destroy, cleanup] ;
   //
   // If parts of the coroutine were outlined to protect against undesirable
   // code motion, these functions will be stored in a struct literal referred to
Info getInfo() const {
     return Result;
   }
   Constant *getRawInfo() const {
     return cast<Constant>(getArgOperand(InfoArg)->stripPointerCasts());
   }
 
   void setInfo(Constant *C) { setArgOperand(InfoArg, C); }
 
+  Function *getCoroutine() const {
+    return cast<Function>(getArgOperand(CoroutineArg)->stripPointerCasts());
+  }
+  void setCoroutineSelf() {
+    assert(isa<ConstantPointerNull>(getArgOperand(CoroutineArg)) &&
+           "Coroutine argument is already assigned");
+    auto *const Int8PtrTy = Type::getInt8PtrTy(getContext());
+    setArgOperand(CoroutineArg,
+                  ConstantExpr::getBitCast(getFunction(), Int8PtrTy));
+  }
 
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_id;
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
 /// This represents the llvm.coro.frame instruction.
 class LLVM_LIBRARY_VISIBILITY CoroFrameInst : 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_frame;
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
 /// This represents the llvm.coro.free instruction.
 class LLVM_LIBRARY_VISIBILITY CoroFreeInst : public IntrinsicInst {
+  enum { IdArg, FrameArg };
+
 public:
+  Value *getFrame() const { return getArgOperand(FrameArg); }
+
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_free;
   }
   static inline bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
 /// This class represents the llvm.coro.begin instruction.
 class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst {
   enum { IdArg, MemArg };
 
 public:
-  CoroIdInst *getId() const {
-    return cast<CoroIdInst>(getArgOperand(IdArg));
-  }
+  CoroIdInst *getId() const { return cast<CoroIdInst>(getArgOperand(IdArg)); }
 
   Value *getMem() const { return getArgOperand(MemArg); }
 
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_begin;
   }
   static inline bool classof(const Value *V) {

lib/Transforms/Coroutines/CoroInternal.h

 
 #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 replaceCoroFree(CoroIdInst *CoroId, bool Elide);
 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;

 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,
+    ResumeField,
+    DestroyField,
+    IndexField,
+    LastKnownField = IndexField
   };
 
   StructType *FrameTy;
   Instruction *FramePtr;
   BasicBlock* AllocaSpillBlock;
   SwitchInst* ResumeSwitch;
   bool HasFinalSuspend;
 
+  IntegerType* getIndexType() const {
+    assert(FrameTy && "frame type not assigned");
+    return cast<IntegerType>(FrameTy->getElementType(IndexField));
+  }
+  ConstantInt *getIndex(uint64_t Value) const {
+    return ConstantInt::get(getIndexType(), Value);
+  }
+
   Shape() = default;
   explicit Shape(Function &F) { buildFrom(F); }
   void buildFrom(Function &F);
 };
 
 void buildCoroutineFrame(Function& F, Shape& Shape);
 
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

lib/Transforms/Coroutines/CoroSplit.cpp

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);
+    ConstantInt *IndexVal = Shape.getIndex(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");
static void replaceFallthroughCoroEnd(IntrinsicInst *End,
   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);
static Function *createClone(Function &F, Twine Suffix, coro::Shape &Shape,
 
   // 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);
+  // Replacing coro.suspend with (0) will result in control flow proceeding to
+  // a resume label associated with a suspend point, replacing it with (1) will
+  // result in control flow proceeding to a cleanup label associated with this
+  // suspend point.
+  auto *NewValue = Builder.getInt8(FnIndex ? 1 : 0);
   for (CoroSuspendInst *CS : Shape.CoroSuspends) {
     auto *MappedCS = cast<CoroSuspendInst>(VMap[CS]);
     MappedCS->replaceAllUsesWith(NewValue);
     MappedCS->eraseFromParent();
   }
 
   // Remove coro.end intrinsics.
   replaceFallthroughCoroEnd(Shape.CoroEnds.front(), VMap);
   // FIXME: coming in upcoming patches:
   // replaceUnwindCoroEnds(Shape.CoroEnds, VMap);
 
+  // We only store resume(0) and destroy(1) addresses in the coroutine frame.
+  // The cleanup(2) clone is only used during devirtualization when coroutine is
+  // eligible for heap elision and thus does not participate in indirect calls
+  // and does not need its address to be stored in the coroutine frame.
+  if (FnIndex < 2) {
     // 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");
+    auto *G = Builder.CreateConstInBoundsGEP2_32(Shape.FrameTy, Shape.FramePtr,
+                                                 0, FnIndex, "fn.addr");
     Builder.CreateStore(NewF, G);
+  }
+
+  // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
+  // to suppress deallocation code.
+  coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
+                        /*Elide=*/FnIndex == 2);
+
   NewF->setCallingConv(CallingConv::Fast);
 
   return NewF;
 }
 
 static void removeCoroEnds(coro::Shape &Shape) {
   for (CoroEndInst *CE : Shape.CoroEnds)
     CE->eraseFromParent();

 }
 
 static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
   coro::Shape Shape(F);
   if (!Shape.CoroBegin)
     return;
 
   buildCoroutineFrame(F, Shape);
+  replaceFrameSize(Shape);
 
   auto *ResumeEntry = createResumeEntryBlock(F, Shape);
-  auto *ResumeClone = createClone(F, ".resume", Shape, ResumeEntry, 0);
-  auto *DestroyClone = createClone(F, ".destroy", Shape, ResumeEntry, 1);
+  auto ResumeClone = createClone(F, ".resume", Shape, ResumeEntry, 0);
+  auto DestroyClone = createClone(F, ".destroy", Shape, ResumeEntry, 1);
+  auto CleanupClone = createClone(F, ".cleanup", Shape, ResumeEntry, 2);
 
   // We no longer need coro.end in F.
   removeCoroEnds(Shape);
 
   postSplitCleanup(F);
   postSplitCleanup(*ResumeClone);
   postSplitCleanup(*DestroyClone);
+  postSplitCleanup(*CleanupClone);
 
-  replaceFrameSize(Shape);
-
-  setCoroInfo(F, Shape.CoroBegin, {ResumeClone, DestroyClone});
-  coro::updateCallGraph(F, {ResumeClone, DestroyClone}, CG, SCC);
+  setCoroInfo(F, Shape.CoroBegin, {ResumeClone, DestroyClone, CleanupClone});
+  coro::updateCallGraph(F, {ResumeClone, DestroyClone, CleanupClone}, 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

lib/Transforms/Coroutines/Coroutines.cpp

for (StringRef Name : List) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
     if (M.getNamedValue(Name))
       return true;
   }
 
   return false;
 }
 
+// Replace all coro.frees associated with the provided CoroId either with 'null'
+// if Elide is true and with its frame parameter otherwise.
+void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
+  SmallVector<CoroFreeInst *, 4> CoroFrees;
+  for (User *U : CoroId->users())
+    if (auto CF = dyn_cast<CoroFreeInst>(U))
+      CoroFrees.push_back(CF);
+
+  if (CoroFrees.empty())
+    return;
+
+  Value *Replacement =
+      Elide ? ConstantPointerNull::get(Type::getInt8PtrTy(CoroId->getContext()))
+            : CoroFrees.front()->getFrame();
+
+  for (CoroFreeInst *CF : CoroFrees) {
+    CF->replaceAllUsesWith(Replacement);
+    CF->eraseFromParent();
+  }
+}
+
 // 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))
void coro::Shape::buildFrom(Function &F) {
   for (CoroFrameInst *CF : CoroFrames) {
     CF->replaceAllUsesWith(CoroBegin);
     CF->eraseFromParent();
   }
 
   // Canonicalize coro.suspend by inserting a coro.save if needed.
   for (CoroSuspendInst *CS : CoroSuspends)
     if (!CS->getCoroSave())
-      createCoroSave(CoroBegin, CoroSuspends.back());
+      createCoroSave(CoroBegin, CS);
 }

test/Transforms/Coroutines/coro-elide.ll

 
 @f.resumers = internal constant [2 x void (i8*)*] [void (i8*)* @f.resume,
                                                    void (i8*)* @f.destroy]
 
 ; a coroutine start function
 define i8* @f() {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
+                          i8* bitcast (i8*()* @f to i8*),
                           i8* bitcast ([2 x void (i8*)*]* @f.resumers to i8*))
   %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @callResume(
 define void @callResume() {
 entry:
-; CHECK: call i8* @llvm.coro.begin
   %hdl = call i8* @f()
 
-; CHECK-NEXT: call void @print(i32 0)
+; CHECK: call void @print(i32 0)
   %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %1 = bitcast i8* %0 to void (i8*)*
   call fastcc void %1(i8* %hdl)
 
 ; CHECK-NEXT: call void @print(i32 1)
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %3 = bitcast i8* %2 to void (i8*)*
   call fastcc void %3(i8* %hdl)
 
 ; CHECK-NEXT: ret void
   ret void
 }
 
 ; CHECK-LABEL: @eh(
 define void @eh() personality i8* null {
 entry:
-; CHECK: call i8* @llvm.coro.begin
   %hdl = call i8* @f()
 
-; CHECK-NEXT: call void @print(i32 0)
+; CHECK: call void @print(i32 0)
   %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %1 = bitcast i8* %0 to void (i8*)*
   invoke void %1(i8* %hdl)
           to label %cont unwind label %ehcleanup
 cont:
   ret void
 
 ehcleanup:
   %tok = cleanuppad within none []
   cleanupret from %tok unwind to caller
 }
 
 ; CHECK-LABEL: @no_devirt_info_null(
 ; no devirtualization here, since coro.begin info parameter is null
 define void @no_devirt_info_null() {
 entry:
-  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
 
 ; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %1 = bitcast i8* %0 to void (i8*)*
   call fastcc void %1(i8* %hdl)
 
 ; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %3 = bitcast i8* %2 to void (i8*)*
   call fastcc void %3(i8* %hdl)
 
 ; CHECK: ret void
   ret void
 }
 
-declare token @llvm.coro.id(i32, i8*, i8*)
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i8* @llvm.coro.frame()
 declare i8* @llvm.coro.subfn.addr(i8*, i8)

test/Transforms/Coroutines/coro-heap-elide.ll

 ; Tests that the dynamic allocation and deallocation of the coroutine frame is
 ; elided and any tail calls referencing the coroutine frame has the tail 
 ; call attribute removed.
 ; RUN: opt < %s -S -inline -coro-elide -instsimplify -simplifycfg | FileCheck %s
 
 declare void @print(i32) nounwind
 
 %f.frame = type {i32}
 
 declare void @bar(i8*)
 
 declare fastcc void @f.resume(%f.frame*)
 declare fastcc void @f.destroy(%f.frame*)
+declare fastcc void @f.cleanup(%f.frame*)
 
 declare void @may_throw()
 declare i8* @CustomAlloc(i32)
 declare void @CustomFree(i8*)
 
[e-mail] = internal constant 
-  [2 x void (%f.frame*)*] [void (%f.frame*)* @f.resume, void (%f.frame*)* @f.destroy]
[e-mail] = internal constant [3 x void (%f.frame*)*] 
+  [void (%f.frame*)* @f.resume, void (%f.frame*)* @f.destroy, void (%f.frame*)* @f.cleanup]
 
 ; a coroutine start function
 define i8* @f() personality i8* null {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
-                      i8* bitcast ([2 x void (%f.frame*)*]* @f.resumers to i8*))
+                      i8* bitcast (i8*()* @f to i8*),
+                      i8* bitcast ([3 x void (%f.frame*)*]* @f.resumers to i8*))
   %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
 dyn.alloc:
   %alloc = call i8* @CustomAlloc(i32 4)
   br label %coro.begin
 coro.begin:
   %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
   invoke void @may_throw() 
     to label %ret unwind label %ehcleanup
 ret:          
   ret i8* %hdl
 
 ehcleanup:
   %tok = cleanuppad within none []
-  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
   %need.dyn.free = icmp ne i8* %mem, null
   br i1 %need.dyn.free, label %dyn.free, label %if.end
 dyn.free:
   call void @CustomFree(i8* %mem)
   br label %if.end
 if.end:
   cleanupret from %tok unwind to caller
 }
 
 ; CHECK-LABEL: @callResume(
 define void @callResume() {
 entry:
 ; CHECK: alloca %f.frame
 ; CHECK-NOT: coro.begin
 ; CHECK-NOT: CustomAlloc
 ; CHECK: call void @may_throw()
   %hdl = call i8* @f()
 
 ; Need to remove 'tail' from the first call to @bar
 ; CHECK-NOT: tail call void @bar(
 ; CHECK: call void @bar(
   tail call void @bar(i8* %hdl)
-; CHECK: tail call void @bar(  
+; CHECK: tail call void @bar(
   tail call void @bar(i8* null)
 
 ; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8* %vFrame)
   %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %1 = bitcast i8* %0 to void (i8*)*
   call fastcc void %1(i8* %hdl)
 
-; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.destroy to void (i8*)*)(i8* %vFrame)
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %3 = bitcast i8* %2 to void (i8*)*
   call fastcc void %3(i8* %hdl)
 
 ; CHECK-NEXT: ret void
   ret void
 }
 
 ; a coroutine start function (cannot elide heap alloc, due to second argument to
 ; coro.begin not pointint to coro.alloc)
 define i8* @f_no_elision() personality i8* null {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
-                      i8* bitcast ([2 x void (%f.frame*)*]* @f.resumers to i8*))
+                      i8* bitcast (i8*()* @f_no_elision to i8*),
+                      i8* bitcast ([3 x void (%f.frame*)*]* @f.resumers to i8*))
   %alloc = call i8* @CustomAlloc(i32 4)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @callResume_no_elision(
 define void @callResume_no_elision() {
 entry:
; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.destroy to void (i8*)*)(i8*
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
   %3 = bitcast i8* %2 to void (i8*)*
   call fastcc void %3(i8* %hdl)
 
 ; CHECK-NEXT: ret void
   ret void
 }
 
-declare token @llvm.coro.id(i32, i8*, i8*)
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i1 @llvm.coro.alloc(token)
-declare i8* @llvm.coro.free(i8*)
+declare i8* @llvm.coro.free(token, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i8* @llvm.coro.frame(token)
 declare i8* @llvm.coro.subfn.addr(i8*, i8)

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)
+  %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)
   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)
+  %mem = call i8* @llvm.coro.free(token %id, 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: 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 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*)
+declare token @llvm.coro.id(i32, i8*, 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)
+  %id = call token @llvm.coro.id(i32 0, i8* null, 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)
+  %mem = call i8* @llvm.coro.free(token %id, 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 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*)
+declare token @llvm.coro.id(i32, i8*, 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/ex0.ll

 ; First example from Doc/Coroutines.rst (two block loop)
 ; RUN: opt < %s -enable-coroutines -O2 -S | FileCheck %s
 
 define i8* @f(i32 %n) {
 entry:
-  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+  %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)
   br label %loop
 
 loop:
   %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
   call void @print(i32 %n.val)
   %0 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %0, label %suspend [i8 0, label %resume 
                                 i8 1, label %cleanup]
 resume:
   %inc = add i32 %n.val, 1
   br label %loop
 
 cleanup:
-  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  %mem = call i8* @llvm.coro.free(token %id, 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: @main(
 define i32 @main() {
 entry:
   %hdl = call i8* @f(i32 4)
   call void @llvm.coro.resume(i8* %hdl)
   call void @llvm.coro.resume(i8* %hdl)
   call void @llvm.coro.destroy(i8* %hdl)
   ret i32 0
 ; CHECK: entry:
 ; CHECK:      call void @print(i32 4)
 ; CHECK:      call void @print(i32 5)
 ; CHECK:      call void @print(i32 6)
 ; CHECK:      ret i32 0
 }
 
-declare token @llvm.coro.id(i32, i8*, i8*)
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.alloc(token)
-declare i8* @llvm.coro.free(i8*)
+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 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/ex1.ll

 ; First example from Doc/Coroutines.rst (one block loop)
 ; 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)
+  %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 noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
   br label %loop
 loop:
   %n.val = phi i32 [ %n, %entry ], [ %inc, %loop ]
   %inc = add nsw i32 %n.val, 1
   call void @print(i32 %n.val)
   %0 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %0, label %suspend [i8 0, label %loop
                                 i8 1, label %cleanup]
 cleanup:
-  %mem = call i8* @llvm.coro.free(i8* %hdl)
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
   call void @free(i8* %mem)
   br label %suspend
 suspend:
   call void @llvm.coro.end(i8* %hdl, i1 false)
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @main(

 ; CHECK:      call void @print(i32 6)
 ; CHECK:      ret i32 0
 }
 
 declare i8* @malloc(i32)
 declare void @free(i8*)
 declare void @print(i32)
 
-declare token @llvm.coro.id(i32, i8*, i8*)
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i32 @llvm.coro.size.i32()
 declare i8* @llvm.coro.begin(token, i8*)
 declare i8 @llvm.coro.suspend(token, i1)
-declare i8* @llvm.coro.free(i8*)
+declare i8* @llvm.coro.free(token, i8*)
 declare void @llvm.coro.end(i8*, i1)
 
 declare void @llvm.coro.resume(i8*)
 declare void @llvm.coro.destroy(i8*)

test/Transforms/Coroutines/ex2.ll

+; Second example from Doc/Coroutines.rst (custom alloc and free functions)
+; 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)
+  %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* @CustomAlloc(i32 %size)
+  br label %coro.begin
+coro.begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %phi)
+  br label %loop
+loop:
+  %n.val = phi i32 [ %n, %coro.begin ], [ %inc, %loop ]
+  %inc = add nsw i32 %n.val, 1
+  call void @print(i32 %n.val)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %loop
+                                i8 1, label %cleanup]
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  %need.dyn.free = icmp ne i8* %mem, null
+  br i1 %need.dyn.free, label %dyn.free, label %suspend
+dyn.free:
+  call void @CustomFree(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 false)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @main
+define i32 @main() {
+entry:
+  %hdl = call i8* @f(i32 4)
+  call void @llvm.coro.resume(i8* %hdl)
+  call void @llvm.coro.resume(i8* %hdl)
+  call void @llvm.coro.destroy(i8* %hdl)
+  ret i32 0
+; CHECK:      call void @print(i32 4)
+; CHECK-NEXT: call void @print(i32 5)
+; CHECK-NEXT: call void @print(i32 6)
+; CHECK-NEXT: ret i32 0
+}
+
+declare i8* @CustomAlloc(i32)
+declare void @CustomFree(i8*)
+declare void @print(i32)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i32 @llvm.coro.size.i32()
+declare i8* @llvm.coro.begin(token, i8*)
+declare i8 @llvm.coro.suspend(token, i1)
+declare i8* @llvm.coro.free(token, i8*)
+declare void @llvm.coro.end(i8*, i1)
+
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)

test/Transforms/Coroutines/ex3.ll

+; Third example from Doc/Coroutines.rst (two suspend points)
+; 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 noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
+  br label %loop
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %loop.resume ]
+  call void @print(i32 %n.val) #4
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %loop.resume
+                                i8 1, label %cleanup]
+loop.resume:
+  %inc = add nsw i32 %n.val, 1
+  %sub = xor i32 %n.val, -1
+  call void @print(i32 %sub)
+  %1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %1, label %suspend [i8 0, label %loop
+                                i8 1, label %cleanup]
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call void @llvm.coro.end(i8* %hdl, i1 false)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @main
+define i32 @main() {
+entry:
+  %hdl = call i8* @f(i32 4)
+  call void @llvm.coro.resume(i8* %hdl)
+  call void @llvm.coro.resume(i8* %hdl)
+  call void @llvm.coro.destroy(i8* %hdl)
+  ret i32 0
+; CHECK:      call void @print(i32 4)
+; CHECK-NEXT: call void @print(i32 -5)
+; CHECK-NEXT: call void @print(i32 5)
+; CHECK:      ret i32 0
+}
+
+declare i8* @malloc(i32)
+declare void @free(i8*)
+declare void @print(i32)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i32 @llvm.coro.size.i32()
+declare i8* @llvm.coro.begin(token, i8*)
+declare i8 @llvm.coro.suspend(token, i1)
+declare i8* @llvm.coro.free(token, i8*)
+declare void @llvm.coro.end(i8*, i1)
+
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(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
 
 define void @f() {
-  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null)
+  %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)
   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)
+  %mem = call i8* @llvm.coro.free(token %id, 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 token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
-declare i8* @llvm.coro.free(i8*)
+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 void @llvm.coro.end(i8*, i1) 
 
 declare noalias i8* @malloc(i32)
 declare void @print(i32)
 declare void @free(i8*)