Skip to content

Commit 43ae2cb

Browse files
yuxuanchen1997yuxuanchen1997
committed
Implement noalloc in CoroSplit
1 parent 5578fde commit 43ae2cb

File tree

9 files changed

+153
-33
lines changed

9 files changed

+153
-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: 106 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,73 @@ struct SwitchCoroutineSplitter {
14551462
setCoroInfo(F, Shape, Clones);
14561463
}
14571464

1465+
// Create a variant of ramp function that does not perform heap allocation
1466+
// for a switch ABI coroutine.
1467+
//
1468+
// The newly split `.noalloc` ramp function has the following differences:
1469+
// - Has one additional frame pointer parameter in lieu of dynamic allocation.
1470+
// - Suppressed allocations by replacing coro.alloc and coro.free.
1471+
static Function *createNoAllocVariant(Function &F, coro::Shape &Shape,
1472+
SmallVectorImpl<Function *> &Clones) {
1473+
auto *OrigFnTy = F.getFunctionType();
1474+
auto OldParams = OrigFnTy->params();
1475+
1476+
SmallVector<Type *> NewParams;
1477+
NewParams.reserve(OldParams.size() + 1);
1478+
NewParams.append(OldParams.begin(), OldParams.end());
1479+
NewParams.push_back(PointerType::getUnqual(Shape.FrameTy));
1480+
1481+
auto *NewFnTy = FunctionType::get(OrigFnTy->getReturnType(), NewParams,
1482+
OrigFnTy->isVarArg());
1483+
Function *NoAllocF =
1484+
Function::Create(NewFnTy, F.getLinkage(), F.getName() + ".noalloc");
1485+
1486+
ValueToValueMapTy VMap;
1487+
unsigned int Idx = 0;
1488+
for (const auto &I : F.args()) {
1489+
VMap[&I] = NoAllocF->getArg(Idx++);
1490+
}
1491+
SmallVector<ReturnInst *, 4> Returns;
1492+
CloneFunctionInto(NoAllocF, &F, VMap,
1493+
CloneFunctionChangeType::LocalChangesOnly, Returns);
1494+
1495+
if (Shape.CoroBegin) {
1496+
auto *NewCoroBegin =
1497+
cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);
1498+
auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());
1499+
coro::replaceCoroFree(NewCoroId, /*Elide=*/true);
1500+
coro::suppressCoroAllocs(NewCoroId);
1501+
NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(Idx));
1502+
NewCoroBegin->eraseFromParent();
1503+
}
1504+
1505+
Module *M = F.getParent();
1506+
M->getFunctionList().insert(M->end(), NoAllocF);
1507+
1508+
removeUnreachableBlocks(*NoAllocF);
1509+
auto NewAttrs = NoAllocF->getAttributes();
1510+
// We just appended the frame pointer as the last argument of the new
1511+
// function.
1512+
auto FrameIdx = NoAllocF->arg_size() - 1;
1513+
// When we elide allocation, we read these attributes to determine the
1514+
// frame size and alignment.
1515+
addFramePointerAttrs(NewAttrs, NoAllocF->getContext(), FrameIdx,
1516+
Shape.FrameSize, Shape.FrameAlign,
1517+
/*NoAlias=*/false);
1518+
1519+
NoAllocF->setAttributes(NewAttrs);
1520+
1521+
Clones.push_back(NoAllocF);
1522+
// Reset the original function's coro info, make the new noalloc variant
1523+
// connected to the original ramp function.
1524+
setCoroInfo(F, Shape, Clones);
1525+
// After copying, set the linkage to internal linkage. Original function
1526+
// may have different linkage, but optimization dependent on this function
1527+
// generally relies on LTO.
1528+
NoAllocF->setLinkage(llvm::GlobalValue::InternalLinkage);
1529+
return NoAllocF;
1530+
}
1531+
14581532
private:
14591533
// Create a resume clone by cloning the body of the original function, setting
14601534
// new entry block and replacing coro.suspend an appropriate value to force
@@ -1913,6 +1987,21 @@ class PrettyStackTraceFunction : public PrettyStackTraceEntry {
19131987
};
19141988
} // namespace
19151989

