[Coroutines] Use allocator overload when available

[modocache]

Summary

Depends on https://reviews.llvm.org/D42605.

An implementation of the behavior described in [dcl.fct.def.coroutine]/7:
when a promise type overloads operator new using a "placement new"
that takes the same argument types as the coroutine function, that
overload is used when allocating the coroutine frame.

Simply passing references to the coroutine function parameters directly
to operator new results in invariant violations in LLVM's coroutine
splitting pass, so this implementation modifies Clang codegen to
produce allocator-specific alloc/store/loads for each parameter being
forwarded to the allocator.

Test Plan: check-clang

[modocache]

Add some diagnostics tests to SemaCXX/coroutines.cpp.

[GorNishanov]

lib/CodeGen/CGCoroutine.cpp

526–578

First, thank you for doing this change!

Second, I am not sure that this part (CGCoroutine.cpp change) belongs in clang.
llvm CoroFrame is doing an unsafe transformation (it was safe until we got the arguments to operator new :-) ).

Moving the stores after the new that loads from those stores is an incorrect transformation. I think it needs to be addressed in llvm. getNotRelocatableInstructions function in CoroSplit.cpp needs to add special handling for AllocaInst and freeze the stores to that Alloca in the blocks preceeding the one with CoroBegin (getCoroBeginPredBlocks).

Also, for this to work for cases where parameter is used both in allocating function and in the body of the coroutine we need to have a copy. Currently, front-end does not create copies for scalar types (see CoroutineStmtBuilder::makeParamMoves() in SemaCoroutine.cpp). I think if we always create copies for all parameters, it will make this change more straightforward.

Third, this code does not handle cases where scalar values passed by reference to an allocation function or a struct passed by value:

Try this code on these:

void *operator new(unsigned long, promise_matching_placement_new_tag,
                   int&, float, double);

or

struct Dummy { int x, y, z; ~Dummy(); };

template<>
struct std::experimental::coroutine_traits<void, promise_matching_placement_new_tag, Dummy&, float, double> {
   struct promise_type {
       void *operator new(unsigned long, promise_matching_placement_new_tag,
                   Dummy, float, double);

I think if this code is changed according to my earlier suggestion of doing copies in clang and freezing stores in the llvm, it should take care the cases above.

These cases need to be added as tests to llvm\tests\Transforms\Coroutines

[GorNishanov]

lib/CodeGen/CGCoroutine.cpp

526–578

Alternatively, we can keep current clang behavior with respect to not doing copies for scalars and instead create a copy in llvm if we decided to freeze the store AND that alloca is used after a suspend point (CoroFrame decided that it needs to spill it into the coroutine frame).

[modocache]

Prevent coroutine parameter stores from being moved, rely on https://reviews.llvm.org/D43000 instead.

[modocache]

Update the tests to work with scalar parameter copies.

[GorNishanov]

lib/Sema/SemaCoroutine.cpp

1071

This does not implement TS specified behavior for non static member functions:

[dcl.fct.def.coroutine]/7 states that an argument list is build up as follows:

The first argument is the amount of space requested, and has type std::size_t. The lvalues p1 ... pn are the succeeding arguments.

Where p1 ... pn are defined earlier in

[dcl.fct.def.coroutine]/3 as:

For a coroutine f that is a non-static member function, let P1 denote the type of the implicit object parameter (13.3.1) and P2 ... Pn be the types of the function parameters; otherwise let P1 ... Pn be the types of the function parameters.

Essentially for non-static member functions, we need to insert implicit object parameter.

Note that lookupPromiseType in SemaCoroutine.cpp when building specialization of std::experimental::coroutine_traits<...> includes implicit object parameter:

// If the function is a non-static member function, add the type
// of the implicit object parameter before the formal parameters.
if (auto *MD = dyn_cast<CXXMethodDecl>(FD)) {
  if (MD->isInstance()) {
    // [over.match.funcs]4
    // For non-static member functions, the type of the implicit object
    // parameter is
    //  -- "lvalue reference to cv X" for functions declared without a
    //      ref-qualifier or with the & ref-qualifier
    //  -- "rvalue reference to cv X" for functions declared with the &&
    //      ref-qualifier
    QualType T =
        MD->getThisType(S.Context)->getAs<PointerType>()->getPointeeType();
    T = FnType->getRefQualifier() == RQ_RValue
            ? S.Context.getRValueReferenceType(T)
            : S.Context.getLValueReferenceType(T, /*SpelledAsLValue*/ true);
    AddArg(T);
  }
}

[GorNishanov]

lib/Sema/SemaCoroutine.cpp

1071

I think promise constructor argument preview is also missing the implicit object parameter.

[modocache]

Thanks for the review, @GorNishanov, and for pointing out that part of the standard! I added a reference to the implicit object parameter, as well as some tests for that behavior. And thanks for pointing out the flaw in https://reviews.llvm.org/D41820 -- I'll submit a separate diff to forward the implicit object parameter to promise type constructors.

[GorNishanov]

LGTM

[modocache]

Hooray, [dcl.fct.def.coroutine]/7 is *in!*

Thank you for all the reviews, @GorNishanov!

lib/CodeGen/CGCoroutine.cpp

void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   // Backend is allowed to elide memory allocations, to help it, emit
   // auto mem = coro.alloc() ? 0 : ... allocation code ...;
   auto *CoroAlloc = Builder.CreateCall(
       CGM.getIntrinsic(llvm::Intrinsic::coro_alloc), {CoroId});
 
