[coro async] Cap the alignment of spilled values (vs. allocas) at the max frame alignment

[aschwaighofer]

Summary

Before this patch we would normally use the ABI alignment which can be
to high for the context alginment.

For spilled values we don't need ABI alignment, since the frame entry's
address is not escaped.

rdar://79664965

[rjmccall]

Shouldn't there be a test case showing that a spilled value with high ABI alignment is placed in an appropriately-underaligned slot, and that the loads and stores to that slot have the necessary alignment? Because the alignment in this test is just coming from the alloca's explicit alignment.

[aschwaighofer]

The spilled value from this test comes from trying to spill %vector_spill which is a value life across a suspend point. Because it has a value it has no alignment. It is misleading that we load the value from an alloca. It could also be the result of a computation. This test case crashes before the patch because we are assigning the abi alignment for the value '%vector_spill`.

[aschwaighofer]

Oh my. You are right I should have verified that when we spill that value though we use the right alignment. Because we don't. :(.

%vector_spill = load <4 x double>, <4 x double>* %vector, align 16
%vector_spill.spill.addr = getelementptr inbounds %my_async_function.Frame, %my_async_function.Frame* %FramePtr, i32 0, i32 1, !dbg !6
 store <4 x double> %vector_spill, <4 x double>* %vector_spill.spill.addr, align 32, !dbg !6

The code that spills/reloads does not use the alignment we have computed for the spill slot.
Rather it just does:

Builder.CreateStore(Def, G);

[aschwaighofer]

Use the compute type alignment for spilled values when spilling/restoring

[rjmccall]

Minor style request, but otherwise LGTM

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1565

I would assume from the name that this is something like the alignment of the type. Maybe SpillAlignment?

[aschwaighofer]

Change variable name from TyAlignment to SpillAlignment

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

private:
   };
 
   const DataLayout &DL;
   LLVMContext &Context;
   uint64_t StructSize = 0;
   Align StructAlign;
   bool IsFinished = false;
 
+  Optional<Align> MaxFrameAlignment;
+
   SmallVector<Field, 8> Fields;
   DenseMap<Value*, unsigned> FieldIndexByKey;
 
 public:
