Skip to content

Commit f8bb0f4

Browse files
yuxuanchen1997yuxuanchen1997
committed
Implement noalloc in CoroSplit
1 parent 8d9c2cc commit f8bb0f4

File tree

9 files changed

+140
-33
lines changed

9 files changed

+140
-33
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ bool declaresIntrinsics(const Module &M,
2626
const std::initializer_list<StringRef>);
2727
void replaceCoroFree(CoroIdInst *CoroId, bool Elide);
2828

29+
void suppressCoroAllocs(CoroIdInst *CoroId);
30+
void suppressCoroAllocs(LLVMContext &Context,
31+
ArrayRef<CoroAllocInst *> CoroAllocs);
32+
2933
/// Attempts to rewrite the location operand of debug intrinsics in terms of
3034
/// the coroutine frame pointer, folding pointer offsets into the DIExpression
3135
/// of the intrinsic.

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

Lines changed: 95 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
#include "llvm/ADT/PriorityWorklist.h"
2626
#include "llvm/ADT/SmallPtrSet.h"
2727
#include "llvm/ADT/SmallVector.h"
28+
#include "llvm/ADT/StringExtras.h"
2829
#include "llvm/ADT/StringRef.h"
2930
#include "llvm/ADT/Twine.h"
3031
#include "llvm/Analysis/CFG.h"
@@ -1179,6 +1180,14 @@ static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
11791180
Shape.AsyncLowering.AsyncFuncPointer->setInitializer(NewFuncPtrStruct);
11801181
}
11811182

1183+
static TypeSize getFrameSizeForShape(coro::Shape &Shape) {
1184+
// In the same function all coro.sizes should have the same result type.
1185+
auto *SizeIntrin = Shape.CoroSizes.back();
1186+
Module *M = SizeIntrin->getModule();
1187+
const DataLayout &DL = M->getDataLayout();
1188+
return DL.getTypeAllocSize(Shape.FrameTy);
1189+
}
1190+
11821191
static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
11831192
if (Shape.ABI == coro::ABI::Async)
11841193
updateAsyncFuncPointerContextSize(Shape);
@@ -1194,10 +1203,8 @@ static void replaceFrameSizeAndAlignment(coro::Shape &Shape) {
11941203

11951204
// In the same function all coro.sizes should have the same result type.
11961205
auto *SizeIntrin = Shape.CoroSizes.back();
1197-
Module *M = SizeIntrin->getModule();
1198-
const DataLayout &DL = M->getDataLayout();
1199-
auto Size = DL.getTypeAllocSize(Shape.FrameTy);
1200-
auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
1206+
auto *SizeConstant =
1207+
ConstantInt::get(SizeIntrin->getType(), getFrameSizeForShape(Shape));
12011208

12021209
for (CoroSizeInst *CS : Shape.CoroSizes) {
12031210
CS->replaceAllUsesWith(SizeConstant);
@@ -1455,6 +1462,62 @@ struct SwitchCoroutineSplitter {
14551462
setCoroInfo(F, Shape, Clones);
14561463
}
14571464

1465+
static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,
1466+
SmallVectorImpl<Function *> &Clones) {
1467+
auto *OrigFnTy = F.getFunctionType();
1468+
auto OldParams = OrigFnTy->params();
1469+
1470+
SmallVector<Type *> NewParams;
1471+
NewParams.reserve(OldParams.size() + 1);
1472+
NewParams.append(OldParams.begin(), OldParams.end());
1473+
NewParams.push_back(PointerType::getUnqual(Shape.FrameTy));
1474+
1475+
auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,
1476+
OrigFnTy->isVarArg());
1477+
Function *NoAllocF =
1478+
Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");
1479+
ValueToValueMapTy VMap;
1480+
unsigned int Idx = 0;
1481+
for (const auto &I : F.args()) {
1482+
VMap[&I] = NoAllocF->getArg(Idx++);
1483+
}
1484+
SmallVector<ReturnInst *, 4> Returns;
1485+
CloneFunctionInto(NoAllocF, &F, VMap,
1486+
CloneFunctionChangeType::LocalChangesOnly, Returns);
1487+
1488+
if (Shape.CoroBegin) {
1489+
auto *NewCoroBegin =
1490+
cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);
1491+
auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());
1492+
coro::replaceCoroFree(NewCoroId, /*Elide=*/true);
1493+
coro::suppressCoroAllocs(NewCoroId);
1494+
NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(Idx));
1495+
NewCoroBegin->eraseFromParent();
1496+
}
1497+
1498+
Module *M = F.getParent();
1499+
M->getFunctionList().insert(M->end(), NoAllocF);
1500+
1501+
removeUnreachableBlocks(*NoAllocF);
1502+
auto NewAttrs = NoAllocF->getAttributes();
1503+
// We just appended the frame pointer as the last argument of the new
1504+
// function.
1505+
auto FrameIdx = NoAllocF->arg_size() - 1;
1506+
// When we elide allocation, we read these attributes to determine the
1507+
// frame size and alignment.
1508+
addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,
1509+
Shape.FrameSize, Shape.FrameAlign,
1510+
/*NoAlias=*/false);
1511+
1512+
NoAllocF->setAttributes(NewAttrs);
1513+
1514+
Clones.push_back(NoAllocF);
1515+
// Reset the original function's coro info, make the new noalloc variant
1516+
// connected to the original ramp function.
1517+
setCoroInfo(F, Shape, Clones);
1518+
return NoAllocF;
1519+
}
1520+
14581521
private:
14591522
// Create a resume clone by cloning the body of the original function, setting
14601523
// new entry block and replacing coro.suspend an appropriate value to force
@@ -1913,6 +1976,21 @@ class PrettyStackTraceFunction : public PrettyStackTraceEntry {
19131976
};
19141977
} // namespace
19151978