1990+
/// Remove calls to llvm.coro.end in the original function.
1991+
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
1992+
if (Shape.ABI != coro::ABI::Switch) {
1993+
for (auto *End : Shape.CoroEnds) {
1994+
replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
1995+
}
1996+
} else {
1997+
for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
1998+
auto &Context = End->getContext();
1999+
End->replaceAllUsesWith(ConstantInt::getFalse(Context));
2000+
End->eraseFromParent();
2001+
}
2002+
}
2003+
}
2004+
19162005
static coro::Shape
19172006
splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19182007
TargetTransformInfo &TTI, bool OptimizeFrame,
@@ -1932,10 +2021,10 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19322021
simplifySuspendPoints(Shape);
19332022
buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
19342023
replaceFrameSizeAndAlignment(Shape);
1935-
2024+
bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
19362025
// If there are no suspend points, no split required, just remove
19372026
// the allocation and deallocation blocks, they are not needed.
1938-
if (Shape.CoroSuspends.empty()) {
2027+
if (isNoSuspendCoroutine) {
19392028
handleNoSuspendCoroutine(Shape);
19402029
} else {
19412030
switch (Shape.ABI) {
@@ -1967,22 +2056,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19672056
for (DbgVariableRecord *DVR : DbgVariableRecords)
19682057
coro::salvageDebugInfo(ArgToAllocaMap, *DVR, Shape.OptimizeFrame,
19692058
false /*UseEntryValue*/);
1970-
return Shape;
1971-
}
19722059

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-
}
2060+
removeCoroEndsFromRampFunction(Shape);
2061+
2062+
if (!isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch) {
2063+
SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
19852064
}
2065+
return Shape;
19862066
}
19872067

19882068
static void updateCallGraphAfterCoroutineSplit(
@@ -2108,18 +2188,18 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21082188
// Split all the coroutines.
21092189
for (LazyCallGraph::Node *N : Coroutines) {
21102190
Function &F = N->getFunction();
2191+
21112192
LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
21122193
<< "\n");
21132194
F.setSplittedCoroutine();
21142195

21152196
SmallVector<Function *, 4> Clones;
2116-
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
2117-
const coro::Shape Shape =
2197+
coro::Shape Shape =
21182198
splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
21192199
OptimizeFrame, MaterializableCallback);
2120-
removeCoroEndsFromRampFunction(Shape);
21212200
updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
21222201

2202+
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
21232203
ORE.emit([&]() {
21242204
return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)
21252205
<< "Split '" << ore::NV("function", F.getName())
@@ -2135,9 +2215,9 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21352215
}
21362216
}
21372217

2138-
for (auto *PrepareFn : PrepareFns) {
2139-
replaceAllPrepares(PrepareFn, CG, C);
2140-
}
2218+
for (auto *PrepareFn : PrepareFns) {
2219+
replaceAllPrepares(PrepareFn, CG, C);
2220+
}
21412221

21422222
return PreservedAnalyses::none();
21432223
}

‎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: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22
; coro.begin.
33
; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
44

5+
; CHECK: @[[RESUMER:.+]] = private constant [4 x ptr] [ptr @f.resume, ptr @f.destroy, ptr @f.cleanup, ptr @f.noalloc]
6+
57
define nonnull ptr @f(i32 %n) presplitcoroutine {
68
; CHECK-LABEL: @f(
79
; CHECK-NEXT: entry:
8-
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @f.resumers)
10+
; CHECK-NEXT: [[ID:%.*]] = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr @[[RESUMER]])
911
; CHECK-NEXT: [[N_ADDR:%.*]] = alloca i32, align 4
1012
; CHECK-NEXT: store i32 [[N:%.*]], ptr [[N_ADDR]], align 4
1113
; 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(ptr noundef nonnull align 8 dereferenceable(24) %{{.*}})
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)