[Coroutine] Properly deal with byval and noalias parameters

[lxfind]

Summary

This patch is to address https://bugs.llvm.org/show_bug.cgi?id=48857.
Previous attempts can be found in D104007 and D101980.
A lot of discussions can be found in those two patches.
To summarize the bug:
When Clang emits IR for coroutines, the first thing it does is to make a copy of every argument to the local stack, so that uses of the arguments in the function will all refer to the local copies instead of the arguments directly.
However, in some cases we find that arguments are still directly used:
When Clang emits IR for a function that has pass-by-value arguments, sometimes it emits an argument with byval attribute. A byval attribute is considered to be local to the function (just like alloca) and hence it can be easily determined that it does not alias other values. If in the IR there exists a memcpy from a byval argument to a local alloca, and then from that local alloca to another alloca, MemCpyOpt will optimize out the first memcpy because byval argument's content will not change. This causes issues because after a coroutine suspension, the byval argument may die outside of the function, and latter uses will lead to memory use-after-free.
This is only a problem for arguments with either byval attribute or noalias attribute, because only these two kinds are considered local. Arguments without these two attributes will be considered to alias coro_suspend and hence we won't have this problem. So we need to be able to deal with these two attributes in coroutines properly.
For noalias arguments, since coro_suspend may potentially change the value of any argument outside of the function, we simply shouldn't mark any argument in a coroutiune as noalias. This can be taken care of in CoroEarly pass.
For byval arguments, if such an argument needs to live across suspensions, we will have to copy their value content to the frame, not just the pointer.

[lxfind]

I just realized that this is now getting very similar to D101980, and we are facing the same problem of not always be able to tell whether a param is byval in the midend. Let me continue the discussion on the BasicAA patch.

[lxfind]

only check byval attribute

[efriedma]

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

224 ↗

(On Diff #352016)

Hmm...

It might make sense to fix this in clang, rather than here, but I guess this is okay?

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1605

Use memcpy to copy large values; load+store results in very inefficient code.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1605

Good to know. We actually use load+store for all the frame rewrites. We can look at changing all of that in a separate patch.

[ChuanqiXu]

This comment was removed by ChuanqiXu.

[lxfind]

I updated the description to reflex the summary of discussions from previous patches. Based on the discussion, we believe that this should be the right approach

[ChuanqiXu]

The overall patch looks good to me.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1141–1142

not just the pointer itself.

'instead of the pointer itself' since we wouldn't store the pointer anymore, right?

1605

Use memcpy to copy large values; load+store results in very inefficient code.

Would the successive pass convert load+store to memcpy?

[lxfind]

update comment

[lxfind]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1605

I think it would

[efriedma]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1605

At -O0, we probably end up with a mess. At higher optimization levels, I think we have some code in instcombine to transform weird load/store operations; not sure if we reliably end up with a memcpy. In any case, no reason to rely on that transform here.

[ChuanqiXu]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

1605

Thanks, I got it.

[ChuanqiXu]

llvm/test/Transforms/Coroutines/coro-byval-param.ll

7

Since this patch deals with 'byval' and 'noalias' arguments, it'd better to add a test case for 'noalias' arguments.

[lxfind]

llvm/test/Transforms/Coroutines/coro-byval-param.ll

7

I am actually not sure what would be a reasonable case that involves noalias (I could just duplicate foo with a noalias arg, but I am hoping it can be more realistic)

[ChuanqiXu]

llvm/test/Transforms/Coroutines/coro-byval-param.ll

7

Yeah, I spent a little time and still don't find realistic examples. If we can't get it in short time, I think the covering the codes may be important too.

[lxfind]

Add test case for noalias argument

[efriedma]

Looking at clang a bit more, I think the only way you end up with noalias arguments at the moment is via "__restrict". Which ends up looking basically like your testcase. Probably not too important in practice.

Patch looks fine from my side, but I'd like to leave final approval to someone more familiar with the coroutine code.

[ChuanqiXu]

LGTM.

[thakis]

Looks like this breaks tests on windows: http://45.33.8.238/win/40255/step_11.txt

Please take a look and revert for now if it takes a while to fix.

[ChuanqiXu]

In D104184#2826151, @thakis wrote:

Looks like this breaks tests on windows: http://45.33.8.238/win/40255/step_11.txt

Please take a look and revert for now if it takes a while to fix.

Although I only test this on Linux, it's hard to believe that this patch would affect that test case. Since this one only touched coroutine module, it shouldn't affect IR without llvm.coro.* intrinsics.

[jeroen.dobbelaere]

In D104184#2825350, @efriedma wrote:

Looking at clang a bit more, I think the only way you end up with noalias arguments at the moment is via "__restrict". Which ends up looking basically like your testcase. Probably not too important in practice.

Patch looks fine from my side, but I'd like to leave final approval to someone more familiar with the coroutine code.

Clang can also introduce noalias arguments when a function returns a struct.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   // 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);
+    Type *FieldType = S.first->getType();
+    // For byval arguments, we need to store the pointed value in the frame,
+    // not just the pointer itself.
+    if (auto *Arg = dyn_cast<Argument>(S.first))
+      if (Arg->hasPassPointeeByValueCopyAttr())
+        FieldType = Arg->getType()->getPointerElementType();
+    FieldIDType Id = B.addField(FieldType, None);
     FrameData.setFieldIndex(S.first, Id);
   }
 
   B.finish(FrameTy);
   FrameData.updateLayoutIndex(B);
   Shape.FrameAlign = B.getStructAlign();
   Shape.FrameSize = B.getStructSize();
 
