Skip to content

Commit 7b89a7c

Browse files
yuxuanchen1997yuxuanchen1997
committed
add CoroAnnotationElidePass
Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60250514
1 parent 2b358ba commit 7b89a7c

14 files changed

+257
-1
lines changed

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//===- CoroAnnotationElide.h - Optimizing a coro_must_elide call ----------===//
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+
//
9+
// \file
10+
// This pass transforms all Call or Invoke instructions that are annotated
11+
// "coro_must_elide" to call the `.noalloc` variant of coroutine instead.
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
#ifndef LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
16+
#define LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
17+
18+
#include "llvm/Analysis/CGSCCPassManager.h"
19+
#include "llvm/Analysis/LazyCallGraph.h"
20+
#include "llvm/IR/PassManager.h"
21+
22+
namespace llvm {
23+
24+
struct CoroAnnotationElidePass : PassInfoMixin<CoroAnnotationElidePass> {
25+
CoroAnnotationElidePass() {}
26+
27+
PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
28+
LazyCallGraph &CG, CGSCCUpdateResult &UR);
29+
30+
static bool isRequired() { return false; }
31+
};
32+
} // end namespace llvm
33+
34+
#endif // LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H

‎llvm/lib/Passes/PassBuilder.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@
135135
#include "llvm/Target/TargetMachine.h"
136136
#include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
137137
#include "llvm/Transforms/CFGuard.h"
138+
#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
138139
#include "llvm/Transforms/Coroutines/CoroCleanup.h"
139140
#include "llvm/Transforms/Coroutines/CoroConditionalWrapper.h"
140141
#include "llvm/Transforms/Coroutines/CoroEarly.h"

‎llvm/lib/Passes/PassBuilderPipelines.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "llvm/Support/VirtualFileSystem.h"
3333
#include "llvm/Target/TargetMachine.h"
3434
#include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
35+
#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
3536
#include "llvm/Transforms/Coroutines/CoroCleanup.h"
3637
#include "llvm/Transforms/Coroutines/CoroConditionalWrapper.h"
3738
#include "llvm/Transforms/Coroutines/CoroEarly.h"
@@ -978,8 +979,8 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
978979
// it's been modified since.
979980
MainCGPipeline.addPass(createCGSCCToFunctionPassAdaptor(
980981
RequireAnalysisPass<ShouldNotRunFunctionPassesAnalysis, Function>()));
981-
982982
MainCGPipeline.addPass(CoroSplitPass(Level != OptimizationLevel::O0));
983+
MainCGPipeline.addPass(CoroAnnotationElidePass());
983984

