From da426afe8557fe714cf88af097a7e4f615232bf7 Mon Sep 17 00:00:00 2001 From: tnowicki Date: Thu, 12 Sep 2024 18:59:24 -0400 Subject: [PATCH 1/4] [Coroutines] Move OptimizeFrame out of Shape * OptimizeFrame is not really a part of the Coroutine Shape info, rather it is specifically for the addFieldForAllocas method called indirectly by buildCoroutineFrame. * This patch passes OptimizeFrame directly to buildCoroutineFrame so it can be provided to addFieldForAllocas instead of keeping it in the Shape. --- llvm/lib/Transforms/Coroutines/ABI.h | 2 +- llvm/lib/Transforms/Coroutines/CoroFrame.cpp | 16 +++++++++------- llvm/lib/Transforms/Coroutines/CoroShape.h | 6 +----- llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 9 +++++---- 4 files changed, 16 insertions(+), 17 deletions(-) diff --git a/llvm/lib/Transforms/Coroutines/ABI.h b/llvm/lib/Transforms/Coroutines/ABI.h index c94bf7d356b65..7fa835e84ca33 100644 --- a/llvm/lib/Transforms/Coroutines/ABI.h +++ b/llvm/lib/Transforms/Coroutines/ABI.h @@ -41,7 +41,7 @@ class LLVM_LIBRARY_VISIBILITY BaseABI { virtual void init() = 0; // Allocate the coroutine frame and do spill/reload as needed. - virtual void buildCoroutineFrame(); + virtual void buildCoroutineFrame(bool OptimizeFrame); // Perform the function splitting according to the ABI. virtual void splitCoroutine(Function &F, coro::Shape &Shape, diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp index 021b1f7a4156b..91530503a7e1e 100644 --- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp @@ -234,7 +234,7 @@ class FrameTypeBuilder { /// Side Effects: Because We sort the allocas, the order of allocas in the /// frame may be different with the order in the source code. void addFieldForAllocas(const Function &F, FrameDataInfo &FrameData, - coro::Shape &Shape); + coro::Shape &Shape, bool OptimizeFrame); /// Add a field to this structure. [[nodiscard]] FieldIDType addField(Type *Ty, MaybeAlign MaybeFieldAlignment, @@ -336,7 +336,8 @@ void FrameDataInfo::updateLayoutIndex(FrameTypeBuilder &B) { void FrameTypeBuilder::addFieldForAllocas(const Function &F, FrameDataInfo &FrameData, - coro::Shape &Shape) { + coro::Shape &Shape, + bool OptimizeFrame) { using AllocaSetType = SmallVector; SmallVector NonOverlapedAllocas; @@ -350,7 +351,7 @@ void FrameTypeBuilder::addFieldForAllocas(const Function &F, } }); - if (!Shape.OptimizeFrame) { + if (!OptimizeFrame) { for (const auto &A : FrameData.Allocas) { AllocaInst *Alloca = A.Alloca; NonOverlapedAllocas.emplace_back(AllocaSetType(1, Alloca)); @@ -860,7 +861,8 @@ static void buildFrameDebugInfo(Function &F, coro::Shape &Shape, // ... spills ... // }; static StructType *buildFrameType(Function &F, coro::Shape &Shape, - FrameDataInfo &FrameData) { + FrameDataInfo &FrameData, + bool OptimizeFrame) { LLVMContext &C = F.getContext(); const DataLayout &DL = F.getDataLayout(); StructType *FrameTy = [&] { @@ -905,7 +907,7 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape, // Because multiple allocas may own the same field slot, // we add allocas to field here. - B.addFieldForAllocas(F, FrameData, Shape); + B.addFieldForAllocas(F, FrameData, Shape, OptimizeFrame); // Add PromiseAlloca to Allocas list so that // 1. updateLayoutIndex could update its index after // `performOptimizedStructLayout` @@ -2056,7 +2058,7 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape, rewritePHIs(F); } -void coro::BaseABI::buildCoroutineFrame() { +void coro::BaseABI::buildCoroutineFrame(bool OptimizeFrame) { SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds); doRematerializations(F, Checker, IsMaterializable); @@ -2087,7 +2089,7 @@ void coro::BaseABI::buildCoroutineFrame() { // Build frame FrameDataInfo FrameData(Spills, Allocas); - Shape.FrameTy = buildFrameType(F, Shape, FrameData); + Shape.FrameTy = buildFrameType(F, Shape, FrameData, OptimizeFrame); Shape.FramePtr = Shape.CoroBegin; // For now, this works for C++ programs only. buildFrameDebugInfo(F, Shape, FrameData); diff --git a/llvm/lib/Transforms/Coroutines/CoroShape.h b/llvm/lib/Transforms/Coroutines/CoroShape.h index f4fb4baa6df31..7daa03beb2542 100644 --- a/llvm/lib/Transforms/Coroutines/CoroShape.h +++ b/llvm/lib/Transforms/Coroutines/CoroShape.h @@ -112,9 +112,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape { Value *FramePtr = nullptr; BasicBlock *AllocaSpillBlock = nullptr; - /// This would only be true if optimization are enabled. - bool OptimizeFrame; - struct SwitchLoweringStorage { SwitchInst *ResumeSwitch; AllocaInst *PromiseAlloca; @@ -265,8 +262,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape { void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const; Shape() = default; - explicit Shape(Function &F, bool OptimizeFrame = false) - : OptimizeFrame(OptimizeFrame) { + explicit Shape(Function &F) { SmallVector CoroFrames; SmallVector UnusedCoroSaves; diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index 4fbda077129fa..9aed4f6522a3f 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -2053,7 +2053,8 @@ void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape, } static void doSplitCoroutine(Function &F, SmallVectorImpl &Clones, - coro::BaseABI &ABI, TargetTransformInfo &TTI) { + coro::BaseABI &ABI, TargetTransformInfo &TTI, + bool OptimizeFrame) { PrettyStackTraceFunction prettyStackTrace(F); auto &Shape = ABI.Shape; @@ -2064,7 +2065,7 @@ static void doSplitCoroutine(Function &F, SmallVectorImpl &Clones, simplifySuspendPoints(Shape); normalizeCoroutine(F, Shape, TTI); - ABI.buildCoroutineFrame(); + ABI.buildCoroutineFrame(OptimizeFrame); replaceFrameSizeAndAlignment(Shape); bool isNoSuspendCoroutine = Shape.CoroSuspends.empty(); @@ -2273,7 +2274,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C, // unreachable blocks before collecting intrinsics into Shape. removeUnreachableBlocks(F); - coro::Shape Shape(F, OptimizeFrame); + coro::Shape Shape(F); if (!Shape.CoroBegin) continue; @@ -2283,7 +2284,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C, SmallVector Clones; auto &TTI = FAM.getResult(F); - doSplitCoroutine(F, Clones, *ABI, TTI); + doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame); CurrentSCC = &updateCallGraphAfterCoroutineSplit( *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM); From a62f918ab018ead4dc771fee7c04d6dfd673e9e7 Mon Sep 17 00:00:00 2001 From: tnowicki Date: Tue, 3 Sep 2024 15:10:54 -0400 Subject: [PATCH 2/4] [Coroutines] Move various util headers to include/llvm for plugin libraries * Move utilities to include/llvm/Transform/Coroutines --- .../llvm}/Transforms/Coroutines/ABI.h | 17 ++-- .../llvm}/Transforms/Coroutines/CoroInstr.h | 95 ++++++++++--------- .../llvm}/Transforms/Coroutines/CoroShape.h | 8 +- .../llvm/Transforms/Coroutines/CoroSplit.h | 1 + .../Coroutines/MaterializationUtils.h | 9 +- .../llvm}/Transforms/Coroutines/SpillUtils.h | 5 +- .../Coroutines/SuspendCrossingInfo.h | 6 +- llvm/lib/Transforms/Coroutines/CoroEarly.cpp | 2 +- llvm/lib/Transforms/Coroutines/CoroFrame.cpp | 9 +- llvm/lib/Transforms/Coroutines/CoroInternal.h | 4 +- llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 6 +- llvm/lib/Transforms/Coroutines/Coroutines.cpp | 6 +- .../Coroutines/MaterializationUtils.cpp | 5 +- llvm/lib/Transforms/Coroutines/SpillUtils.cpp | 2 +- .../Coroutines/SuspendCrossingInfo.cpp | 2 +- .../Transforms/Coroutines/ExtraRematTest.cpp | 1 + 16 files changed, 95 insertions(+), 83 deletions(-) rename llvm/{lib => include/llvm}/Transforms/Coroutines/ABI.h (89%) rename llvm/{lib => include/llvm}/Transforms/Coroutines/CoroInstr.h (90%) rename llvm/{lib => include/llvm}/Transforms/Coroutines/CoroShape.h (96%) rename llvm/{lib => include/llvm}/Transforms/Coroutines/MaterializationUtils.h (76%) rename llvm/{lib => include/llvm}/Transforms/Coroutines/SpillUtils.h (93%) rename llvm/{lib => include/llvm}/Transforms/Coroutines/SuspendCrossingInfo.h (97%) diff --git a/llvm/lib/Transforms/Coroutines/ABI.h b/llvm/include/llvm/Transforms/Coroutines/ABI.h similarity index 89% rename from llvm/lib/Transforms/Coroutines/ABI.h rename to llvm/include/llvm/Transforms/Coroutines/ABI.h index 7fa835e84ca33..e7568d275c161 100644 --- a/llvm/lib/Transforms/Coroutines/ABI.h +++ b/llvm/include/llvm/Transforms/Coroutines/ABI.h @@ -12,12 +12,13 @@ // ABI enum and ABI class are used by the Coroutine passes when lowering. //===----------------------------------------------------------------------===// -#ifndef LIB_TRANSFORMS_COROUTINES_ABI_H -#define LIB_TRANSFORMS_COROUTINES_ABI_H +#ifndef LLVM_TRANSFORMS_COROUTINES_ABI_H +#define LLVM_TRANSFORMS_COROUTINES_ABI_H -#include "CoroShape.h" -#include "SuspendCrossingInfo.h" #include "llvm/Analysis/TargetTransformInfo.h" +#include "llvm/Transforms/Coroutines/CoroShape.h" +#include "llvm/Transforms/Coroutines/MaterializationUtils.h" +#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" namespace llvm { @@ -30,7 +31,7 @@ namespace coro { // ABI operations. The ABIs (e.g. Switch, Async, Retcon{Once}) are the common // ABIs. -class LLVM_LIBRARY_VISIBILITY BaseABI { +class BaseABI { public: BaseABI(Function &F, coro::Shape &S, std::function IsMaterializable) @@ -56,7 +57,7 @@ class LLVM_LIBRARY_VISIBILITY BaseABI { std::function IsMaterializable; }; -class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI { +class SwitchABI : public BaseABI { public: SwitchABI(Function &F, coro::Shape &S, std::function IsMaterializable) @@ -69,7 +70,7 @@ class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI { TargetTransformInfo &TTI) override; }; -class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI { +class AsyncABI : public BaseABI { public: AsyncABI(Function &F, coro::Shape &S, std::function IsMaterializable) @@ -82,7 +83,7 @@ class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI { TargetTransformInfo &TTI) override; }; -class LLVM_LIBRARY_VISIBILITY AnyRetconABI : public BaseABI { +class AnyRetconABI : public BaseABI { public: AnyRetconABI(Function &F, coro::Shape &S, std::function IsMaterializable) diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h similarity index 90% rename from llvm/lib/Transforms/Coroutines/CoroInstr.h rename to llvm/include/llvm/Transforms/Coroutines/CoroInstr.h index a31703fe01304..1aa1996a4c009 100644 --- a/llvm/lib/Transforms/Coroutines/CoroInstr.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h @@ -22,8 +22,8 @@ // the Coroutine library. //===----------------------------------------------------------------------===// -#ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H -#define LLVM_LIB_TRANSFORMS_COROUTINES_COROINSTR_H +#ifndef LLVM_TRANSFORMS_COROUTINES_COROINSTR_H +#define LLVM_TRANSFORMS_COROUTINES_COROINSTR_H #include "llvm/IR/GlobalVariable.h" #include "llvm/IR/IntrinsicInst.h" @@ -32,7 +32,7 @@ namespace llvm { /// This class represents the llvm.coro.subfn.addr instruction. -class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst { +class CoroSubFnInst : public IntrinsicInst { enum { FrameArg, IndexArg }; public: @@ -67,7 +67,7 @@ class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst { }; /// This represents the llvm.coro.alloc instruction. -class LLVM_LIBRARY_VISIBILITY CoroAllocInst : public IntrinsicInst { +class CoroAllocInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -82,7 +82,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocInst : public IntrinsicInst { // FIXME: add callback metadata // FIXME: make a proper IntrinisicInst. Currently this is not possible, // because llvm.coro.await.suspend.* can be invoked. -class LLVM_LIBRARY_VISIBILITY CoroAwaitSuspendInst : public CallBase { +class CoroAwaitSuspendInst : public CallBase { enum { AwaiterArg, FrameArg, WrapperArg }; public: @@ -112,7 +112,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAwaitSuspendInst : public CallBase { }; /// This represents a common base class for llvm.coro.id instructions. -class LLVM_LIBRARY_VISIBILITY AnyCoroIdInst : public IntrinsicInst { +class AnyCoroIdInst : public IntrinsicInst { public: CoroAllocInst *getCoroAlloc() { for (User *U : users()) @@ -143,7 +143,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdInst : public IntrinsicInst { }; /// This represents the llvm.coro.id instruction. -class LLVM_LIBRARY_VISIBILITY CoroIdInst : public AnyCoroIdInst { +class CoroIdInst : public AnyCoroIdInst { enum { AlignArg, PromiseArg, CoroutineArg, InfoArg }; public: @@ -232,7 +232,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdInst : public AnyCoroIdInst { /// This represents either the llvm.coro.id.retcon or /// llvm.coro.id.retcon.once instruction. -class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst { +class AnyCoroIdRetconInst : public AnyCoroIdInst { enum { SizeArg, AlignArg, StorageArg, PrototypeArg, AllocArg, DeallocArg }; public: @@ -246,9 +246,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst { return cast(getArgOperand(AlignArg))->getAlignValue(); } - Value *getStorage() const { - return getArgOperand(StorageArg); - } + Value *getStorage() const { return getArgOperand(StorageArg); } /// Return the prototype for the continuation function. The type, /// attributes, and calling convention of the continuation function(s) @@ -270,8 +268,8 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst { // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { auto ID = I->getIntrinsicID(); - return ID == Intrinsic::coro_id_retcon - || ID == Intrinsic::coro_id_retcon_once; + return ID == Intrinsic::coro_id_retcon || + ID == Intrinsic::coro_id_retcon_once; } static bool classof(const Value *V) { return isa(V) && classof(cast(V)); @@ -279,8 +277,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroIdRetconInst : public AnyCoroIdInst { }; /// This represents the llvm.coro.id.retcon instruction. -class LLVM_LIBRARY_VISIBILITY CoroIdRetconInst - : public AnyCoroIdRetconInst { +class CoroIdRetconInst : public AnyCoroIdRetconInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -292,8 +289,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconInst }; /// This represents the llvm.coro.id.retcon.once instruction. -class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceInst - : public AnyCoroIdRetconInst { +class CoroIdRetconOnceInst : public AnyCoroIdRetconInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -305,7 +301,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceInst }; /// This represents the llvm.coro.id.async instruction. -class LLVM_LIBRARY_VISIBILITY CoroIdAsyncInst : public AnyCoroIdInst { +class CoroIdAsyncInst : public AnyCoroIdInst { enum { SizeArg, AlignArg, StorageArg, AsyncFuncPtrArg }; public: @@ -356,7 +352,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdAsyncInst : public AnyCoroIdInst { }; /// This represents the llvm.coro.context.alloc instruction. -class LLVM_LIBRARY_VISIBILITY CoroAsyncContextAllocInst : public IntrinsicInst { +class CoroAsyncContextAllocInst : public IntrinsicInst { enum { AsyncFuncPtrArg }; public: @@ -375,8 +371,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAsyncContextAllocInst : public IntrinsicInst { }; /// This represents the llvm.coro.context.dealloc instruction. -class LLVM_LIBRARY_VISIBILITY CoroAsyncContextDeallocInst - : public IntrinsicInst { +class CoroAsyncContextDeallocInst : public IntrinsicInst { enum { AsyncContextArg }; public: @@ -396,7 +391,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAsyncContextDeallocInst /// This represents the llvm.coro.async.resume instruction. /// During lowering this is replaced by the resume function of a suspend point /// (the continuation function). -class LLVM_LIBRARY_VISIBILITY CoroAsyncResumeInst : public IntrinsicInst { +class CoroAsyncResumeInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -408,7 +403,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAsyncResumeInst : public IntrinsicInst { }; /// This represents the llvm.coro.async.size.replace instruction. -class LLVM_LIBRARY_VISIBILITY CoroAsyncSizeReplace : public IntrinsicInst { +class CoroAsyncSizeReplace : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -420,7 +415,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAsyncSizeReplace : public IntrinsicInst { }; /// This represents the llvm.coro.frame instruction. -class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst { +class CoroFrameInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -432,7 +427,7 @@ class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst { }; /// This represents the llvm.coro.free instruction. -class LLVM_LIBRARY_VISIBILITY CoroFreeInst : public IntrinsicInst { +class CoroFreeInst : public IntrinsicInst { enum { IdArg, FrameArg }; public: @@ -447,8 +442,8 @@ class LLVM_LIBRARY_VISIBILITY CoroFreeInst : public IntrinsicInst { } }; -/// This class represents the llvm.coro.begin instruction. -class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst { +/// This class represents the llvm.coro.begin instructions. +class CoroBeginInst : public IntrinsicInst { enum { IdArg, MemArg }; public: @@ -468,7 +463,7 @@ class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst { }; /// This represents the llvm.coro.save instruction. -class LLVM_LIBRARY_VISIBILITY CoroSaveInst : public IntrinsicInst { +class CoroSaveInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -480,7 +475,7 @@ class LLVM_LIBRARY_VISIBILITY CoroSaveInst : public IntrinsicInst { }; /// This represents the llvm.coro.promise instruction. -class LLVM_LIBRARY_VISIBILITY CoroPromiseInst : public IntrinsicInst { +class CoroPromiseInst : public IntrinsicInst { enum { FrameArg, AlignArg, FromArg }; public: @@ -505,7 +500,7 @@ class LLVM_LIBRARY_VISIBILITY CoroPromiseInst : public IntrinsicInst { } }; -class LLVM_LIBRARY_VISIBILITY AnyCoroSuspendInst : public IntrinsicInst { +class AnyCoroSuspendInst : public IntrinsicInst { public: CoroSaveInst *getCoroSave() const; @@ -521,7 +516,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroSuspendInst : public IntrinsicInst { }; /// This represents the llvm.coro.suspend instruction. -class LLVM_LIBRARY_VISIBILITY CoroSuspendInst : public AnyCoroSuspendInst { +class CoroSuspendInst : public AnyCoroSuspendInst { enum { SaveArg, FinalArg }; public: @@ -553,7 +548,7 @@ inline CoroSaveInst *AnyCoroSuspendInst::getCoroSave() const { } /// This represents the llvm.coro.suspend.async instruction. -class LLVM_LIBRARY_VISIBILITY CoroSuspendAsyncInst : public AnyCoroSuspendInst { +class CoroSuspendAsyncInst : public AnyCoroSuspendInst { public: enum { StorageArgNoArg, @@ -594,7 +589,7 @@ class LLVM_LIBRARY_VISIBILITY CoroSuspendAsyncInst : public AnyCoroSuspendInst { }; /// This represents the llvm.coro.suspend.retcon instruction. -class LLVM_LIBRARY_VISIBILITY CoroSuspendRetconInst : public AnyCoroSuspendInst { +class CoroSuspendRetconInst : public AnyCoroSuspendInst { public: op_iterator value_begin() { return arg_begin(); } const_op_iterator value_begin() const { return arg_begin(); } @@ -619,7 +614,7 @@ class LLVM_LIBRARY_VISIBILITY CoroSuspendRetconInst : public AnyCoroSuspendInst }; /// This represents the llvm.coro.size instruction. -class LLVM_LIBRARY_VISIBILITY CoroSizeInst : public IntrinsicInst { +class CoroSizeInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -631,7 +626,7 @@ class LLVM_LIBRARY_VISIBILITY CoroSizeInst : public IntrinsicInst { }; /// This represents the llvm.coro.align instruction. -class LLVM_LIBRARY_VISIBILITY CoroAlignInst : public IntrinsicInst { +class CoroAlignInst : public IntrinsicInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -643,7 +638,7 @@ class LLVM_LIBRARY_VISIBILITY CoroAlignInst : public IntrinsicInst { }; /// This represents the llvm.end.results instruction. -class LLVM_LIBRARY_VISIBILITY CoroEndResults : public IntrinsicInst { +class CoroEndResults : public IntrinsicInst { public: op_iterator retval_begin() { return arg_begin(); } const_op_iterator retval_begin() const { return arg_begin(); } @@ -671,7 +666,7 @@ class LLVM_LIBRARY_VISIBILITY CoroEndResults : public IntrinsicInst { } }; -class LLVM_LIBRARY_VISIBILITY AnyCoroEndInst : public IntrinsicInst { +class AnyCoroEndInst : public IntrinsicInst { enum { FrameArg, UnwindArg, TokenArg }; public: @@ -700,7 +695,7 @@ class LLVM_LIBRARY_VISIBILITY AnyCoroEndInst : public IntrinsicInst { }; /// This represents the llvm.coro.end instruction. -class LLVM_LIBRARY_VISIBILITY CoroEndInst : public AnyCoroEndInst { +class CoroEndInst : public AnyCoroEndInst { public: // Methods to support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { @@ -712,7 +707,7 @@ class LLVM_LIBRARY_VISIBILITY CoroEndInst : public AnyCoroEndInst { }; /// This represents the llvm.coro.end instruction. -class LLVM_LIBRARY_VISIBILITY CoroAsyncEndInst : public AnyCoroEndInst { +class CoroAsyncEndInst : public AnyCoroEndInst { enum { FrameArg, UnwindArg, MustTailCallFuncArg }; public: @@ -736,12 +731,11 @@ class LLVM_LIBRARY_VISIBILITY CoroAsyncEndInst : public AnyCoroEndInst { }; /// This represents the llvm.coro.alloca.alloc instruction. -class LLVM_LIBRARY_VISIBILITY CoroAllocaAllocInst : public IntrinsicInst { +class CoroAllocaAllocInst : public IntrinsicInst { enum { SizeArg, AlignArg }; + public: - Value *getSize() const { - return getArgOperand(SizeArg); - } + Value *getSize() const { return getArgOperand(SizeArg); } Align getAlignment() const { return cast(getArgOperand(AlignArg))->getAlignValue(); } @@ -756,8 +750,9 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocaAllocInst : public IntrinsicInst { }; /// This represents the llvm.coro.alloca.get instruction. -class LLVM_LIBRARY_VISIBILITY CoroAllocaGetInst : public IntrinsicInst { +class CoroAllocaGetInst : public IntrinsicInst { enum { AllocArg }; + public: CoroAllocaAllocInst *getAlloc() const { return cast(getArgOperand(AllocArg)); @@ -773,8 +768,9 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocaGetInst : public IntrinsicInst { }; /// This represents the llvm.coro.alloca.free instruction. -class LLVM_LIBRARY_VISIBILITY CoroAllocaFreeInst : public IntrinsicInst { +class CoroAllocaFreeInst : public IntrinsicInst { enum { AllocArg }; + public: CoroAllocaAllocInst *getAlloc() const { return cast(getArgOperand(AllocArg)); @@ -789,6 +785,13 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocaFreeInst : public IntrinsicInst { } }; +namespace coro { + +// Check if BasicBlock suspends the coroutine +bool isSuspendBlock(BasicBlock *BB); + +} // End namespace coro. + } // End namespace llvm. -#endif +#endif // LLVM_TRANSFORMS_COROUTINES_COROINSTR_H diff --git a/llvm/lib/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h similarity index 96% rename from llvm/lib/Transforms/Coroutines/CoroShape.h rename to llvm/include/llvm/Transforms/Coroutines/CoroShape.h index 7daa03beb2542..2c0e9257fcf92 100644 --- a/llvm/lib/Transforms/Coroutines/CoroShape.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h @@ -12,9 +12,9 @@ #ifndef LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H #define LLVM_TRANSFORMS_COROUTINES_COROSHAPE_H -#include "CoroInstr.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/PassManager.h" +#include "llvm/Transforms/Coroutines/CoroInstr.h" namespace llvm { @@ -22,6 +22,10 @@ class CallGraph; namespace coro { +// This enum is used by CoroSplitPass's internal methods for implementing ABI +// specific behavior. Ideally, this would not be necessary and all behavior +// would be implemented by the ABI objects. However, it is not practical and +// not necessary to move such code to the ABI objects. enum class ABI { /// The "resume-switch" lowering, where there are separate resume and /// destroy functions that are shared between all suspend points. The @@ -49,7 +53,7 @@ enum class ABI { // Holds structural Coroutine Intrinsics for a particular function and other // values used during CoroSplit pass. -struct LLVM_LIBRARY_VISIBILITY Shape { +struct Shape { CoroBeginInst *CoroBegin = nullptr; SmallVector CoroEnds; SmallVector CoroSizes; diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h index 52b6c4918ada9..a5fd57f8f9dfa 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h @@ -18,6 +18,7 @@ #include "llvm/Analysis/CGSCCPassManager.h" #include "llvm/Analysis/LazyCallGraph.h" #include "llvm/IR/PassManager.h" +#include "llvm/Transforms/Coroutines/ABI.h" namespace llvm { diff --git a/llvm/lib/Transforms/Coroutines/MaterializationUtils.h b/llvm/include/llvm/Transforms/Coroutines/MaterializationUtils.h similarity index 76% rename from llvm/lib/Transforms/Coroutines/MaterializationUtils.h rename to llvm/include/llvm/Transforms/Coroutines/MaterializationUtils.h index f391851c97b3b..d8fc0c86a6fb5 100644 --- a/llvm/lib/Transforms/Coroutines/MaterializationUtils.h +++ b/llvm/include/llvm/Transforms/Coroutines/MaterializationUtils.h @@ -6,11 +6,10 @@ // //===----------------------------------------------------------------------===// -#include "SuspendCrossingInfo.h" -#include "llvm/IR/Instruction.h" +#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" -#ifndef LIB_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H -#define LIB_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H +#ifndef LLVM_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H +#define LLVM_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H namespace llvm { @@ -27,4 +26,4 @@ void doRematerializations(Function &F, SuspendCrossingInfo &Checker, } // namespace llvm -#endif // LIB_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H +#endif // LLVM_TRANSFORMS_COROUTINES_MATERIALIZATIONUTILS_H diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.h b/llvm/include/llvm/Transforms/Coroutines/SpillUtils.h similarity index 93% rename from llvm/lib/Transforms/Coroutines/SpillUtils.h rename to llvm/include/llvm/Transforms/Coroutines/SpillUtils.h index 8843b611e0842..6cdf83c0603f4 100644 --- a/llvm/lib/Transforms/Coroutines/SpillUtils.h +++ b/llvm/include/llvm/Transforms/Coroutines/SpillUtils.h @@ -6,8 +6,9 @@ // //===----------------------------------------------------------------------===// -#include "CoroInternal.h" -#include "SuspendCrossingInfo.h" +#include "llvm/IR/Dominators.h" +#include "llvm/Transforms/Coroutines/CoroShape.h" +#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" #ifndef LLVM_TRANSFORMS_COROUTINES_SPILLINGINFO_H #define LLVM_TRANSFORMS_COROUTINES_SPILLINGINFO_H diff --git a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.h b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h similarity index 97% rename from llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.h rename to llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h index db889966dcf1d..49cae6dde47e5 100644 --- a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.h +++ b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h @@ -12,16 +12,16 @@ // ptrs in the BlockToIndexMapping. //===----------------------------------------------------------------------===// -#ifndef LLVM_LIB_TRANSFORMS_COROUTINES_SUSPENDCROSSINGINFO_H -#define LLVM_LIB_TRANSFORMS_COROUTINES_SUSPENDCROSSINGINFO_H +#ifndef LLVM_TRANSFORMS_COROUTINES_SUSPENDCROSSINGINFO_H +#define LLVM_TRANSFORMS_COROUTINES_SUSPENDCROSSINGINFO_H -#include "CoroInstr.h" #include "llvm/ADT/BitVector.h" #include "llvm/ADT/PostOrderIterator.h" #include "llvm/ADT/SmallVector.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Function.h" #include "llvm/IR/Instruction.h" +#include "llvm/Transforms/Coroutines/CoroInstr.h" namespace llvm { diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp index 093a81648c92e..a3674306f3e10 100644 --- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp @@ -8,12 +8,12 @@ #include "llvm/Transforms/Coroutines/CoroEarly.h" #include "CoroInternal.h" -#include "CoroShape.h" #include "llvm/IR/DIBuilder.h" #include "llvm/IR/Function.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/InstIterator.h" #include "llvm/IR/Module.h" +#include "llvm/Transforms/Coroutines/CoroShape.h" using namespace llvm; diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp index 91530503a7e1e..a83c9e2638b48 100644 --- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp @@ -15,11 +15,7 @@ // the value into the coroutine frame. //===----------------------------------------------------------------------===// -#include "ABI.h" #include "CoroInternal.h" -#include "MaterializationUtils.h" -#include "SpillUtils.h" -#include "SuspendCrossingInfo.h" #include "llvm/ADT/BitVector.h" #include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/SmallString.h" @@ -33,6 +29,11 @@ #include "llvm/IR/Module.h" #include "llvm/Support/Debug.h" #include "llvm/Support/OptimizedStructLayout.h" +#include "llvm/Transforms/Coroutines/ABI.h" +#include "llvm/Transforms/Coroutines/CoroInstr.h" +#include "llvm/Transforms/Coroutines/MaterializationUtils.h" +#include "llvm/Transforms/Coroutines/SpillUtils.h" +#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" #include "llvm/Transforms/Utils/BasicBlockUtils.h" #include "llvm/Transforms/Utils/Local.h" #include "llvm/Transforms/Utils/PromoteMemToReg.h" diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h index 88d0f83c98c9e..a0b52063aca10 100644 --- a/llvm/lib/Transforms/Coroutines/CoroInternal.h +++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h @@ -11,10 +11,10 @@ #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H -#include "CoroInstr.h" -#include "CoroShape.h" #include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/IR/IRBuilder.h" +#include "llvm/Transforms/Coroutines/CoroInstr.h" +#include "llvm/Transforms/Coroutines/CoroShape.h" namespace llvm { diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index 9aed4f6522a3f..ef1f27118bc14 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -19,10 +19,7 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Coroutines/CoroSplit.h" -#include "ABI.h" -#include "CoroInstr.h" #include "CoroInternal.h" -#include "MaterializationUtils.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/PriorityWorklist.h" #include "llvm/ADT/SmallPtrSet.h" @@ -64,6 +61,9 @@ #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" diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp index 184efbfe903d2..f4d9a7a8aa856 100644 --- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -10,10 +10,7 @@ // //===----------------------------------------------------------------------===// -#include "ABI.h" -#include "CoroInstr.h" #include "CoroInternal.h" -#include "CoroShape.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Analysis/CallGraph.h" @@ -29,6 +26,9 @@ #include "llvm/IR/Type.h" #include "llvm/Support/Casting.h" #include "llvm/Support/ErrorHandling.h" +#include "llvm/Transforms/Coroutines/ABI.h" +#include "llvm/Transforms/Coroutines/CoroInstr.h" +#include "llvm/Transforms/Coroutines/CoroShape.h" #include "llvm/Transforms/Utils/Local.h" #include #include diff --git a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp index 708e8734175f9..c3ea0977d4211 100644 --- a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp +++ b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp @@ -9,12 +9,13 @@ // This file contains classes used to materialize insts after suspends points. //===----------------------------------------------------------------------===// -#include "MaterializationUtils.h" -#include "SpillUtils.h" +#include "llvm/Transforms/Coroutines/MaterializationUtils.h" +#include "CoroInternal.h" #include "llvm/ADT/PostOrderIterator.h" #include "llvm/IR/Dominators.h" #include "llvm/IR/InstIterator.h" #include "llvm/IR/Instruction.h" +#include "llvm/Transforms/Coroutines/SpillUtils.h" #include using namespace llvm; diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp index 96b5c8440e5f9..dc06d9c4cd1c7 100644 --- a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp +++ b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp @@ -6,7 +6,7 @@ // //===----------------------------------------------------------------------===// -#include "SpillUtils.h" +#include "llvm/Transforms/Coroutines/SpillUtils.h" #include "llvm/Analysis/CFG.h" #include "llvm/Analysis/PtrUseVisitor.h" #include "llvm/IR/CFG.h" diff --git a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp index 84699e653db60..f18f23306befb 100644 --- a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp +++ b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp @@ -12,7 +12,7 @@ // ptrs in the BlockToIndexMapping. //===----------------------------------------------------------------------===// -#include "SuspendCrossingInfo.h" +#include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" // The "coro-suspend-crossing" flag is very noisy. There is another debug type, // "coro-frame", which results in leaner debug spew. diff --git a/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp b/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp index 9aab3cfd9cf10..1d55889a32d7a 100644 --- a/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp +++ b/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp @@ -11,6 +11,7 @@ #include "llvm/Passes/PassBuilder.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Testing/Support/Error.h" +#include "llvm/Transforms/Coroutines/ABI.h" #include "llvm/Transforms/Coroutines/CoroSplit.h" #include "gtest/gtest.h" From 52688e8836cd9d27b90478d59c15b121e8af5f0d Mon Sep 17 00:00:00 2001 From: tnowicki Date: Tue, 3 Sep 2024 15:36:01 -0400 Subject: [PATCH 3/4] [Coroutines] Support Custom ABIs and plugin libraries * Add the llvm.coro.begin.custom intrinsic * Add constructors to CoroSplit that take a list of generators that create the custom ABI object. * Add has/getCustomABI methods to CoroBeginInst class. * Add a unittest for the custom ABI. --- .../llvm/Analysis/TargetTransformInfoImpl.h | 1 + llvm/include/llvm/IR/Intrinsics.td | 3 +- llvm/include/llvm/Transforms/Coroutines/ABI.h | 8 +- .../llvm/Transforms/Coroutines/CoroInstr.h | 19 +++- .../llvm/Transforms/Coroutines/CoroSplit.h | 15 +++- .../lib/Transforms/Coroutines/CoroCleanup.cpp | 4 +- llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 36 +++++++- llvm/lib/Transforms/Coroutines/Coroutines.cpp | 4 +- .../Transforms/Coroutines/ExtraRematTest.cpp | 87 +++++++++++++++++++ 9 files changed, 164 insertions(+), 13 deletions(-) diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h index eca8818cc25e6..72af1201b7d1f 100644 --- a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h +++ b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h @@ -769,6 +769,7 @@ class TargetTransformInfoImplBase { case Intrinsic::experimental_gc_relocate: case Intrinsic::coro_alloc: case Intrinsic::coro_begin: + case Intrinsic::coro_begin_custom_abi: case Intrinsic::coro_free: case Intrinsic::coro_end: case Intrinsic::coro_frame: diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td index 1e243eb27b312..0896c2b820928 100644 --- a/llvm/include/llvm/IR/Intrinsics.td +++ b/llvm/include/llvm/IR/Intrinsics.td @@ -1716,7 +1716,8 @@ def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty], [IntrNoMem]>; def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty], [WriteOnly>]>; - +def int_coro_begin_custom_abi : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty, llvm_i32_ty], + [WriteOnly>]>; def int_coro_free : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty], [IntrReadMem, IntrArgMemOnly, ReadOnly>, diff --git a/llvm/include/llvm/Transforms/Coroutines/ABI.h b/llvm/include/llvm/Transforms/Coroutines/ABI.h index e7568d275c161..c1392fdcda41d 100644 --- a/llvm/include/llvm/Transforms/Coroutines/ABI.h +++ b/llvm/include/llvm/Transforms/Coroutines/ABI.h @@ -29,7 +29,13 @@ 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. +// ABIs. However, specific users often need to modify the behavior of these, +// such as for C++20 or Swift. This can be accomplished by inheriting one of +// the common ABIs and overriding one or more of the methods to create a custom +// ABI. The custom ABI is specified with the coro.begin.custom.abi intrinsic +// instead of the coro.begin intrinsic by providing an i32 in the last arg. +// This is used to lookup a generator for the custom ABI from a set of +// generators provided to the CoroSplitPass constructor. class BaseABI { public: diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h index 1aa1996a4c009..656edbdd6fdbe 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h @@ -124,7 +124,8 @@ class AnyCoroIdInst : public IntrinsicInst { IntrinsicInst *getCoroBegin() { for (User *U : users()) if (auto *II = dyn_cast(U)) - if (II->getIntrinsicID() == Intrinsic::coro_begin) + if (II->getIntrinsicID() == Intrinsic::coro_begin || + II->getIntrinsicID() == Intrinsic::coro_begin_custom_abi) return II; llvm_unreachable("no coro.begin associated with coro.id"); } @@ -442,20 +443,30 @@ class CoroFreeInst : public IntrinsicInst { } }; -/// This class represents the llvm.coro.begin instructions. +/// This class represents the llvm.coro.begin or llvm.coro.begin.custom.abi +/// instructions. class CoroBeginInst : public IntrinsicInst { - enum { IdArg, MemArg }; + enum { IdArg, MemArg, CustomABIArg }; public: AnyCoroIdInst *getId() const { return cast(getArgOperand(IdArg)); } + bool hasCustomABI() const { + return getIntrinsicID() == Intrinsic::coro_begin_custom_abi; + } + + int getCustomABI() const { + return cast(getArgOperand(CustomABIArg))->getZExtValue(); + } + Value *getMem() const { return getArgOperand(MemArg); } // Methods for support type inquiry through isa, cast, and dyn_cast: static bool classof(const IntrinsicInst *I) { - return I->getIntrinsicID() == Intrinsic::coro_begin; + return I->getIntrinsicID() == Intrinsic::coro_begin || + I->getIntrinsicID() == Intrinsic::coro_begin_custom_abi; } static bool classof(const Value *V) { return isa(V) && classof(cast(V)); diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h index a5fd57f8f9dfa..888e924c15973 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h @@ -28,17 +28,28 @@ struct Shape; } // namespace coro struct CoroSplitPass : PassInfoMixin { + using BaseABITy = + std::function(Function &, coro::Shape &)>; CoroSplitPass(bool OptimizeFrame = false); + + CoroSplitPass(SmallVector GenCustomABIs, + bool OptimizeFrame = false); + + // For back compatibility, constructor takes a materiaizlable callback. CoroSplitPass(std::function MaterializableCallback, bool OptimizeFrame = false); + // For back compatibility, constructor takes a materiaizlable callback. + CoroSplitPass(std::function MaterializableCallback, + SmallVector GenCustomABIs, + bool OptimizeFrame = false); + PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM, LazyCallGraph &CG, CGSCCUpdateResult &UR); + static bool isRequired() { return true; } - using BaseABITy = - std::function(Function &, coro::Shape &)>; // Generator for an ABI transformer BaseABITy CreateAndInitABI; diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp index dd92b3593af92..1cda7f93f72a2 100644 --- a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp @@ -53,6 +53,7 @@ bool Lowerer::lower(Function &F) { default: continue; case Intrinsic::coro_begin: + case Intrinsic::coro_begin_custom_abi: II->replaceAllUsesWith(II->getArgOperand(1)); break; case Intrinsic::coro_free: @@ -112,7 +113,8 @@ static bool declaresCoroCleanupIntrinsics(const Module &M) { M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr", "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon", "llvm.coro.id.async", "llvm.coro.id.retcon.once", - "llvm.coro.async.size.replace", "llvm.coro.async.resume"}); + "llvm.coro.async.size.replace", "llvm.coro.async.resume", + "llvm.coro.begin.custom.abi"}); } PreservedAnalyses CoroCleanupPass::run(Module &M, diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index ef1f27118bc14..57ce0a9d2c0b9 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -2200,7 +2200,15 @@ static void addPrepareFunction(const Module &M, static std::unique_ptr CreateNewABI(Function &F, coro::Shape &S, - std::function IsMatCallback) { + std::function IsMatCallback, + const SmallVector GenCustomABIs) { + if (S.CoroBegin->hasCustomABI()) { + unsigned CustomABI = S.CoroBegin->getCustomABI(); + if (CustomABI >= GenCustomABIs.size()) + llvm_unreachable("Custom ABI not found amoung those specified"); + return GenCustomABIs[CustomABI](F, S); + } + switch (S.ABI) { case coro::ABI::Switch: return std::unique_ptr( @@ -2221,7 +2229,17 @@ CreateNewABI(Function &F, coro::Shape &S, CoroSplitPass::CoroSplitPass(bool OptimizeFrame) : CreateAndInitABI([](Function &F, coro::Shape &S) { std::unique_ptr ABI = - CreateNewABI(F, S, coro::isTriviallyMaterializable); + CreateNewABI(F, S, coro::isTriviallyMaterializable, {}); + ABI->init(); + return ABI; + }), + OptimizeFrame(OptimizeFrame) {} + +CoroSplitPass::CoroSplitPass( + SmallVector GenCustomABIs, bool OptimizeFrame) + : CreateAndInitABI([=](Function &F, coro::Shape &S) { + std::unique_ptr ABI = + CreateNewABI(F, S, coro::isTriviallyMaterializable, GenCustomABIs); ABI->init(); return ABI; }), @@ -2232,7 +2250,19 @@ CoroSplitPass::CoroSplitPass(bool OptimizeFrame) CoroSplitPass::CoroSplitPass(std::function IsMatCallback, bool OptimizeFrame) : CreateAndInitABI([=](Function &F, coro::Shape &S) { - std::unique_ptr ABI = CreateNewABI(F, S, IsMatCallback); + std::unique_ptr ABI = CreateNewABI(F, S, IsMatCallback, {}); + ABI->init(); + return ABI; + }), + OptimizeFrame(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 IsMatCallback, + SmallVector GenCustomABIs, bool OptimizeFrame) + : CreateAndInitABI([=](Function &F, coro::Shape &S) { + std::unique_ptr ABI = CreateNewABI(F, S, IsMatCallback, GenCustomABIs); ABI->init(); return ABI; }), diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp index f4d9a7a8aa856..1c45bcd7f6a83 100644 --- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp +++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp @@ -73,6 +73,7 @@ static const char *const CoroIntrinsics[] = { "llvm.coro.await.suspend.handle", "llvm.coro.await.suspend.void", "llvm.coro.begin", + "llvm.coro.begin.custom.abi", "llvm.coro.destroy", "llvm.coro.done", "llvm.coro.end", @@ -247,7 +248,8 @@ void coro::Shape::analyze(Function &F, } break; } - case Intrinsic::coro_begin: { + case Intrinsic::coro_begin: + case Intrinsic::coro_begin_custom_abi: { auto CB = cast(II); // Ignore coro id's that aren't pre-split. diff --git a/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp b/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp index 1d55889a32d7a..c3394fdaa940b 100644 --- a/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp +++ b/llvm/unittests/Transforms/Coroutines/ExtraRematTest.cpp @@ -182,4 +182,91 @@ TEST_F(ExtraRematTest, TestCoroRematWithCallback) { CallInst *CI = getCallByName(Resume1, "should.remat"); ASSERT_TRUE(CI); } + +StringRef TextCoroBeginCustomABI = R"( + define ptr @f(i32 %n) presplitcoroutine { + entry: + %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null) + %size = call i32 @llvm.coro.size.i32() + %alloc = call ptr @malloc(i32 %size) + %hdl = call ptr @llvm.coro.begin.custom.abi(token %id, ptr %alloc, i32 0) + + %inc1 = add i32 %n, 1 + %val2 = call i32 @should.remat(i32 %inc1) + %sp1 = call i8 @llvm.coro.suspend(token none, i1 false) + switch i8 %sp1, label %suspend [i8 0, label %resume1 + i8 1, label %cleanup] + resume1: + %inc2 = add i32 %val2, 1 + %sp2 = call i8 @llvm.coro.suspend(token none, i1 false) + switch i8 %sp1, label %suspend [i8 0, label %resume2 + i8 1, label %cleanup] + + resume2: + call void @print(i32 %val2) + call void @print(i32 %inc2) + br label %cleanup + + cleanup: + %mem = call ptr @llvm.coro.free(token %id, ptr %hdl) + call void @free(ptr %mem) + br label %suspend + suspend: + call i1 @llvm.coro.end(ptr %hdl, i1 0) + ret ptr %hdl + } + + declare ptr @llvm.coro.free(token, ptr) + declare i32 @llvm.coro.size.i32() + declare i8 @llvm.coro.suspend(token, i1) + declare void @llvm.coro.resume(ptr) + declare void @llvm.coro.destroy(ptr) + + declare token @llvm.coro.id(i32, ptr, ptr, ptr) + declare i1 @llvm.coro.alloc(token) + declare ptr @llvm.coro.begin.custom.abi(token, ptr, i32) + declare i1 @llvm.coro.end(ptr, i1) + + declare i32 @should.remat(i32) + + declare noalias ptr @malloc(i32) + declare void @print(i32) + declare void @free(ptr) + )"; + +// SwitchABI with overridden isMaterializable +class ExtraCustomABI : public coro::SwitchABI { +public: + ExtraCustomABI(Function &F, coro::Shape &S) + : coro::SwitchABI(F, S, ExtraMaterializable) {} +}; + +TEST_F(ExtraRematTest, TestCoroRematWithCustomABI) { + ParseAssembly(TextCoroBeginCustomABI); + + ASSERT_TRUE(M); + + CoroSplitPass::BaseABITy GenCustomABI = [](Function &F, coro::Shape &S) { + return std::unique_ptr(new ExtraCustomABI(F, S)); + }; + + CGSCCPassManager CGPM; + CGPM.addPass(CoroSplitPass({GenCustomABI})); + MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM))); + MPM.run(*M, MAM); + + // Verify that extra rematerializable instruction has been rematerialized + Function *F = M->getFunction("f.resume"); + ASSERT_TRUE(F) << "could not find split function f.resume"; + + BasicBlock *Resume1 = getBasicBlockByName(F, "resume1"); + ASSERT_TRUE(Resume1) + << "could not find expected BB resume1 in split function"; + + // With callback the extra rematerialization of the function should have + // happened + CallInst *CI = getCallByName(Resume1, "should.remat"); + ASSERT_TRUE(CI); +} + } // namespace From e6c1978a84d11ba7dccef863b806b763689fd3bf Mon Sep 17 00:00:00 2001 From: tnowicki Date: Mon, 9 Sep 2024 15:15:02 -0400 Subject: [PATCH 4/4] [Coroutines] Documentation for plugin libraries and custom ABIs --- llvm/docs/Coroutines.rst | 90 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst index 5679aefcb421d..d53f4ed2a4e70 100644 --- a/llvm/docs/Coroutines.rst +++ b/llvm/docs/Coroutines.rst @@ -312,6 +312,7 @@ 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 second parameter of `coro.begin` is given a block of memory to be used if the coroutine frame needs to be allocated dynamically. + The `coro.id`_ intrinsic serves as coroutine identity useful in cases when the `coro.begin`_ intrinsic get duplicated by optimization passes such as jump-threading. @@ -749,6 +750,65 @@ and python iterator `__next__` would look like: return *(int*)coro.promise(hdl, 4, false); } +Custom ABIs and Plugin Libraries +-------------------------------- + +Plugin libraries can extend coroutine lowering enabling a wide variety of users +to utilize the coroutine transformation passes. An existing coroutine lowering +is extended by: 1. defining custom ABIs that inherit from the existing ABIs, +2. give a list of generators for the custom ABIs when constructing the +`CoroSplit`_ pass, and 3. use `coro.begin.custom.abi` in place of `coro.begin` +with an additional parameter for the index of the generator/ABI to be used for +the coroutine. + +A custom ABI overriding the SwitchABI's materialization looks like: + +.. code-block:: c++ + + class CustomSwitchABI : public coro::SwitchABI { + public: + CustomSwitchABI(Function &F, coro::Shape &S) + : coro::SwitchABI(F, S, ExtraMaterializable) {} + }; + +Giving a list of custom ABI generators while constructing the `CoroSplit` +pass looks like: + +.. code-block:: c++ + + CoroSplitPass::BaseABITy GenCustomABI = [](Function &F, coro::Shape &S) { + return new CustomSwitchABI(F, S); + }; + + CGSCCPassManager CGPM; + CGPM.addPass(CoroSplitPass({GenCustomABI})); + +The LLVM IR for a coroutine using a Coroutine with a custom ABI looks like: + +.. code-block:: llvm + + define ptr @f(i32 %n) presplitcoroutine_custom_abi { + entry: + %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null) + %size = call i32 @llvm.coro.size.i32() + %alloc = call ptr @malloc(i32 %size) + %hdl = call noalias ptr @llvm.coro.begin.custom.abi(token %id, ptr %alloc, i32 0) + 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 ptr @llvm.coro.free(token %id, ptr %hdl) + call void @free(ptr %mem) + br label %suspend + suspend: + %unused = call i1 @llvm.coro.end(ptr %hdl, i1 false, token none) + ret ptr %hdl + } Intrinsics ========== @@ -1007,6 +1067,36 @@ with small positive and negative offsets). A frontend should emit exactly one `coro.begin` intrinsic per coroutine. +.. _coro.begin.custom.abi: + +'llvm.coro.begin.custom.abi' Intrinsic +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:: + + declare ptr @llvm.coro.begin.custom.abi(token , ptr , i32) + +Overview: +""""""""" + +The '``llvm.coro.begin.custom.abi``' intrinsic is used in place of the +`coro.begin` intrinsic with an additional parameter to specify the custom ABI +for the coroutine. The return is identical to that of the `coro.begin` +intrinsic. + +Arguments: +"""""""""" + +The first and second arguments are identical to those of the `coro.begin` +intrinsic. + +The third argument is an i32 index of the generator list given to the +`CoroSplit` pass specifying the custom ABI generator lor this coroutine. + +Semantics: +"""""""""" + +The semantics are identical to those of the `coro.begin` intrinsic. + .. _coro.free: 'llvm.coro.free' Intrinsic