Skip to content

Commit

Permalink
[Coroutines] Move OptimizeFrame out of Shape (#111017)
Browse files Browse the repository at this point in the history
* 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.

Co-authored-by: tnowicki <[email protected]>
  • Loading branch information
TylerNowicki and tnowicki authored Oct 8, 2024
1 parent c0a2915 commit 6e5d612
Show file tree
Hide file tree
Showing 4 changed files with 16 additions and 17 deletions.
2 changes: 1 addition & 1 deletion llvm/lib/Transforms/Coroutines/ABI.h
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
16 changes: 9 additions & 7 deletions llvm/lib/Transforms/Coroutines/CoroFrame.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<AllocaInst *, 4>;
SmallVector<AllocaSetType, 4> NonOverlapedAllocas;

Expand All @@ -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));
Expand Down Expand Up @@ -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 = [&] {
Expand Down Expand Up @@ -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`
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);
Expand Down
6 changes: 1 addition & 5 deletions llvm/lib/Transforms/Coroutines/CoroShape.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<CoroFrameInst *, 8> CoroFrames;
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;

Expand Down
9 changes: 5 additions & 4 deletions llvm/lib/Transforms/Coroutines/CoroSplit.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2053,7 +2053,8 @@ void coro::SwitchABI::splitCoroutine(Function &F, coro::Shape &Shape,
}

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

auto &Shape = ABI.Shape;
Expand All @@ -2064,7 +2065,7 @@ static void doSplitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
simplifySuspendPoints(Shape);

normalizeCoroutine(F, Shape, TTI);
ABI.buildCoroutineFrame();
ABI.buildCoroutineFrame(OptimizeFrame);
replaceFrameSizeAndAlignment(Shape);

bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
Expand Down Expand Up @@ -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;

Expand All @@ -2283,7 +2284,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,

SmallVector<Function *, 4> Clones;
auto &TTI = FAM.getResult<TargetIRAnalysis>(F);
doSplitCoroutine(F, Clones, *ABI, TTI);
doSplitCoroutine(F, Clones, *ABI, TTI, OptimizeFrame);
CurrentSCC = &updateCallGraphAfterCoroutineSplit(
*N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);

Expand Down

0 comments on commit 6e5d612

Please sign in to comment.