[coroutines] Respect alloca alignment requirements when building coroutine frame

[GorNishanov]

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 }

[modocache]

Thank you, @GorNishanov, this is excellent! I've pulled this down and confirmed it fixes some nasty errors in a codebase I work on, related to [[ https://github.com/facebook/folly/blob/3cec1c6b4ccfa4b80effad376f6321ce9b70cec1/folly/Function.h#L253 | this usage of std::aligned_storage ]]. The problem also reproduced with this much smaller example: https://reviews.llvm.org/P8076 -- and I've confirmed this change of yours fixes that, too. So, thank you!

The code itself looks good to me!

lib/Transforms/Coroutines/CoroFrame.cpp

 #undef DEBUG_TYPE // "coro-suspend-crossing"
 #define DEBUG_TYPE "coro-frame"
 
 // We build up the list of spills for every case where a use is separated
 // from the definition by a suspend point.
 
 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<Instruction>(U)) {}
 
   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
 
 // Note that there may be more than one record with the same value of Def in
 // the SpillInfo vector.
 using SpillInfo = SmallVector<Spill, 8>;
 
 #ifndef NDEBUG
 static void dump(StringRef Title, SpillInfo const &Spills) {
   dbgs() << "------------- " << Title << "--------------\n";
   Value *CurrentValue = nullptr;
   for (auto const &E : Spills) {
     if (CurrentValue != E.def()) {
       CurrentValue = E.def();
       CurrentValue->dump();
     }
     dbgs() << "   user: ";
     E.user()->dump();
   }
 }
 #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<Type *> 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;
 //     ResumeFnTy DestroyFnAddr;
 //     int ResumeIndex;
 //     ... promise (if present) ...
 //     ... 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);
   auto *FramePtrTy = FrameTy->getPointerTo();
   auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
                                  /*IsVarArgs=*/false);
   auto *FnPtrTy = FnTy->getPointerTo();
 
   // Figure out how wide should be an integer type storing the suspend index.
   unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
   Type *PromiseType = Shape.PromiseAlloca
                           ? Shape.PromiseAlloca->getType()->getElementType()
                           : Type::getInt1Ty(C);
   SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, PromiseType,
                                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;
 
     CurrentDef = S.def();
     // PromiseAlloca was already added to Types array earlier.
     if (CurrentDef == Shape.PromiseAlloca)
       continue;
 
     Type *Ty = nullptr;
-    if (auto *AI = dyn_cast<AllocaInst>(CurrentDef))
+    if (auto *AI = dyn_cast<AllocaInst>(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);
 
   return FrameTy;
 }
 
 // We need to make room to insert a spill after initial PHIs, but before
 // catchswitch instruction. Placing it before violates the requirement that
static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
   PointerType *FramePtrTy = Shape.FrameTy->getPointerTo();
   auto *FramePtr =
       cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
   Type *FrameTy = FramePtrTy->getElementType();
 
   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
   // 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 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"));
     return isa<AllocaInst>(CurrentValue)
                ? G
                : Builder.CreateLoad(G,
                                     CurrentValue->getName() + Twine(".reload"));
   };
 
   for (auto const &E : Spills) {
     // If we have not seen the value, generate a spill.
     if (CurrentValue != E.def()) {
       CurrentValue = E.def();
       CurrentBlock = nullptr;
       CurrentReload = nullptr;
 
-      ++Index;
+      Index = E.fieldIndex();
 
       if (auto *AI = dyn_cast<AllocaInst>(CurrentValue)) {
         // Spilled AllocaInst will be replaced with GEP from the coroutine frame
         // there is no spill required.
         Allocas.emplace_back(AI, Index);
         if (!AI->isStaticAlloca())
           report_fatal_error("Coroutines cannot handle non static allocas yet");
       } else {

lib/Transforms/Coroutines/CoroInternal.h

struct LLVM_LIBRARY_VISIBILITY Shape {
   SmallVector<CoroSuspendInst *, 4> CoroSuspends;
 
   // Field Indexes for known coroutine frame fields.
   enum {
     ResumeField,
     DestroyField,
     PromiseField,
     IndexField,
-    LastKnownField = IndexField
   };
 
   StructType *FrameTy;
   Instruction *FramePtr;
   BasicBlock *AllocaSpillBlock;
   SwitchInst *ResumeSwitch;
   AllocaInst *PromiseAlloca;
   bool HasFinalSuspend;

test/Transforms/Coroutines/coro-padding.ll

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