[Coroutines] Add `llvm.coro.align`, `llvm.coro.raw.frame.ptr.offset` and `llvm.coro.raw.frame.ptr.addr` intrinsics

[ychen]

Summary

llvm.coro.align returns the alignment of the coroutine frame in bytes.

llvm.coro.raw.frame.ptr.offset returns the byte offset of the raw memory block address (returned by the allocator) in the over-aligned coroutine frame. This is for switched-resume scheme only. The returned value is undefined if the coroutine frame is not over-aligned.

llvm.coro.raw.frame.ptr.addr returns the address storing the raw frame address. The returned address is either an alloca (before llvm.coro.begin) or a coroutine frame field (after llvm.coro.begin).

[ychen]

• Fix doc.

[ChuanqiXu]

The overall patch looks good to me. It seems like this patch did nothing without D97915. And it is different from normal NFC changes. I am not sure if it is more proper to merge this with D97915.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1059

It looks like Shape.SwitchLowering.FramePtrOffset is used but never set other than initializing.

[ychen]

In D102145#2747094, @ChuanqiXu wrote:

The overall patch looks good to me. It seems like this patch did nothing without D97915. And it is different from normal NFC changes. I am not sure if it is more proper to merge this with D97915.

It could be merged with D97915. I was afraid the patch size would be too big to make reviewing inconvenient. I'll add a standalone test for added intrinsics to make the patch self-ccontain.

[ychen]

• Add llvm.coro.raw.frame.ptr.addr to help storing raw frame address during allocation. llvm.coro.raw.frame.ptr.offset cannot be used during allocation because coroutine lowering would overwrite the raw frame address field with %alloc.frame.ptr alloca after llvm.coro.begin. The malloc returned address should be stored to %alloc.frame.ptr instead of the raw frame address field directly.
• Add tests for added intrinsics.

[ychen]

• Fix a test typo.

[ChuanqiXu]

llvm/docs/Coroutines.rst

964

If this is only valid under switch-resume coroutines, we need to mark them in the doc.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

796–825

If the new added intrinsic is intended for all APIs, we should move them out of the if blocks.

820

If we add a new alloca simply, it shouldn't work properly. For example,

void foo() {
     %a = alloca i32
     %b = alloca i32
     %promise = alloca PromiseType 
}

The frame we want may be:

%FrameType = type { [i8 x num of padding], resume_fn, destroy_fn, %promise, i32, i32, index}

However, if we simply add a new alloca with type i8*, we may get the FrameType with:

%FrameType = type { resume_fn, destroy_fn, %promise, i32, i32, i32*, index}

I don't think this could solve our problems. Or does it work only with the successor patches?

[ychen]

Thanks for the review!

llvm/docs/Coroutines.rst

964

It's general, i.e not specific to switch-resume coroutines.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

796–825

It is only supported for switched-resume coroutines. (mentioned in the doc)

820

Yes, following patch D97915 would do the overallocate&adjust frame start address. Here checks if (1) the frame is overaligned and (2) if llvm.coro.raw.frame.ptr.* intrinsics are used. If both are true, create an alloca that goes to the frame. coro-frame-overalign.ll test has the details.

[ChuanqiXu]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

796–825

hmmm, doesn't it conflict with the replies above? Or does it mean these intrinsics could be used generally by design but we just don't support them for other ABI now?

820

So does it mean this code would only work correctly if D97915 applies? It is a little bit weird for me. I understand that it would be clear er if we split a patch into separate ones. But it would be odd if these patches depends on each other.

