Skip to content

Commit 66227bf

Browse files
TylerNowickiTylerNowicki
authored
[Coroutines] ABI Objects to improve code separation between different ABIs, users and utilities. (#109713)
This patch re-lands #109338 and fixes the various test failures. --- Original description --- * Adds an ABI object class hierarchy to implement the coroutine ABIs (Switch, Asyc, and Retcon{Once}) * The ABI object improves the separation of the code related to users, ABIs and utilities. * No code changes are required by any existing users. * Each ABI overrides delegate methods for initialization, building the coroutine frame and splitting the coroutine, other methods may be added later. * CoroSplit invokes a generator lambda to instantiate the ABI object and calls the ABI object to carry out its primary operations. In a follow-up change this will be used to instantiated customized ABIs according to a new intrinsic llvm.coro.begin.custom.abi. * Note, in a follow-up change additional constructors will be added to the ABI objects (Switch, Asyc, and AnyRetcon) to allow a set of generator lambdas to be passed in from a CoroSplit constructor that will also be added. The init() method is separate from the constructor to avoid duplication of its code. It is a virtual method so it can be invoked by CreateAndInitABI or potentially CoroSplit::run(). I wasn't sure if we should call init() from within CoroSplit::run(), CreateAndInitABI, or perhaps hard code a call in each constructor. One consideration is that init() can change the IR, so perhaps it should appear in CoroSplit::run(), but this looks a bit odd. Invoking init() in the constructor would result in the call appearing 4 times per ABI (after adding the new constructors). Invoking it in CreateAndInitABI seems to be a balance between these. See RFC for more info: https://discourse.llvm.org/t/rfc-abi-objects-for-coroutines/81057
1 parent 66f846d commit 66227bf

File tree

7 files changed

+209
-69
lines changed

7 files changed

+209
-69
lines changed

‎llvm/include/llvm/Transforms/Coroutines/CoroSplit.h

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,19 +21,26 @@
2121

2222
namespace llvm {
2323

24+
namespace coro {
25+
class BaseABI;
26+
class Shape;
27+
} // namespace coro
28+
2429
struct CoroSplitPass : PassInfoMixin<CoroSplitPass> {
25-
const std::function<bool(Instruction &)> MaterializableCallback;
2630

2731
CoroSplitPass(bool OptimizeFrame = false);
2832
CoroSplitPass(std::function<bool(Instruction &)> MaterializableCallback,
29-
bool OptimizeFrame = false)
30-
: MaterializableCallback(MaterializableCallback),
31-
OptimizeFrame(OptimizeFrame) {}
33+
bool OptimizeFrame = false);
3234

3335
PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
3436
LazyCallGraph &CG, CGSCCUpdateResult &UR);
3537
static bool isRequired() { return true; }
3638

39+
using BaseABITy =
40+
std::function<std::unique_ptr<coro::BaseABI>(Function &, coro::Shape &)>;
41+
// Generator for an ABI transformer
42+
BaseABITy CreateAndInitABI;
43+
3744
// Would be true if the Optimization level isn't O0.
3845
bool OptimizeFrame;
3946
};