auto GetFramePointer = [&](Value *Orig) -> Value * {
     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 IsByValArg = 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();
 
       // If we're spilling an Argument, make sure we clear 'nocapture'
       // from the coroutine function.
       Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::NoCapture);
 
+      if (Arg->hasPassPointeeByValueCopyAttr())
+        IsByValArg = true;
+
     } else if (auto *CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
       // Don't spill immediately after a suspend; splitting assumes
       // that the suspend will be followed by a branch.
       InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHI();
     } else {
       auto *I = cast<Instruction>(Def);
       if (!DT.dominates(CB, I)) {
         // If it is not dominated by CoroBegin, then spill should be
if (auto *Arg = dyn_cast<Argument>(Def)) {
         InsertPt = I->getNextNode();
       }
     }
 
     auto Index = FrameData.getFieldIndex(Def);
     Builder.SetInsertPoint(InsertPt);
     auto *G = Builder.CreateConstInBoundsGEP2_32(
         FrameTy, FramePtr, 0, Index, Def->getName() + Twine(".spill.addr"));
+    if (IsByValArg) {
+      // For byval arguments, we need to store the pointed value in the frame,
+      // not just the pointer itself.
+      auto *Value =
+          Builder.CreateLoad(Def->getType()->getPointerElementType(), Def);
+      Builder.CreateStore(Value, G);
+    } else {
       Builder.CreateStore(Def, G);
+    }
 
     BasicBlock *CurrentBlock = nullptr;
     Value *CurrentReload = nullptr;
     for (auto *U : E.second) {
       // If we have not seen the use block, create a load instruction to reload
       // the spilled value from the coroutine frame. Populates the Value pointer
       // reference provided with the frame GEP.
       if (CurrentBlock != U->getParent()) {
         CurrentBlock = U->getParent();
         Builder.SetInsertPoint(&*CurrentBlock->getFirstInsertionPt());
 
         auto *GEP = GetFramePointer(E.first);
         GEP->setName(E.first->getName() + Twine(".reload.addr"));
+        if (IsByValArg)
+          CurrentReload = GEP;
+        else
           CurrentReload = Builder.CreateLoad(
               FrameTy->getElementType(FrameData.getFieldIndex(E.first)), GEP,
               E.first->getName() + Twine(".reload"));
 
         TinyPtrVector<DbgDeclareInst *> DIs = FindDbgDeclareUses(Def);
         for (DbgDeclareInst *DDI : DIs) {
           bool AllowUnresolved = false;
           // This dbg.declare is preserved for all coro-split function
           // fragments. It will be unreachable in the main function, and
           // processed by coro::salvageDebugInfo() by CoroCloner.
           DIBuilder(*CurrentBlock->getParent()->getParent(), AllowUnresolved)

llvm/test/Transforms/Coroutines/coro-byval-param.ll

+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+%promise_type = type { i8 }
+%struct.A = type <{ i64, i64, i32, [4 x i8] }>
+
+; Function Attrs: noinline ssp uwtable mustprogress
+define %promise_type* @foo(%struct.A* nocapture readonly byval(%struct.A) align 8 %a1) #0 {
+entry:
+  %__promise = alloca %promise_type, align 1
+  %a2 = alloca %struct.A, align 8
+  %0 = getelementptr inbounds %promise_type, %promise_type* %__promise, i64 0, i32 0
+  %1 = call token @llvm.coro.id(i32 16, i8* nonnull %0, i8* bitcast (%promise_type* (%struct.A*)* @foo to i8*), i8* null)
+  %2 = call i1 @llvm.coro.alloc(token %1)
+  br i1 %2, label %coro.alloc, label %coro.init
+
+coro.alloc:                                       ; preds = %entry
+  %3 = call i64 @llvm.coro.size.i64()
+  %call = call noalias nonnull i8* @_Znwm(i64 %3) #9
+  br label %coro.init
+
+coro.init:                                        ; preds = %coro.alloc, %entry
+  %4 = phi i8* [ null, %entry ], [ %call, %coro.alloc ]
+  %5 = call i8* @llvm.coro.begin(token %1, i8* %4) #10
+  %6 = bitcast %struct.A* %a1 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* nonnull %0) #2
+  %call2 = call %promise_type* @_ZN4task12promise_type17get_return_objectEv(%promise_type* nonnull dereferenceable(1) %__promise)
+  call void @initial_suspend(%promise_type* nonnull dereferenceable(1) %__promise)
+  %7 = call token @llvm.coro.save(i8* null)
+  call fastcc void @_ZNSt12experimental13coroutines_v116coroutine_handleIN4task12promise_typeEE12from_addressEPv(i8* %5) #2
+  %8 = call i8 @llvm.coro.suspend(token %7, i1 false)
+  switch i8 %8, label %coro.ret [
+    i8 0, label %init.ready
+    i8 1, label %cleanup33
+  ]
+
+init.ready:                                       ; preds = %coro.init
+  %9 = bitcast %struct.A* %a2 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 24, i8* nonnull %9) #2
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %9, i8* align 8 %6, i64 24, i1 false)
+  call void @llvm.lifetime.end.p0i8(i64 24, i8* nonnull %9) #2
+  call void @_ZN4task12promise_type13final_suspendEv(%promise_type* nonnull dereferenceable(1) %__promise) #2
+  %10 = call token @llvm.coro.save(i8* null)
+  call fastcc void @_ZNSt12experimental13coroutines_v116coroutine_handleIN4task12promise_typeEE12from_addressEPv(i8* %5) #2
+  %11 = call i8 @llvm.coro.suspend(token %10, i1 true) #10
+  %switch = icmp ult i8 %11, 2
+  br i1 %switch, label %cleanup33, label %coro.ret
+
+cleanup33:                                        ; preds = %init.ready, %coro.init
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* nonnull %0) #2
+  %12 = call i8* @llvm.coro.free(token %1, i8* %5)
+  %.not = icmp eq i8* %12, null
+  br i1 %.not, label %coro.ret, label %coro.free
+
+coro.free:                                        ; preds = %cleanup33
+  call void @_ZdlPv(i8* nonnull %12) #2
+  br label %coro.ret
+
+coro.ret:                                         ; preds = %coro.free, %cleanup33, %init.ready, %coro.init
+  %13 = call i1 @llvm.coro.end(i8* null, i1 false) #10
+  ret %promise_type* %call2
+}
+
+; check that the frame contains the entire struct, instead of just the struct pointer
+; CHECK: %foo.Frame = type { void (%foo.Frame*)*, void (%foo.Frame*)*, %promise_type, %struct.A, i1 }
+
+; Function Attrs: argmemonly nounwind readonly
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
+
+; Function Attrs: nounwind
+declare i1 @llvm.coro.alloc(token) #2
+
+; Function Attrs: nobuiltin nofree allocsize(0)
+declare nonnull i8* @_Znwm(i64) local_unnamed_addr #3
+
+; Function Attrs: nounwind readnone
+declare i64 @llvm.coro.size.i64() #4
+
+; Function Attrs: nounwind
+declare i8* @llvm.coro.begin(token, i8* writeonly) #2
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #5
+
+; Function Attrs: argmemonly nofree nounwind willreturn
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg) #6
+
+; Function Attrs: noinline nounwind ssp uwtable willreturn mustprogress
+declare %promise_type* @_ZN4task12promise_type17get_return_objectEv(%promise_type* nonnull dereferenceable(1)) local_unnamed_addr #7 align 2
+
+; Function Attrs: noinline nounwind ssp uwtable willreturn mustprogress
+declare void @initial_suspend(%promise_type* nonnull dereferenceable(1)) local_unnamed_addr #7 align 2
+
+; Function Attrs: nounwind
+declare token @llvm.coro.save(i8*) #2
+
+; Function Attrs: noinline nounwind ssp uwtable willreturn mustprogress
+declare hidden fastcc void @_ZNSt12experimental13coroutines_v116coroutine_handleIN4task12promise_typeEE12from_addressEPv(i8*) unnamed_addr #7 align 2
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #5
+
+; Function Attrs: nounwind
+declare i8 @llvm.coro.suspend(token, i1) #2
+
+; Function Attrs: noinline nounwind ssp uwtable willreturn mustprogress
+declare void @_ZN4task12promise_type13final_suspendEv(%promise_type* nonnull dereferenceable(1)) local_unnamed_addr #7 align 2
+
+; Function Attrs: nounwind
+declare i1 @llvm.coro.end(i8*, i1) #2
+
+; Function Attrs: nobuiltin nounwind
+declare void @_ZdlPv(i8*) local_unnamed_addr #8
+
+; Function Attrs: argmemonly nounwind readonly
+declare i8* @llvm.coro.free(token, i8* nocapture readonly) #1
+
+attributes #0 = { noinline ssp uwtable mustprogress "coroutine.presplit"="1" "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="penryn" "target-features"="+cx16,+cx8,+fxsr,+mmx,+sahf,+sse,+sse2,+sse3,+sse4.1,+ssse3,+x87" "tune-cpu"="generic" }
+attributes #1 = { argmemonly nounwind readonly }
+attributes #2 = { nounwind }
+attributes #3 = { nobuiltin nofree allocsize(0) "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="penryn" "target-features"="+cx16,+cx8,+fxsr,+mmx,+sahf,+sse,+sse2,+sse3,+sse4.1,+ssse3,+x87" "tune-cpu"="generic" }
+attributes #4 = { nounwind readnone }
+attributes #5 = { argmemonly nofree nosync nounwind willreturn }
+attributes #6 = { argmemonly nofree nounwind willreturn }
+attributes #7 = { noinline nounwind ssp uwtable willreturn mustprogress "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="penryn" "target-features"="+cx16,+cx8,+fxsr,+mmx,+sahf,+sse,+sse2,+sse3,+sse4.1,+ssse3,+x87" "tune-cpu"="generic" }
+attributes #8 = { nobuiltin nounwind "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="penryn" "target-features"="+cx16,+cx8,+fxsr,+mmx,+sahf,+sse,+sse2,+sse3,+sse4.1,+ssse3,+x87" "tune-cpu"="generic" }
+attributes #9 = { allocsize(0) }
+attributes #10 = { noduplicate }
+