[ychen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

820

So does it mean this code would only work correctly if D97915 applies?

No. This patch only lowers the new intrinsics under switched-resume ABI. Without D97915, there are no uses of these new intrinsics, so this patch basically no-op before D97915 applies.

[ChuanqiXu]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

820

It is still weird for me. It looks like these codes are redundant if we don't applies D97915. I am still prefer to merge this with D97915.

[ychen]

ping?

[ChuanqiXu]

It looks like that other reviewers are not available at this time. I am curious about the alignment problem for coroutine frame since I met confusing regression problem due to alignment mismatch in the past. Although the condition to trigger the alignment mismatch problem is not common (it requires the user to specify the alignment for a variable with alignment large than 16), I am desired to complete the (series) patch(es). So again, it looks better for me to merge this with D97915 to review and test. How do you think about it @ychen ?

[ychen]

In D102145#2809809, @ChuanqiXu wrote:

It looks like that other reviewers are not available at this time. I am curious about the alignment problem for coroutine frame since I met confusing regression problem due to alignment mismatch in the past. Although the condition to trigger the alignment mismatch problem is not common (it requires the user to specify the alignment for a variable with alignment large than 16), I am desired to complete the (series) patch(es). So again, it looks better for me to merge this with D97915 to review and test. How do you think about it @ychen ?

Thanks for taking another look. Sure, I'll update the patch tomorrow.

[ychen]

327141f was landed as an alternative.

llvm/docs/Coroutines.rst

 None
 
 Semantics:
 """"""""""
 
 The `coro.size` intrinsic is lowered to a constant representing the size of
 the coroutine frame. 
 
+.. _coro.align:
+
+'llvm.coro.align' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+    declare i32 @llvm.coro.align.i32()
+    declare i64 @llvm.coro.align.i64()
+
+Overview:
+"""""""""
+
+The '``llvm.coro.align``' intrinsic returns the alignment of the coroutine frame
+in bytes.
+
+Arguments:
+""""""""""
+
+None
+
+Semantics:
+""""""""""
+
+The `coro.align` intrinsic is lowered to a constant representing the alignment
+of the coroutine frame.
+
+.. _coro.raw.frame.ptr.offset:
+
+'llvm.coro.raw.frame.ptr.offset' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+    declare i32 @llvm.coro.raw.frame.ptr.offset.i32()
+    declare i64 @llvm.coro.raw.frame.ptr.offset.i64()
+
+Overview:
+"""""""""
+
+The '``llvm.coro.raw.frame.ptr.offset``' intrinsic returns the byte offset of
+the raw memory block address (returned by the allocator) in coroutine frame.
+This is only supported for switched-resume coroutines. The return value is
+undefined when the coroutine frame is not overaligned.
+
+Arguments:
+""""""""""
+
+None
+
+Semantics:
+""""""""""
+
+The `coro.raw.frame.ptr.offset` intrinsic is lowered to a constant representing
+the byte offset of the raw memory block address in coroutine frame.
+
+
+.. _coro.raw.frame.ptr.addr:
+
+'llvm.coro.raw.frame.ptr.addr' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+::
+
+    declare i8** @llvm.coro.raw.frame.ptr.addr()
+
+Overview:
+"""""""""
+
+The '``llvm.coro.raw.frame.ptr.addr``' intrinsic returns the address storing the raw
+frame address. The returned address is either an alloca or a coroutine frame
+field. This is only supported for switched-resume coroutines. The return value
+is undefined when the coroutine frame is not overaligned.
+
+Arguments:
+""""""""""
+
+None
+
+Semantics:
+""""""""""
+
+The `coro.raw.frame.ptr.offset.addr` intrinsic is lowered to either an alloca
+or a coroutine frame field storing the raw frame address.
+
 .. _coro.begin:
 
 'llvm.coro.begin' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare i8* @llvm.coro.begin(token <id>, i8* <mem>)
 

llvm/include/llvm/IR/Intrinsics.td

def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
                                NoCapture<ArgIndex<1>>]>;
 def int_coro_end : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty], []>;
 def int_coro_end_async
     : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty, llvm_vararg_ty], []>;
 
 def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_noop : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
