diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h index a2be1099ff68fc..944ce602d4c108 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h @@ -21,19 +21,26 @@ namespace llvm { +namespace coro { +class BaseABI; +class Shape; +} // namespace coro + struct CoroSplitPass : PassInfoMixin { - const std::function MaterializableCallback; + // BaseABITy generates an instance of a coro ABI. + using BaseABITy = std::function; CoroSplitPass(bool OptimizeFrame = false); CoroSplitPass(std::function MaterializableCallback, - bool OptimizeFrame = false) - : MaterializableCallback(MaterializableCallback), - OptimizeFrame(OptimizeFrame) {} + bool OptimizeFrame = false); PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM, LazyCallGraph &CG, CGSCCUpdateResult &UR); static bool isRequired() { return true; } + // Generator for an ABI transformer + BaseABITy CreateAndInitABI; + // Would be true if the Optimization level isn't O0. bool OptimizeFrame; }; diff --git a/llvm/lib/Transforms/Coroutines/ABI.h b/llvm/lib/Transforms/Coroutines/ABI.h new file mode 100644 index 00000000000000..20ef06f50c2faa --- /dev/null +++ b/llvm/lib/Transforms/Coroutines/ABI.h @@ -0,0 +1,109 @@ +//===- ABI.h - Coroutine ABI Transformers ---------------------*- C++ -*---===// +// +// 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 file declares the pass that analyzes a function for coroutine intrs and +// a transformer class that contains methods for handling different steps of +// coroutine lowering. +//===----------------------------------------------------------------------===// + +#ifndef LIB_TRANSFORMS_COROUTINES_ABI_H +#define LIB_TRANSFORMS_COROUTINES_ABI_H + +#include "CoroShape.h" +#include "MaterializationUtils.h" +#include "SuspendCrossingInfo.h" +#include "llvm/Analysis/TargetTransformInfo.h" + +namespace llvm { + +class Function; + +namespace coro { + +// This interface/API is to provide an object oriented way to implement ABI +// functionality. This is intended to replace use of the ABI enum to perform +// ABI operations. The ABIs (e.g. Switch, Async, Retcon{Once}) are the common +// ABIs. + +class LLVM_LIBRARY_VISIBILITY BaseABI { +public: + BaseABI(Function &F, Shape &S) + : F(F), Shape(S), IsMaterializable(coro::isTriviallyMaterializable) {} + + BaseABI(Function &F, coro::Shape &S, + std::function IsMaterializable) + : F(F), Shape(S), IsMaterializable(IsMaterializable) {} + + // Initialize the coroutine ABI + virtual void init() = 0; + + // Allocate the coroutine frame and do spill/reload as needed. + virtual void buildCoroutineFrame(); + + // Perform the function splitting according to the ABI. + virtual void splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) = 0; + + Function &F; + coro::Shape &Shape; + + // Callback used by coro::BaseABI::buildCoroutineFrame for rematerialization. + // It is provided to coro::doMaterializations(..). + std::function IsMaterializable; +}; + +class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI { +public: + SwitchABI(Function &F, coro::Shape &S) : BaseABI(F, S) {} + + SwitchABI(Function &F, coro::Shape &S, + std::function IsMaterializable) + : BaseABI(F, S, IsMaterializable) {} + + void init() override; + + void splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) override; +}; + +class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI { +public: + AsyncABI(Function &F, coro::Shape &S) : BaseABI(F, S) {} + + AsyncABI(Function &F, coro::Shape &S, + std::function IsMaterializable) + : BaseABI(F, S, IsMaterializable) {} + + void init() override; + + void splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) override; +}; + +class LLVM_LIBRARY_VISIBILITY AnyRetconABI : public BaseABI { +public: + AnyRetconABI(Function &F, coro::Shape &S) : BaseABI(F, S) {} + + AnyRetconABI(Function &F, coro::Shape &S, + std::function IsMaterializable) + : BaseABI(F, S, IsMaterializable) {} + + void init() override; + + void splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) override; +}; + +} // end namespace coro + +} // end namespace llvm + +#endif // LLVM_TRANSFORMS_COROUTINES_ABI_H diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp index c08f56b024dfc7..021b1f7a4156b9 100644 --- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp @@ -15,6 +15,7 @@ // the value into the coroutine frame. //===----------------------------------------------------------------------===// +#include "ABI.h" #include "CoroInternal.h" #include "MaterializationUtils.h" #include "SpillUtils.h" @@ -2055,11 +2056,9 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape, rewritePHIs(F); } -void coro::buildCoroutineFrame( - Function &F, Shape &Shape, - const std::function &MaterializableCallback) { +void coro::BaseABI::buildCoroutineFrame() { SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds); - doRematerializations(F, Checker, MaterializableCallback); + doRematerializations(F, Checker, IsMaterializable); const DominatorTree DT(F); if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon && diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h index fcbd31878bdea7..88d0f83c98c9ec 100644 --- a/llvm/lib/Transforms/Coroutines/CoroInternal.h +++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h @@ -62,9 +62,6 @@ struct LowererBase { bool defaultMaterializable(Instruction &V); void normalizeCoroutine(Function &F, coro::Shape &Shape, TargetTransformInfo &TTI); -void buildCoroutineFrame( - Function &F, Shape &Shape, - const std::function &MaterializableCallback); CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn, TargetTransformInfo &TTI, ArrayRef Arguments, IRBuilder<> &); diff --git a/llvm/lib/Transforms/Coroutines/CoroShape.h b/llvm/lib/Transforms/Coroutines/CoroShape.h index dcfe94ca076bd8..f4fb4baa6df314 100644 --- a/llvm/lib/Transforms/Coroutines/CoroShape.h +++ b/llvm/lib/Transforms/Coroutines/CoroShape.h @@ -275,7 +275,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape { invalidateCoroutine(F, CoroFrames); return; } - initABI(); cleanCoroutine(CoroFrames, UnusedCoroSaves); } }; diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index 382bdfff1926f7..a17e16cc8e2233 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -19,6 +19,7 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Coroutines/CoroSplit.h" +#include "ABI.h" #include "CoroInstr.h" #include "CoroInternal.h" #include "llvm/ADT/DenseMap.h" @@ -1779,9 +1780,9 @@ CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn, return TailCall; } -static void splitAsyncCoroutine(Function &F, coro::Shape &Shape, - SmallVectorImpl &Clones, - TargetTransformInfo &TTI) { +void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) { assert(Shape.ABI == coro::ABI::Async); assert(Clones.empty()); // Reset various things that the optimizer might have decided it @@ -1874,9 +1875,9 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape, } } -static void splitRetconCoroutine(Function &F, coro::Shape &Shape, - SmallVectorImpl &Clones, - TargetTransformInfo &TTI) { +void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) { assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce); assert(Clones.empty()); @@ -2044,26 +2045,27 @@ static bool hasSafeElideCaller(Function &F) { return false; } -static coro::Shape -splitCoroutine(Function &F, SmallVectorImpl &Clones, - TargetTransformInfo &TTI, bool OptimizeFrame, - std::function MaterializableCallback) { - PrettyStackTraceFunction prettyStackTrace(F); +void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape, + SmallVectorImpl &Clones, + TargetTransformInfo &TTI) { + SwitchCoroutineSplitter::split(F, Shape, Clones, TTI); +} - // The suspend-crossing algorithm in buildCoroutineFrame get tripped - // up by uses in unreachable blocks, so remove them as a first pass. - removeUnreachableBlocks(F); +static void doSplitCoroutine(Function &F, SmallVectorImpl &Clones, + coro::BaseABI &ABI, TargetTransformInfo &TTI) { + PrettyStackTraceFunction prettyStackTrace(F); - coro::Shape Shape(F, OptimizeFrame); - if (!Shape.CoroBegin) - return Shape; + auto &Shape = ABI.Shape; + assert(Shape.CoroBegin); lowerAwaitSuspends(F, Shape); simplifySuspendPoints(Shape); + normalizeCoroutine(F, Shape, TTI); - buildCoroutineFrame(F, Shape, MaterializableCallback); + ABI.buildCoroutineFrame(); replaceFrameSizeAndAlignment(Shape); + bool isNoSuspendCoroutine = Shape.CoroSuspends.empty(); bool shouldCreateNoAllocVariant = !isNoSuspendCoroutine && @@ -2075,18 +2077,7 @@ splitCoroutine(Function &F, SmallVectorImpl &Clones, if (isNoSuspendCoroutine) { handleNoSuspendCoroutine(Shape); } else { - switch (Shape.ABI) { - case coro::ABI::Switch: - SwitchCoroutineSplitter::split(F, Shape, Clones, TTI); - break; - case coro::ABI::Async: - splitAsyncCoroutine(F, Shape, Clones, TTI); - break; - case coro::ABI::Retcon: - case coro::ABI::RetconOnce: - splitRetconCoroutine(F, Shape, Clones, TTI); - break; - } + ABI.splitCoroutine(F, Shape, Clones, TTI); } // Replace all the swifterror operations in the original function. @@ -2107,8 +2098,6 @@ splitCoroutine(Function &F, SmallVectorImpl &Clones, if (shouldCreateNoAllocVariant) SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones); - - return Shape; } static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit( @@ -2207,8 +2196,53 @@ static void addPrepareFunction(const Module &M, Fns.push_back(PrepareFn); } +static coro::BaseABI *CreateNewABI(Function &F, coro::Shape &S) { + switch (S.ABI) { + case coro::ABI::Switch: + return new coro::SwitchABI(F, S); + case coro::ABI::Async: + return new coro::AsyncABI(F, S); + case coro::ABI::Retcon: + return new coro::AnyRetconABI(F, S); + case coro::ABI::RetconOnce: + return new coro::AnyRetconABI(F, S); + } + llvm_unreachable("Unknown ABI"); +} + CoroSplitPass::CoroSplitPass(bool OptimizeFrame) - : MaterializableCallback(coro::defaultMaterializable), + : CreateAndInitABI([](Function &F, coro::Shape &S) { + coro::BaseABI *ABI = CreateNewABI(F, S); + ABI->init(); + return ABI; + }), + OptimizeFrame(OptimizeFrame) {} + +static coro::BaseABI * +CreateNewABIIsMat(Function &F, coro::Shape &S, + std::function IsMatCallback) { + switch (S.ABI) { + case coro::ABI::Switch: + return new coro::SwitchABI(F, S, IsMatCallback); + case coro::ABI::Async: + return new coro::AsyncABI(F, S, IsMatCallback); + case coro::ABI::Retcon: + return new coro::AnyRetconABI(F, S, IsMatCallback); + case coro::ABI::RetconOnce: + return new coro::AnyRetconABI(F, S, IsMatCallback); + } + llvm_unreachable("Unknown ABI"); +} + +// For back compatibility, constructor takes a materializable callback and +// creates a generator for an ABI with a modified materializable callback. +CoroSplitPass::CoroSplitPass(std::function IsMatCallback, + bool OptimizeFrame) + : CreateAndInitABI([=](Function &F, coro::Shape &S) { + coro::BaseABI *ABI = CreateNewABIIsMat(F, S, IsMatCallback); + ABI->init(); + return ABI; + }), OptimizeFrame(OptimizeFrame) {} PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C, @@ -2241,12 +2275,23 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C, Function &F = N->getFunction(); LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName() << "\n"); + + // The suspend-crossing algorithm in buildCoroutineFrame gets tripped up + // by unreachable blocks, so remove them as a first pass. Remove the + // unreachable blocks before collecting intrinsics into Shape. + removeUnreachableBlocks(F); + + coro::Shape Shape(F, OptimizeFrame); + if (!Shape.CoroBegin) + continue; + F.setSplittedCoroutine(); + std::unique_ptr ABI(CreateAndInitABI(F, Shape)); + SmallVector Clones; - coro::Shape Shape = - splitCoroutine(F, Clones, FAM.getResult(F), - OptimizeFrame, MaterializableCallback); + auto &TTI = FAM.getResult(F); + doSplitCoroutine(F, Clones, *ABI, TTI); CurrentSCC = &updateCallGraphAfterCoroutineSplit( *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM); diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp index ef0abd2d823bbd..c7a04b8fd965b7 100644 --- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -10,6 +10,7 @@ // //===----------------------------------------------------------------------===// +#include "ABI.h" #include "CoroInstr.h" #include "CoroInternal.h" #include "CoroShape.h" @@ -372,11 +373,10 @@ void coro::Shape::invalidateCoroutine( } } -// Perform semantic checking and initialization of the ABI -void coro::Shape::initABI() { - switch (ABI) { - case coro::ABI::Switch: { - for (auto *AnySuspend : CoroSuspends) { +void coro::SwitchABI::init() { + assert(Shape.ABI == coro::ABI::Switch); + { + for (auto *AnySuspend : Shape.CoroSuspends) { auto Suspend = dyn_cast(AnySuspend); if (!Suspend) { #ifndef NDEBUG @@ -386,21 +386,22 @@ void coro::Shape::initABI() { } if (!Suspend->getCoroSave()) - createCoroSave(CoroBegin, Suspend); + createCoroSave(Shape.CoroBegin, Suspend); } - break; } - case coro::ABI::Async: { - break; - }; - case coro::ABI::Retcon: - case coro::ABI::RetconOnce: { +} + +void coro::AsyncABI::init() { assert(Shape.ABI == coro::ABI::Async); } + +void coro::AnyRetconABI::init() { + assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce); + { // Determine the result value types, and make sure they match up with // the values passed to the suspends. - auto ResultTys = getRetconResultTypes(); - auto ResumeTys = getRetconResumeTypes(); + auto ResultTys = Shape.getRetconResultTypes(); + auto ResumeTys = Shape.getRetconResumeTypes(); - for (auto *AnySuspend : CoroSuspends) { + for (auto *AnySuspend : Shape.CoroSuspends) { auto Suspend = dyn_cast(AnySuspend); if (!Suspend) { #ifndef NDEBUG @@ -427,7 +428,7 @@ void coro::Shape::initABI() { #ifndef NDEBUG Suspend->dump(); - RetconLowering.ResumePrototype->getFunctionType()->dump(); + Shape.RetconLowering.ResumePrototype->getFunctionType()->dump(); #endif report_fatal_error("argument to coro.suspend.retcon does not " "match corresponding prototype function result"); @@ -436,7 +437,7 @@ void coro::Shape::initABI() { if (SI != SE || RI != RE) { #ifndef NDEBUG Suspend->dump(); - RetconLowering.ResumePrototype->getFunctionType()->dump(); + Shape.RetconLowering.ResumePrototype->getFunctionType()->dump(); #endif report_fatal_error("wrong number of arguments to coro.suspend.retcon"); } @@ -455,7 +456,7 @@ void coro::Shape::initABI() { if (SuspendResultTys.size() != ResumeTys.size()) { #ifndef NDEBUG Suspend->dump(); - RetconLowering.ResumePrototype->getFunctionType()->dump(); + Shape.RetconLowering.ResumePrototype->getFunctionType()->dump(); #endif report_fatal_error("wrong number of results from coro.suspend.retcon"); } @@ -463,17 +464,13 @@ void coro::Shape::initABI() { if (SuspendResultTys[I] != ResumeTys[I]) { #ifndef NDEBUG Suspend->dump(); - RetconLowering.ResumePrototype->getFunctionType()->dump(); + Shape.RetconLowering.ResumePrototype->getFunctionType()->dump(); #endif report_fatal_error("result from coro.suspend.retcon does not " "match corresponding prototype function param"); } } } - break; - } - default: - llvm_unreachable("coro.begin is not dependent on a coro.id call"); } }