[SimplifyCFG] Skip hoisting common instructions that return token type

[ychen]

Summary

By LangRef, hoisting token-returning instructions obsures the origin
so it should be skipped. Found this issue while investigating a
CoroSplit pass crash.

[ychen]

run update_test_checks.py on llvm/test/Transforms/LowerMatrixIntrinsics/transpose-opts.ll

[nikic]

Don't think we have a clear understanding of the transforms that are valid for token types, but this generally looks reasonable to me. We should avoid making token values flow across control flow edges, even if it does not require introducing a phi.

llvm/test/Transforms/LowerMatrixIntrinsics/transpose-opts.ll

975 ↗

(On Diff #441817)

Unrelated change? Feel free to commit separately.

llvm/test/Transforms/SimplifyCFG/hoist-skip-token.ll

5

Use update_test_checks please.

37

It should be possible to simplify this test further, e.g. the outer loop structure doesn't really seem relevant here.

[ychen]

• rebase
• simplify the test case
• use update_test_checks.py

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1477

Ugh, I could have sworn I left a comment here, but apparently I didn't: This check shouldn't be here, it should be in the do/while loop below. Otherwise you'll only check that the first instruction doesn't return a token type, but may still hoist it if it's a later instruction.

[ychen]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1477

Ah, thanks for the remainder. After giving it another thought, I'm curious if the check should be in isIdenticalToWhenDefined to say that any two token-returning instructions are distinct? WDYT?

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1477

I think isIdenticalToWhenDefined() is a too generic API for this check. I don't think there's anything generally wrong with merging token-returning calls -- the problem here is that we want to either hoist both the token-producing and token-consuming calls, or neither, and the simplest way to do that is to never hoist them.

[ychen]

Address Nikita's feeback

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

1477

Agreed. Thanks.

[nikic]

LGTM

[efriedma]

I think we need to clarify LangRef if this is illegal in general for tokens in general. If it is, we need to explain why. If it isn't, we need an attribute the express the right semantics here.

All LangRef really says at the moment is that a token can't be put behind a phi or select; it doesn't really describe if/when it's legal to merge instructions that produce tokens.

[ychen]

In D129025#3631057, @efriedma wrote:

I think we need to clarify LangRef if this is illegal in general for tokens in general. If it is, we need to explain why. If it isn't, we need an attribute the express the right semantics here.

All LangRef really says at the moment is that a token can't be put behind a phi or select; it doesn't really describe if/when it's legal to merge instructions that produce tokens.

Yes. I also found the token type description could be more clear. I'll send a follow-up LangRef patch to clarify the intent.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

bool SimplifyCFGOpt::HoistThenElseCodeToIf(BranchInst *BI,
   DbgInfoIntrinsic *DBI2 = dyn_cast<DbgInfoIntrinsic>(I2);
   if (!DBI1 || !DBI2 || !DBI1->isIdenticalToWhenDefined(DBI2)) {
     while (isa<DbgInfoIntrinsic>(I1))
       I1 = &*BB1_Itr++;
     while (isa<DbgInfoIntrinsic>(I2))
       I2 = &*BB2_Itr++;
   }
   if (isa<PHINode>(I1) || !I1->isIdenticalToWhenDefined(I2))
     return false;
 
   BasicBlock *BIParent = BI->getParent();
 
   bool Changed = false;
 
   auto _ = make_scope_exit([&]() {
     if (Changed)
       ++NumHoistCommonCode;
bool SimplifyCFGOpt::HoistThenElseCodeToIf(BranchInst *BI,
   }
 
   do {
     // If we are hoisting the terminator instruction, don't move one (making a
     // broken BB), instead clone it, and remove BI.
     if (I1->isTerminator())
       goto HoistTerminator;
 
+    // Hoisting token-returning instructions would obscure the origin.
+    if (I1->getType()->isTokenTy())
+      return Changed;
+
     // If we're going to hoist a call, make sure that the two instructions we're
     // commoning/hoisting are both marked with musttail, or neither of them is
     // marked as such. Otherwise, we might end up in a situation where we hoist
     // from a block where the terminator is a `ret` to a block where the terminator
     // is a `br`, and `musttail` calls expect to be followed by a return.
     auto *C1 = dyn_cast<CallInst>(I1);
     auto *C2 = dyn_cast<CallInst>(I2);
     if (C1 && C2)

llvm/test/Transforms/SimplifyCFG/hoist-skip-token.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes='simplifycfg<hoist-common-insts>' -S | FileCheck %s
+
+declare token @llvm.coro.save(ptr)
+declare i8 @llvm.coro.suspend(token, i1)
+
+define void @f(i32 %x) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[X:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[AWAIT_SUSPEND:%.*]], label [[FINAL_SUSPEND:%.*]]
+; CHECK:       await.suspend:
+; CHECK-NEXT:    [[TMP0:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.coro.suspend(token [[TMP0]], i1 false)
+; CHECK-NEXT:    br label [[CORO_RET:%.*]]
+; CHECK:       final.suspend:
+; CHECK-NEXT:    [[TMP2:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP3:%.*]] = call i8 @llvm.coro.suspend(token [[TMP2]], i1 true)
+; CHECK-NEXT:    br label [[CORO_RET]]
+; CHECK:       coro.ret:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp slt i32 %x, 0
+  br i1 %cmp, label %await.suspend, label %final.suspend
+
+await.suspend:
+  %0 = call token @llvm.coro.save(ptr null)
+  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
+  br label %coro.ret
+
+final.suspend:
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token %2, i1 true)
+  br label %coro.ret
+
+coro.ret:
+  ret void
+}