984985
// Make sure we don't affect potential future NoRerun CGSCC adaptors.
985986
MIWP.addLateModulePass(createModuleToFunctionPassAdaptor(

‎llvm/lib/Passes/PassRegistry.def

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,7 @@ CGSCC_PASS("attributor-light-cgscc", AttributorLightCGSCCPass())
234234
CGSCC_PASS("invalidate<all>", InvalidateAllAnalysesPass())
235235
CGSCC_PASS("no-op-cgscc", NoOpCGSCCPass())
236236
CGSCC_PASS("openmp-opt-cgscc", OpenMPOptCGSCCPass())
237+
CGSCC_PASS("coro-annotation-elide", CoroAnnotationElidePass())
237238
#undef CGSCC_PASS
238239

239240
#ifndef CGSCC_PASS_WITH_PARAMS

‎llvm/lib/Transforms/Coroutines/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
add_llvm_component_library(LLVMCoroutines
22
Coroutines.cpp
3+
CoroAnnotationElide.cpp
34
CoroCleanup.cpp
45
CoroConditionalWrapper.cpp
56
CoroEarly.cpp

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

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
//===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
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+
9+
//===----------------------------------------------------------------------===//
10+
11+
#include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
12+
13+
#include "llvm/Analysis/LazyCallGraph.h"
14+
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
15+
#include "llvm/IR/Analysis.h"
16+
#include "llvm/IR/IRBuilder.h"
17+
#include "llvm/IR/InstIterator.h"
18+
#include "llvm/IR/Instruction.h"
19+
#include "llvm/IR/Module.h"
20+
#include "llvm/IR/PassManager.h"
21+
#include "llvm/Transforms/Utils/CallGraphUpdater.h"
22+
23+
#include <cassert>
24+
25+
using namespace llvm;
26+
27+
#define DEBUG_TYPE "coro-annotation-elide"
28+
29+
#define CORO_MUST_ELIDE_ANNOTATION "coro_must_elide"
30+
31+
static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
32+
for (Instruction &I : F->getEntryBlock())
33+
if (!isa<AllocaInst>(&I))
34+
return &I;
35+
llvm_unreachable("no terminator in the entry block");
36+
}
37+
38+
static Value *allocateFrameInCaller(Function *Caller, uint64_t FrameSize,
39+
Align FrameAlign) {
40+
LLVMContext &C = Caller->getContext();
41+
BasicBlock::iterator InsertPt =
42+
getFirstNonAllocaInTheEntryBlock(Caller)->getIterator();
43+
const DataLayout &DL = Caller->getDataLayout();
44+
auto FrameTy = ArrayType::get(Type::getInt8Ty(C), FrameSize);
45+
auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);
46+
Frame->setAlignment(FrameAlign);
47+
return new BitCastInst(Frame, PointerType::getUnqual(C), "vFrame", InsertPt);
48+
}
49+
50+
static void processCall(CallBase *CB, Function *Caller, Function *NewCallee,
51+
uint64_t FrameSize, Align FrameAlign) {
52+
auto *FramePtr = allocateFrameInCaller(Caller, FrameSize, FrameAlign);
53+
auto NewCBInsertPt = CB->getIterator();
54+
llvm::CallBase *NewCB = nullptr;
55+
SmallVector<Value *, 4> NewArgs;
56+
NewArgs.append(CB->arg_begin(), CB->arg_end());
57+
NewArgs.push_back(FramePtr);
58+
59+
if (auto *CI = dyn_cast<CallInst>(CB)) {
60+
auto *NewCI = CallInst::Create(NewCallee->getFunctionType(), NewCallee,
61+
NewArgs, "", NewCBInsertPt);
62+
NewCI->setTailCallKind(CI->getTailCallKind());
63+
NewCB = NewCI;
64+
} else if (auto *II = dyn_cast<InvokeInst>(CB)) {
65+
NewCB = InvokeInst::Create(NewCallee->getFunctionType(), NewCallee,
66+
II->getNormalDest(), II->getUnwindDest(),
67+
NewArgs, std::nullopt, "", NewCBInsertPt);
68+
} else {
69+
llvm_unreachable("CallBase should either be Call or Invoke!");
70+
}
71+
72+
NewCB->setCalledFunction(NewCallee->getFunctionType(), NewCallee);
73+
NewCB->setCallingConv(CB->getCallingConv());
74+
NewCB->setAttributes(CB->getAttributes());
75+
NewCB->setDebugLoc(CB->getDebugLoc());
76+
std::copy(CB->bundle_op_info_begin(), CB->bundle_op_info_end(),
77+
NewCB->bundle_op_info_begin());
78+
79+
NewCB->removeFnAttr(llvm::Attribute::CoroMustElide);
80+
CB->replaceAllUsesWith(NewCB);
81+
CB->eraseFromParent();
82+
}
83+
84+
PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
85+
CGSCCAnalysisManager &AM,
86+
LazyCallGraph &CG,
87+
CGSCCUpdateResult &UR) {
88+
bool Changed = false;
89+
CallGraphUpdater CGUpdater;
90+
CGUpdater.initialize(CG, C, AM, UR);
91+
92+
auto &FAM =
93+
AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
94+
95+
for (LazyCallGraph::Node &N : C) {
96+
Function *Callee = &N.getFunction();
97+
Function *NewCallee = Callee->getParent()->getFunction(
98+
(Callee->getName() + ".noalloc").str());
99+
if (!NewCallee) {
100+
continue;
101+
}
102+
103+
auto FramePtrArgPosition = NewCallee->arg_size() - 1;
104+
auto FrameSize =
105+
NewCallee->getParamDereferenceableBytes(FramePtrArgPosition);
106+
auto FrameAlign =
107+
NewCallee->getParamAlign(FramePtrArgPosition).valueOrOne();
108+
109+
SmallVector<CallBase *, 4> Users;
110+
for (auto *U : Callee->users()) {
111+
if (auto *CB = dyn_cast<CallBase>(U)) {
112+
Users.push_back(CB);
113+
}
114+
}
115+
116+
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(*Callee);
117+
118+
for (auto *CB : Users) {
119+
auto *Caller = CB->getFunction();
120+
if (Caller && Caller->isPresplitCoroutine() &&
121+
CB->hasFnAttr(llvm::Attribute::CoroMustElide)) {
122+
processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
123+
CGUpdater.reanalyzeFunction(*Caller);
124+
125+
ORE.emit([&]() {
126+
return OptimizationRemark(DEBUG_TYPE, "CoroAnnotationElide", Caller)
127+
<< "'" << ore::NV("callee", Callee->getName())
128+
<< "' elided in '" << ore::NV("caller", Caller->getName());
129+
});
130+
Changed = true;
131+
}
132+
}
133+
}
134+
return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
135+
}

