[PATCH] D86190: [LICM][Coroutine] Don't sink stores from loops with coro.suspend instructions.

[JunMa via Phabricator via llvm-commits]

junparser created this revision.
junparser added reviewers: wenlei, reames, fhahn, efriedma.
junparser added a project: LLVM.
Herald added subscribers: llvm-commits, asbirlea, modocache, hiraditya.
junparser requested review of this revision.

See pr46990(https://bugs.llvm.org/show_bug.cgi?id=46990). LICM should not sink store instructions to loop exit blocks which cross coro.suspend intrinsics. This breaks semantic of coro.suspend intrinsic which return to caller directly. Also this leads to use-after-free if the coroutine is freed before control returns to the caller in multithread environment.

This patch disable promotion by check whether loop contains coro.suspend intrinsics. And promotion optimization will run in under CGSCC pipeline after coroutine function been splitted.

Test Plan: check-llvm, pr46990


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D86190

Files:
  llvm/lib/Transforms/Scalar/LICM.cpp
  llvm/test/Transforms/LICM/sink-with-coroutine.ll


Index: llvm/test/Transforms/LICM/sink-with-coroutine.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/LICM/sink-with-coroutine.ll
@@ -0,0 +1,51 @@
+; RUN: opt -S < %s -licm | FileCheck %s
+
+; LICM across a @coro.suspend.
+
+define i64 @licm(i64 %n) #0 {
+; CHECK-LABEL: @licm
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    br label [[BB0:%.*]]
+; CHECK:       bb0:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, [[BB0]] ], [ [[T5:%.*]], [[AWAIT_READY:%.*]] ]
+; CHECK-NEXT:    [[T5]] = add i64 [[I]], 1
+; CHECK-NEXT:    [[SUSPEND:%.*]] = call i8 @llvm.coro.suspend(token none, i1 false)
+; CHECK-NEXT:    switch i8 [[SUSPEND]], label [[BB2:%.*]] [
+; CHECK-NEXT:    i8 0, label [[AWAIT_READY]]
+; CHECK-NEXT:    ]
+; CHECK:       await.ready:
+; CHECK-NEXT:    store i64 1, i64* [[P]], align 4
+; CHECK-NEXT:    [[T6:%.*]] = icmp ult i64 [[T5]], [[N:%.*]]
+; CHECK-NEXT:    br i1 [[T6]], label [[LOOP]], label [[BB2]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %p = alloca i64
+  br label %bb0
+
+bb0:
+  br label %loop
+
+loop:
+  %i = phi i64 [ 0, %bb0 ], [ %t5, %await.ready ]
+  %t5 = add i64 %i, 1
+  %suspend = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %suspend, label %bb2 [
+  i8 0, label %await.ready
+  ]
+
+
+await.ready:
+  store i64 1, i64* %p
+  %t6 = icmp ult i64 %t5, %n
+  br i1 %t6, label %loop, label %bb2
+
+bb2:
+  ret i64 0
+}
+
+declare i8  @llvm.coro.suspend(token, i1)
Index: llvm/lib/Transforms/Scalar/LICM.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/LICM.cpp
+++ llvm/lib/Transforms/Scalar/LICM.cpp
@@ -301,6 +301,7 @@
   std::unique_ptr<AliasSetTracker> CurAST;
   std::unique_ptr<MemorySSAUpdater> MSSAU;
   bool NoOfMemAccTooLarge = false;
+  bool HasCoroSuspendInst = false;
   unsigned LicmMssaOptCounter = 0;
 
   if (!MSSA) {
@@ -312,6 +313,13 @@
 
     unsigned AccessCapCount = 0;
     for (auto *BB : L->getBlocks()) {
+      // Don't sink stores from loops with coroutine suspend instructions.
+      HasCoroSuspendInst |= llvm::any_of(*BB, [](Instruction &I) {
+        IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
+        return II && (II->getIntrinsicID() == Intrinsic::coro_suspend ||
+                      II->getIntrinsicID() == Intrinsic::coro_suspend_retcon);
+      });
+
       if (auto *Accesses = MSSA->getBlockAccesses(BB)) {
         for (const auto &MA : *Accesses) {
           (void)MA;
@@ -363,7 +371,7 @@
   // preheader for SSA updater, so also avoid sinking when no preheader
   // is available.
   if (!DisablePromotion && Preheader && L->hasDedicatedExits() &&
-      !NoOfMemAccTooLarge) {
+      !NoOfMemAccTooLarge && !HasCoroSuspendInst) {
     // Figure out the loop exits and their insertion points
     SmallVector<BasicBlock *, 8> ExitBlocks;
     L->getUniqueExitBlocks(ExitBlocks);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D86190.286457.patch
Type: text/x-patch
Size: 3023 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200819/0ab8e3f5/attachment.bin>