[coroutine] update promise object's final layout index

[ychen]

Summary

promise is a header field but it is not guaranteed that it would be the third
field of the frame due to performOptimizedStructLayout..

[ychen]

ping?

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

746

PromiseAlloca is added here as a header field, which means it should have a fixed index?
Could it be there are bugs in performOptimizedStructLayout that moved it around somehow?

[ychen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

746

I'm inclined to say it should be fixed offset which is not changed after performOptimizedStructLayout.

1. The use of IsHeader in addFieldForAllocas suggests that.
2. The comment below in OptimizedStructLayout.h

/// - Fields may be assigned a fixed offset in the layout.  If there are
///   gaps among the fixed-offset fields, the algorithm may attempt
///   to allocate flexible-offset fields into those gaps.  If that's
///   undesirable, the caller should "block out" those gaps by e.g.
///   just creating a single fixed-offset field that represents the
///   entire "header".

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

746

You are right.
In that case, I wonder if it is still necessary to add PromiseAlloca as a header field?

[ychen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

746

Yeah, that's an option too. TBH, I'm not quite sure about the importance or implication of PromiseAlloca using flexible offset. Or maybe it would make llvm.coro.promise implementation hard or impossible?

[lxfind]

Thanks. I agree your fix is correct.
Some small suggestions commented

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

You can move this code to right after "B.addFieldForAllocas(F, FrameData, Shape);" (comments will need to be updated)
Then you won't need to modify "updateLayoutIndex".
That way I think it's actually cleaner.

[ychen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

My first attempt was here. It seems FrameTypeBuilder object is not available here to do the update.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

Not update here. But move these code to right after "B.addFieldForAllocas(F, FrameData, Shape);"

[ychen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

I may misunderstand this. Right after "B.addFieldForAllocas(F, FrameData, Shape);", performOptimizedStructLayout is not yet done, how do I update the index?

2299–2300

Or you mean after B.finish(FrameTy);?

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

You won't need to update the index manually, if you append the promise alloca to the alloca list after addFieldForAllocas. updateLayoutIndex will take care of it naturally.

[ychen]

Address comments

[ychen]

rebase

[ychen]

rebase

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

2299–2300

Thanks for the explanation!

[lxfind]

Thank you!

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

if (Shape.ABI == coro::ABI::Switch) {
     (void)B.addField(FnPtrTy, None, /*header*/ true);
     (void)B.addField(FnPtrTy, None, /*header*/ true);
 
     // PromiseAlloca field needs to be explicitly added here because it's
     // a header field with a fixed offset based on its alignment. Hence it
     // needs special handling and cannot be added to FrameData.Allocas.
     if (PromiseAlloca)
       FrameData.setFieldIndex(
           PromiseAlloca, B.addFieldForAlloca(PromiseAlloca, /*header*/ true));
 
     // Add a field to store the suspend index.  This doesn't need to
     // be in the header.
     unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
     Type *IndexType = Type::getIntNTy(C, IndexBits);
 
     SwitchIndexFieldId = B.addField(IndexType, None);
   } else {
     assert(PromiseAlloca == nullptr && "lowering doesn't support promises");
   }
 
   // Because multiple allocas may own the same field slot,
   // we add allocas to field here.
   B.addFieldForAllocas(F, FrameData, Shape);
+  // Add PromiseAlloca to Allocas list so that
+  // 1. updateLayoutIndex could update its index after
+  // `performOptimizedStructLayout`
+  // 2. it is processed in insertSpills.
+  if (Shape.ABI == coro::ABI::Switch && PromiseAlloca)
+    // We assume that the promise alloca won't be modified before
+    // CoroBegin and no alias will be create before CoroBegin.
+    FrameData.Allocas.emplace_back(
+        PromiseAlloca, DenseMap<Instruction *, llvm::Optional<APInt>>{}, false);
   // Create an entry for every spilled value.
   for (auto &S : FrameData.Spills) {
     FieldIDType Id = B.addField(S.first->getType(), None);
     FrameData.setFieldIndex(S.first, Id);
   }
 
   B.finish(FrameTy);
   FrameData.updateLayoutIndex(B);
for (User *U : I.users())
               "token definition is separated from the use by a suspend point");
         FrameData.Spills[&I].push_back(cast<Instruction>(U));
       }
   }
   LLVM_DEBUG(dumpSpills("Spills", FrameData.Spills));
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
       Shape.ABI == coro::ABI::Async)
     sinkSpillUsesAfterCoroBegin(F, FrameData, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, FrameData);
-  // Add PromiseAlloca to Allocas list so that it is processed in insertSpills.
-  if (Shape.ABI == coro::ABI::Switch && Shape.SwitchLowering.PromiseAlloca)
-    // We assume that the promise alloca won't be modified before
-    // CoroBegin and no alias will be create before CoroBegin.
-    FrameData.Allocas.emplace_back(
-        Shape.SwitchLowering.PromiseAlloca,
-        DenseMap<Instruction *, llvm::Optional<APInt>>{}, false);
   Shape.FramePtr = insertSpills(FrameData, Shape);
   lowerLocalAllocas(LocalAllocas, DeadInstructions);
 
   for (auto I : DeadInstructions)
     I->eraseFromParent();
 }

llvm/test/Transforms/Coroutines/coro-spill-promise.ll

+; Check that promise object is reloaded from the correct index of the coro frame.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+%"class.task::promise_type" = type { [64 x i8] }
+
+declare void @consume(i32*)
+declare void @consume2(%"class.task::promise_type"*)
+
+define i8* @f() "coroutine.presplit"="1" {
+entry:
+  %data = alloca i32, align 4
+  %__promise = alloca %"class.task::promise_type", align 64
+  %pv = bitcast %"class.task::promise_type"* %__promise to i8*
+  %id = call token @llvm.coro.id(i32 0, i8* %pv, 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(i32* %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(i32* %data)
+  call void @consume2(%"class.task::promise_type"* %__promise)
+  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
+}
+
+; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i32, i1, [43 x i8], %"class.task::promise_type" }
+
+; CHECK-LABEL: @f.resume(
+; CHECK: %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
+; CHECK: call void @consume2(%"class.task::promise_type"* %[[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*)