[coroutine] Respect alloca alignment requirements.

[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 }

Also, removed useless std::sort for spills that was not required and was causing needless non-determinism of builds since it was sorting based on addresses of the instruction that may result in different sorted order between the builds

[lewissbaker]

All looks fine to me.

[GorNishanov]

Put out an updated version of this patch in

https://reviews.llvm.org/D45221

lib/Transforms/Coroutines/CoroFrame.cpp

struct Spill : std::pair<Value *, Instruction *> {
   using base = std::pair<Value *, Instruction *>;
 
   Spill(Value *Def, User *U) : base(Def, cast<Instruction>(U)) {}
 
   Value *def() const { return first; }
   Instruction *user() const { return second; }
   BasicBlock *userBlock() const { return second->getParent(); }
 
-  std::pair<Value *, BasicBlock *> getKey() const {
-    return {def(), userBlock()};
+  // 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 field() const {
+    assert(FieldNo && "Accessing unassigned field");
+    return FieldNo;
+  }
+  void setField(unsigned FieldNumber) {
+    assert(!FieldNo && "Reassigning field number");
+    FieldNo = FieldNumber;
   }
 
-  bool operator<(Spill const &rhs) const { return getKey() < rhs.getKey(); }
+private:
+  unsigned FieldNo = 0;
 };
 
 // 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 instruciton
+// 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.setField(Types.size());
+
     Types.push_back(Ty);
+    Padder.addType(Ty);
   }
   FrameTy->setBody(Types);
 
   return FrameTy;
 }
 
 // Replace all alloca and SSA values that are accessed across suspend points
 // with GetElementPointer from coroutine frame + loads and stores. Create an
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;
 
   // 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.field();
 
       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 {

 
 static void rewritePHIs(BasicBlock &BB) {
   // For every incoming edge we will create a block holding all
   // incoming values in a single PHI nodes.
   //
   // loop:
   //    %n.val = phi i32[%n, %entry], [%inc, %loop]
   //
-  // It will create:  
+  // It will create:
   //
   // loop.from.entry:
   //    %n.loop.pre = phi i32 [%n, %entry]
   //    br %label loop
   // loop.from.loop:
   //    %inc.loop.pre = phi i32 [%inc, %loop]
   //    br %label loop
   //
for (Instruction &I : instructions(F))
         for (User *U : I.users())
           if (Checker.isDefinitionAcrossSuspend(I, U))
             Spills.emplace_back(&I, U);
 
     if (Spills.empty())
       break;
 
     // Rewrite materializable instructions to be materialized at the use point.
-    std::sort(Spills.begin(), Spills.end());
     DEBUG(dump("Materializations", Spills));
     rewriteMaterializableInstructions(Builder, Spills);
     Spills.clear();
   }
 
   // Collect the spills for arguments and other not-materializable values.
   for (Argument &A : F.getArgumentList())
     for (User *U : A.users())
for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {
         // We cannot spill a token.
         if (I.getType()->isTokenTy())
           report_fatal_error(
               "token definition is separated from the use by a suspend point");
         Spills.emplace_back(&I, U);
       }
   }
-  std::sort(Spills.begin(), Spills.end());
   DEBUG(dump("Spills", Spills));
   moveSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);
 }

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