CoroSplit.cpp | Explore in Territory

//===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// This pass builds the coroutine frame and outlines resume and destroy parts
// of the coroutine into separate functions.
//
// We present a coroutine to an LLVM as an ordinary function with suspension
// points marked up with intrinsics. We let the optimizer party on the coroutine
// as a single function for as long as possible. Shortly before the coroutine is
// eligible to be inlined into its callers, we split up the coroutine into parts
// corresponding to an initial, resume and destroy invocations of the coroutine,
// add them to the current SCC and restart the IPO pipeline to optimize the
// coroutine subfunctions we extracted before proceeding to the caller of the
// coroutine.
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Coroutines/CoroSplit.h"
#include "CoroInternal.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/PriorityWorklist.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Analysis/CFG.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/Analysis/LazyCallGraph.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/IR/Argument.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/CallingConv.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Coroutines/ABI.h"
#include "llvm/Transforms/Coroutines/CoroInstr.h"
#include "llvm/Transforms/Coroutines/MaterializationUtils.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/CallGraphUpdater.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <initializer_list>
#include <iterator>

usingnamespacellvm;

#define DEBUG_TYPE …

namespace {

/// A little helper class for building
class CoroCloner { … };

} // end anonymous namespace

// FIXME:
// Lower the intrinisc in CoroEarly phase if coroutine frame doesn't escape
// and it is known that other transformations, for example, sanitizers
// won't lead to incorrect code.
static void lowerAwaitSuspend(IRBuilder<> &Builder, CoroAwaitSuspendInst *CB,
                              coro::Shape &Shape) { … }

static void lowerAwaitSuspends(Function &F, coro::Shape &Shape) { … }

static void maybeFreeRetconStorage(IRBuilder<> &Builder,
                                   const coro::Shape &Shape, Value *FramePtr,
                                   CallGraph *CG) { … }

/// Replace an llvm.coro.end.async.
/// Will inline the must tail call function call if there is one.
/// \returns true if cleanup of the coro.end block is needed, false otherwise.
static bool replaceCoroEndAsync(AnyCoroEndInst *End) { … }

/// Replace a non-unwind call to llvm.coro.end.
static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
                                      const coro::Shape &Shape, Value *FramePtr,
                                      bool InResume, CallGraph *CG) { … }

// Mark a coroutine as done, which implies that the coroutine is finished and
// never get resumed.
//
// In resume-switched ABI, the done state is represented by storing zero in
// ResumeFnAddr.
//
// NOTE: We couldn't omit the argument `FramePtr`. It is necessary because the
// pointer to the frame in splitted function is not stored in `Shape`.
static void markCoroutineAsDone(IRBuilder<> &Builder, const coro::Shape &Shape,
                                Value *FramePtr) { … }

/// Replace an unwind call to llvm.coro.end.
static void replaceUnwindCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                                 Value *FramePtr, bool InResume,
                                 CallGraph *CG) { … }

static void replaceCoroEnd(AnyCoroEndInst *End, const coro::Shape &Shape,
                           Value *FramePtr, bool InResume, CallGraph *CG) { … }

// In the resume function, we remove the last case  (when coro::Shape is built,
// the final suspend point (if present) is always the last element of
// CoroSuspends array) since it is an undefined behavior to resume a coroutine
// suspended at the final suspend point.
// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL
// and the coroutine doesn't suspend at the final suspend point actually (this
// is possible since the coroutine is considered suspended at the final suspend
// point if promise.unhandled_exception() exits via an exception), we can
// remove the last case.
void CoroCloner::handleFinalSuspend() { … }

static FunctionType *
getFunctionTypeFromAsyncSuspend(AnyCoroSuspendInst *Suspend) { … }

static Function *createCloneDeclaration(Function &OrigF, coro::Shape &Shape,
                                        const Twine &Suffix,
                                        Module::iterator InsertBefore,
                                        AnyCoroSuspendInst *ActiveSuspend) { … }

/// Replace uses of the active llvm.coro.suspend.retcon/async call with the
/// arguments to the continuation function.
///
/// This assumes that the builder has a meaningful insertion point.
void CoroCloner::replaceRetconOrAsyncSuspendUses() { … }

void CoroCloner::replaceCoroSuspends() { … }

void CoroCloner::replaceCoroEnds() { … }

static void replaceSwiftErrorOps(Function &F, coro::Shape &Shape,
                                 ValueToValueMapTy *VMap) { … }

/// Returns all DbgVariableIntrinsic in F.
static std::pair<SmallVector<DbgVariableIntrinsic *, 8>,
                 SmallVector<DbgVariableRecord *>>
collectDbgVariableIntrinsics(Function &F) { … }

void CoroCloner::replaceSwiftErrorOps() { … }

void CoroCloner::salvageDebugInfo() { … }

void CoroCloner::replaceEntryBlock() { … }

