[coroutines] Fix spills of static array allocas

[ben-clayton]

Summary

CoroFrame was not considering static array allocas, and was only ever reserving a single element in the coroutine frame.
This meant that stores to the non-zero'th element would corrupt later frame data.

Store static array allocas as field arrays in the coroutine frame.

Added test.

[Orlando]

Hi, I added a few inline comments, mostly just nits on style.

lib/Transforms/Coroutines/CoroFrame.cpp

405

count -> Count

507

indices -> Indices

514

count -> Count

test/Transforms/Coroutines/coro-frame-arrayalloca.ll

4

Is this second RUN redundant?

[ben-clayton]

Fixed capitalization of locals 'Count' and 'Indices', removed duplicate '; RUN:' line in test.

[ben-clayton]

test/Transforms/Coroutines/coro-frame-arrayalloca.ll

4

Good spot! Yes, removed.

[GorNishanov]

LGTM! Thank you for doing the fix.
What is the original source that synthesized allocas with count?

I tried to reproduce it from c++ with char a[5], but it got mapped to llvm array types.

[ben-clayton]

Thank you GorNishanov.

What is the original source that synthesized allocas with count?

SwiftShader (a CPU implementation of OpenGL ES and Vulkan) which uses Reactor for JIT code generation of shaders and compute / vertex / fragment pipelines.
I'm adding coroutine support for handling SPIR-V memory barriers which bring all shader invocations up to the same barrier before continuing.
The change is here if you're curious to see the coroutine implementation.

I do not believe I have permissions to submit this. Please can you do the honours?

Thanks again,
Ben

[ben-clayton]

I do not believe I have permissions to submit this. Please can you do the honours?

Gentle ping!

[ben-clayton]

Another friendly ping.

[lewissbaker]

llvm/trunk/lib/Transforms/Coroutines/CoroFrame.cpp

429 ↗

(On Diff #199349)

Does this call need to be changed to Padder.addType(Types.back())?

I can imagine a case where Ty is char and we have a local variable of type char[3] with the next local variable having type int.

Couldn't this result in the Padding object thinking it needs to pad to the next int value by adding char[3] whereas it actually only needs to pad with char?

lib/Transforms/Coroutines/CoroFrame.cpp

for (auto &S : Spills) {
     if (CurrentDef == S.def())
       continue;
 
     CurrentDef = S.def();
     // PromiseAlloca was already added to Types array earlier.
     if (CurrentDef == Shape.PromiseAlloca)
       continue;
 
+    uint64_t Count = 1;
     Type *Ty = nullptr;
     if (auto *AI = dyn_cast<AllocaInst>(CurrentDef)) {
       Ty = AI->getAllocatedType();
       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);
         }
       }
+      if (auto *CI = dyn_cast<ConstantInt>(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());
+    if (Count == 1)
       Types.push_back(Ty);
+    else
+      Types.push_back(ArrayType::get(Ty, Count));
     Padder.addType(Ty);
   }
   FrameTy->setBody(Types);
 
   return FrameTy;
 }
 
 // We need to make room to insert a spill after initial PHIs, but before

 //    br label PostSpill
 //
 //  PostSpill:
 //    whatever
 //
 //
 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();
   auto *FramePtr =
       cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
 
   Value *CurrentValue = nullptr;
   BasicBlock *CurrentBlock = nullptr;
   Value *CurrentReload = nullptr;
   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
   // need to be changed, not just the access across suspend points
   // we remember allocas and their indices to be handled once we processed
   // all the spills.
   SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
   // Promise alloca (if present) has a fixed field number (Shape::PromiseField)
   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<Value *, 3> Indices = {
+        ConstantInt::get(Type::getInt32Ty(C), 0),
+        ConstantInt::get(Type::getInt32Ty(C), Index),
+    };
+
+    if (auto *AI = dyn_cast<AllocaInst>(Orig)) {
+      if (auto *CI = dyn_cast<ConstantInt>(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<AllocaInst>(CurrentValue)
                ? G
                : Builder.CreateLoad(FrameTy->getElementType(Index), G,
                                     CurrentValue->getName() + Twine(".reload"));
   };
 
   for (auto const &E : Spills) {
     // If we have not seen the value, generate a spill.
static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
   Shape.AllocaSpillBlock =
       FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
   Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),
                                           "PostSpill");
 
   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<Instruction>(G)) here,
     // as we are changing location of the instruction.
     G->takeName(P.first);
     P.first->replaceAllUsesWith(G);
     P.first->eraseFromParent();
   }
   return FramePtr;
 }

test/Transforms/Coroutines/coro-frame-arrayalloca.ll

+; 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*)