Skip to content

Commit c641497

Browse files
yuxuanchen1997yuxuanchen1997
authored
[Coroutines] Inline the .noalloc ramp function marked coro_safe_elide (#114004)
1 parent 37b4df4 commit c641497

File tree

7 files changed

+159
-58
lines changed

7 files changed

+159
-58
lines changed

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,18 @@
1717
#ifndef LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
1818
#define LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
1919

20+
#include "llvm/Analysis/CGSCCPassManager.h"
21+
#include "llvm/Analysis/LazyCallGraph.h"
2022
#include "llvm/IR/PassManager.h"
2123

2224
namespace llvm {
2325

24-
class Function;
25-
2626
struct CoroAnnotationElidePass : PassInfoMixin<CoroAnnotationElidePass> {
27-
PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
27+
CoroAnnotationElidePass() {}
28+
29+
PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
30+
LazyCallGraph &CG, CGSCCUpdateResult &UR);
31+
2832
static bool isRequired() { return false; }
2933
};
3034
} // end namespace llvm

‎llvm/lib/Passes/PassBuilderPipelines.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -984,8 +984,7 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
984984

985985
if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink) {
986986
MainCGPipeline.addPass(CoroSplitPass(Level != OptimizationLevel::O0));
987-
MainCGPipeline.addPass(
988-
createCGSCCToFunctionPassAdaptor(CoroAnnotationElidePass()));
987+
MainCGPipeline.addPass(CoroAnnotationElidePass());
989988
}
990989

991990
// Make sure we don't affect potential future NoRerun CGSCC adaptors.
@@ -1036,7 +1035,8 @@ PassBuilder::buildModuleInlinerPipeline(OptimizationLevel Level,
10361035
if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink) {
10371036
MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(
10381037
CoroSplitPass(Level != OptimizationLevel::O0)));
1039-
MPM.addPass(createModuleToFunctionPassAdaptor(CoroAnnotationElidePass()));
1038+
MPM.addPass(
1039+
createModuleToPostOrderCGSCCPassAdaptor(CoroAnnotationElidePass()));
10401040
}
10411041

10421042
return MPM;

‎llvm/lib/Passes/PassRegistry.def‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ CGSCC_PASS("attributor-light-cgscc", AttributorLightCGSCCPass())
245245
CGSCC_PASS("invalidate<all>", InvalidateAllAnalysesPass())
246246
CGSCC_PASS("no-op-cgscc", NoOpCGSCCPass())
247247
CGSCC_PASS("openmp-opt-cgscc", OpenMPOptCGSCCPass())
248+
CGSCC_PASS("coro-annotation-elide", CoroAnnotationElidePass())
248249
#undef CGSCC_PASS
249250

250251
#ifndef CGSCC_PASS_WITH_PARAMS
@@ -346,7 +347,6 @@ FUNCTION_PASS("complex-deinterleaving", ComplexDeinterleavingPass(TM))
346347
FUNCTION_PASS("consthoist", ConstantHoistingPass())
347348
FUNCTION_PASS("constraint-elimination", ConstraintEliminationPass())
348349
FUNCTION_PASS("coro-elide", CoroElidePass())
349-
FUNCTION_PASS("coro-annotation-elide", CoroAnnotationElidePass())
350350
FUNCTION_PASS("correlated-propagation", CorrelatedValuePropagationPass())
351351
FUNCTION_PASS("count-visits", CountVisitsPass())
352352
FUNCTION_PASS("dce", DCEPass())

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

Lines changed: 76 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,16 @@
1616

1717
#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
1818

19+
#include "llvm/Analysis/CGSCCPassManager.h"
1920
#include "llvm/Analysis/LazyCallGraph.h"
2021
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
2122
#include "llvm/IR/Analysis.h"
2223
#include "llvm/IR/IRBuilder.h"
2324
#include "llvm/IR/Instruction.h"
2425
#include "llvm/IR/Module.h"
2526
#include "llvm/IR/PassManager.h"
27+
#include "llvm/Transforms/Utils/CallGraphUpdater.h"
28+
#include "llvm/Transforms/Utils/Cloning.h"
2629

2730
#include <cassert>
2831

@@ -90,61 +93,86 @@ static void processCall(CallBase *CB, Function *Caller, Function *NewCallee,
9093

9194
NewCB->removeFnAttr(llvm::Attribute::CoroElideSafe);
9295
CB->replaceAllUsesWith(NewCB);
93-
CB->eraseFromParent();
96+
97+
InlineFunctionInfo IFI;
98+
InlineResult IR = InlineFunction(*NewCB, IFI);
99+
if (IR.isSuccess()) {
100+
CB->eraseFromParent();
101+
} else {
102+
NewCB->replaceAllUsesWith(CB);
103+
NewCB->eraseFromParent();
104+
}
94105
}
95106

