From d4712715dd3dfc7c387e7e97a5eee92a3ecfcee3 Mon Sep 17 00:00:00 2001 From: Gor Nishanov Date: Tue, 3 Apr 2018 20:54:20 +0000 Subject: [PATCH] [coroutines] Respect alloca alignment requirements when building coroutine frame Summary: If an alloca need to be stored in the coroutine frame and it has an alignment specified and the alignment does not match the natural alignment of the alloca type. Insert appropriate padding into the coroutine frame to make sure that it gets requested alignment. For example for a packet type (which natural alignment is 1), but alloca alignment is 8, we may need to insert a padding field with required number of bytes to make sure it is properly aligned. ``` %PackedStruct = type <{ i64 }> ... %data = alloca %PackedStruct, align 8 ``` If the previous field in the coroutine frame had alignment 2, we would have [6 x i8] inserted before %PackedStruct in the coroutine frame: ``` %f.Frame = type { ..., i16, [6 x i8], %PackedStruct } ``` Reviewers: rnk, lewissbaker, modocache Reviewed By: modocache Subscribers: EricWF, llvm-commits Differential Revision: https://reviews.llvm.org/D45221 llvm-svn: 329112 --- llvm/lib/Transforms/Coroutines/CoroFrame.cpp | 95 ++++++++++++++++++++++--- llvm/lib/Transforms/Coroutines/CoroInternal.h | 1 - llvm/test/Transforms/Coroutines/coro-padding.ll | 61 ++++++++++++++++ 3 files changed, 148 insertions(+), 9 deletions(-) create mode 100644 llvm/test/Transforms/Coroutines/coro-padding.ll diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp index 2bf142a7..85423b3 100644 --- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp @@ -263,8 +263,9 @@ SuspendCrossingInfo::SuspendCrossingInfo(Function &F, coro::Shape &Shape) namespace { class Spill { - Value *Def; - Instruction *User; + Value *Def = nullptr; + Instruction *User = nullptr; + unsigned FieldNo = 0; public: Spill(Value *Def, llvm::User *U) : Def(Def), User(cast(U)) {} @@ -272,6 +273,20 @@ public: Value *def() const { return Def; } Instruction *user() const { return User; } BasicBlock *userBlock() const { return User->getParent(); } + + // Note that field index is stored in the first SpillEntry for a particular + // definition. Subsequent mentions of a defintion do not have fieldNo + // assigned. This works out fine as the users of Spills capture the info about + // the definition the first time they encounter it. Consider refactoring + // SpillInfo into two arrays to normalize the spill representation. + unsigned fieldIndex() const { + assert(FieldNo && "Accessing unassigned field"); + return FieldNo; + } + void setFieldIndex(unsigned FieldNumber) { + assert(!FieldNo && "Reassigning field number"); + FieldNo = FieldNumber; + } }; } // namespace @@ -294,6 +309,55 @@ static void dump(StringRef Title, SpillInfo const &Spills) { } #endif +// We cannot rely solely on natural alignment of a type when building a +// coroutine frame and if the alignment specified on the Alloca instruction +// differs from the natural alignment of the alloca type we will need to insert +// padding. +struct PaddingCalculator { + const DataLayout &DL; + LLVMContext &Context; + unsigned StructSize = 0; + + PaddingCalculator(LLVMContext &Context, DataLayout const &DL) + : DL(DL), Context(Context) {} + + // Replicate the logic from IR/DataLayout.cpp to match field offset + // computation for LLVM structs. + void addType(Type *Ty) { + unsigned TyAlign = DL.getABITypeAlignment(Ty); + if ((StructSize & (TyAlign - 1)) != 0) + StructSize = alignTo(StructSize, TyAlign); + + StructSize += DL.getTypeAllocSize(Ty); // Consume space for this data item. + } + + void addTypes(SmallVectorImpl const &Types) { + for (auto *Ty : Types) + addType(Ty); + } + + unsigned computePadding(Type *Ty, unsigned ForcedAlignment) { + unsigned TyAlign = DL.getABITypeAlignment(Ty); + auto Natural = alignTo(StructSize, TyAlign); + auto Forced = alignTo(StructSize, ForcedAlignment); + + // Return how many bytes of padding we need to insert. + if (Natural != Forced) + return std::max(Natural, Forced) - StructSize; + + // Rely on natural alignment. + return 0; + } + + // If padding required, return the padding field type to insert. + ArrayType *getPaddingType(Type *Ty, unsigned ForcedAlignment) { + if (auto Padding = computePadding(Ty, ForcedAlignment)) + return ArrayType::get(Type::getInt8Ty(Context), Padding); + + return nullptr; + } +}; + // Build a struct that will keep state for an active coroutine. // struct f.frame { // ResumeFnTy ResumeFnAddr; @@ -305,6 +369,8 @@ static void dump(StringRef Title, SpillInfo const &Spills) { static StructType *buildFrameType(Function &F, coro::Shape &Shape, SpillInfo &Spills) { LLVMContext &C = F.getContext(); + const DataLayout &DL = F.getParent()->getDataLayout(); + PaddingCalculator Padder(C, DL); SmallString<32> Name(F.getName()); Name.append(".Frame"); StructType *FrameTy = StructType::create(C, Name); @@ -322,8 +388,10 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape, Type::getIntNTy(C, IndexBits)}; Value *CurrentDef = nullptr; + Padder.addTypes(Types); + // Create an entry for every spilled value. - for (auto const &S : Spills) { + for (auto &S : Spills) { if (CurrentDef == S.def()) continue; @@ -333,12 +401,22 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape, continue; Type *Ty = nullptr; - if (auto *AI = dyn_cast(CurrentDef)) + if (auto *AI = dyn_cast(CurrentDef)) { Ty = AI->getAllocatedType(); - else + if (unsigned AllocaAlignment = AI->getAlignment()) { + // If alignment is specified in alloca, see if we need to insert extra + // padding. + if (auto PaddingTy = Padder.getPaddingType(Ty, AllocaAlignment)) { + Types.push_back(PaddingTy); + Padder.addType(PaddingTy); + } + } + } else { Ty = CurrentDef->getType(); - + } + S.setFieldIndex(Types.size()); Types.push_back(Ty); + Padder.addType(Ty); } FrameTy->setBody(Types); @@ -399,7 +477,7 @@ static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) { Value *CurrentValue = nullptr; BasicBlock *CurrentBlock = nullptr; Value *CurrentReload = nullptr; - unsigned Index = coro::Shape::LastKnownField; + unsigned Index = 0; // Proper field number will be read from field definition. // We need to keep track of any allocas that need "spilling" // since they will live in the coroutine frame now, all access to them @@ -414,6 +492,7 @@ static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) { // Create a load instruction to reload the spilled value from the coroutine // frame. auto CreateReload = [&](Instruction *InsertBefore) { + assert(Index && "accessing unassigned field number"); Builder.SetInsertPoint(InsertBefore); auto *G = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index, CurrentValue->getName() + @@ -431,7 +510,7 @@ static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) { CurrentBlock = nullptr; CurrentReload = nullptr; - ++Index; + Index = E.fieldIndex(); if (auto *AI = dyn_cast(CurrentValue)) { // Spilled AllocaInst will be replaced with GEP from the coroutine frame diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h index 1eac88d..8e690d6 100644 --- a/llvm/lib/Transforms/Coroutines/CoroInternal.h +++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h @@ -76,7 +76,6 @@ struct LLVM_LIBRARY_VISIBILITY Shape { DestroyField, PromiseField, IndexField, - LastKnownField = IndexField }; StructType *FrameTy; diff --git a/llvm/test/Transforms/Coroutines/coro-padding.ll b/llvm/test/Transforms/Coroutines/coro-padding.ll new file mode 100644 index 0000000..87b5bf7 --- /dev/null +++ b/llvm/test/Transforms/Coroutines/coro-padding.ll @@ -0,0 +1,61 @@ +; Check that we will insert the correct padding if natural alignment of the +; spilled data does not match the alignment specified in alloca instruction. +; RUN: opt < %s -coro-split -S | FileCheck %s + +%PackedStruct = type <{ i64 }> + +declare void @consume(%PackedStruct*) + +define i8* @f() "coroutine.presplit"="1" { +entry: + %data = alloca %PackedStruct, align 8 + %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null) + %size = call i32 @llvm.coro.size.i32() + %alloc = call i8* @malloc(i32 %size) + %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc) + call void @consume(%PackedStruct* %data) + %0 = call i8 @llvm.coro.suspend(token none, i1 false) + switch i8 %0, label %suspend [i8 0, label %resume + i8 1, label %cleanup] +resume: + call void @consume(%PackedStruct* %data) + br label %cleanup + +cleanup: + %mem = call i8* @llvm.coro.free(token %id, i8* %hdl) + call void @free(i8* %mem) + br label %suspend +suspend: + call i1 @llvm.coro.end(i8* %hdl, i1 0) + ret i8* %hdl +} + +; See if the padding was inserted before PackedStruct +; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i1, i1, [6 x i8], %PackedStruct } + +; See if we used correct index to access packed struct (padding is field 4) +; CHECK-LABEL: @f( +; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5 +; CHECK-NEXT: call void @consume(%PackedStruct* %[[DATA]]) +; CHECK: ret i8* + +; See if we used correct index to access packed struct (padding is field 4) +; CHECK-LABEL: @f.resume( +; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5 +; CHECK-NEXT: call void @consume(%PackedStruct* %[[DATA]]) +; CHECK: ret void + +declare i8* @llvm.coro.free(token, i8*) +declare i32 @llvm.coro.size.i32() +declare i8 @llvm.coro.suspend(token, i1) +declare void @llvm.coro.resume(i8*) +declare void @llvm.coro.destroy(i8*) + +declare token @llvm.coro.id(i32, i8*, i8*, i8*) +declare i1 @llvm.coro.alloc(token) +declare i8* @llvm.coro.begin(token, i8*) +declare i1 @llvm.coro.end(i8*, i1) + +declare noalias i8* @malloc(i32) +declare double @print(double) +declare void @free(i8*) -- 2.7.4