‎llvm/lib/Transforms/Coroutines/ABI.h

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
//===- ABI.h - Coroutine lowering class definitions (ABIs) ----*- C++ -*---===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
// This file defines coroutine lowering classes. The interface for coroutine
9+
// lowering is defined by BaseABI. Each lowering method (ABI) implements the
10+
// interface. Note that the enum class ABI, such as ABI::Switch, determines
11+
// which ABI class, such as SwitchABI, is used to lower the coroutine. Both the
12+
// ABI enum and ABI class are used by the Coroutine passes when lowering.
13+
//===----------------------------------------------------------------------===//
14+
15+
#ifndef LIB_TRANSFORMS_COROUTINES_ABI_H
16+
#define LIB_TRANSFORMS_COROUTINES_ABI_H
17+
18+
#include "CoroShape.h"
19+
#include "SuspendCrossingInfo.h"
20+
#include "llvm/Analysis/TargetTransformInfo.h"
21+
22+
namespace llvm {
23+
24+
class Function;
25+
26+
namespace coro {
27+
28+
// This interface/API is to provide an object oriented way to implement ABI
29+
// functionality. This is intended to replace use of the ABI enum to perform
30+
// ABI operations. The ABIs (e.g. Switch, Async, Retcon{Once}) are the common
31+
// ABIs.
32+
33+
class LLVM_LIBRARY_VISIBILITY BaseABI {
34+
public:
35+
BaseABI(Function &F, coro::Shape &S,
36+
std::function<bool(Instruction &)> IsMaterializable)
37+
: F(F), Shape(S), IsMaterializable(IsMaterializable) {}
38+
virtual ~BaseABI() = default;
39+
40+
// Initialize the coroutine ABI
41+
virtual void init() = 0;
42+
43+
// Allocate the coroutine frame and do spill/reload as needed.
44+
virtual void buildCoroutineFrame();
45+
46+
// Perform the function splitting according to the ABI.
47+
virtual void splitCoroutine(Function &F, coro::Shape &Shape,
48+
SmallVectorImpl<Function *> &Clones,
49+
TargetTransformInfo &TTI) = 0;
50+
51+
Function &F;
52+
coro::Shape &Shape;
53+
54+
// Callback used by coro::BaseABI::buildCoroutineFrame for rematerialization.
55+
// It is provided to coro::doMaterializations(..).
56+
std::function<bool(Instruction &I)> IsMaterializable;
57+
};
58+
59+
class LLVM_LIBRARY_VISIBILITY SwitchABI : public BaseABI {
60+
public:
61+
SwitchABI(Function &F, coro::Shape &S,
62+
std::function<bool(Instruction &)> IsMaterializable)
63+
: BaseABI(F, S, IsMaterializable) {}
64+
65+
void init() override;
66+
67+
void splitCoroutine(Function &F, coro::Shape &Shape,
68+
SmallVectorImpl<Function *> &Clones,
69+
TargetTransformInfo &TTI) override;
70+
};
71+
72+
class LLVM_LIBRARY_VISIBILITY AsyncABI : public BaseABI {
73+
public:
74+
AsyncABI(Function &F, coro::Shape &S,
75+
std::function<bool(Instruction &)> IsMaterializable)
76+
: BaseABI(F, S, IsMaterializable) {}
77+
78+
void init() override;
79+
80+
void splitCoroutine(Function &F, coro::Shape &Shape,
81+
SmallVectorImpl<Function *> &Clones,
82+
TargetTransformInfo &TTI) override;
83+
};
84+
85+
class LLVM_LIBRARY_VISIBILITY AnyRetconABI : public BaseABI {
86+
public:
87+
AnyRetconABI(Function &F, coro::Shape &S,
88+
std::function<bool(Instruction &)> IsMaterializable)
89+
: BaseABI(F, S, IsMaterializable) {}
90+
91+
void init() override;
92+
93+
void splitCoroutine(Function &F, coro::Shape &Shape,
94+
SmallVectorImpl<Function *> &Clones,
95+
TargetTransformInfo &TTI) override;
96+
};
97+
98+
} // end namespace coro
99+
100+
} // end namespace llvm
101+
102+
#endif // LLVM_TRANSFORMS_COROUTINES_ABI_H

‎llvm/lib/Transforms/Coroutines/CoroFrame.cpp

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
// the value into the coroutine frame.
1616
//===----------------------------------------------------------------------===//
1717

18+
#include "ABI.h"
1819
#include "CoroInternal.h"
1920
#include "MaterializationUtils.h"
2021
#include "SpillUtils.h"
@@ -2055,11 +2056,9 @@ void coro::normalizeCoroutine(Function &F, coro::Shape &Shape,
20552056
rewritePHIs(F);
20562057
}
20572058

2058-
void coro::buildCoroutineFrame(
2059-
Function &F, Shape &Shape,
2060-
const std::function<bool(Instruction &)> &MaterializableCallback) {
2059+
void coro::BaseABI::buildCoroutineFrame() {
20612060
SuspendCrossingInfo Checker(F, Shape.CoroSuspends, Shape.CoroEnds);
2062-
doRematerializations(F, Checker, MaterializableCallback);
2061+
doRematerializations(F, Checker, IsMaterializable);
20632062

20642063
const DominatorTree DT(F);
20652064
if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&

‎llvm/lib/Transforms/Coroutines/CoroInternal.h

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,6 @@ struct LowererBase {
6262
bool defaultMaterializable(Instruction &V);
6363
void normalizeCoroutine(Function &F, coro::Shape &Shape,
6464
TargetTransformInfo &TTI);
65-
void buildCoroutineFrame(
66-
Function &F, Shape &Shape,
67-
const std::function<bool(Instruction &)> &MaterializableCallback);
6865
CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
6966
TargetTransformInfo &TTI,
7067
ArrayRef<Value *> Arguments, IRBuilder<> &);

‎llvm/lib/Transforms/Coroutines/CoroShape.h

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
275275
invalidateCoroutine(F, CoroFrames);
276276
return;
277277
}
278-
initABI();
279278
cleanCoroutine(CoroFrames, UnusedCoroSaves);
280279
}
281280
};