96-
PreservedAnalyses CoroAnnotationElidePass::run(Function &F,
97-
FunctionAnalysisManager &FAM) {
107+
PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
108+
CGSCCAnalysisManager &AM,
109+
LazyCallGraph &CG,
110+
CGSCCUpdateResult &UR) {
98111
bool Changed = false;
112+
CallGraphUpdater CGUpdater;
113+
CGUpdater.initialize(CG, C, AM, UR);
99114

100-
Function *NewCallee =
101-
F.getParent()->getFunction((F.getName() + ".noalloc").str());
102-
103-
if (!NewCallee)
104-
return PreservedAnalyses::all();
105-
106-
auto FramePtrArgPosition = NewCallee->arg_size() - 1;
107-
auto FrameSize = NewCallee->getParamDereferenceableBytes(FramePtrArgPosition);
108-
auto FrameAlign = NewCallee->getParamAlign(FramePtrArgPosition).valueOrOne();
115+
auto &FAM =
116+
AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
109117

110-
SmallVector<CallBase *, 4> Users;
111-
for (auto *U : F.users()) {
112-
if (auto *CB = dyn_cast<CallBase>(U)) {
113-
if (CB->getCalledFunction() == &F)
114-
Users.push_back(CB);
115-
}
116-
}
117-
118-
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
119-
120-
for (auto *CB : Users) {
121-
auto *Caller = CB->getFunction();
122-
if (!Caller)
118+
for (LazyCallGraph::Node &N : C) {
119+
Function *Callee = &N.getFunction();
120+
Function *NewCallee = Callee->getParent()->getFunction(
121+
(Callee->getName() + ".noalloc").str());
122+
if (!NewCallee)
123123
continue;
124124

125-
bool IsCallerPresplitCoroutine = Caller->isPresplitCoroutine();
126-
bool HasAttr = CB->hasFnAttr(llvm::Attribute::CoroElideSafe);
127-
if (IsCallerPresplitCoroutine && HasAttr) {
128-
processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
129-
130-
ORE.emit([&]() {
131-
return OptimizationRemark(DEBUG_TYPE, "CoroAnnotationElide", Caller)
132-
<< "'" << ore::NV("callee", F.getName()) << "' elided in '"
133-
<< ore::NV("caller", Caller->getName()) << "'";
134-
});
135-
136-
FAM.invalidate(*Caller, PreservedAnalyses::none());
137-
Changed = true;
138-
} else {
139-
ORE.emit([&]() {
140-
return OptimizationRemarkMissed(DEBUG_TYPE, "CoroAnnotationElide",
141-
Caller)
142-
<< "'" << ore::NV("callee", F.getName()) << "' not elided in '"
143-
<< ore::NV("caller", Caller->getName()) << "' (caller_presplit="
144-
<< ore::NV("caller_presplit", IsCallerPresplitCoroutine)
145-
<< ", elide_safe_attr=" << ore::NV("elide_safe_attr", HasAttr)
146-
<< ")";
147-
});
125+
SmallVector<CallBase *, 4> Users;
126+
for (auto *U : Callee->users()) {
127+
if (auto *CB = dyn_cast<CallBase>(U)) {
128+
if (CB->getCalledFunction() == Callee)
129+
Users.push_back(CB);
130+
}
131+
}
132+
auto FramePtrArgPosition = NewCallee->arg_size() - 1;
133+
auto FrameSize =
134+
NewCallee->getParamDereferenceableBytes(FramePtrArgPosition);
135+
auto FrameAlign =
136+
NewCallee->getParamAlign(FramePtrArgPosition).valueOrOne();
137+
138+
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(*Callee);
139+
140+
for (auto *CB : Users) {
141+
auto *Caller = CB->getFunction();
142+
if (!Caller)
143+
continue;
144+
145+
bool IsCallerPresplitCoroutine = Caller->isPresplitCoroutine();
146+
bool HasAttr = CB->hasFnAttr(llvm::Attribute::CoroElideSafe);
147+
if (IsCallerPresplitCoroutine && HasAttr) {
148+
auto *CallerN = CG.lookup(*Caller);
149+
auto *CallerC = CG.lookupSCC(*CallerN);
150+
processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
151+
152+
ORE.emit([&]() {
153+
return OptimizationRemark(DEBUG_TYPE, "CoroAnnotationElide", Caller)
154+
<< "'" << ore::NV("callee", Callee->getName())
155+
<< "' elided in '" << ore::NV("caller", Caller->getName())
156+
<< "'";
157+
});
158+
159+
FAM.invalidate(*Caller, PreservedAnalyses::none());
160+
Changed = true;
161+
updateCGAndAnalysisManagerForCGSCCPass(CG, *CallerC, *CallerN, AM, UR,
162+
FAM);
163+
164+
} else {
165+
ORE.emit([&]() {
166+
return OptimizationRemarkMissed(DEBUG_TYPE, "CoroAnnotationElide",
167+
Caller)
168+
<< "'" << ore::NV("callee", Callee->getName())
169+
<< "' not elided in '" << ore::NV("caller", Caller->getName())
170+
<< "' (caller_presplit="
171+
<< ore::NV("caller_presplit", IsCallerPresplitCoroutine)
172+
<< ", elide_safe_attr=" << ore::NV("elide_safe_attr", HasAttr)
173+
<< ")";
174+
});
175+
}
148176
}
149177
}
150178

‎llvm/test/Transforms/Coroutines/coro-transform-must-elide.ll‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,13 @@ define ptr @caller() #0 {
5959
entry:
6060
%task = call ptr @callee(i8 0) #1
6161
ret ptr %task
62-
63-
; CHECK: %[[FRAME:.+]] = alloca [32 x i8], align 8
64-
; CHECK-NEXT: %[[TASK:.+]] = call ptr @callee.noalloc(i8 0, ptr %[[FRAME]])
62+
; CHECK: %[[TASK:.+]] = alloca %struct.Task, align 8
63+
; CHECK-NEXT: %[[FRAME:.+]] = alloca [32 x i8], align 8
64+
; CHECK-NEXT: call void @llvm.lifetime.start.p0(i64 8, ptr %[[TASK]])
65+
; CHECK-NEXT: %[[ID:.+]] = call token @llvm.coro.id(i32 0, ptr null, ptr @callee, ptr @callee.resumers)
66+
; CHECK-NEXT: %[[HDL:.+]] = call ptr @llvm.coro.begin(token %[[ID]], ptr null)
67+
; CHECK-NEXT: store ptr %[[HDL]], ptr %[[TASK]], align 8
68+
; CHECK-NEXT: call void @llvm.lifetime.end.p0(i64 8, ptr %[[TASK]])
6569
; CHECK-NEXT: ret ptr %[[TASK]]
6670
}
6771

‎llvm/test/Transforms/Coroutines/gh114487-crash-in-cgscc.ll‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
; Verify that we don't crash when eliding coro_elide_safe callsites.
2+
; RUN: opt < %s -passes='cgscc(function<>(simplifycfg<>),function-attrs,coro-split,coro-annotation-elide)' -S | FileCheck %s
3+
4+
; CHECK-LABEL: define void @foo()
5+
define void @foo() presplitcoroutine personality ptr null {
6+
entry:
7+
%0 = call token @llvm.coro.save(ptr null)
8+
br label %branch
9+
10+
branch:
11+
; Check that we don't call bar at all.
12+
; CHECK-NOT: call void @bar{{.*}}
13+
call void @bar() coro_elide_safe
14+
; CHECK: call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @bar.resumers)
15+
ret void
16+
}
17+
18+
; CHECK-LABEL: define void @bar()
19+
define void @bar() presplitcoroutine personality ptr null {
20+
entry:
21+
%0 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
22+
%1 = call ptr @llvm.coro.begin(token %0, ptr null)
23+
%2 = call token @llvm.coro.save(ptr null)
24+
%3 = call i8 @llvm.coro.suspend(token none, i1 false)
25+
ret void
26+
}
27+
28+
declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) nounwind
29+
declare ptr @llvm.coro.begin(token, ptr writeonly) nounwind
30+
declare token @llvm.coro.save(ptr) nomerge nounwind
31+
declare i8 @llvm.coro.suspend(token, i1) nounwind
32+