-  FrameTypeBuilder(LLVMContext &Context, DataLayout const &DL)
-      : DL(DL), Context(Context) {}
+  FrameTypeBuilder(LLVMContext &Context, DataLayout const &DL,
+                   Optional<Align> MaxFrameAlignment)
+      : DL(DL), Context(Context), MaxFrameAlignment(MaxFrameAlignment) {}
 
   /// Add a field to this structure for the storage of an `alloca`
   /// instruction.
   LLVM_NODISCARD FieldIDType addFieldForAlloca(AllocaInst *AI,
                                                bool IsHeader = false) {
     Type *Ty = AI->getAllocatedType();
 
     // Make an array type if this is a static array allocation.
public:
   ///
   /// Side Effects: Because We sort the allocas, the order of allocas in the
   /// frame may be different with the order in the source code.
   void addFieldForAllocas(const Function &F, FrameDataInfo &FrameData,
                           coro::Shape &Shape);
 
   /// Add a field to this structure.
   LLVM_NODISCARD FieldIDType addField(Type *Ty, MaybeAlign FieldAlignment,
-                                      bool IsHeader = false) {
+                                      bool IsHeader = false,
+                                      bool IsSpillOfValue = false) {
     assert(!IsFinished && "adding fields to a finished builder");
     assert(Ty && "must provide a type for a field");
 
     // The field size is always the alloc size of the type.
     uint64_t FieldSize = DL.getTypeAllocSize(Ty);
 
     // For an alloca with size=0, we don't need to add a field and they
     // can just point to any index in the frame. Use index 0.
     if (FieldSize == 0) {
       return 0;
     }
 
     // The field alignment might not be the type alignment, but we need
     // to remember the type alignment anyway to build the type.
-    Align TyAlignment = DL.getABITypeAlign(Ty);
-    if (!FieldAlignment) FieldAlignment = TyAlignment;
+    // If we are spilling values we don't need to worry about ABI alignment
+    // concerns.
+    auto ABIAlign = DL.getABITypeAlign(Ty);
+    Align TyAlignment =
+        (IsSpillOfValue && MaxFrameAlignment)
+            ? (*MaxFrameAlignment < ABIAlign ? *MaxFrameAlignment : ABIAlign)
+            : ABIAlign;
+    if (!FieldAlignment) {
+      FieldAlignment = TyAlignment;
+    }
 
     // Lay out header fields immediately.
     uint64_t Offset;
     if (IsHeader) {
       Offset = alignTo(StructSize, FieldAlignment);
       StructSize = Offset + FieldSize;
 
     // Everything else has a flexible offset.
static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   LLVMContext &C = F.getContext();
   const DataLayout &DL = F.getParent()->getDataLayout();
   StructType *FrameTy = [&] {
     SmallString<32> Name(F.getName());
     Name.append(".Frame");
     return StructType::create(C, Name);
   }();
 
-  FrameTypeBuilder B(C, DL);
+  // We will use this value to cap the alignment of spilled values.
+  Optional<Align> MaxFrameAlignment;
+  if (Shape.ABI == coro::ABI::Async)
+    MaxFrameAlignment = Shape.AsyncLowering.getContextAlignment();
+  FrameTypeBuilder B(C, DL, MaxFrameAlignment);
 
   AllocaInst *PromiseAlloca = Shape.getPromiseAlloca();
   Optional<FieldIDType> SwitchIndexFieldId;
 
   if (Shape.ABI == coro::ABI::Switch) {
     auto *FramePtrTy = FrameTy->getPointerTo();
     auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
                                    /*IsVarArg=*/false);
static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   // Create an entry for every spilled value.
   for (auto &S : FrameData.Spills) {
     Type *FieldType = S.first->getType();
     // For byval arguments, we need to store the pointed value in the frame,
     // instead of the pointer itself.
     if (const Argument *A = dyn_cast<Argument>(S.first))
       if (A->hasByValAttr())
         FieldType = FieldType->getPointerElementType();
-    FieldIDType Id = B.addField(FieldType, None);
+    FieldIDType Id =
+        B.addField(FieldType, None, false /*header*/, true /*IsSpillOfValue*/);
     FrameData.setFieldIndex(S.first, Id);
   }
 
   B.finish(FrameTy);
   FrameData.updateLayoutIndex(B);
   Shape.FrameAlign = B.getStructAlign();
   Shape.FrameSize = B.getStructSize();
 
if (isa<AllocaInst>(Orig)) {
         return Builder.CreateBitCast(GEP, Orig->getType(),
                                      Orig->getName() + Twine(".cast"));
     }
     return GEP;
   };
 
   for (auto const &E : FrameData.Spills) {
     Value *Def = E.first;
     // Create a store instruction storing the value into the
     // coroutine frame.
     Instruction *InsertPt = nullptr;
     bool NeedToCopyArgPtrValue = false;
     if (auto *Arg = dyn_cast<Argument>(Def)) {
       // For arguments, we will place the store instruction right after
       // the coroutine frame pointer instruction, i.e. bitcast of
       // coro.begin from i8* to %f.frame*.
       InsertPt = FramePtr->getNextNode();

llvm/test/Transforms/Coroutines/coro-async.ll

entry:
   %resume_ctxt = load i8*, i8** %resume_ctxt_addr, align 8
   ret i8* %resume_ctxt
 }
 
 
 define swiftcc void @my_async_function(i8* swiftasync %async.ctxt, %async.task* %task, %async.actor* %actor) !dbg !1 {
 entry:
   %tmp = alloca { i64, i64 }, align 8
+  %vector = alloca <4 x double>, align 16
   %proj.1 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 0
   %proj.2 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 1
 
   %id = call token @llvm.coro.id.async(i32 128, i32 16, i32 0,
           i8* bitcast (<{i32, i32}>* @my_async_function_fp to i8*))
   %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
   store i64 0, i64* %proj.1, align 8
   store i64 1, i64* %proj.2, align 8
entry:
   %resume.func_ptr = call i8* @llvm.coro.async.resume()
   store i8* %resume.func_ptr, i8** %return_to_caller.addr
 
   ; store caller context into callee context
   %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr
   %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
   %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
+  %vector_spill = load <4 x double>, <4 x double>* %vector, align 16
   %res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
                                                   i8* %resume.func_ptr,
                                                   i8* %resume_proj_fun,
                                                   void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
                                                   i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor), !dbg !5
 
   call void @llvm.coro.async.context.dealloc(i8* %callee_context)
   %continuation_task_arg = extractvalue {i8*, i8*, i8*} %res, 1
   %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
   %val = load i64, i64* %proj.1
   call void @some_user(i64 %val)
   %val.2 = load i64, i64* %proj.2
   call void @some_user(i64 %val.2)
-
+  store <4 x double> %vector_spill, <4 x double>* %vector, align 16
   tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor)
   call i1 (i8*, i1, ...) @llvm.coro.end.async(i8* %hdl, i1 0)
   unreachable
 }
 
 define void @my_async_function_pa(i8* %ctxt, %async.task* %task, %async.actor* %actor) {
   call void @llvm.coro.async.size.replace(i8* bitcast (<{i32, i32}>* @my_async_function_pa_fp to i8*), i8* bitcast (<{i32, i32}>* @my_async_function_fp to i8*))
   call swiftcc void @my_async_function(i8* %ctxt, %async.task* %task, %async.actor* %actor)