[Coroutines] Part 6: Elide dynamic allocation of a coroutine frame when possible

[GorNishanov]

Summary

A particular coroutine usage pattern, where a coroutine is created, manipulated and
destroyed by the same calling function, is common for coroutines implementing
RAII idiom and is suitable for allocation elision optimization which avoid
dynamic allocation by storing the coroutine frame as a static alloca in its
caller.

coro.free and coro.alloc intrinsics are used to indicate which code needs to be suppressed
when dynamic allocation elision happens:

entry:
  %elide = call i8* @llvm.coro.alloc()
  %need.dyn.alloc = icmp ne i8* %elide, null
  br i1 %need.dyn.alloc, label %coro.begin, label %dyn.alloc
dyn.alloc:
  %alloc = call i8* @CustomAlloc(i32 4)
  br label %coro.begin
coro.begin:
  %phi = phi i8* [ %elide, %entry ], [ %alloc, %dyn.alloc ]
  %hdl = call i8* @llvm.coro.begin(i8* %phi, i32 0, i8* null,
                          i8* bitcast ([2 x void (%f.frame*)*]* @f.resumers to i8*))

and

  %mem = call i8* @llvm.coro.free(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:
  ...

If heap allocation elision is performed, we replace coro.alloc with a static alloca on the caller frame and coro.free with null constant.

Also, we need to make sure that if there are any tail calls referencing the coroutine frame, we need to remote tail call attribute, since now coroutine frame lives on the stack.

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)
3.Add empty coroutine passes. (https://reviews.llvm.org/D22847)
4.Add coroutine devirtualization + tests.
ab) Lower coro.resume and coro.destroy (https://reviews.llvm.org/D22998)
c) Do devirtualization (https://reviews.llvm.org/D23229)
5.Add CGSCC restart trigger + tests. (https://reviews.llvm.org/D23234)
6.Add coroutine heap elision + tests. <= we are here
7.Add the rest of the logic (split into more patches)

[GorNishanov]

clang-format the changes. Fix the title.

[majnemer]

lib/Transforms/Coroutines/CoroElide.cpp

91–95

This doesn't seem conservative enough because you are only check if the call's operands alias the alloca. I could imagine the call referencing the coroutine frame via other means (load/store via global).

I think getModRefInfo() != MRI_NoModRef would capture this better.

104–106

Should we report_fatal_error if the CallSite is must_tail?

111

auto *

130

Do we have to worry about the alignment of this alloca?

210–211

Is this still safe if there are multiple coro.begin calls and some of them refer to llvm.coro.alloc and others do not?

210–219

We would typically sink this assignment into the if.

220

auto *

[GorNishanov]

Address the feedback

lib/Transforms/Coroutines/CoroElide.cpp

130

Added comment:

// 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.

210–211

Addressed by making coro.begin return the token and coro.frame taken that token returning the coroutine frame address. This should prevent coro.begin being duplicated.

[GorNishanov]

modRef check was too conservative. Reverting back to original behavior.

[GorNishanov]

1. coro.begin now has a parameter pointing directly at coro.alloc if it is present or null, otherwise
2. documentation + test + code updated to support this tweak
3. clarified the FIXME comment explaining that we need to add a safety check when heap allocation elision is possible
4. added a no_elision unit test for the cases when coro.alloc is not provided

[majnemer]

LGTM with nits addressed.

lib/Transforms/Coroutines/CoroInstr.h

106–112

I'd phrase this another way to make it more conservative:

CoroAllocInst *getAlloc() const {
  if (auto *CAI = dyn_cast<CoroAllocInst>getArgOperand(ElideArg->stripPointerCasts()))
    return CAI;

  return nullptr;
}

This way, we won't crash if undef or a bitcast or null or some other weird operation is sitting in the elide arg.

[GorNishanov]

Nits addressed, Rebased to top of the tree, will land soon

docs/Coroutines.rst

 The LLVM IR for this coroutine looks like this:
 
 .. code-block:: none
 
   define i8* @f(i32 %n) {
   entry:
     %size = call i32 @llvm.coro.size.i32()
     %alloc = call i8* @malloc(i32 %size)
-    %hdl = call noalias i8* @llvm.coro.begin(i8* %alloc, i32 0, i8* null, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %alloc, i32 0, i8* null, i8* null)
+    %hdl = call noalias i8* @llvm.coro.frame(token %beg)
     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)
     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
 lowered to a constant representing the size required for the coroutine frame. 
-The `coro.begin`_ intrinsic initializes the coroutine frame and returns the 
-coroutine handle. The first parameter of `coro.begin` is given a block of memory 
-to be used if the coroutine frame needs to be allocated dynamically.
+The `coro.begin`_ intrinsic initializes the coroutine frame and returns the a
+token that is used to obtain the coroutine handle via `coro.frame` intrinsic.
+The first parameter of `coro.begin` is given a block of memory to be used if the
+coroutine frame needs to be allocated dynamically.
 
 The `cleanup` block destroys the coroutine frame. The `coro.free`_ intrinsic, 
 given the coroutine handle, returns a pointer of the memory block to be freed or
 `null` if the coroutine frame was not allocated dynamically. The `cleanup` 
 block is entered when coroutine runs to completion by itself or destroyed via
 call to the `coro.destroy`_ intrinsic.
 
 The `suspend` block contains code to be executed when coroutine runs to 

 .. code-block:: text
 
   %f.frame = type { void (%f.frame*)*, void (%f.frame*)*, i32 }
 
 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:: llvm
+.. code-block:: none
 
   define i8* @f(i32 %n) {
   entry:
     %alloc = call noalias i8* @malloc(i32 24)
-    %0 = call noalias i8* @llvm.coro.begin(i8* %alloc, i32 0, i8* null, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %alloc, i32 0, i8* null, i8* null)
+    %0 = call i8* @llvm.coro.frame(token %beg)
     %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
    
     %inc = add nsw i32 %n, 1
     %inc.spill.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2

 RAII idiom and is suitable for allocation elision optimization which avoid 
 dynamic allocation by storing the coroutine frame as a static `alloca` in its 
 caller.
 
 In the entry block, we will call `coro.alloc`_ intrinsic that will return `null`
 when dynamic allocation is required, and an address of an alloca on the caller's
 frame where coroutine frame can be stored if dynamic allocation is elided.
 
-.. code-block:: llvm
+.. code-block:: none
 
   entry:
     %elide = call i8* @llvm.coro.alloc()
     %need.dyn.alloc = icmp ne i8* %elide, null
     br i1 %need.dyn.alloc, label %coro.begin, label %dyn.alloc
   dyn.alloc:
     %size = call i32 @llvm.coro.size.i32()
     %alloc = call i8* @CustomAlloc(i32 %size)
     br label %coro.begin
   coro.begin:
     %phi = phi i8* [ %elide, %entry ], [ %alloc, %dyn.alloc ]
-    %hdl = call noalias i8* @llvm.coro.begin(i8* %phi, i32 0, i8* null, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %phi, i32 0, i8* null, i8* null)
 
 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:
entry:
     %need.dyn.alloc = icmp ne i8* %elide, null
     br i1 %need.dyn.alloc, label %coro.begin, label %dyn.alloc
   dyn.alloc:
     %size = call i32 @llvm.coro.size.i32()
     %alloc = call i8* @malloc(i32 %size)
     br label %coro.begin
   coro.begin:
     %phi = phi i8* [ %elide, %entry ], [ %alloc, %dyn.alloc ]
-    %hdl = call noalias i8* @llvm.coro.begin(i8* %phi, i32 0, i8* %pv, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %phi, i32 0, i8* %pv, i8* null)
+    %hdl = call i8* @llvm.coro.frame(token %beg)
     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]

 Using this intrinsic on a coroutine that does not have a coroutine promise
 leads to undefined behavior. It is possible to read and modify coroutine
 promise of the coroutine which is currently executing. The coroutine author and
 a coroutine user are responsible to makes sure there is no data races.
 
 Example:
 """"""""
 
-.. code-block:: llvm
+.. code-block:: text
 
   define i8* @f(i32 %n) {
   entry:
     %promise = alloca i32
     %pv = bitcast i32* %promise to i8*
     ...
     ; the third argument to coro.begin points to the coroutine promise.
-    %hdl = call noalias i8* @llvm.coro.begin(i8* %alloc, i32 0, i8* %pv, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %alloc, i32 0, i8* %pv, i8* null)
+    %hdl = call noalias i8* @llvm.coro.frame(token %beg)
     ...
     store i32 42, i32* %promise ; store something into the promise
     ...
     ret i8* %hdl
   }
 
   define i32 @main() {
   entry:

 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i8* @llvm.coro.begin(i8* <mem>, i32 <align>, i8* <promise>, i8* <fnaddr>)
 
 Overview:
 """""""""
 
-The '``llvm.coro.begin``' intrinsic returns an address of the coroutine frame.
+The '``llvm.coro.begin``' intrinsic captures coroutine initialization 
+information and returns a token that can be used by `coro.frame` intrinsic to
+return an address of the coroutine frame.
 
 Arguments:
 """"""""""
 
 The first argument is a pointer to a block of memory where coroutine frame
 will be stored.
 
 The second 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 third argument, if not `null`, designates a particular alloca instruction to
 be a `coroutine promise`_.
 
 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:
 """"""""""
 
 Depending on the alignment requirements of the objects in the coroutine frame
-and/or on the codegen compactness reasons the pointer returned from `coro.begin` 
-may be at offset to the `%mem` argument. (This could be beneficial if 
-instructions that express relative access to data can be more compactly encoded 
-with small positive and negative offsets).
+and/or on the codegen compactness reasons the pointer returned from `coro.frame`
+associated witha a particular `coro.begin` may be at offset to the `%mem` 
+argument. (This could be beneficial if instructions that express relative access
+to data can be more compactly encoded with small positive and negative offsets).
 
 A frontend should emit exactly one `coro.begin` intrinsic per coroutine.
 
 .. _coro.free:
 
 'llvm.coro.free' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i8* @llvm.coro.free(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.
+returned by prior `coro.frame` call.
 
 Example (custom deallocation function):
 """""""""""""""""""""""""""""""""""""""
 
 .. code-block:: llvm
 
   cleanup:
     %mem = call i8* @llvm.coro.free(i8* %frame)

 If the coroutine is eligible for heap elision, this intrinsic is lowered to an 
 alloca storing the coroutine frame. Otherwise, it is lowered to constant `null`.
 
 A frontend should emit at most one `coro.alloc` intrinsic per coroutine.
 
 Example:
 """"""""
 
-.. code-block:: llvm
+.. code-block:: text
 
   entry:
     %elide = call i8* @llvm.coro.alloc()
     %0 = icmp ne i8* %elide, null
     br i1 %0, label %coro.begin, label %coro.alloc
 
   coro.alloc:
     %frame.size = call i32 @llvm.coro.size()
     %alloc = call i8* @MyAlloc(i32 %frame.size)
     br label %coro.begin
 
   coro.begin:
     %phi = phi i8* [ %elide, %entry ], [ %alloc, %coro.alloc ]
-    %frame = call i8* @llvm.coro.begin(i8* %phi, i32 0, i8* null, i8* null)
+    %beg = call token @llvm.coro.begin(i8* %phi, i32 0, i8* null, i8* null)
+    %frame = call i8* @llvm.coro.frame(token %beg)
 
 .. _coro.frame:
 
 'llvm.coro.frame' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i8* @llvm.coro.frame()
 
 Overview:
 """""""""
 
 The '``llvm.coro.frame``' intrinsic returns an address of the coroutine frame of
 the enclosing coroutine.
 
 Arguments:
 """"""""""
 
-None
+A token that refers to `coro.begin` instruction.
 
 Semantics:
 """"""""""
 
-This intrinsic is lowered to refer to the `coro.begin`_ instruction. This is
-a frontend convenience intrinsic that makes it easier to refer to the
-coroutine frame.
+This intrinsic is lowered to refer to address of the coroutine frame. 
 
 .. _coro.end:
 
 'llvm.coro.end' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare void @llvm.coro.end(i8* <handle>, i1 <unwind>)

 ---------
 The pass CoroSplit buides coroutine frame and outlines resume and destroy parts 
 into separate functions.
 
 CoroElide
 ---------
 The pass CoroElide examines if the inlined coroutine is eligible for heap 
 allocation elision optimization. If so, it replaces `coro.alloc` and 
-`coro.begin` intrinsic with an address of a coroutine frame placed on its caller
+`coro.frame` intrinsic with an address of a coroutine frame placed on its caller
 and replaces `coro.free` intrinsics with `null` to remove the deallocation code. 
 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.

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_alloc : Intrinsic<[llvm_ptr_ty], [], []>;
-def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty, llvm_i32_ty,
+def int_coro_begin : Intrinsic<[llvm_token_ty], [llvm_ptr_ty, llvm_i32_ty,
                                 llvm_ptr_ty, llvm_ptr_ty],
                                [WriteOnly<0>, ReadNone<2>, ReadOnly<3>,
                                 NoCapture<3>]>;
 
 def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
                               [IntrArgMemOnly, ReadOnly<0>, NoCapture<0>]>;
 def int_coro_end : Intrinsic<[], [llvm_ptr_ty, llvm_i1_ty], []>;
 
-def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
+def int_coro_frame : Intrinsic<[llvm_ptr_ty], [llvm_token_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], []>;
 
 def int_coro_param : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_ptr_ty],
                                [IntrNoMem, ReadNone<0>, ReadNone<1>]>;
 

lib/Transforms/Coroutines/CoroElide.cpp

 // to coroutine sub-functions.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/ErrorHandling.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-elide"
 
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct CoroElide : FunctionPass {
   static char ID;
   CoroElide() : FunctionPass(ID) {}
 
   bool NeedsToRun = false;
 
   bool doInitialization(Module &M) override {
     NeedsToRun = coro::declaresIntrinsics(M, {"llvm.coro.begin"});
     return false;
   }
 
   bool runOnFunction(Function &F) override;
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
     AU.setPreservesCFG();
   }
 };
 }
 
+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_END(
+    CoroElide, "coro-elide",
+    "Coroutine frame allocation elision and indirect calls replacement", false,
+    false)
+
+Pass *llvm::createCoroElidePass() { return new CoroElide(); }
+
+//===----------------------------------------------------------------------===//
+//                              Implementation
+//===----------------------------------------------------------------------===//
+
 // 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())
     return;
 
   // See if we need to bitcast the constant to match the type of the intrinsic
   // being replaced. Note: All coro.subfn.addr intrinsics return the same type,
   // so we only need to examine the type of the first one in the list.
   Type *IntrTy = Users.front()->getType();
   Type *ValueTy = Value->getType();
   if (ValueTy != IntrTy) {
     // May need to tweak the function type to match the type expected at the
     // use site.
     assert(ValueTy->isPointerTy() && IntrTy->isPointerTy());
     Value = ConstantExpr::getBitCast(Value, IntrTy);
   }
 
   // Now the value type matches the type of the intrinsic. Replace them all!
   for (CoroSubFnInst *I : Users)
     replaceAndRecursivelySimplify(I, Value);
 }
 
+// See if any operand of the call instruction references the coroutine frame.
+static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) {
+  for (Value *Op : CI->operand_values())
+    if (AA.alias(Op, Frame) != NoAlias)
+      return true;
+  return false;
+}
+
+// Look for any tail calls referencing the coroutine frame and remove tail
+// attribute from them, since now coroutine frame resides on the stack and tail
+// call implies that the function does not references anything on the stack.
+static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) {
+  Function &F = *Frame->getFunction();
+  MemoryLocation Mem(Frame);
+  for (Instruction &I : instructions(F))
+    if (auto *Call = dyn_cast<CallInst>(&I))
+      if (Call->isTailCall() && (operandReferences(Call, Frame, AA) ||
+                                 AA.getModRefInfo(Call, Mem) != MRI_NoModRef)) {
+        // FIXME: If we ever hit this check. Evaluate whether it is more
+        // appropriate to retain musttail and allow the code to compile.
+        if (Call->isMustTailCall())
+          report_fatal_error("Call referring to the coroutine frame cannot be "
+                             "marked as musttail");
+        Call->setTailCall(false);
+      }
+}
+
+// Given a resume function @f.resume(%f.frame* %frame), returns %f.frame type.
+static Type *getFrameType(Function *Resume) {
+  auto *ArgType = Resume->getArgumentList().front().getType();
+  return cast<PointerType>(ArgType)->getElementType();
+}
+
+// Finds first non alloca instruction in the entry block of a function.
+static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
+  for (Instruction &I : F->getEntryBlock())
+    if (!isa<AllocaInst>(&I))
+      return &I;
+  llvm_unreachable("no terminator in the entry block");
+}
+
+// To elide heap allocations we need to suppress code blocks guarded by
+// llvm.coro.alloc and llvm.coro.free instructions.
+static void elideHeapAllocations(CoroBeginInst *CoroBegin, Type *FrameTy,
+                                 CoroAllocInst *AllocInst, AAResults &AA) {
+  LLVMContext &C = CoroBegin->getContext();
+  auto *InsertPt = getFirstNonAllocaInTheEntryBlock(CoroBegin->getFunction());
+
+  // 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);
+
+  // Replacing llvm.coro.alloc with non-null value will suppress dynamic
+  // allocation as it is expected for the frontend to generate the code that
+  // looks like:
+  //   mem = coro.alloc();
+  //   if (!mem) mem = malloc(coro.size());
+  //   coro.begin(mem, ...)
+  AllocInst->replaceAllUsesWith(FrameVoidPtr);
+  AllocInst->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));
+  coro::replaceAllCoroFrees(CoroBegin, NullPtr);
+  CoroBegin->lowerTo(FrameVoidPtr);
+
+  // 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);
+}
+
 // See if there are any coro.subfn.addr intrinsics directly referencing
 // the coro.begin. If found, replace them with an appropriate coroutine
 // subfunction associated with that coro.begin.
-static bool replaceIndirectCalls(CoroBeginInst *CoroBegin) {
+static bool replaceIndirectCalls(CoroBeginInst *CoroBegin, AAResults &AA) {
   SmallVector<CoroSubFnInst *, 8> ResumeAddr;
   SmallVector<CoroSubFnInst *, 8> DestroyAddr;
 
-  for (User *U : CoroBegin->users()) {
+  for (User *CF : CoroBegin->users()) {
+    assert(isa<CoroFrameInst>(CF) &&
+           "CoroBegin can be only used by coro.frame instructions");
+    for (User *U : CF->users()) {
       if (auto *II = dyn_cast<CoroSubFnInst>(U)) {
         switch (II->getIndex()) {
         case CoroSubFnInst::ResumeIndex:
           ResumeAddr.push_back(II);
           break;
         case CoroSubFnInst::DestroyIndex:
           DestroyAddr.push_back(II);
           break;
         default:
           llvm_unreachable("unexpected coro.subfn.addr constant");
         }
       }
     }
+  }
   if (ResumeAddr.empty() && DestroyAddr.empty())
     return false;
 
   // PostSplit coro.begin refers to an array of subfunctions in its Info
   // argument.
   ConstantArray *Resumers = CoroBegin->getInfo().Resumers;
   assert(Resumers && "PostSplit coro.begin 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;
+
   auto *DestroyAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::DestroyIndex);
-
-  replaceWithConstant(ResumeAddrConstant, ResumeAddr);
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
+
+  // If llvm.coro.begin refers to llvm.coro.alloc, we can elide the allocation.
+  if (auto *AllocInst = CoroBegin->getAlloc()) {
+    // FIXME: Do more sophisticated check for when we can do heap elision.
+    // Something like: for every exit from the function where coro.begin is
+    // live, there is a coro.free or coro.destroy dominating that exit block.
+    // At the moment we simply assume that if we found at least one coro.destroy
+    // referencing the coro.begin, we can elide the heap allocation.
+    auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
+    elideHeapAllocations(CoroBegin, FrameTy, AllocInst, AA);
+  }
+
   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))
bool CoroElide::runOnFunction(Function &F) {
   for (auto &I : instructions(F))
     if (auto *CB = dyn_cast<CoroBeginInst>(&I))
       if (CB->getInfo().isPostSplit())
         CoroBegins.push_back(CB);
 
   if (CoroBegins.empty())
     return Changed;
 
+  AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
   for (auto *CB : CoroBegins)
-    Changed |= replaceIndirectCalls(CB);
+    Changed |= replaceIndirectCalls(CB, AA);
 
   return Changed;
 }
-
-char CoroElide::ID = 0;
-INITIALIZE_PASS_BEGIN(
-    CoroElide, "coro-elide",
-    "Coroutine frame allocation elision and indirect calls replacement", false,
-    false)
-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/CoroInstr.h

public:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_subfn_addr;
   }
   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 CoroAllocInst : 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_alloc;
+  }
+  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 {
+public:
+  // 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 { MemArg, AlignArg, PromiseArg, InfoArg };
 
 public:
+  // See if there is a coro.alloc alternative to dynamic memory allocation.
+  CoroAllocInst *getAlloc() const {
+    if (auto PN = dyn_cast<PHINode>(getMem()))
+      for (Value *V : PN->incoming_values())
+        if (auto *CA = dyn_cast<CoroAllocInst>(V))
+          return CA;
+    return nullptr;
+  }
+
+  Value *getMem() const { return getArgOperand(MemArg); }
+
   Constant *getRawInfo() const {
     return cast<Constant>(getArgOperand(InfoArg)->stripPointerCasts());
   }
 
   void setInfo(Constant *C) { setArgOperand(InfoArg, C); }
 
   // Info argument of coro.begin is
   //   fresh out of the frontend: null ;
Info getInfo() const {
     Constant *Initializer = GV->getInitializer();
     if ((Result.OutlinedParts = dyn_cast<ConstantStruct>(Initializer)))
       return Result;
 
     Result.Resumers = cast<ConstantArray>(Initializer);
     return Result;
   }
 
+  // Replaces all coro.frame intrinsics that are associated with this coro.begin
+  // to a replacement value and removes coro.begin and all of the coro.frame
+  // intrinsics.
+  void lowerTo(Value* Replacement) {
+    SmallVector<CoroFrameInst*, 4> FrameInsts;
+    for (auto *CF : this->users())
+      FrameInsts.push_back(cast<CoroFrameInst>(CF));
+
+    for (auto *CF : FrameInsts) {
+      CF->replaceAllUsesWith(Replacement);
+      CF->eraseFromParent();
+    }
+
+    this->eraseFromParent();
+  }
+
   // 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) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
 } // End namespace llvm.

lib/Transforms/Coroutines/CoroInternal.h

 #define UNPREPARED_FOR_SPLIT "0"
 #define PREPARED_FOR_SPLIT "1"
 
 #define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger"
 
 namespace coro {
 
 bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
+void replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement);
 
 // Keeps data and helper functions for lowering coroutine intrinsics.
 struct LowererBase {
   Module &TheModule;
   LLVMContext &Context;
   FunctionType *const ResumeFnType;
 
   LowererBase(Module &M);
   Value *makeSubFnCall(Value *Arg, int Index, Instruction *InsertPt);
 };
 
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

lib/Transforms/Coroutines/Coroutines.cpp

bool coro::declaresIntrinsics(Module &M,
   for (StringRef Name : List) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
     if (M.getNamedValue(Name))
       return true;
   }
 
   return false;
 }
+
+// Find all llvm.coro.free instructions associated with the provided coro.begin
+// and replace them with the provided replacement value.
+void coro::replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement) {
+  SmallVector<CoroFreeInst *, 4> CoroFrees;
+  for (User *FramePtr: CB->users())
+    for (User *U : FramePtr->users())
+      if (auto *CF = dyn_cast<CoroFreeInst>(U))
+        CoroFrees.push_back(CF);
+
+  if (CoroFrees.empty())
+    return;
+
+  for (CoroFreeInst *CF : CoroFrees) {
+    CF->replaceAllUsesWith(Replacement);
+    CF->eraseFromParent();
+  }
+}

test/Transforms/Coroutines/coro-elide.ll

 ; Tests that the coro.destroy and coro.resume are devirtualized where possible,
 ; SCC pipeline restarts and inlines the direct calls.
-; RUN: opt < %s -S -inline -coro-elide | FileCheck %s
+; RUN: opt < %s -S -inline -coro-elide -dce | FileCheck %s
 
 declare void @print(i32) nounwind
 
 ; resume part of the coroutine
 define fastcc void @f.resume(i8*) {
   tail call void @print(i32 0)
   ret void
 }
 
 ; destroy part of the coroutine
 define fastcc void @f.destroy(i8*) {
   tail call void @print(i32 1)
   ret void
 }
 
 @f.resumers = internal constant [2 x void (i8*)*] [void (i8*)* @f.resume,
                                                    void (i8*)* @f.destroy]
 
 ; a coroutine start function
 define i8* @f() {
 entry:
-  %hdl = call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
+  %tok = call token @llvm.coro.begin(i8* null, i32 0, i8* null,
                           i8* bitcast ([2 x void (i8*)*]* @f.resumers to i8*))
+  %hdl = call i8* @llvm.coro.frame(token %tok)
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @callResume(
 define void @callResume() {
 entry:
-; CHECK: call i8* @llvm.coro.begin
+; CHECK: call token @llvm.coro.begin
   %hdl = call i8* @f()
 
 ; CHECK-NEXT: 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
+; CHECK: call token @llvm.coro.begin
   %hdl = call i8* @f()
 
 ; CHECK-NEXT: 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:
-  %hdl = call i8* @llvm.coro.begin(i8* null, i32 0, i8* null, i8* null)
+  %tok = call token @llvm.coro.begin(i8* null, i32 0, i8* null, i8* null)
+  %hdl = call i8* @llvm.coro.frame(token %tok)
 
 ; 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)
   %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
; CHECK: 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 i8* @llvm.coro.begin(i8*, i32, i8*, i8*)
+declare token @llvm.coro.begin(i8*, i32, i8*, i8*)
+declare i8* @llvm.coro.frame(token)
 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 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]
+
+; a coroutine start function
+define i8* @f() personality i8* null {
+entry:
+  %elide = call i8* @llvm.coro.alloc()
+  %need.dyn.alloc = icmp ne i8* %elide, null
+  br i1 %need.dyn.alloc, label %coro.begin, label %dyn.alloc
+dyn.alloc:
+  %alloc = call i8* @CustomAlloc(i32 4)
+  br label %coro.begin
+coro.begin:
+  %phi = phi i8* [ %elide, %entry ], [ %alloc, %dyn.alloc ]
+  %beg = call token @llvm.coro.begin(i8* %phi, i32 0, i8* null,
+                          i8* bitcast ([2 x void (%f.frame*)*]* @f.resumers to i8*))
+  %hdl = call i8* @llvm.coro.frame(token %beg)
+  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)
+  %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(  
+  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)
+  %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 i8* @llvm.coro.alloc()
+declare i8* @llvm.coro.free(i8*)
+declare token @llvm.coro.begin(i8*, i32, i8*, i8*)
+declare i8* @llvm.coro.frame(token)
+declare i8* @llvm.coro.subfn.addr(i8*, 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.begin(i8*, i32, i8*, i8*)
+
+; a coroutine start function
+define void @f() {
+  call token @llvm.coro.begin(i8* null, i32 0, i8* null, i8* null)
+  ret void
+}