‎llvm/test/Transforms/Coroutines/gh114487-non-inlinable.ll‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
; Verify that we don't crash when eliding coro_elide_safe callsites.
2+
; RUN: opt < %s -passes='cgscc(function<>(simplifycfg<>),function-attrs,coro-annotation-elide)' -S | FileCheck %s
3+
4+
; CHECK-LABEL: define void @foo()
5+
define void @foo() presplitcoroutine personality ptr null {
6+
entry:
7+
%0 = call token @llvm.coro.save(ptr null)
8+
br label %branch
9+
10+
branch:
11+
; Check that we still call bar() because we can't inline bar.noalloc()
12+
; CHECK: call void @bar()
13+
call void @bar() coro_elide_safe
14+
ret void
15+
}
16+
17+
; CHECK-LABEL: define void @bar()
18+
define void @bar() personality ptr null {
19+
entry:
20+
%0 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
21+
%1 = call ptr @llvm.coro.begin(token %0, ptr null)
22+
%2 = call token @llvm.coro.save(ptr null)
23+
%3 = call i8 @llvm.coro.suspend(token none, i1 false)
24+
ret void
25+
}
26+
27+
declare void @bar.noalloc(ptr)
28+
29+
declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) nounwind
30+
declare ptr @llvm.coro.begin(token, ptr writeonly) nounwind
31+
declare token @llvm.coro.save(ptr) nomerge nounwind
32+
declare i8 @llvm.coro.suspend(token, i1) nounwind
33+

0 commit comments

Comments
 (0)