1979+
/// Remove calls to llvm.coro.end in the original function.
1980+
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
1981+
if (Shape.ABI != coro::ABI::Switch) {
1982+
for (auto *End : Shape.CoroEnds) {
1983+
replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
1984+
}
1985+
} else {
1986+
for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
1987+
auto &Context = End->getContext();
1988+
End->replaceAllUsesWith(ConstantInt::getFalse(Context));
1989+
End->eraseFromParent();
1990+
}
1991+
}
1992+
}
1993+
19161994
static coro::Shape
19171995
splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19181996
TargetTransformInfo &TTI, bool OptimizeFrame,
@@ -1932,10 +2010,10 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19322010
simplifySuspendPoints(Shape);
19332011
buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
19342012
replaceFrameSizeAndAlignment(Shape);
1935-
2013+
bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
19362014
// If there are no suspend points, no split required, just remove
19372015
// the allocation and deallocation blocks, they are not needed.
1938-
if (Shape.CoroSuspends.empty()) {
2016+
if (isNoSuspendCoroutine) {
19392017
handleNoSuspendCoroutine(Shape);
19402018
} else {
19412019
switch (Shape.ABI) {
@@ -1967,22 +2045,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19672045
for (DbgVariableRecord *DVR : DbgVariableRecords)
19682046
coro::salvageDebugInfo(ArgToAllocaMap, *DVR, Shape.OptimizeFrame,
19692047
false /*UseEntryValue*/);
1970-
return Shape;
1971-
}
19722048

1973-
/// Remove calls to llvm.coro.end in the original function.
1974-
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
1975-
if (Shape.ABI != coro::ABI::Switch) {
1976-
for (auto *End : Shape.CoroEnds) {
1977-
replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
1978-
}
1979-
} else {
1980-
for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
1981-
auto &Context = End->getContext();
1982-
End->replaceAllUsesWith(ConstantInt::getFalse(Context));
1983-
End->eraseFromParent();
1984-
}
2049+
removeCoroEndsFromRampFunction(Shape);
2050+
2051+
if (!isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch) {
2052+
SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
19852053
}
2054+
return Shape;
19862055
}
19872056

19882057
static void updateCallGraphAfterCoroutineSplit(
@@ -2108,18 +2177,18 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21082177
// Split all the coroutines.
21092178
for (LazyCallGraph::Node *N : Coroutines) {
21102179
Function &F = N->getFunction();
2180+
21112181
LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
21122182
<< "\n");
21132183
F.setSplittedCoroutine();
21142184

21152185
SmallVector<Function *, 4> Clones;
2116-
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
2117-
const coro::Shape Shape =
2186+
coro::Shape Shape =
21182187
splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
21192188
OptimizeFrame, MaterializableCallback);
2120-
removeCoroEndsFromRampFunction(Shape);
21212189
updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
21222190

2191+
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
21232192
ORE.emit([&]() {
21242193
return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)
21252194
<< "Split '" << ore::NV("function", F.getName())
@@ -2135,9 +2204,9 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21352204
}
21362205
}
21372206

2138-
for (auto *PrepareFn : PrepareFns) {
2139-
replaceAllPrepares(PrepareFn, CG, C);
2140-
}
2207+
for (auto *PrepareFn : PrepareFns) {
2208+
replaceAllPrepares(PrepareFn, CG, C);
2209+
}
21412210

21422211
return PreservedAnalyses::none();
21432212
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,33 @@ void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
145145
}
146146
}
147147