/// Derive the value of the new frame pointer.
Value *CoroCloner::deriveNewFramePointer() { … }

/// Adjust the scope line of the funclet to the first line number after the
/// suspend point. This avoids a jump in the line table from the function
/// declaration (where prologue instructions are attributed to) to the suspend
/// point.
/// Only adjust the scope line when the files are the same.
/// If no candidate line number is found, fallback to the line of ActiveSuspend.
static void updateScopeLine(Instruction *ActiveSuspend,
                            DISubprogram &SPToUpdate) { … }

static void addFramePointerAttrs(AttributeList &Attrs, LLVMContext &Context,
                                 unsigned ParamIndex, uint64_t Size,
                                 Align Alignment, bool NoAlias) { … }

static void addAsyncContextAttrs(AttributeList &Attrs, LLVMContext &Context,
                                 unsigned ParamIndex) { … }

static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context,
                              unsigned ParamIndex) { … }

/// Clone the body of the original function into a resume function of
/// some sort.
void CoroCloner::create() { … }

static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) { … }

static TypeSize getFrameSizeForShape(coro::Shape &Shape) { … }

static void replaceFrameSizeAndAlignment(coro::Shape &Shape) { … }

static void postSplitCleanup(Function &F) { … }

// Coroutine has no suspend points. Remove heap allocation for the coroutine
// frame if possible.
static void handleNoSuspendCoroutine(coro::Shape &Shape) { … }

// SimplifySuspendPoint needs to check that there is no calls between
// coro_save and coro_suspend, since any of the calls may potentially resume
// the coroutine and if that is the case we cannot eliminate the suspend point.
static bool hasCallsInBlockBetween(Instruction *From, Instruction *To) { … }

static bool hasCallsInBlocksBetween(BasicBlock *SaveBB, BasicBlock *ResDesBB) { … }

static bool hasCallsBetween(Instruction *Save, Instruction *ResumeOrDestroy) { … }

// If a SuspendIntrin is preceded by Resume or Destroy, we can eliminate the
// suspend point and replace it with nornal control flow.
static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
                                 CoroBeginInst *CoroBegin) { … }

// Remove suspend points that are simplified.
static void simplifySuspendPoints(coro::Shape &Shape) { … }

namespace {

struct SwitchCoroutineSplitter { … };

} // namespace

static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,
                                       Value *Continuation) { … }

/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.
static void coerceArguments(IRBuilder<> &Builder, FunctionType *FnTy,
                            ArrayRef<Value *> FnArgs,
                            SmallVectorImpl<Value *> &CallArgs) { … }

CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                                   TargetTransformInfo &TTI,
                                   ArrayRef<Value *> Arguments,
                                   IRBuilder<> &Builder) { … }

void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                    SmallVectorImpl<Function *> &Clones,
                                    TargetTransformInfo &TTI) { … }

void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                        SmallVectorImpl<Function *> &Clones,
                                        TargetTransformInfo &TTI) { … }

namespace {
class PrettyStackTraceFunction : public PrettyStackTraceEntry { … };
} // namespace

/// Remove calls to llvm.coro.end in the original function.
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) { … }

static bool hasSafeElideCaller(Function &F) { … }

void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
                                     SmallVectorImpl<Function *> &Clones,
                                     TargetTransformInfo &TTI) { … }

static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
                             coro::BaseABI &ABI, TargetTransformInfo &TTI,
                             bool OptimizeFrame) { … }

static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
    LazyCallGraph::Node &N, const coro::Shape &Shape,
    const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
    LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
    FunctionAnalysisManager &FAM) { … }

/// Replace a call to llvm.coro.prepare.retcon.
static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,
                           LazyCallGraph::SCC &C) { … }

static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,
                               LazyCallGraph::SCC &C) { … }

static void addPrepareFunction(const Module &M,
                               SmallVectorImpl<Function *> &Fns,
                               StringRef Name) { … }

static std::unique_ptr<coro::BaseABI>
CreateNewABI(Function &F, coro::Shape &S,
             std::function<bool(Instruction &)> IsMatCallback,
             const SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs) { … }

CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
    : … { … }

CoroSplitPass::CoroSplitPass(
    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
    : … { … }

// For back compatibility, constructor takes a materializable callback and
// creates a generator for an ABI with a modified materializable callback.
CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,
                             bool OptimizeFrame)
    : … { … }

// For back compatibility, constructor takes a materializable callback and
// creates a generator for an ABI with a modified materializable callback.
CoroSplitPass::CoroSplitPass(
    std::function<bool(Instruction &)> IsMatCallback,
    SmallVector<CoroSplitPass::BaseABITy> GenCustomABIs, bool OptimizeFrame)
    : … { … }

PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
                                     CGSCCAnalysisManager &AM,
                                     LazyCallGraph &CG, CGSCCUpdateResult &UR) { … }
llvm/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
