[Coroutines] Don't move stores for allocator args

[modocache]

Summary

The behavior described in Coroutines TS [dcl.fct.def.coroutine]/7
allows coroutine parameters to be passed into allocator functions.
The instructions to store values into the alloca'd parameters must not
be moved past the frame allocation, otherwise uninitialized values are
passed to the allocator.

Test Plan: check-llvm

[compnerd]

lib/Transforms/Coroutines/CoroSplit.cpp

669–674 ↗

(On Diff #133128)

Why not write this as:

for (const auto &User : Alloca->users())
  if (auto *SI = dyn_cast<llvm::StoreInst>(&User))
    DoNotRelocate.insert(SI);

[modocache]

Thanks, @compnerd! Updated to use your suggestion.

[modocache]

Added a test.

[GorNishanov]

lib/Transforms/Coroutines/CoroSplit.cpp

668 ↗

(On Diff #133317)

if (RelocBlocks.count(A->getParent()) != 0) {

I am wondering if this check necessary?

RelocBlocks contains the block with CoroBegin and all preceeding blocks.
We prime the work queue with CoroBegin and terminators of the preceeding blocks
All of the operands of instructions in the Work queue should be in the preceeding blocks.

I can imagine doing this check for SI, so that we will not place in DoNotRelocate StoreInst that may occur after Coro.Begin.

671 ↗

(On Diff #133317)

I do not think just DoNotRelocate.insert(SI) is sufficient. If we are freezing the store, we need to freeze the operands as well. So this part should look something like:

if (DoNotRelocate.count(I) == 0) {
  Work.push_back(I);
  DoNotRelocate.insert(I);
}

Consider this fragment (you probably won't be able to get it straight from clang, but, with some creative optimizer passes you may get something like this:

define i8* @f(i32 %argument) "coroutine.presplit"="1" {
entry:
  %argument.addr = alloca i32, align 4
  %x = add i32 %argument, 5 ; will get moved 
  store i32 %x, i32* %argument.addr, align 4 ; << frozen
  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)

To fix, I believe your loop looking for stores needs to look like:

if (auto *A = dyn_cast<AllocaInst>(I)) {
  // Stores to alloca instructions that occur before coroutine frame
  // is allocated should not be moved; the stored values may be used
  // by the coroutine frame allocator.
  for (const auto &User : A->users())
    if (auto *SI = dyn_cast<llvm::StoreInst>(User))
      if (RelocBlocks.count(SI->getParent()) != 0 &&
          DoNotRelocate.count(SI) == 0) {
        Work.push_back(SI);
        DoNotRelocate.insert(SI);
      }
  continue;
}

[modocache]

Thanks for the review, @GorNishanov! I added something similar to your example of a store operand to the test case. Indeed, without the changes you proposed, the existence of the operand caused an invariant to be violated in LLVM. With your suggested changes, the assertion doesn't fire and the new test passes.

[GorNishanov]

LGTM!

[modocache]

Thank you, @GorNishanov!

llvm/trunk/lib/Transforms/Coroutines/CoroSplit.cpp

getNotRelocatableInstructions(CoroBeginInst *CoroBegin,
   for (BasicBlock *BB : RelocBlocks)
     if (BB != CoroBeginBB)
       Work.push_back(BB->getTerminator());
 
   // For every instruction in the Work list, place its operands in DoNotRelocate
   // set.
   do {
     Instruction *Current = Work.pop_back_val();
+    DEBUG(dbgs() << "CoroSplit: Will not relocate: " << *Current << "\n");
     DoNotRelocate.insert(Current);
     for (Value *U : Current->operands()) {
       auto *I = dyn_cast<Instruction>(U);
       if (!I)
         continue;
-      if (isa<AllocaInst>(U))
+
+      if (auto *A = dyn_cast<AllocaInst>(I)) {
+        // Stores to alloca instructions that occur before the coroutine frame
+        // is allocated should not be moved; the stored values may be used by
+        // the coroutine frame allocator. The operands to those stores must also
+        // remain in place.
+        for (const auto &User : A->users())
+          if (auto *SI = dyn_cast<llvm::StoreInst>(User))
+            if (RelocBlocks.count(SI->getParent()) != 0 &&
+                DoNotRelocate.count(SI) == 0) {
+              Work.push_back(SI);
+              DoNotRelocate.insert(SI);
+            }
         continue;
+      }
+
       if (DoNotRelocate.count(I) == 0) {
         Work.push_back(I);
         DoNotRelocate.insert(I);
       }
     }
   } while (!Work.empty());
   return DoNotRelocate;
 }

llvm/trunk/test/Transforms/Coroutines/coro-split-alloc.ll

+; Tests that coro-split passes initialized values to coroutine frame allocator.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+define i8* @f(i32 %argument) "coroutine.presplit"="1" {
+entry:
+  %argument.addr = alloca i32, align 4
+  %incremented = add i32 %argument, 1
+  store i32 %incremented, i32* %argument.addr, align 4
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %need.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.alloc, label %dyn.alloc, label %begin
+
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %allocator_argument = load i32, i32* %argument.addr, align 4
+  %alloc = call i8* @custom_alloctor(i32 %size, i32 %allocator_argument)
+  br label %begin
+
+begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %phi)
+  %print_argument = load i32, i32* %argument.addr, align 4
+  call void @print(i32 %print_argument)
+  %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 @print(i32 1)
+  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(
+; CHECK: %argument.addr = alloca i32
+; CHECK: %incremented = add i32 %argument, 1
+; CHECK-NEXT: store i32 %incremented, i32* %argument.addr
+; CHECK-LABEL: dyn.alloc:
+; CHECK: %allocator_argument = load i32, i32* %argument.addr
+; CHECK: %alloc = call i8* @custom_alloctor(i32 24, i32 %allocator_argument)
+; CHECK-LABEL: begin:
+; CHECK: %print_argument = load i32, i32* %argument.addr
+; CHECK: call void @print(i32 %print_argument)
+
+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* @custom_alloctor(i32, i32)
+declare void @print(i32)
+declare void @free(i8*)