From 162db89815eb7e3e7f574f9af6a737f9ff137526 Mon Sep 17 00:00:00 2001 From: tnowicki Date: Thu, 12 Sep 2024 17:13:46 -0400 Subject: [PATCH 1/5] [Coroutines] Define ABI objects for each type of lowering * Define an ABI object for Switch, Retcon, and Async lowerings * Perform initialization of each type of lowering as part of ABI initialization. * Make buildCoroutineFrame and splitCoroutine interfaces to the ABI. --- .../llvm/Transforms/Coroutines/CoroSplit.h | 15 ++- llvm/lib/Transforms/Coroutines/ABI.h | 109 ++++++++++++++++ llvm/lib/Transforms/Coroutines/CoroFrame.cpp | 7 +- llvm/lib/Transforms/Coroutines/CoroInternal.h | 3 - llvm/lib/Transforms/Coroutines/CoroShape.h | 1 - llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 117 ++++++++++++------ llvm/lib/Transforms/Coroutines/Coroutines.cpp | 43 +++---- 7 files changed, 224 insertions(+), 71 deletions(-) create mode 100644 llvm/lib/Transforms/Coroutines/ABI.h diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h index a2be1099ff68f..944ce602d4c10 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 0000000000000..20ef06f50c2fa --- /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 c08f56b024dfc..021b1f7a4156b 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 fcbd31878bdea..88d0f83c98c9e 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 dcfe94ca076bd..f4fb4baa6df31 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 382bdfff1926f..a17e16cc8e223 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 ef0abd2d823bb..c7a04b8fd965b 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"); } } From cb38006066883cb240b527788418fcb9ea081ad3 Mon Sep 17 00:00:00 2001 From: tnowicki Date: Thu, 12 Sep 2024 18:59:24 -0400 Subject: [PATCH 2/5] [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 provied 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 20ef06f50c2fa..0aa3d7bc6f728 100644 --- a/llvm/lib/Transforms/Coroutines/ABI.h +++ b/llvm/lib/Transforms/Coroutines/ABI.h @@ -42,7 +42,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 a17e16cc8e223..7c0de35319bf6 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -2052,7 +2052,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; @@ -2063,7 +2064,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(); @@ -2281,7 +2282,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; @@ -2291,7 +2292,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 c1c55a5048ffd64a8b668979a3e44078ae79c1ab Mon Sep 17 00:00:00 2001 From: tnowicki Date: Tue, 3 Sep 2024 15:10:54 -0400 Subject: [PATCH 3/5] [Coroutines] Move various util headers to include/llvm for plugin libraries * Move utilities to include/llvm/Transform/Coroutines --- .../llvm}/Transforms/Coroutines/ABI.h | 18 ++-- .../llvm}/Transforms/Coroutines/CoroInstr.h | 99 ++++++++++--------- .../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 | 4 +- 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, 98 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 0aa3d7bc6f728..9f896e9ad235c 100644 --- a/llvm/lib/Transforms/Coroutines/ABI.h +++ b/llvm/include/llvm/Transforms/Coroutines/ABI.h @@ -10,13 +10,13 @@ // coroutine 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 "MaterializationUtils.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 { @@ -29,7 +29,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, Shape &S) : F(F), Shape(S), IsMaterializable(coro::isTriviallyMaterializable) {} @@ -57,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) : BaseABI(F, S) {} @@ -72,7 +72,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) : BaseABI(F, S) {} @@ -87,7 +87,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) : BaseABI(F, S) {} 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..4e63cbd08a938 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,17 @@ class LLVM_LIBRARY_VISIBILITY CoroAllocaFreeInst : public IntrinsicInst { } }; +namespace coro { + +// Check for structural coroutine intrinsics that should not be spilled into +// the coroutine frame. +bool isCoroutineStructureIntrinsic(Instruction &I); + +// 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 944ce602d4c10..d37b6da512143 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 5f8efd1a8f32e..f36ae2274c79f 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 7c0de35319bf6..1bef5f697f150 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -19,8 +19,6 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Coroutines/CoroSplit.h" -#include "ABI.h" -#include "CoroInstr.h" #include "CoroInternal.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/PriorityWorklist.h" @@ -63,6 +61,8 @@ #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/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 c7a04b8fd965b..d591f1c6988c6 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 f213ac1c8d7d5..e85c43b3fbc95 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 0d20cfe7fabe9ccd0d9c7db25b489fa35195a98c Mon Sep 17 00:00:00 2001 From: tnowicki Date: Tue, 3 Sep 2024 15:36:01 -0400 Subject: [PATCH 4/5] [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 | 10 +++ .../lib/Transforms/Coroutines/CoroCleanup.cpp | 4 +- llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 47 +++++++++- llvm/lib/Transforms/Coroutines/Coroutines.cpp | 4 +- .../Transforms/Coroutines/ExtraRematTest.cpp | 87 +++++++++++++++++++ 9 files changed, 171 insertions(+), 12 deletions(-) diff --git a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h index c592bc8f6ba2a..d44b1dff9b4ae 100644 --- a/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h +++ b/llvm/include/llvm/Analysis/TargetTransformInfoImpl.h @@ -767,6 +767,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 0a74a217a5f01..7a9124d4ead3c 100644 --- a/llvm/include/llvm/IR/Intrinsics.td +++ b/llvm/include/llvm/IR/Intrinsics.td @@ -1709,7 +1709,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 9f896e9ad235c..986ab14b5f5a1 100644 --- a/llvm/include/llvm/Transforms/Coroutines/ABI.h +++ b/llvm/include/llvm/Transforms/Coroutines/ABI.h @@ -27,7 +27,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 4e63cbd08a938..982e61c3f1c0c 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 d37b6da512143..997c15e76593e 100644 --- a/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h +++ b/llvm/include/llvm/Transforms/Coroutines/CoroSplit.h @@ -32,7 +32,17 @@ struct CoroSplitPass : PassInfoMixin { using BaseABITy = std::function; 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, 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 1bef5f697f150..a22d6b981807c 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -2197,7 +2197,16 @@ static void addPrepareFunction(const Module &M, Fns.push_back(PrepareFn); } -static coro::BaseABI *CreateNewABI(Function &F, coro::Shape &S) { +static coro::BaseABI * +CreateNewABI(Function &F, coro::Shape &S, + 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 new coro::SwitchABI(F, S); @@ -2213,7 +2222,16 @@ static coro::BaseABI *CreateNewABI(Function &F, coro::Shape &S) { CoroSplitPass::CoroSplitPass(bool OptimizeFrame) : CreateAndInitABI([](Function &F, coro::Shape &S) { - coro::BaseABI *ABI = CreateNewABI(F, S); + coro::BaseABI *ABI = CreateNewABI(F, S, {}); + ABI->init(); + return ABI; + }), + OptimizeFrame(OptimizeFrame) {} + +CoroSplitPass::CoroSplitPass( + SmallVector GenCustomABIs, bool OptimizeFrame) + : CreateAndInitABI([=](Function &F, coro::Shape &S) { + coro::BaseABI *ABI = CreateNewABI(F, S, GenCustomABIs); ABI->init(); return ABI; }), @@ -2221,7 +2239,15 @@ CoroSplitPass::CoroSplitPass(bool OptimizeFrame) static coro::BaseABI * CreateNewABIIsMat(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 new coro::SwitchABI(F, S, IsMatCallback); @@ -2240,7 +2266,20 @@ CreateNewABIIsMat(Function &F, coro::Shape &S, CoroSplitPass::CoroSplitPass(std::function IsMatCallback, bool OptimizeFrame) : CreateAndInitABI([=](Function &F, coro::Shape &S) { - coro::BaseABI *ABI = CreateNewABIIsMat(F, S, IsMatCallback); + coro::BaseABI *ABI = CreateNewABIIsMat(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) { + coro::BaseABI *ABI = + CreateNewABIIsMat(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 d591f1c6988c6..16c6468665c30 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", @@ -246,7 +247,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..e52c803b59abd 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 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 225685b065faf82492d221606d19c46ad1dc37f2 Mon Sep 17 00:00:00 2001 From: tnowicki Date: Wed, 4 Sep 2024 18:44:27 -0400 Subject: [PATCH 5/5] [Coroutines] Improve dump of BB label to avoid str copies --- .../Coroutines/SuspendCrossingInfo.h | 5 ++- .../Coroutines/MaterializationUtils.cpp | 25 +++++++----- .../Coroutines/SuspendCrossingInfo.cpp | 39 +++++++++++-------- 3 files changed, 43 insertions(+), 26 deletions(-) diff --git a/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h index 49cae6dde47e5..88cbf88acc4cd 100644 --- a/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h +++ b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h @@ -25,6 +25,8 @@ namespace llvm { +class ModuleSlotTracker; + // Provides two way mapping between the blocks and numbers. class BlockToIndexMapping { SmallVector V; @@ -96,7 +98,8 @@ class SuspendCrossingInfo { // Print order is in RPO void dump() const; void dump(StringRef Label, BitVector const &BV, - const ReversePostOrderTraversal &RPOT) const; + const ReversePostOrderTraversal &RPOT, + ModuleSlotTracker &MST) const; #endif SuspendCrossingInfo(Function &F, diff --git a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp index c3ea0977d4211..6327cea64c0de 100644 --- a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp +++ b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp @@ -15,6 +15,7 @@ #include "llvm/IR/Dominators.h" #include "llvm/IR/InstIterator.h" #include "llvm/IR/Instruction.h" +#include "llvm/IR/ModuleSlotTracker.h" #include "llvm/Transforms/Coroutines/SpillUtils.h" #include @@ -104,19 +105,25 @@ struct RematGraph { } #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) - static std::string getBasicBlockLabel(const BasicBlock *BB) { - if (BB->hasName()) - return BB->getName().str(); - - std::string S; - raw_string_ostream OS(S); - BB->printAsOperand(OS, false); - return OS.str().substr(1); + static void dumpBasicBlockLabel(const BasicBlock *BB, + ModuleSlotTracker &MST) { + if (BB->hasName()) { + dbgs() << BB->getName(); + return; + } + + dbgs() << MST.getLocalSlot(BB); } void dump() const { + BasicBlock *BB = EntryNode->Node->getParent(); + Function *F = BB->getParent(); + + ModuleSlotTracker MST(F->getParent()); + MST.incorporateFunction(*F); + dbgs() << "Entry ("; - dbgs() << getBasicBlockLabel(EntryNode->Node->getParent()); + dumpBasicBlockLabel(BB, MST); dbgs() << ") : " << *EntryNode->Node << "\n"; for (auto &E : Remats) { dbgs() << *(E.first) << "\n"; diff --git a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp index f18f23306befb..c9bb3395a9949 100644 --- a/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp +++ b/llvm/lib/Transforms/Coroutines/SuspendCrossingInfo.cpp @@ -13,6 +13,7 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Coroutines/SuspendCrossingInfo.h" +#include "llvm/IR/ModuleSlotTracker.h" // The "coro-suspend-crossing" flag is very noisy. There is another debug type, // "coro-frame", which results in leaner debug spew. @@ -20,24 +21,26 @@ namespace llvm { #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) -static std::string getBasicBlockLabel(const BasicBlock *BB) { - if (BB->hasName()) - return BB->getName().str(); - - std::string S; - raw_string_ostream OS(S); - BB->printAsOperand(OS, false); - return OS.str().substr(1); +static void dumpBasicBlockLabel(const BasicBlock *BB, ModuleSlotTracker &MST) { + if (BB->hasName()) { + dbgs() << BB->getName(); + return; + } + + dbgs() << MST.getLocalSlot(BB); } -LLVM_DUMP_METHOD void SuspendCrossingInfo::dump( - StringRef Label, BitVector const &BV, - const ReversePostOrderTraversal &RPOT) const { +LLVM_DUMP_METHOD void +SuspendCrossingInfo::dump(StringRef Label, BitVector const &BV, + const ReversePostOrderTraversal &RPOT, + ModuleSlotTracker &MST) const { dbgs() << Label << ":"; for (const BasicBlock *BB : RPOT) { auto BBNo = Mapping.blockToIndex(BB); - if (BV[BBNo]) - dbgs() << " " << getBasicBlockLabel(BB); + if (BV[BBNo]) { + dbgs() << " "; + dumpBasicBlockLabel(BB, MST); + } } dbgs() << "\n"; } @@ -49,12 +52,16 @@ LLVM_DUMP_METHOD void SuspendCrossingInfo::dump() const { BasicBlock *const B = Mapping.indexToBlock(0); Function *F = B->getParent(); + ModuleSlotTracker MST(F->getParent()); + MST.incorporateFunction(*F); + ReversePostOrderTraversal RPOT(F); for (const BasicBlock *BB : RPOT) { auto BBNo = Mapping.blockToIndex(BB); - dbgs() << getBasicBlockLabel(BB) << ":\n"; - dump(" Consumes", Block[BBNo].Consumes, RPOT); - dump(" Kills", Block[BBNo].Kills, RPOT); + dumpBasicBlockLabel(BB, MST); + dbgs() << ":\n"; + dump(" Consumes", Block[BBNo].Consumes, RPOT, MST); + dump(" Kills", Block[BBNo].Kills, RPOT, MST); } dbgs() << "\n"; }