   Builder.CreateCondBr(CoroAlloc, AllocBB, InitBB);
 
   EmitBlock(AllocBB);
-  auto *AllocateCall = EmitScalarExpr(S.getAllocate());
+  // Emit the call to the coroutine frame allocation function.
+  auto *AllocateCall = cast<llvm::CallInst>(EmitScalarExpr(S.getAllocate()));
+
+  // The backend coroutine split transform will move stores and loads for the
+  // coroutine function's arguments down past the first suspend point.
+  //
+  // However, in the case that the coroutine frame is, as per
+  // [dcl.fct.def.coroutine]/7, being allocated with an allocation function
+  // matching the coroutine function's arguments, we need to ensure that the
+  // allocation function is passed arguments that have values stored in them.
+  //
+  // Here, we generate instructions to store the coroutine function's arguments
+  // separately, and then pass them into the allocation function. First, we
+  // search for the coroutine function argument allocas that correspond to the
+  // arguments passed into the allocation function.
+  for (unsigned OpIdx = 0, OpEnd = AllocateCall->getNumArgOperands();
+       OpIdx != OpEnd; ++OpIdx) {
+    if (auto *AllocateOp =
+            dyn_cast<llvm::Instruction>(AllocateCall->getArgOperand(OpIdx))) {
+      for (auto &AllocateOpOp : AllocateOp->operands()) {
+        if (auto *Alloca = dyn_cast<llvm::AllocaInst>(AllocateOpOp)) {
+          // We've found the alloca instruction. Now we search for the store
+          // instruction that stores the coroutine function argument into that
+          // alloca's address.
+          for (auto I = Alloca->user_begin(), E = Alloca->user_end(); I != E;
+               ++I) {
+            if (auto *Store = dyn_cast<llvm::StoreInst>(*I)) {
+              // Now we generate an alloca, store, and a load, to replace the
+              // allocation function call instruction operands.
+              llvm::BasicBlock::iterator InsertPt = Builder.GetInsertPoint();
+              Builder.SetInsertPoint(AllocateCall);
+
+              llvm::AllocaInst *NewAlloca = Builder.CreateAlloca(
+                  Alloca->getAllocatedType(), Alloca->getArraySize(),
+                  "coro.allocate." + Alloca->getName());
+              NewAlloca->setAlignment(Alloca->getAlignment());
+
+              Address NewAllocaAddress = {
+                  NewAlloca,
+                  CharUnits::fromQuantity(NewAlloca->getAlignment())};
+              Builder.CreateStore(Store->getOperand(0), NewAllocaAddress);
+
+              llvm::LoadInst *NewLoad = Builder.CreateLoad(NewAllocaAddress);
+              AllocateCall->setOperand(OpIdx, NewLoad);
+
+              Builder.SetInsertPoint(AllocBB, InsertPt);
+            }
+          }
+        }
+      }
+    }
+  }
+
   auto *AllocOrInvokeContBB = Builder.GetInsertBlock();
 