‎llvm/lib/Transforms/Coroutines/CoroSplit.cpp

Lines changed: 73 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@
1919
//===----------------------------------------------------------------------===//
2020

2121
#include "llvm/Transforms/Coroutines/CoroSplit.h"
22+
#include "ABI.h"
2223
#include "CoroInstr.h"
2324
#include "CoroInternal.h"
25+
#include "MaterializationUtils.h"
2426
#include "llvm/ADT/DenseMap.h"
2527
#include "llvm/ADT/PriorityWorklist.h"
2628
#include "llvm/ADT/SmallPtrSet.h"
@@ -1779,9 +1781,9 @@ CallInst *coro::createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
17791781
return TailCall;
17801782
}
17811783

1782-
static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
1783-
SmallVectorImpl<Function *> &Clones,
1784-
TargetTransformInfo &TTI) {
1784+
void coro::AsyncABI::splitCoroutine(Function &F, coro::Shape &Shape,
1785+
SmallVectorImpl<Function *> &Clones,
1786+
TargetTransformInfo &TTI) {
17851787
assert(Shape.ABI == coro::ABI::Async);
17861788
assert(Clones.empty());
17871789
// Reset various things that the optimizer might have decided it
@@ -1874,9 +1876,9 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
18741876
}
18751877
}
18761878

1877-
static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
1878-
SmallVectorImpl<Function *> &Clones,
1879-
TargetTransformInfo &TTI) {
1879+
void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,
1880+
SmallVectorImpl<Function *> &Clones,
1881+
TargetTransformInfo &TTI) {
18801882
assert(Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce);
18811883
assert(Clones.empty());
18821884

@@ -2044,26 +2046,27 @@ static bool hasSafeElideCaller(Function &F) {
20442046
return false;
20452047
}
20462048

2047-
static coro::Shape
2048-
splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
2049-
TargetTransformInfo &TTI, bool OptimizeFrame,
2050-
std::function<bool(Instruction &)> MaterializableCallback) {
2051-
PrettyStackTraceFunction prettyStackTrace(F);
2049+
void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
2050+
SmallVectorImpl<Function *> &Clones,
2051+
TargetTransformInfo &TTI) {
2052+
SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
2053+
}
20522054

2053-
// The suspend-crossing algorithm in buildCoroutineFrame get tripped
2054-
// up by uses in unreachable blocks, so remove them as a first pass.
2055-
removeUnreachableBlocks(F);
2055+
static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
2056+
coro::BaseABI &ABI, TargetTransformInfo &TTI) {
2057+
PrettyStackTraceFunction prettyStackTrace(F);
20562058

2057-
coro::Shape Shape(F, OptimizeFrame);
2058-
if (!Shape.CoroBegin)
2059-
return Shape;
2059+
auto &Shape = ABI.Shape;
2060+
assert(Shape.CoroBegin);
20602061

20612062
lowerAwaitSuspends(F, Shape);
20622063

20632064
simplifySuspendPoints(Shape);
2065+
20642066
normalizeCoroutine(F, Shape, TTI);
2065-
buildCoroutineFrame(F, Shape, MaterializableCallback);
2067+
ABI.buildCoroutineFrame();
20662068
replaceFrameSizeAndAlignment(Shape);
2069+
20672070
bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
20682071

20692072
bool shouldCreateNoAllocVariant = !isNoSuspendCoroutine &&
@@ -2075,18 +2078,7 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
20752078
if (isNoSuspendCoroutine) {
20762079
handleNoSuspendCoroutine(Shape);
20772080
} else {
2078-
switch (Shape.ABI) {
2079-
case coro::ABI::Switch:
2080-
SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
2081-
break;
2082-
case coro::ABI::Async:
2083-
splitAsyncCoroutine(F, Shape, Clones, TTI);
2084-
break;
2085-
case coro::ABI::Retcon:
2086-
case coro::ABI::RetconOnce:
2087-
splitRetconCoroutine(F, Shape, Clones, TTI);
2088-
break;
2089-
}
2081+
ABI.splitCoroutine(F, Shape, Clones, TTI);
20902082
}
20912083

20922084
// Replace all the swifterror operations in the original function.
@@ -2107,8 +2099,6 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
21072099

21082100
if (shouldCreateNoAllocVariant)
21092101
SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
2110-
2111-
return Shape;
21122102
}
21132103

21142104
static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
@@ -2207,8 +2197,44 @@ static void addPrepareFunction(const Module &M,
22072197
Fns.push_back(PrepareFn);
22082198
}
22092199

2200+
static std::unique_ptr<coro::BaseABI>
2201+
CreateNewABI(Function &F, coro::Shape &S,
2202+
std::function<bool(Instruction &)> IsMatCallback) {
2203+
switch (S.ABI) {
2204+
case coro::ABI::Switch:
2205+
return std::unique_ptr<coro::BaseABI>(
2206+
new coro::SwitchABI(F, S, IsMatCallback));
2207+
case coro::ABI::Async:
2208+
return std::unique_ptr<coro::BaseABI>(
2209+
new coro::AsyncABI(F, S, IsMatCallback));
2210+
case coro::ABI::Retcon:
2211+
return std::unique_ptr<coro::BaseABI>(
2212+
new coro::AnyRetconABI(F, S, IsMatCallback));
2213+
case coro::ABI::RetconOnce:
2214+
return std::unique_ptr<coro::BaseABI>(
2215+
new coro::AnyRetconABI(F, S, IsMatCallback));
2216+
}
2217+
llvm_unreachable("Unknown ABI");
2218+
}
2219+
22102220
CoroSplitPass::CoroSplitPass(bool OptimizeFrame)
2211-
: MaterializableCallback(coro::defaultMaterializable),
2221+
: CreateAndInitABI([](Function &F, coro::Shape &S) {
2222+
std::unique_ptr<coro::BaseABI> ABI =
2223+
CreateNewABI(F, S, coro::isTriviallyMaterializable);
2224+
ABI->init();
2225+
return std::move(ABI);
2226+
}),
2227+
OptimizeFrame(OptimizeFrame) {}
2228+
2229+
// For back compatibility, constructor takes a materializable callback and
2230+
// creates a generator for an ABI with a modified materializable callback.
2231+
CoroSplitPass::CoroSplitPass(std::function<bool(Instruction &)> IsMatCallback,
2232+
bool OptimizeFrame)
2233+
: CreateAndInitABI([=](Function &F, coro::Shape &S) {
2234+
std::unique_ptr<coro::BaseABI> ABI = CreateNewABI(F, S, IsMatCallback);
2235+
ABI->init();
2236+
return std::move(ABI);
2237+
}),
22122238
OptimizeFrame(OptimizeFrame) {}
22132239

22142240
PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
@@ -2241,12 +2267,23 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
22412267
Function &F = N->getFunction();
22422268
LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
22432269
<< "\n");
2270+
2271+
// The suspend-crossing algorithm in buildCoroutineFrame gets tripped up
2272+
// by unreachable blocks, so remove them as a first pass. Remove the
2273+
// unreachable blocks before collecting intrinsics into Shape.
2274+
removeUnreachableBlocks(F);
2275+
2276+
coro::Shape Shape(F, OptimizeFrame);
2277+
if (!Shape.CoroBegin)
2278+
continue;
2279+
22442280
F.setSplittedCoroutine();
22452281

2282+
std::unique_ptr<coro::BaseABI> ABI = CreateAndInitABI(F, Shape);
2283+
22462284
SmallVector<Function *, 4> Clones;
2247-
coro::Shape Shape =
2248-
splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
2249-
OptimizeFrame, MaterializableCallback);
2285+
auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
2286+
doSplitCoroutine(F, Clones, *ABI, TTI);
22502287
CurrentSCC = &updateCallGraphAfterCoroutineSplit(
22512288
*N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);
22522289

0 commit comments

Comments
 (0)