Index: lib/Transforms/Coroutines/CoroFrame.cpp =================================================================== --- lib/Transforms/Coroutines/CoroFrame.cpp +++ lib/Transforms/Coroutines/CoroFrame.cpp @@ -402,6 +402,7 @@ if (CurrentDef == Shape.PromiseAlloca) continue; + uint64_t Count = 1; Type *Ty = nullptr; if (auto *AI = dyn_cast(CurrentDef)) { Ty = AI->getAllocatedType(); @@ -413,11 +414,18 @@ Padder.addType(PaddingTy); } } + if (auto *CI = dyn_cast(AI->getArraySize())) + Count = CI->getValue().getZExtValue(); + else + report_fatal_error("Coroutines cannot handle non static allocas yet"); } else { Ty = CurrentDef->getType(); } S.setFieldIndex(Types.size()); - Types.push_back(Ty); + if (Count == 1) + Types.push_back(Ty); + else + Types.push_back(ArrayType::get(Ty, Count)); Padder.addType(Ty); } FrameTy->setBody(Types); @@ -470,6 +478,7 @@ // static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) { auto *CB = Shape.CoroBegin; + LLVMContext &C = CB->getContext(); IRBuilder<> Builder(CB->getNextNode()); StructType *FrameTy = Shape.FrameTy; PointerType *FramePtrTy = FrameTy->getPointerTo(); @@ -491,14 +500,38 @@ if (Shape.PromiseAlloca) Allocas.emplace_back(Shape.PromiseAlloca, coro::Shape::PromiseField); + // Create a GEP with the given index into the coroutine frame for the original + // value Orig. Appends an extra 0 index for array-allocas, preserving the + // original type. + auto GetFramePointer = [&](uint32_t Index, Value *Orig) -> Value * { + SmallVector Indices = { + ConstantInt::get(Type::getInt32Ty(C), 0), + ConstantInt::get(Type::getInt32Ty(C), Index), + }; + + if (auto *AI = dyn_cast(Orig)) { + if (auto *CI = dyn_cast(AI->getArraySize())) { + auto Count = CI->getValue().getZExtValue(); + if (Count > 1) { + Indices.push_back(ConstantInt::get(Type::getInt32Ty(C), 0)); + } + } else { + report_fatal_error("Coroutines cannot handle non static allocas yet"); + } + } + + return Builder.CreateInBoundsGEP(FrameTy, FramePtr, Indices); + }; + // 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() + - Twine(".reload.addr")); + + auto *G = GetFramePointer(Index, CurrentValue); + G->setName(CurrentValue->getName() + Twine(".reload.addr")); + return isa(CurrentValue) ? G : Builder.CreateLoad(FrameTy->getElementType(Index), G, @@ -588,8 +621,8 @@ Builder.SetInsertPoint(&Shape.AllocaSpillBlock->front()); // If we found any allocas, replace all of their remaining uses with Geps. for (auto &P : Allocas) { - auto *G = - Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, P.second); + auto *G = GetFramePointer(P.second, P.first); + // We are not using ReplaceInstWithInst(P.first, cast(G)) here, // as we are changing location of the instruction. G->takeName(P.first); Index: test/Transforms/Coroutines/coro-frame-arrayalloca.ll =================================================================== --- /dev/null +++ test/Transforms/Coroutines/coro-frame-arrayalloca.ll @@ -0,0 +1,72 @@ +; Check that we can handle spills of array allocas +; RUN: opt < %s -coro-split -S | FileCheck %s + +declare void @consume.double.ptr(double*) +declare void @consume.i32.ptr(i32*) + +define i8* @f() "coroutine.presplit"="1" { +entry: + %prefix = alloca double + %data = alloca i32, i32 4 + %suffix = alloca double + %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.double.ptr(double* %prefix) + call void @consume.i32.ptr(i32* %data) + call void @consume.double.ptr(double* %suffix) + %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.double.ptr(double* %prefix) + call void @consume.i32.ptr(i32* %data) + call void @consume.double.ptr(double* %suffix) + 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 array alloca was stored as an array field. +; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i1, i1, double, [4 x i32], double } + +; See if we used correct index to access prefix, data, suffix (@f) +; CHECK-LABEL: @f( +; CHECK: %prefix = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4 +; CHECK-NEXT: %data = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5 +; CHECK-NEXT: %suffix = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 6 +; CHECK-NEXT: call void @consume.double.ptr(double* %prefix) +; CHECK-NEXT: call void @consume.i32.ptr(i32* %data) +; CHECK-NEXT: call void @consume.double.ptr(double* %suffix) +; CHECK: ret i8* + +; See if we used correct index to access prefix, data, suffix (@f.resume) +; CHECK-LABEL: @f.resume( +; CHECK: %[[SUFFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 6 +; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5 +; CHECK: %[[PREFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4 +; CHECK: call void @consume.double.ptr(double* %[[PREFIX]]) +; CHECK-NEXT: call void @consume.i32.ptr(i32* %[[DATA]]) +; CHECK-NEXT: call void @consume.double.ptr(double* %[[SUFFIX]]) + +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*)