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

[junparser]

Summary

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.
Also Promotion optimization will run in CGSCC pipeline after coroutine function been splitted.

Test Plan: check-llvm, pr46990

[efriedma]

Are you sure this is the right fix? At first glance, ensuring %p lives long enough seems like the sort of thing coroutine lowering should be able to deal with. If it isn't, at the very least we need documentation describing what optimizations are legal.

[lxfind]

llvm/test/Transforms/LICM/sink-with-coroutine.ll

37

It seems to me the problem is that the LICM generated code is mixing the loop exit basic block vs the return basic block. It just happens that both the suspend switch and loop exit condition jumps to the same block %bb2, but they mean different things.
I think we should still be able to do LICM when there is coroutine, but it's just that we need to make sure the instructions are moved to the loop exit block, while the suspend switch should remain jumping to the return block?

[junparser]

In D86190#2227340, @efriedma wrote:

Are you sure this is the right fix? At first glance, ensuring %p lives long enough seems like the sort of thing coroutine lowering should be able to deal with. If it isn't, at the very least we need documentation describing what optimizations are legal.

I just disable promotion before coroutine lowering and count on promotion after lowering which might not be the right fix. @lxfind's comment seems to be the right solution. Need some investigation

[junparser]

llvm/test/Transforms/LICM/sink-with-coroutine.ll

37

Thanks for the suggestion. This might be the right fix, I'll check it later.

[lxfind]

@junparser I have been re-thinking about this patch, and now I actually think this is the right patch to go.
Other than the correctness reasons, I realize that it is almost always harmful to do LICM for a loop that contains a coroutine suspension. The reason is that in most cases LICM moves memory operations out of the loop. (LICM does move constant calls as well but it is relatively rare to have constant calls that's not inlined, furthermore there are other complications such as TLS access and etc.).
For memory operations, moving them out of the loop does not make the loop faster if there is a coroutine suspend in the loop, because the value moved out needs to be put on the coroutine frame anyway. So LICM not only does not eliminate memory access but increases the frame size.
I propose to bring this patch back. (I can help send it too)
What do you think?

llvm/lib/Transforms/Scalar/LICM.cpp

bool LoopInvariantCodeMotion::runOnLoop(
   // just exit.
   if (hasDisableLICMTransformsHint(L)) {
     return false;
   }
 
   std::unique_ptr<AliasSetTracker> CurAST;
   std::unique_ptr<MemorySSAUpdater> MSSAU;
   bool NoOfMemAccTooLarge = false;
+  bool HasCoroSuspendInst = false;
   unsigned LicmMssaOptCounter = 0;
 
   if (!MSSA) {
     LLVM_DEBUG(dbgs() << "LICM: Using Alias Set Tracker.\n");
     CurAST = collectAliasInfoForLoop(L, LI, AA);
   } else {
     LLVM_DEBUG(dbgs() << "LICM: Using MemorySSA.\n");
     MSSAU = std::make_unique<MemorySSAUpdater>(MSSA);
 
     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;
           AccessCapCount++;
           if (AccessCapCount > LicmMssaNoAccForPromotionCap) {
             NoOfMemAccTooLarge = true;
             break;
           }
bool LoopInvariantCodeMotion::runOnLoop(
   // Now that all loop invariants have been removed from the loop, promote any
   // memory references to scalars that we can.
   // Don't sink stores from loops without dedicated block exits. Exits
   // containing indirect branches are not transformed by loop simplify,
   // make sure we catch that. An additional load may be generated in the
   // 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);
 
     // We can't insert into a catchswitch.
     bool HasCatchSwitch = llvm::any_of(ExitBlocks, [](BasicBlock *Exit) {
       return isa<CatchSwitchInst>(Exit->getTerminator());
     });

llvm/test/Transforms/LICM/sink-with-coroutine.ll

+; 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)