   // Handle allocation failure if 'ReturnStmtOnAllocFailure' was provided.
   if (auto *RetOnAllocFailure = S.getReturnStmtOnAllocFailure()) {
     auto *RetOnFailureBB = createBasicBlock("coro.ret.on.failure");
 
     // See if allocation was successful.
     auto *NullPtr = llvm::ConstantPointerNull::get(Int8PtrTy);

lib/Sema/SemaCoroutine.cpp

assert(!IsPromiseDependentType &&
          "cannot make statement while the promise type is dependent");
   QualType PromiseType = Fn.CoroutinePromise->getType();
 
   if (S.RequireCompleteType(Loc, PromiseType, diag::err_incomplete_type))
     return false;
 
   const bool RequiresNoThrowAlloc = ReturnStmtOnAllocFailure != nullptr;
 
-  // FIXME: Add support for stateful allocators.
+  // [dcl.fct.def.coroutine]/7
+  // Lookup allocation functions using a parameter list composed of the
+  // requested size of the coroutine state being allocated, followed by
+  // the coroutine function's arguments. If a matching allocation function
+  // exists, use it. Otherwise, use an allocation function that just takes
+  // the requested size.
 
   FunctionDecl *OperatorNew = nullptr;
   FunctionDecl *OperatorDelete = nullptr;
   FunctionDecl *UnusedResult = nullptr;
   bool PassAlignment = false;
   SmallVector<Expr *, 1> PlacementArgs;
 
+  // [dcl.fct.def.coroutine]/7
+  // "The allocation function’s name is looked up in the scope of P.
+  // [...] If the lookup finds an allocation function in the scope of P,
+  // overload resolution is performed on a function call created by assembling
+  // an argument list."
+  for (auto *PD : FD.parameters()) {
+    if (PD->getType()->isDependentType())
+      continue;
+
+    // Build a reference to the parameter.
+    auto PDLoc = PD->getLocation();
+    ExprResult PDRefExpr =
+        S.BuildDeclRefExpr(PD, PD->getOriginalType().getNonReferenceType(),
+                           ExprValueKind::VK_LValue, PDLoc);
+    if (PDRefExpr.isInvalid())
+      return false;
+
+    PlacementArgs.push_back(PDRefExpr.get());
+  }
   S.FindAllocationFunctions(Loc, SourceRange(),
                             /*UseGlobal*/ false, PromiseType,
                             /*isArray*/ false, PassAlignment, PlacementArgs,
-                            OperatorNew, UnusedResult);
+                            OperatorNew, UnusedResult, /*Diagnose*/ false);
+
+  // [dcl.fct.def.coroutine]/7
+  // "If no matching function is found, overload resolution is performed again
+  // on a function call created by passing just the amount of space required as
+  // an argument of type std::size_t."
+  if (!OperatorNew && !PlacementArgs.empty()) {
+    PlacementArgs.clear();
+    S.FindAllocationFunctions(Loc, SourceRange(),
+                              /*UseGlobal*/ false, PromiseType,
+                              /*isArray*/ false, PassAlignment,
+                              PlacementArgs, OperatorNew, UnusedResult);
+  }
 
   bool IsGlobalOverload =
       OperatorNew && !isa<CXXRecordDecl>(OperatorNew->getDeclContext());
   // If we didn't find a class-local new declaration and non-throwing new
   // was is required then we need to lookup the non-throwing global operator
   // instead.
   if (RequiresNoThrowAlloc && (!OperatorNew || IsGlobalOverload)) {
     auto *StdNoThrow = buildStdNoThrowDeclRef(S, Loc);
     if (!StdNoThrow)
       return false;
     PlacementArgs = {StdNoThrow};
     OperatorNew = nullptr;
     S.FindAllocationFunctions(Loc, SourceRange(),
                               /*UseGlobal*/ true, PromiseType,
                               /*isArray*/ false, PassAlignment, PlacementArgs,
                               OperatorNew, UnusedResult);
   }
 
-  assert(OperatorNew && "expected definition of operator new to be found");
+  if (!OperatorNew)
+    return false;
 
   if (RequiresNoThrowAlloc) {
     const auto *FT = OperatorNew->getType()->getAs<FunctionProtoType>();
     if (!FT->isNothrow(S.Context, /*ResultIfDependent*/ false)) {
       S.Diag(OperatorNew->getLocation(),
              diag::err_coroutine_promise_new_requires_nothrow)
           << OperatorNew;
       S.Diag(Loc, diag::note_coroutine_promise_call_implicitly_required)

test/CodeGenCoroutines/coro-alloc.cpp

extern "C" void f1(promise_new_tag ) {
   // CHECK: call i8* @_ZNSt12experimental16coroutine_traitsIJv15promise_new_tagEE12promise_typenwEm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
   co_return;
 }
 
+struct promise_matching_placement_new_tag {};
+
+template<>
+struct std::experimental::coroutine_traits<void, promise_matching_placement_new_tag, int, float, double> {
+  struct promise_type {
+    void *operator new(unsigned long, promise_matching_placement_new_tag,
+                       int, float, double);
+    void get_return_object() {}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+  };
+};
+
+// CHECK-LABEL: f1a(
+extern "C" void f1a(promise_matching_placement_new_tag, int x, float y , double z) {
+  // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
+  // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
+  // CHECK: store i32 %x, i32* %coro.allocate.x.addr, align 4
+  // CHECK: %[[INT:.+]] = load i32, i32* %coro.allocate.x.addr, align 4
+  // CHECK: store float %y, float* %coro.allocate.y.addr, align 4
+  // CHECK: %[[FLOAT:.+]] = load float, float* %coro.allocate.y.addr, align 4
+  // CHECK: store double %z, double* %coro.allocate.z.addr, align 8
+  // CHECK: %[[DOUBLE:.+]] = load double, double* %coro.allocate.z.addr, align 8
+  // CHECK: call i8* @_ZNSt12experimental16coroutine_traitsIJv34promise_matching_placement_new_tagifdEE12promise_typenwEmS1_ifd(i64 %[[SIZE]], i32 %[[INT]], float %[[FLOAT]], double %[[DOUBLE]])
+  co_return;
+}
+
 struct promise_delete_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_delete_tag> {
   struct promise_type {
     void operator delete(void*);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }