Skip to content

Commit 19bd17b

Browse files
yuxuanchen1997yuxuanchen1997
committed
Implement LLVM bits
1 parent d38f060 commit 19bd17b

File tree

9 files changed

+142
-33
lines changed

9 files changed

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

1981+
/// Remove calls to llvm.coro.end in the original function.
1982+
static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
1983+
if (Shape.ABI != coro::ABI::Switch) {
1984+
for (auto *End : Shape.CoroEnds) {
1985+
replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
1986+
}
1987+
} else {
1988+
for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
1989+
auto &Context = End->getContext();
1990+
End->replaceAllUsesWith(ConstantInt::getFalse(Context));
1991+
End->eraseFromParent();
1992+
}
1993+
}
1994+
}
1995+
19161996
static coro::Shape
19171997
splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19181998
TargetTransformInfo &TTI, bool OptimizeFrame,
@@ -1932,10 +2012,10 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19322012
simplifySuspendPoints(Shape);
19332013
buildCoroutineFrame(F, Shape, TTI, MaterializableCallback);
19342014
replaceFrameSizeAndAlignment(Shape);
1935-
2015+
bool isNoSuspendCoroutine = Shape.CoroSuspends.empty();
19362016
// If there are no suspend points, no split required, just remove
19372017
// the allocation and deallocation blocks, they are not needed.
1938-
if (Shape.CoroSuspends.empty()) {
2018+
if (isNoSuspendCoroutine) {
19392019
handleNoSuspendCoroutine(Shape);
19402020
} else {
19412021
switch (Shape.ABI) {
@@ -1967,22 +2047,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
19672047
for (DbgVariableRecord *DVR : DbgVariableRecords)
19682048
coro::salvageDebugInfo(ArgToAllocaMap, *DVR, Shape.OptimizeFrame,
19692049
false /*UseEntryValue*/);
1970-
return Shape;
1971-
}
19722050

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-
}
2051+
removeCoroEndsFromRampFunction(Shape);
2052+
2053+
if (!isNoSuspendCoroutine && Shape.ABI == coro::ABI::Switch) {
2054+
SwitchCoroutineSplitter::createNoAllocVariant(F, Shape, Clones);
19852055
}
2056+
return Shape;
19862057
}
19872058

19882059
static void updateCallGraphAfterCoroutineSplit(
@@ -2108,18 +2179,18 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21082179
// Split all the coroutines.
21092180
for (LazyCallGraph::Node *N : Coroutines) {
21102181
Function &F = N->getFunction();
2182+
21112183
LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
21122184
<< "\n");
21132185
F.setSplittedCoroutine();
21142186

21152187
SmallVector<Function *, 4> Clones;
2116-
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
2117-
const coro::Shape Shape =
2188+
coro::Shape Shape =
21182189
splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
21192190
OptimizeFrame, MaterializableCallback);
2120-
removeCoroEndsFromRampFunction(Shape);
21212191
updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
21222192

2193+
auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
21232194
ORE.emit([&]() {
21242195
return OptimizationRemark(DEBUG_TYPE, "CoroSplit", &F)
21252196
<< "Split '" << ore::NV("function", F.getName())
@@ -2135,9 +2206,9 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
21352206
}
21362207
}
21372208

2138-
for (auto *PrepareFn : PrepareFns) {
2139-
replaceAllPrepares(PrepareFn, CG, C);
2140-
}
2209+
for (auto *PrepareFn : PrepareFns) {
2210+
replaceAllPrepares(PrepareFn, CG, C);
2211+
}
21412212

21422213
return PreservedAnalyses::none();
21432214
}

‎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)