Transforms/Coroutines/CoroShape.h

0001 //===- CoroShape.h - Coroutine info for lowering --------------*- C++ -*---===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 // This file declares the shape info struct that is required by many coroutine
0009 // utility methods.
0010 //===----------------------------------------------------------------------===//
0011
0012 #ifndef LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H
0013 #define LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H
0014
0015 #include "llvm/IR/IRBuilder.h"
0016 #include "llvm/IR/PassManager.h"
0017 #include "llvm/Transforms/Coroutines/CoroInstr.h"
0018
0019 namespace llvm {
0020
0021 class CallGraph;
0022
0023 namespace coro {
0024
0025 enum class ABI {
0026   /// The "resume-switch" lowering, where there are separate resume and
0027   /// destroy functions that are shared between all suspend points.  The
0028   /// coroutine frame implicitly stores the resume and destroy functions,
0029   /// the current index, and any promise value.
0030   Switch,
0031
0032   /// The "returned-continuation" lowering, where each suspend point creates a
0033   /// single continuation function that is used for both resuming and
0034   /// destroying.  Does not support promises.
0035   Retcon,
0036
0037   /// The "unique returned-continuation" lowering, where each suspend point
0038   /// creates a single continuation function that is used for both resuming
0039   /// and destroying.  Does not support promises.  The function is known to
0040   /// suspend at most once during its execution, and the return value of
0041   /// the continuation is void.
0042   RetconOnce,
0043
0044   /// The "async continuation" lowering, where each suspend point creates a
0045   /// single continuation function. The continuation function is available as an
0046   /// intrinsic.
0047   Async,
0048 };
0049
0050 // Holds structural Coroutine Intrinsics for a particular function and other
0051 // values used during CoroSplit pass.
0052 struct Shape {
0053   CoroBeginInst *CoroBegin = nullptr;
0054   SmallVector<AnyCoroEndInst *, 4> CoroEnds;
0055   SmallVector<CoroSizeInst *, 2> CoroSizes;
0056   SmallVector<CoroAlignInst *, 2> CoroAligns;
0057   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
0058   SmallVector<CoroAwaitSuspendInst *, 4> CoroAwaitSuspends;
0059   SmallVector<CallInst *, 2> SymmetricTransfers;
0060
0061   // Values invalidated by replaceSwiftErrorOps()
0062   SmallVector<CallInst *, 2> SwiftErrorOps;
0063
0064   void clear() {
0065     CoroBegin = nullptr;
0066     CoroEnds.clear();
0067     CoroSizes.clear();
0068     CoroAligns.clear();
0069     CoroSuspends.clear();
0070     CoroAwaitSuspends.clear();
0071     SymmetricTransfers.clear();
0072
0073     SwiftErrorOps.clear();
0074
0075     FrameTy = nullptr;
0076     FramePtr = nullptr;
0077     AllocaSpillBlock = nullptr;
0078   }
0079
0080   // Scan the function and collect the above intrinsics for later processing
0081   void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
0082                SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
0083   // If for some reason, we were not able to find coro.begin, bailout.
0084   void invalidateCoroutine(Function &F,
0085                            SmallVectorImpl<CoroFrameInst *> &CoroFrames);
0086   // Perform ABI related initial transformation
0087   void initABI();
0088   // Remove orphaned and unnecessary intrinsics
0089   void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
0090                       SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
0091
0092   // Field indexes for special fields in the switch lowering.
0093   struct SwitchFieldIndex {
0094     enum {
0095       Resume,
0096       Destroy
0097
0098       // The promise field is always at a fixed offset from the start of
0099       // frame given its type, but the index isn't a constant for all
0100       // possible frames.
0101
0102       // The switch-index field isn't at a fixed offset or index, either;
0103       // we just work it in where it fits best.
0104     };
0105   };
0106
0107   coro::ABI ABI;
0108
0109   StructType *FrameTy = nullptr;
0110   Align FrameAlign;
0111   uint64_t FrameSize = 0;
0112   Value *FramePtr = nullptr;
0113   BasicBlock *AllocaSpillBlock = nullptr;
0114
0115   struct SwitchLoweringStorage {
0116     SwitchInst *ResumeSwitch;
0117     AllocaInst *PromiseAlloca;
0118     BasicBlock *ResumeEntryBlock;
0119     unsigned IndexField;
0120     unsigned IndexAlign;
0121     unsigned IndexOffset;
0122     bool HasFinalSuspend;
0123     bool HasUnwindCoroEnd;
0124   };
0125
0126   struct RetconLoweringStorage {
0127     Function *ResumePrototype;
0128     Function *Alloc;
0129     Function *Dealloc;
0130     BasicBlock *ReturnBlock;
0131     bool IsFrameInlineInStorage;
0132   };
0133
0134   struct AsyncLoweringStorage {
0135     Value *Context;
0136     CallingConv::ID AsyncCC;
0137     unsigned ContextArgNo;
0138     uint64_t ContextHeaderSize;
0139     uint64_t ContextAlignment;
0140     uint64_t FrameOffset; // Start of the frame.
0141     uint64_t ContextSize; // Includes frame size.
0142     GlobalVariable *AsyncFuncPointer;
0143
0144     Align getContextAlignment() const { return Align(ContextAlignment); }
0145   };
0146
0147   union {
0148     SwitchLoweringStorage SwitchLowering;
0149     RetconLoweringStorage RetconLowering;
0150     AsyncLoweringStorage AsyncLowering;
0151   };
0152
0153   CoroIdInst *getSwitchCoroId() const {
0154     assert(ABI == coro::ABI::Switch);
0155     return cast<CoroIdInst>(CoroBegin->getId());
0156   }
0157
0158   AnyCoroIdRetconInst *getRetconCoroId() const {
0159     assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
0160     return cast<AnyCoroIdRetconInst>(CoroBegin->getId());
0161   }
0162
0163   CoroIdAsyncInst *getAsyncCoroId() const {
0164     assert(ABI == coro::ABI::Async);
0165     return cast<CoroIdAsyncInst>(CoroBegin->getId());
0166   }
0167
0168   unsigned getSwitchIndexField() const {
0169     assert(ABI == coro::ABI::Switch);
0170     assert(FrameTy && "frame type not assigned");
0171     return SwitchLowering.IndexField;
0172   }
0173   IntegerType *getIndexType() const {
0174     assert(ABI == coro::ABI::Switch);
0175     assert(FrameTy && "frame type not assigned");
0176     return cast<IntegerType>(FrameTy->getElementType(getSwitchIndexField()));
0177   }
0178   ConstantInt *getIndex(uint64_t Value) const {
0179     return ConstantInt::get(getIndexType(), Value);
0180   }
0181
0182   PointerType *getSwitchResumePointerType() const {
0183     assert(ABI == coro::ABI::Switch);
0184     assert(FrameTy && "frame type not assigned");
0185     return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
0186   }
0187
0188   FunctionType *getResumeFunctionType() const {
0189     switch (ABI) {
0190     case coro::ABI::Switch:
0191       return FunctionType::get(Type::getVoidTy(FrameTy->getContext()),
0192                                PointerType::getUnqual(FrameTy->getContext()),
0193                                /*IsVarArg=*/false);
0194     case coro::ABI::Retcon:
0195     case coro::ABI::RetconOnce:
0196       return RetconLowering.ResumePrototype->getFunctionType();
0197     case coro::ABI::Async:
0198       // Not used. The function type depends on the active suspend.
0199       return nullptr;
0200     }
0201
0202     llvm_unreachable("Unknown coro::ABI enum");
0203   }
0204
0205   ArrayRef<Type *> getRetconResultTypes() const {
0206     assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
0207     auto FTy = CoroBegin->getFunction()->getFunctionType();
0208
0209     // The safety of all this is checked by checkWFRetconPrototype.
0210     if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
0211       return STy->elements().slice(1);
0212     } else {
0213       return ArrayRef<Type *>();
0214     }
0215   }
0216
0217   ArrayRef<Type *> getRetconResumeTypes() const {
0218     assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
0219
0220     // The safety of all this is checked by checkWFRetconPrototype.
0221     auto FTy = RetconLowering.ResumePrototype->getFunctionType();
0222     return FTy->params().slice(1);
0223   }
0224
0225   CallingConv::ID getResumeFunctionCC() const {
0226     switch (ABI) {
0227     case coro::ABI::Switch:
0228       return CallingConv::Fast;
0229
0230     case coro::ABI::Retcon:
0231     case coro::ABI::RetconOnce:
0232       return RetconLowering.ResumePrototype->getCallingConv();
0233     case coro::ABI::Async:
0234       return AsyncLowering.AsyncCC;
0235     }
0236     llvm_unreachable("Unknown coro::ABI enum");
0237   }
0238
0239   AllocaInst *getPromiseAlloca() const {
0240     if (ABI == coro::ABI::Switch)
0241       return SwitchLowering.PromiseAlloca;
0242     return nullptr;
0243   }
0244
0245   BasicBlock::iterator getInsertPtAfterFramePtr() const {
0246     if (auto *I = dyn_cast<Instruction>(FramePtr)) {
0247       BasicBlock::iterator It = std::next(I->getIterator());
0248       It.setHeadBit(true); // Copy pre-RemoveDIs behaviour.
0249       return It;
0250     }
0251     return cast<Argument>(FramePtr)->getParent()->getEntryBlock().begin();
0252   }
0253
0254   /// Allocate memory according to the rules of the active lowering.
0255   ///
0256   /// \param CG - if non-null, will be updated for the new call
0257   Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;
0258
0259   /// Deallocate memory according to the rules of the active lowering.
0260   ///
0261   /// \param CG - if non-null, will be updated for the new call
0262   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
0263
0264   Shape() = default;
0265   explicit Shape(Function &F) {
0266     SmallVector<CoroFrameInst *, 8> CoroFrames;
0267     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
0268
0269     analyze(F, CoroFrames, UnusedCoroSaves);
0270     if (!CoroBegin) {
0271       invalidateCoroutine(F, CoroFrames);
0272       return;
0273     }
0274     cleanCoroutine(CoroFrames, UnusedCoroSaves);
0275   }
0276 };
0277
0278 } // end namespace coro
0279
0280 } // end namespace llvm
0281
0282 #endif // LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H