148+
void coro::suppressCoroAllocs(CoroIdInst *CoroId) {
149+
SmallVector<CoroAllocInst *, 4> CoroAllocs;
150+
for (User *U : CoroId->users())
151+
if (auto *CA = dyn_cast<CoroAllocInst>(U))
152+
CoroAllocs.push_back(CA);
153+
154+
if (CoroAllocs.empty())
155+
return;
156+
157+
coro::suppressCoroAllocs(CoroId->getContext(), CoroAllocs);
158+
}
159+
160+
// Replacing llvm.coro.alloc with false will suppress dynamic
161+
// allocation as it is expected for the frontend to generate the code that
162+
// looks like:
163+
// id = coro.id(...)
164+
// mem = coro.alloc(id) ? malloc(coro.size()) : 0;
165+
// coro.begin(id, mem)
166+
void coro::suppressCoroAllocs(LLVMContext &Context,
167+
ArrayRef<CoroAllocInst *> CoroAllocs) {
168+
auto *False = ConstantInt::getFalse(Context);
169+
for (auto *CA : CoroAllocs) {
170+
CA->replaceAllUsesWith(False);
171+
CA->eraseFromParent();
172+
}
173+
}
174+
148175
static void clear(coro::Shape &Shape) {
149176
Shape.CoroBegin = nullptr;
150177
Shape.CoroEnds.clear();

‎llvm/test/Transforms/Coroutines/ArgAddr.ll

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
define nonnull ptr @f(i32 %n) presplitcoroutine {
66
; CHECK-LABEL: @f(
77
; CHECK-NEXT: entry:
8-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @f.resumers)
8+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
99
; CHECK-NEXT: [[N_ADDR:%.*]] = alloca i32, align 4
1010
; CHECK-NEXT: store i32 [[N:%.*]], ptr [[N_ADDR]], align 4
1111
; CHECK-NEXT: [[CALL:%.*]] = tail call ptr @malloc(i32 24)

‎llvm/test/Transforms/Coroutines/coro-alloca-07.ll

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ declare void @free(ptr)
6262

6363
; CHECK-LABEL: @f(
6464
; CHECK-NEXT: entry:
65-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @f.resumers)
65+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
6666
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i32 48)
6767
; CHECK-NEXT: [[HDL:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
6868
; CHECK-NEXT: store ptr @f.resume, ptr [[HDL]], align 8

‎llvm/test/Transforms/Coroutines/coro-alloca-loop-carried-address.ll

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
define void @foo() presplitcoroutine {
88
; CHECK-LABEL: @foo(
99
; CHECK-NEXT: entry:
10-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @foo.resumers)
10+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
1111
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i64 40)
1212
; CHECK-NEXT: [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
1313
; CHECK-NEXT: store ptr @foo.resume, ptr [[VFRAME]], align 8

‎llvm/test/Transforms/Coroutines/coro-lifetime-end.ll

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ declare void @consume.i8.array(ptr)
1313
define void @HasNoLifetimeEnd() presplitcoroutine {
1414
; CHECK-LABEL: define void @HasNoLifetimeEnd() {
1515
; CHECK-NEXT: entry:
16-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @HasNoLifetimeEnd.resumers)
16+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
1717
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i64 16)
1818
; CHECK-NEXT: [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
1919
; CHECK-NEXT: store ptr @HasNoLifetimeEnd.resume, ptr [[VFRAME]], align 8
@@ -50,7 +50,7 @@ exit:
5050
define void @LifetimeEndAfterCoroEnd() presplitcoroutine {
5151
; CHECK-LABEL: define void @LifetimeEndAfterCoroEnd() {
5252
; CHECK-NEXT: entry:
53-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @LifetimeEndAfterCoroEnd.resumers)
53+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
5454
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i64 16)
5555
; CHECK-NEXT: [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
5656
; CHECK-NEXT: store ptr @LifetimeEndAfterCoroEnd.resume, ptr [[VFRAME]], align 8
@@ -88,7 +88,7 @@ exit:
8888
define void @BranchWithoutLifetimeEnd() presplitcoroutine {
8989
; CHECK-LABEL: define void @BranchWithoutLifetimeEnd() {
9090
; CHECK-NEXT: entry:
91-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @BranchWithoutLifetimeEnd.resumers)
91+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
9292
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i64 16)
9393
; CHECK-NEXT: [[VFRAME:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
9494
; CHECK-NEXT: store ptr @BranchWithoutLifetimeEnd.resume, ptr [[VFRAME]], align 8

‎llvm/test/Transforms/Coroutines/coro-spill-after-phi.ll

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
define ptr @f(i1 %n) presplitcoroutine {
99
; CHECK-LABEL: @f(
1010
; CHECK-NEXT: entry:
11-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @f.resumers)
11+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @{{.*}})
1212
; CHECK-NEXT: [[ALLOC:%.*]] = call ptr @malloc(i32 32)
1313
; CHECK-NEXT: [[HDL:%.*]] = call noalias nonnull ptr @llvm.coro.begin(token [[ID]], ptr [[ALLOC]])
1414
; CHECK-NEXT: store ptr @f.resume, ptr [[HDL]], align 8

‎llvm/test/Transforms/Coroutines/coro-split-00.ll

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,13 @@ suspend:
6363
; CHECK-NOT: call void @free(
6464
; CHECK: ret void
6565

66+
; CHECK-LABEL: @f.noalloc({{.*}})
67+
; CHECK-NOT: call ptr @malloc
68+
; CHECK: call void @print(i32 0)
69+
; CHECK-NOT: call void @print(i32 1)
70+
; CHECK-NOT: call void @free(
71+
; CHECK: ret ptr %{{.*}}
72+
6673
declare ptr @llvm.coro.free(token, ptr)
6774
declare i32 @llvm.coro.size.i32()
6875
declare i8 @llvm.coro.suspend(token, i1)

0 commit comments

Comments
 (0)