‎llvm/test/Other/new-pm-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,7 @@
226226
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
227227
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
228228
; CHECK-O-NEXT: Running pass: CoroSplitPass
229+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
229230
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
230231
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
231232
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@
153153
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
154154
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
155155
; CHECK-O-NEXT: Running pass: CoroSplitPass
156+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
156157
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
157158
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
158159
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@
137137
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
138138
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
139139
; CHECK-O-NEXT: Running pass: CoroSplitPass
140+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
140141
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
141142
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
142143
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@
145145
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
146146
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
147147
; CHECK-O-NEXT: Running pass: CoroSplitPass
148+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
148149
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
149150
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
150151
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-prelink-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,7 @@
185185
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
186186
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
187187
; CHECK-O-NEXT: Running pass: CoroSplitPass
188+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
188189
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
189190
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
190191
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@
184184
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
185185
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
186186
; CHECK-O-NEXT: Running pass: CoroSplitPass
187+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
187188
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
188189
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
189190
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

‎llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@
149149
; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
150150
; CHECK-O-NEXT: Running analysis: ShouldNotRunFunctionPassesAnalysis
151151
; CHECK-O-NEXT: Running pass: CoroSplitPass
152+
; CHECK-O-NEXT: Running pass: CoroAnnotationElidePass
152153
; CHECK-O-NEXT: Running pass: InvalidateAnalysisPass<{{.*}}ShouldNotRunFunctionPassesAnalysis
153154
; CHECK-O-NEXT: Invalidating analysis: ShouldNotRunFunctionPassesAnalysis
154155
; CHECK-O-NEXT: Invalidating analysis: InlineAdvisorAnalysis

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
; Testing elide performed its job for calls to coroutines marked safe.
2+
; RUN: opt < %s -S -passes='cgscc(coro-annotation-elide)' | FileCheck %s
3+
4+
%struct.Task = type { ptr }
5+
6+
declare void @print(i32) nounwind
7+
8+
; resume part of the coroutine
9+
define fastcc void @callee.resume(ptr dereferenceable(1)) {
10+
tail call void @print(i32 0)
11+
ret void
12+
}
13+
14+
; destroy part of the coroutine
15+
define fastcc void @callee.destroy(ptr) {
16+
tail call void @print(i32 1)
17+
ret void
18+
}
19+
20+
; cleanup part of the coroutine
21+
define fastcc void @callee.cleanup(ptr) {
22+
tail call void @print(i32 2)
23+
ret void
24+
}
25+
26+
@callee.resumers = internal constant [3 x ptr] [
27+
ptr @callee.resume, ptr @callee.destroy, ptr @callee.cleanup]
28+
29+
declare void @alloc(i1) nounwind
30+
31+
; CHECK-LABEL: define ptr @callee
32+
define ptr @callee(i8 %arg) {
33+
entry:
34+
%task = alloca %struct.Task, align 8
35+
%id = call token @llvm.coro.id(i32 0, ptr null,
36+
ptr @callee,
37+
ptr @callee.resumers)
38+
%alloc = call i1 @llvm.coro.alloc(token %id)
39+
%hdl = call ptr @llvm.coro.begin(token %id, ptr null)
40+
store ptr %hdl, ptr %task
41+
ret ptr %task
42+
}
43+
44+
; CHECK-LABEL: define ptr @callee.noalloc
45+
define ptr @callee.noalloc(i8 %arg, ptr dereferenceable(32) align(8) %frame) {
46+
entry:
47+
%task = alloca %struct.Task, align 8
48+
%id = call token @llvm.coro.id(i32 0, ptr null,
49+
ptr @callee,
50+
ptr @callee.resumers)
51+
%hdl = call ptr @llvm.coro.begin(token %id, ptr null)
52+
store ptr %hdl, ptr %task
53+
ret ptr %task
54+
}
55+
56+
; CHECK-LABEL: define ptr @caller()
57+
; Function Attrs: presplitcoroutine
58+
define ptr @caller() #0 {
59+
entry:
60+
%task = call ptr @callee(i8 0) #1
61+
ret ptr %task
62+
63+
; CHECK: %[[ALLOCA:.+]] = alloca [32 x i8], align 8
64+
; CHECK-NEXT: %[[FRAME:.+]] = bitcast ptr %[[ALLOCA]] to ptr
65+
; CHECK-NEXT: %[[TASK:.+]] = call ptr @callee.noalloc(i8 0, ptr %[[FRAME]])
66+
; CHECK-NEXT: ret ptr %[[TASK]]
67+
}
68+
69+
declare token @llvm.coro.id(i32, ptr, ptr, ptr)
70+
declare ptr @llvm.coro.begin(token, ptr)
71+
declare ptr @llvm.coro.frame()
72+
declare ptr @llvm.coro.subfn.addr(ptr, i8)
73+
declare i1 @llvm.coro.alloc(token)
74+
75+
attributes #0 = { presplitcoroutine }
76+
attributes #1 = { coro_must_elide }

0 commit comments

Comments
 (0)