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

[Xun Li via Phabricator via llvm-commits]

lxfind created this revision.
lxfind added reviewers: junparser, efriedma, ChuanqiXu.
Herald added subscribers: hoy, modimo, wenlei, asbirlea, hiraditya.
lxfind requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

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.
This is a resubmit of D86190 <https://reviews.llvm.org/D86190>.
Disabling LICM for loops with coroutine suspension is a better option not only for correctness purpose but also for performance purpose.
In most cases LICM sinks memory operations. In the case of coroutine, sinking memory operation out of the loop does not improve performance since coroutien needs to get data from the frame anyway. In fact LICM would hurt coroutine performance since it adds more entries to the frame.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D96928

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,53 @@
+; RUN: opt -S < %s -passes=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:    [[RES:%.*]] = call i1 @llvm.coro.end(i8* null, i1 false)
+; 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:
+  %res = call i1 @llvm.coro.end(i8* null, i1 false)
+  ret i64 0
+}
+
+declare i8  @llvm.coro.suspend(token, i1)
+declare i1  @llvm.coro.end(i8*, i1)
Index: llvm/lib/Transforms/Scalar/LICM.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/LICM.cpp
+++ llvm/lib/Transforms/Scalar/LICM.cpp
@@ -359,6 +359,23 @@
   std::unique_ptr<MemorySSAUpdater> MSSAU;
   std::unique_ptr<SinkAndHoistLICMFlags> Flags;
 
+  // Don't sink stores from loops with coroutine suspend instructions.
+  // There are two reasons for this:
+  // 1. LICM would sink instructions into the default destination of
+  // the coroutine switch. The default destination of the switch is to special
+  // handle the case where the coroutine is destroyed. No instruction can be
+  // sunk there.
+  // 2. In most cases LICM sinks memory operations. In the case of coroutine,
+  // sinking memory operation out of the loop does not improve performance
+  // since coroutien needs to get data from the frame anyway. In fact LICM
+  // would hurt coroutine performance since it adds more entries to the frame.
+  bool HasCoroSuspendInst = llvm::any_of(L->getBlocks(), [](BasicBlock *BB) {
+    return llvm::any_of(*BB, [](Instruction &I) {
+      IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
+      return II && II->getIntrinsicID() == Intrinsic::coro_suspend;
+    });
+  });
+
   if (!MSSA) {
     LLVM_DEBUG(dbgs() << "LICM: Using Alias Set Tracker.\n");
     CurAST = collectAliasInfoForLoop(L, LI, AA);
@@ -405,7 +422,7 @@
   // preheader for SSA updater, so also avoid sinking when no preheader
   // is available.
   if (!DisablePromotion && Preheader && L->hasDedicatedExits() &&
-      !Flags->tooManyMemoryAccesses()) {
+      !Flags->tooManyMemoryAccesses() && !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: D96928.324507.patch
Type: text/x-patch
Size: 3581 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210218/81cfd4ff/attachment.bin>