[LICM][Coroutine] Don't sink stores to coroutine suspend point.

[junparser]

Summary

This patch is the refactor of D86190.
See pr46990(https://bugs.llvm.org/show_bug.cgi?id=46990). Instead of using idea (not sinking store instructions to loop exit blocks which cross coro.suspend intrinsics) in D86190. This patch split exit block of the suspend point alone and ignore it when do promotion.

Test Plan: check-llvm, pr46990

[junparser]

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

[junparser]

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

[efriedma]

If there are transformations that are specifically illegal with corouties, I'd like to see a section in https://llvm.org/docs/Coroutines.html describing what transforms are and are not legal. Just saying that the current version of the coroutine lowering code can't handle it isn't sufficient.

[junparser]

In D87817#2280618, @efriedma wrote:

If there are transformations that are specifically illegal with corouties, I'd like to see a section in https://llvm.org/docs/Coroutines.html describing what transforms are and are not legal. Just saying that the current version of the coroutine lowering code can't handle it isn't sufficient.

Thanks for the suggestion, I'll add it later.
Let's say for this case, we do not wanna move any instructions (maybe this is too restrictive, using !mayReadOrWriteMemory && !mayHaveSideEffects) to path between coroutine suspend point (coro.suspend intrinsic) and return to caller (coro.end intrinsic).

[efriedma]

Okay, I'll wait for the update.

llvm/lib/Transforms/Scalar/LICM.cpp

400

Is there any particular reason to expect the edge from the coro_suspend to the coro_end is a switch?

[junparser]

sorry for the late reply.

llvm/lib/Transforms/Scalar/LICM.cpp

400

The result of this coro.suspend lead to basic blocks where coroutine should proceed when suspended (-1), resumed (0) or destroyed (1). we also generate switch from front end.

Although sometimes it would be optimized to conditional branch because of that both of resumed and destroyed case lead to the same basic block. since we are inside loop, the resumed basic block and destroyed basic block should not be same.

[junparser]

Update the doc. @efriedma is this ok?

[efriedma]

llvm/docs/Coroutines.rst

1561

This isn't what I was looking for.

Say I'm someone who doesn't know anything about coroutines, and I'm writing a new optimization. What do I need to know about coroutines so I don't break them?

llvm/lib/Transforms/Scalar/LICM.cpp

400

the resumed basic block and destroyed basic block should not be same

"should not be the same" here seems to mean "usually won't be the same"; is there any rule that actually precludes it?

[asbirlea]

Can you please extend the testing to cover at least some of these cases: multiple predecesssors, multiple exit blocks, multiple predecessors with a coroutine switch?

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

2

Add new pass manager testing using -passes.

3

Add memoryssa verification with -verify-memoryssa

11

Nit: update test to use the variables generated.
For example here, replace bb0: with [[BB0]], generated one line above.

18

Following up on the question of "is it always a switch", this looks like it could be have optimized to:

 %cond = icmp eq i8  [[SUSPEND]], 0
br i8 [[AWAIT_READY]], label [[BB2_SPLIT_LOOP_EXIT1:%.*]]

For example, simple loop unswitch may do such a change, and licm is in the same LPM, so perhaps this is possible.

[junparser]

llvm/docs/Coroutines.rst

1561

The coroutine framework is designed to not interfere other transformation, Ideally, we do not need know anything about coroutine when write optimization. so this is still temporary fix.

llvm/lib/Transforms/Scalar/LICM.cpp

400

the resumed basic block and destroyed basic block should not be same

"should not be the same" here seems to mean "usually won't be the same"; is there any rule that actually precludes it?

Except we write llvm ir directly with two branches, it is always a switch inside loop.

[junparser]

In D87817#2288579, @asbirlea wrote:

Can you please extend the testing to cover at least some of these cases: multiple predecesssors, multiple exit blocks, multiple predecessors with a coroutine switch?

ok, I'll add later.

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

18

I missed the cleanup label. the switch is

switch i8 %suspend, label %bb2 [
i8 0, label %await.ready
i8 1, label %cleanup
]

[junparser]

In D87817#2278816, @junparser wrote:

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

@lxfind, what do you think about this idea?

[efriedma]

llvm/docs/Coroutines.rst

1561

If you're not sure how it's supposed to work, I'm not sure we want to add temporary hacks before it's settled.

If you want to push forward with this anyway, please add comments to LICM to make it clear that this is a temporary state.

[lxfind]

In D87817#2289365, @junparser wrote:

In D87817#2278816, @junparser wrote:

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

@lxfind, what do you think about this idea?

Are you suggesting to not have an edge from the coro.suspend switch to the coro.end block? Would that make the coro.end BB a dead BB and could potentially get removed?

This is likely the right solution, but we do need to make it more general. In order for this to always work, I wonder if the proper place to insert this transformation is in LoopSimplify, which is a pass always executed before any Loop pass. That might guarantee that any loop pass will be able to handle this properly. I am not too familiar with loop passes, but maybe someone else could comment on that.

[junparser]

In D87817#2291684, @lxfind wrote:

In D87817#2289365, @junparser wrote:

In D87817#2278816, @junparser wrote:

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

@lxfind, what do you think about this idea?

Are you suggesting to not have an edge from the coro.suspend switch to the coro.end block? Would that make the coro.end BB a dead BB and could potentially get removed?

No, There is always path from entry to coro.end BB.

This is likely the right solution, but we do need to make it more general. In order for this to always work, I wonder if the proper place to insert this transformation is in LoopSimplify, which is a pass always executed before any Loop pass. That might guarantee that any loop pass will be able to handle this properly. I am not too familiar with loop passes, but maybe someone else could comment on that.

Loop exit-block insertion in LoopSimplify pass guarantees that all exit blocks from the loop only have predecessors from inside of the loop. It will break the dedicated form of exit blocks when we consider ignore coro.suspend path. Maybe there is other idea i haven't thought about.

[lxfind]

In D87817#2294122, @junparser wrote:

In D87817#2291684, @lxfind wrote:

In D87817#2289365, @junparser wrote:

In D87817#2278816, @junparser wrote:

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

@lxfind, what do you think about this idea?

Are you suggesting to not have an edge from the coro.suspend switch to the coro.end block? Would that make the coro.end BB a dead BB and could potentially get removed?

No, There is always path from entry to coro.end BB.

This is likely the right solution, but we do need to make it more general. In order for this to always work, I wonder if the proper place to insert this transformation is in LoopSimplify, which is a pass always executed before any Loop pass. That might guarantee that any loop pass will be able to handle this properly. I am not too familiar with loop passes, but maybe someone else could comment on that.

Loop exit-block insertion in LoopSimplify pass guarantees that all exit blocks from the loop only have predecessors from inside of the loop. It will break the dedicated form of exit blocks when we consider ignore coro.suspend path. Maybe there is other idea i haven't thought about.

I wasn't referring to ignoring the coro.suspend path, but to move this code in this patch into the LoopSimplify pass, so that all loop passes can benefit, not sure if that makes sense. Anyway, I don't have a better idea.. So don't let that block you here.

[junparser]

In D87817#2294244, @lxfind wrote:

In D87817#2294122, @junparser wrote:

In D87817#2291684, @lxfind wrote:

In D87817#2289365, @junparser wrote:

In D87817#2278816, @junparser wrote:

In D87817#2278717, @junparser wrote:

We try to find an general solution to handle some conservative & mismatch analysis in transformations caused by the change of cfg with coro.suspend intrinsic. However afaik, it seems quiet difficult to deal this without change framework fo coroutine.
This patch just fixes the known issue, it may happens in other passes, FYI.

one of the idea is that we do not generate default case for switch instruction of coro.suspend intrinsic. Instead, we find fall through coro.end for coro.suspend directly in corosplit pass. The only uncertain thing is whether fall through coro.end is moved by passes which may change the semantic.

@lxfind, what do you think about this idea?

Are you suggesting to not have an edge from the coro.suspend switch to the coro.end block? Would that make the coro.end BB a dead BB and could potentially get removed?

No, There is always path from entry to coro.end BB.

This is likely the right solution, but we do need to make it more general. In order for this to always work, I wonder if the proper place to insert this transformation is in LoopSimplify, which is a pass always executed before any Loop pass. That might guarantee that any loop pass will be able to handle this properly. I am not too familiar with loop passes, but maybe someone else could comment on that.

Loop exit-block insertion in LoopSimplify pass guarantees that all exit blocks from the loop only have predecessors from inside of the loop. It will break the dedicated form of exit blocks when we consider ignore coro.suspend path. Maybe there is other idea i haven't thought about.

I wasn't referring to ignoring the coro.suspend path, but to move this code in this patch into the LoopSimplify pass, so that all loop passes can benefit, not sure if that makes sense. Anyway, I don't have a better idea.. So don't let that block you here.

Thanks for the suggestion.

[junparser]

@efriedma Consider that temporary fix maybe not the best solution for such issue. We would try to find other general solution. Before that, I'll keep this patch as plan changes.

llvm/docs/Coroutines.rst

 
 #. Cannot handle coroutines with `inalloca` parameters (used in x86 on Windows).
 
 #. Alignment is ignored by coro.begin and coro.free intrinsics.
 
 #. Make required changes to make sure that coroutine optimizations work with
    LTO.
 
+#. Make required changes to make sure not to insert or move any instructions
+   into edges from suspended(-1) case of coro.suspend intrinsics to their label.
+
 #. More tests, more tests, more tests

llvm/lib/Transforms/Scalar/LICM.cpp

bool LoopInvariantCodeMotion::runOnLoop(
   // preheader for SSA updater, so also avoid sinking when no preheader
   // is available.
   if (!DisablePromotion && Preheader && L->hasDedicatedExits() &&
       !NoOfMemAccTooLarge) {
     // Figure out the loop exits and their insertion points
     SmallVector<BasicBlock *, 8> ExitBlocks;
     L->getUniqueExitBlocks(ExitBlocks);
 
+    bool HasCatchSwitch = false;
+    for (auto *BIter = ExitBlocks.end(); BIter != ExitBlocks.begin();) {
+      BasicBlock *ExitBB = *--BIter;
       // We can't insert into a catchswitch.
-    bool HasCatchSwitch = llvm::any_of(ExitBlocks, [](BasicBlock *Exit) {
-      return isa<CatchSwitchInst>(Exit->getTerminator());
-    });
+      HasCatchSwitch |= isa<CatchSwitchInst>(ExitBB->getTerminator());
+
+      // Find exit edge from coroutine suspend block.
+      SmallSetVector<BasicBlock *, 8> SuspendBBs;
+      SmallSetVector<BasicBlock *, 8> PredBBs(pred_begin(ExitBB),
+                                              pred_end(ExitBB));
+      for (auto *PredBB : PredBBs) {
+
+        if (auto *SWI = dyn_cast<SwitchInst>(PredBB->getTerminator())) {
+          IntrinsicInst *II = dyn_cast<IntrinsicInst>(SWI->getCondition());
+          if (II && (II->getIntrinsicID() == Intrinsic::coro_suspend) &&
+              SWI->getDefaultDest() == ExitBB)
+            SuspendBBs.insert(PredBB);
+        }
+      }
+
+      if (!SuspendBBs.empty()) {
+        // Split all of the predcessor blocks to keep dedicated form.
+        for (auto *PredBB : PredBBs) {
+          BasicBlock *NewPred = SplitBlockPredecessors(
+              ExitBB, PredBB, ".split.loop.exit", DT, LI, MSSAU.get(), true);
+          // Skip exit blocks comes from coroutine suspend block since we
+          // do not want sink anything into them.
+          if (!SuspendBBs.count(PredBB))
+            ExitBlocks.push_back(NewPred);
+        }
+
+        ExitBlocks.erase(BIter);
+      }
+    }
 
     if (!HasCatchSwitch) {
       SmallVector<Instruction *, 8> InsertPts;
       SmallVector<MemoryAccess *, 8> MSSAInsertPts;
       InsertPts.reserve(ExitBlocks.size());
       if (MSSAU)
         MSSAInsertPts.reserve(ExitBlocks.size());
       for (BasicBlock *ExitBlock : ExitBlocks) {

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; 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_SPLIT_LOOP_EXIT1:%.*]] [
+; CHECK-NEXT:    i8 0, label [[AWAIT_READY]]
+; CHECK-NEXT:    ]
+; CHECK:       await.ready:
+; CHECK-NEXT:    [[T6:%.*]] = icmp ult i64 [[T5]], [[N:%.*]]
+; CHECK-NEXT:    br i1 [[T6]], label [[LOOP]], label [[BB2_SPLIT_LOOP_EXIT:%.*]]
+; CHECK:       bb2.split.loop.exit:
+; CHECK-NEXT:    store i64 1, i64* [[P]], align 1
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb2.split.loop.exit1:
+; CHECK-NEXT:    br 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)