+def int_coro_align : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
+def int_coro_raw_frame_ptr_offset : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
+def int_coro_raw_frame_ptr_addr : Intrinsic<[llvm_ptrptr_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_suspend_retcon : Intrinsic<[llvm_any_ty], [llvm_vararg_ty], []>;
 def int_coro_prepare_retcon : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
                                         [IntrNoMem]>;
 def int_coro_alloca_alloc : Intrinsic<[llvm_token_ty],
                                       [llvm_anyint_ty, llvm_i32_ty], []>;

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

 
 #include "CoroInternal.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Analysis/PtrUseVisitor.h"
 #include "llvm/Analysis/StackLifetime.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/CFG.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/OptimizedStructLayout.h"
uint64_t getStructSize() const {
     return StructSize;
   }
 
   Align getStructAlign() const {
     assert(IsFinished && "not yet finished!");
     return StructAlign;
   }
 
+  SmallVector<Field, 8> &getFields() { return Fields; }
+
   FieldIDType getLayoutFieldIndex(FieldIDType Id) const {
     assert(IsFinished && "not yet finished!");
     return Fields[Id].LayoutFieldIndex;
   }
 };
 } // namespace
 
 void FrameDataInfo::updateLayoutIndex(FrameTypeBuilder &B) {
if (Shape.ABI == coro::ABI::Switch) {
     SwitchIndexFieldId = B.addField(IndexType, None);
   } else {
     assert(PromiseAlloca == nullptr && "lowering doesn't support promises");
   }
 
   // Because multiple allocas may own the same field slot,
   // we add allocas to field here.
   B.addFieldForAllocas(F, FrameData, Shape);
+
+  // Create an entry for every spilled value.
+  for (auto &S : FrameData.Spills) {
+    FieldIDType Id = B.addField(S.first->getType(), None);
+    FrameData.setFieldIndex(S.first, Id);
+  }
+
+  Optional<FieldIDType> FramePtrField = None;
+  if (Shape.ABI == coro::ABI::Switch) {
     // Add PromiseAlloca to Allocas list so that
     // 1. updateLayoutIndex could update its index after
     // `performOptimizedStructLayout`
     // 2. it is processed in insertSpills.
-  if (Shape.ABI == coro::ABI::Switch && PromiseAlloca)
+    if (PromiseAlloca)
       // We assume that the promise alloca won't be modified before
       // CoroBegin and no alias will be create before CoroBegin.
       FrameData.Allocas.emplace_back(
-        PromiseAlloca, DenseMap<Instruction *, llvm::Optional<APInt>>{}, false);
-  // Create an entry for every spilled value.
-  for (auto &S : FrameData.Spills) {
-    FieldIDType Id = B.addField(S.first->getType(), None);
-    FrameData.setFieldIndex(S.first, Id);
+          PromiseAlloca, DenseMap<Instruction *, llvm::Optional<APInt>>{},
+          false);
+
+    Align FrameAlign =
+        std::max_element(
+            B.getFields().begin(), B.getFields().end(),
+            [](auto &F1, auto &F2) { return F1.Alignment < F2.Alignment; })
+            ->Alignment;
+
+    // Check for over-alignment.
+    Value *PtrAddr =
+        ConstantPointerNull::get(Type::getInt8PtrTy(C)->getPointerTo());
+    unsigned NewAlign = Shape.getSwitchCoroId()->getAlignment();
+    bool NeedFramePtrField = Shape.CoroRawFramePtrOffsets.size() > 0 ||
+                             Shape.CoroRawFramePtrAddrs.size() > 0;
+    if (NeedFramePtrField && NewAlign && FrameAlign > NewAlign) {
+      BasicBlock &Entry = F.getEntryBlock();
+      IRBuilder<> Builder(&Entry, Entry.getFirstInsertionPt());
+
+      // Reserve frame space for raw frame pointer.
+      Value *Mem = Shape.CoroBegin->getMem();
+      AllocaInst *FramePtrAddr =
+          Builder.CreateAlloca(Mem->getType(), nullptr, "alloc.frame.ptr");
+      PtrAddr = FramePtrAddr;
+      FramePtrField = B.addFieldForAlloca(FramePtrAddr);
+      FrameData.setFieldIndex(FramePtrAddr, *FramePtrField);
+      FrameData.Allocas.emplace_back(
+          FramePtrAddr, DenseMap<Instruction *, llvm::Optional<APInt>>{}, true);
+    }
+
+    for (CoroRawFramePtrAddrInst *C : Shape.CoroRawFramePtrAddrs) {
+      C->replaceAllUsesWith(PtrAddr);
+      C->eraseFromParent();
+    }
   }
 
   B.finish(FrameTy);
   FrameData.updateLayoutIndex(B);
   Shape.FrameAlign = B.getStructAlign();
   Shape.FrameSize = B.getStructSize();
 
   switch (Shape.ABI) {
   case coro::ABI::Switch:
     // In the switch ABI, remember the switch-index field.
     Shape.SwitchLowering.IndexField =
         B.getLayoutFieldIndex(*SwitchIndexFieldId);
 
+    if (FramePtrField) {
+      FieldIDType FieldIdx = B.getLayoutFieldIndex(*FramePtrField);
+      Shape.SwitchLowering.FramePtrOffset =
+          DL.getStructLayout(FrameTy)->getElementOffset(FieldIdx);
+    }
+
     // Also round the frame size up to a multiple of its alignment, as is
     // generally expected in C/C++.
     Shape.FrameSize = alignTo(Shape.FrameSize, Shape.FrameAlign);
     break;
 
   // In the retcon ABI, remember whether the frame is inline in the storage.
   case coro::ABI::Retcon:
   case coro::ABI::RetconOnce: {

llvm/lib/Transforms/Coroutines/CoroInstr.h

 // the Coroutine library.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H
 #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H
 
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 
 /// This class represents the llvm.coro.subfn.addr instruction.
 class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst {
   enum { FrameArg, IndexArg };
 

 public:
   AllocaInst *getPromise() const {
     Value *Arg = getArgOperand(PromiseArg);
     return isa<ConstantPointerNull>(Arg)
                ? nullptr
                : cast<AllocaInst>(Arg->stripPointerCasts());
   }
 
+  unsigned getAlignment() const {
+    return cast<ConstantInt>(getArgOperand(AlignArg))->getZExtValue();
+  }
+
   void clearPromise() {
     Value *Arg = getArgOperand(PromiseArg);
     setArgOperand(PromiseArg,
                   ConstantPointerNull::get(Type::getInt8PtrTy(getContext())));
     if (isa<AllocaInst>(Arg))
       return;
     assert((isa<BitCastInst>(Arg) || isa<GetElementPtrInst>(Arg)) &&
            "unexpected instruction designating the promise");
public:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_size;
   }
   static bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
 };
 
+/// This represents the llvm.coro.align instruction.
+class LLVM_LIBRARY_VISIBILITY CoroAlignInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_align;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.raw.frame.ptr.offset instruction.
+class LLVM_LIBRARY_VISIBILITY CoroRawFramePtrOffsetInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_raw_frame_ptr_offset;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.raw.frame.ptr.addr instruction.
+class LLVM_LIBRARY_VISIBILITY CoroRawFramePtrAddrInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_raw_frame_ptr_addr;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 class LLVM_LIBRARY_VISIBILITY AnyCoroEndInst : public IntrinsicInst {
   enum { FrameArg, UnwindArg };
 
 public:
   bool isFallthrough() const { return !isUnwind(); }
   bool isUnwind() const {
     return cast<Constant>(getArgOperand(UnwindArg))->isOneValue();
   }

llvm/lib/Transforms/Coroutines/CoroInternal.h

 };
 
 // Holds structural Coroutine Intrinsics for a particular function and other
 // values used during CoroSplit pass.
 struct LLVM_LIBRARY_VISIBILITY Shape {
   CoroBeginInst *CoroBegin;
   SmallVector<AnyCoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
+  SmallVector<CoroAlignInst *, 2> CoroAligns;
+  SmallVector<CoroRawFramePtrOffsetInst *, 2> CoroRawFramePtrOffsets;
+  SmallVector<CoroRawFramePtrAddrInst *, 2> CoroRawFramePtrAddrs;
   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
   SmallVector<CallInst*, 2> SwiftErrorOps;
 
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
     enum {
       Resume,
       Destroy
struct LLVM_LIBRARY_VISIBILITY Shape {
   /// This would only be true if optimization are enabled.
   bool ReuseFrameSlot;
 
   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
     BasicBlock *ResumeEntryBlock;
     unsigned IndexField;
+    unsigned FramePtrOffset;
     bool HasFinalSuspend;
   };
 
   struct RetconLoweringStorage {
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
     BasicBlock *ReturnBlock;

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
   auto *NewContextSize = ConstantInt::get(OrigContextSize->getType(),
                                           Shape.AsyncLowering.ContextSize);
   auto *NewFuncPtrStruct = ConstantStruct::get(
       FuncPtrStruct->getType(), OrigRelativeFunOffset, NewContextSize);
 
   Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);
 }
 
-static void replaceFrameSize(coro::Shape &Shape) {
+static void replaceFrameSizeAndAlign(coro::Shape &Shape) {
   if (Shape.ABI == coro::ABI::Async)
     updateAsyncFuncPointerContextSize(Shape);
 
-  if (Shape.CoroSizes.empty())
-    return;
-
+  if (!Shape.CoroSizes.empty()) {
     // In the same function all coro.sizes should have the same result type.
     auto *SizeIntrin = Shape.CoroSizes.back();
     Module *M = SizeIntrin->getModule();
     const DataLayout &DL = M->getDataLayout();
     auto Size = DL.getTypeAllocSize(Shape.FrameTy);
     auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
 
     for (CoroSizeInst *CS : Shape.CoroSizes) {
       CS->replaceAllUsesWith(SizeConstant);
       CS->eraseFromParent();
     }
   }
 
+  if (!Shape.CoroAligns.empty()) {
+    auto *Intrin = Shape.CoroAligns.back();
+    auto *AlignConstant =
+        ConstantInt::get(Intrin->getType(), Shape.FrameAlign.value());
+
+    for (CoroAlignInst *CS : Shape.CoroAligns) {
+      CS->replaceAllUsesWith(AlignConstant);
+      CS->eraseFromParent();
+    }
+  }
+
+  if (!Shape.CoroRawFramePtrOffsets.empty()) {
+    auto *Intrin = Shape.CoroRawFramePtrOffsets.back();
+    auto *FramePtrOffset = ConstantInt::get(
+        Intrin->getType(), Shape.SwitchLowering.FramePtrOffset);
+
+    for (CoroRawFramePtrOffsetInst *CS : Shape.CoroRawFramePtrOffsets) {
+      CS->replaceAllUsesWith(FramePtrOffset);
+      CS->eraseFromParent();
+    }
+  }
+}
+
 // Create a global constant array containing pointers to functions provided and
 // set Info parameter of CoroBegin to point at this constant. Example:
 //
 //   @f.resumers = internal constant [2 x void(%f.frame*)*]
 //                    [void(%f.frame*)* @f.resume, void(%f.frame*)* @f.destroy]
 //   define void @f() {
 //     ...
 //     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
static coro::Shape splitCoroutine(Function &F,
   removeUnreachableBlocks(F);
 
   coro::Shape Shape(F, ReuseFrameSlot);
   if (!Shape.CoroBegin)
     return Shape;
 
   simplifySuspendPoints(Shape);
   buildCoroutineFrame(F, Shape);
-  replaceFrameSize(Shape);
+  replaceFrameSizeAndAlign(Shape);
 
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
     handleNoSuspendCoroutine(Shape);
   } else {
     switch (Shape.ABI) {
     case coro::ABI::Switch:

llvm/lib/Transforms/Coroutines/Coroutines.cpp

void coro::updateCallGraph(Function &ParentFunc, ArrayRef<Function *> NewFuncs,
 
   SCC.initialize(Nodes);
 }
 
 static void clear(coro::Shape &Shape) {
   Shape.CoroBegin = nullptr;
   Shape.CoroEnds.clear();
   Shape.CoroSizes.clear();
+  Shape.CoroAligns.clear();
+  Shape.CoroRawFramePtrOffsets.clear();
+  Shape.CoroRawFramePtrAddrs.clear();
   Shape.CoroSuspends.clear();
 
   Shape.FrameTy = nullptr;
   Shape.FramePtr = nullptr;
   Shape.AllocaSpillBlock = nullptr;
 }
 
 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
void coro::Shape::buildFrom(Function &F) {
   for (Instruction &I : instructions(F)) {
     if (auto II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
       case Intrinsic::coro_size:
         CoroSizes.push_back(cast<CoroSizeInst>(II));
         break;
+      case Intrinsic::coro_align:
+        CoroAligns.push_back(cast<CoroAlignInst>(II));
+        break;
+      case Intrinsic::coro_raw_frame_ptr_offset:
+        CoroRawFramePtrOffsets.push_back(cast<CoroRawFramePtrOffsetInst>(II));
+        break;
+      case Intrinsic::coro_raw_frame_ptr_addr:
+        CoroRawFramePtrAddrs.push_back(cast<CoroRawFramePtrAddrInst>(II));
+        break;
       case Intrinsic::coro_frame:
         CoroFrames.push_back(cast<CoroFrameInst>(II));
         break;
       case Intrinsic::coro_save:
         // After optimizations, coro_suspends using this coro_save might have
         // been removed, remember orphaned coro_saves to remove them later.
         if (II->use_empty())
           UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));
void coro::Shape::buildFrom(Function &F) {
   switch (auto IdIntrinsic = Id->getIntrinsicID()) {
   case Intrinsic::coro_id: {
     auto SwitchId = cast<CoroIdInst>(Id);
     this->ABI = coro::ABI::Switch;
     this->SwitchLowering.HasFinalSuspend = HasFinalSuspend;
     this->SwitchLowering.ResumeSwitch = nullptr;
     this->SwitchLowering.PromiseAlloca = SwitchId->getPromise();
     this->SwitchLowering.ResumeEntryBlock = nullptr;
+    this->SwitchLowering.FramePtrOffset = 0;
 
     for (auto AnySuspend : CoroSuspends) {
       auto Suspend = dyn_cast<CoroSuspendInst>(AnySuspend);
       if (!Suspend) {
 #ifndef NDEBUG
         AnySuspend->dump();
 #endif
         report_fatal_error("coro.id must be paired with coro.suspend");

llvm/test/Transforms/Coroutines/coro-frame-overalign.ll

+; Check that `llvm.coro.align`, `llvm.coro.raw.frame.ptr.offset` and
+; `@llvm.coro.raw.frame.ptr.alloca` are lowered correctly.
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+%PackedStruct = type <{ i64 }>
+
+declare void @consume(%PackedStruct*, i32, i32, i8**)
+declare void @consume2(i32, i32)
+
+define i8* @f() "coroutine.presplit"="1" {
+entry:
+  %data = alloca %PackedStruct, align 32
+  %id = call token @llvm.coro.id(i32 16, 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)
+  %align = call i32 @llvm.coro.align.i32()
+  %offset = call i32 @llvm.coro.raw.frame.ptr.offset.i32()
+  %addr = call i8** @llvm.coro.raw.frame.ptr.addr()
+  call void @consume(%PackedStruct* %data, i32 %align, i32 %offset, i8** %addr)
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  br label %cleanup
+
+cleanup:
+  %align2 = call i32 @llvm.coro.align.i32()
+  %offset2 = call i32 @llvm.coro.raw.frame.ptr.offset.i32()
+  call void @consume2(i32 %align2, i32 %offset2)
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; See if the raw frame address was inserted into the frame.
+; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i8*, i1, [7 x i8], %PackedStruct }
+
+; See if we used correct index to access frame addr field (field 2).
+; CHECK-LABEL: @f(
+; CHECK: %alloc.frame.ptr = alloca i8*, align 8
+; CHECK: %[[FIELD:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2
+; CHECK: %[[ADDR:.+]] = load i8*, i8** %alloc.frame.ptr, align 8
+; CHECK: store i8* %[[ADDR]], i8** %[[FIELD]], align 8
+; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
+; CHECK: call void @consume(%PackedStruct* %[[DATA]], i32 32, i32 16, i8** %[[FIELD]])
+; CHECK: ret i8*
+
+; See if `llvm.coro.align` and `llvm.coro.raw.frame.ptr.offset` are lowered
+; correctly during deallocation.
+; CHECK-LABEL: @f.destroy(
+; CHECK: call void @consume2(i32 32, i32 16)
+; CHECK: call void @free(i8* %{{.*}})
+
+; CHECK-LABEL: @f.cleanup(
+; CHECK: call void @consume2(i32 32, i32 16)
+; CHECK: call void @free(i8*
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i32 @llvm.coro.align.i32()
+declare i32 @llvm.coro.raw.frame.ptr.offset.i32()
+declare i8** @llvm.coro.raw.frame.ptr.addr()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare double @print(double)
+declare void @free(i8*)