[llvm] 8a7f5ad - We can only move static allocas into the resume entry points

[Arnold Schwaighofer via llvm-commits]

Author: Arnold Schwaighofer
Date: 2021-02-05T06:06:10-08:00
New Revision: 8a7f5ad0fdbcef6d945e9aa598c60354caa4f566

URL: https://github.com/llvm/llvm-project/commit/8a7f5ad0fdbcef6d945e9aa598c60354caa4f566
DIFF: https://github.com/llvm/llvm-project/commit/8a7f5ad0fdbcef6d945e9aa598c60354caa4f566.diff

LOG: We can only move static allocas into the resume entry points

Dynamic allocas that still exist have been verified to be only used
'locally' not accross a suspend point.

rdar://73903220

Differential Revision: https://reviews.llvm.org/D96071

Added: 
    

Modified: 
    llvm/lib/Transforms/Coroutines/CoroSplit.cpp
    llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index e5876d612b6c..c64dbd7f3ea6 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -715,15 +715,17 @@ void CoroCloner::replaceEntryBlock() {
   }
   }
 
-  // Any alloca that's still being used but not reachable from the new entry
-  // needs to be moved to the new entry.
+  // Any static alloca that's still being used but not reachable from the new
+  // entry needs to be moved to the new entry.
   Function *F = OldEntry->getParent();
   DominatorTree DT{*F};
   for (auto IT = inst_begin(F), End = inst_end(F); IT != End;) {
     Instruction &I = *IT++;
-    if (!isa<AllocaInst>(&I) || I.use_empty())
+    auto *Alloca = dyn_cast<AllocaInst>(&I);
+    if (!Alloca || I.use_empty())
       continue;
-    if (DT.isReachableFromEntry(I.getParent()))
+    if (DT.isReachableFromEntry(I.getParent()) ||
+        !isa<ConstantInt>(Alloca->getArraySize()))
       continue;
     I.moveBefore(*Entry, Entry->getFirstInsertionPt());
   }

diff  --git a/llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll b/llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll
index 19cc873fb329..53a596518082 100644
--- a/llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll
+++ b/llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll
@@ -230,6 +230,40 @@ end:
   unreachable
 }
 
+declare i32 @getSize()
+define {i8*, i32} @k(i8* %buffer, i32 %n, i1 %cond) {
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, i8* %buffer, i8* bitcast ({i8*, i32} (i8*, i1)* @prototype_g to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  br i1 %cond, label %alloca_block, label %non_alloca_block
+
+suspend:
+  %unwind = call i1 (...) @llvm.coro.suspend.retcon.i1(i32 %n)
+  br i1 %unwind, label %cleanup, label %resume
+
+resume:
+  br label %cleanup
+
+alloca_block:
+  %size = call i32 @getSize()
+  ; This will get lowered to a dynamic alloca.
+  ; Make sure code that runs after that lowering does not hoist the dynamic
+  ; alloca into the entry block of the resume function.
+  %alloca = call token @llvm.coro.alloca.alloc.i32(i32 %size, i32 8)
+  %ptr = call i8* @llvm.coro.alloca.get(token %alloca)
+  call void @use(i8* %ptr)
+  call void @llvm.coro.alloca.free(token %alloca)
+  br label %suspend
+
+non_alloca_block:
+  %ignore = call i32 @getSize()
+  br label %suspend
+
+cleanup:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
+
 declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.suspend.retcon.i1(...)