[IR][coroutine] make final and non-final llvm.coro.saves unique

[ychen]

Summary

An alternative to D129025.

For

define void @f(i32 %x) {
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
}

SimplifyCFG would hoist llvm.coro.save into the entry block. However,
since one is for the final suspend and the other is not, the transformation
is not desired.

I proposed D129025, which requires a update to the token type definition
in LangRef to make merging two token-returning instructions illegal
(which I didn't do yet and I think it may affect other users of the token
type like WinEH, GC intrinsics). D129025 also excludes some valid
transformation: for this test case, if both branches of the if statement
contains non-final suspend, it is ok to hoisting llvm.coro.save.

So this patch, I propose to directly change the IR equality between two
llvm.coro.save to disable the invalid hoisting. I don't like the
special casing of llvm.coro.save in the llvm/IR space however, the
only other choice would be to change llvm.coro.save in some way like:

1. add new flag parameter to llvm.coro.save
2. add llvm.coro.save.final

I'm not a fan of both.

Any thoughts?

[ychen]

update

[ChuanqiXu]

D129025 also excludes some valid
transformation: for this test case, if both branches of the if statement
contains non-final suspend, it is ok to hoisting llvm.coro.save.

It looks not valid to me to hoist non-final suspend llvm.coro.save. From https://github.com/llvm/llvm-project/blob/011d2bf86487520c3515f16e0b1d32994bf2b48f/llvm/lib/Transforms/Coroutines/CoroSplit.cpp#L388-L405 we could know a llvm.coro.save would be lowered into:

%index.addr = getelementptr %f.frame... (index field number)
store i32 %IndexVal, i32* %index.addr1

So the following code:

if.then:
  %0 = call token @llvm.coro.save(ptr null)
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %merge

if.else:
  %2 = call token @llvm.coro.save(ptr null)
  %3 = call i8 @llvm.coro.suspend(token %2, i1 false)
  br label %merge

would be lowered into:

if.then:
  %0 = getelementptr %f.frame... (index field number)
  store i32 0, i32* %0 ; assume this is the first suspend points
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %merge

if.else:
  %2 = getelementptr %f.frame... (index field number)
  store i32 1, i32* %2
  %3 = call i8 @llvm.coro.suspend(token %2, i1 false)
  br label %merge

so it clearly makes no sense to merge the non-final coro saves. D129025 looks good enough to me temporarily.

however, the only other choice would be to change llvm.coro.save in some way like:

1.add new flag parameter to llvm.coro.save
2.add llvm.coro.save.final

I'm not a fan of both.

So option2 is not valid. But I am a fan of option 1. Since it looks naturally to merge 2 @llvm.coro.save(ptr null) but it wouldn't be naturally to merge @llvm.coro.save(ptr null, i32 0) and @llvm.coro.save(ptr null, i32 1). So we solve the problem in the IR level.

[nikic]

Another possibility is to mark llvm.coro.save as IntrNoMerge, though that would prevent any merging (similar to the SimplifyCFG patch, just only for this intrinsic and not other token-returning intrinsics).

I'm not familiar with these coro intrinsics, but if these have two different ways of being lowered, then adding a flag to distinguish them generally makes sense to me.

[ychen]

so it clearly makes no sense to merge the non-final coro saves. D129025 looks good enough to me temporarily.

Agreed.

however, the only other choice would be to change llvm.coro.save in some way like:

1.add new flag parameter to llvm.coro.save
2.add llvm.coro.save.final

I'm not a fan of both.

So option2 is not valid. But I am a fan of option 1. Since it looks naturally to merge 2 @llvm.coro.save(ptr null) but it wouldn't be naturally to merge @llvm.coro.save(ptr null, i32 0) and @llvm.coro.save(ptr null, i32 1). So we solve the problem in the IR level.

it looks naturally to merge 2 @llvm.coro.save(ptr null)

Hmmm, I don't quite understand this. Do you please give an example?

One more complication is that for

define void @f(i32 %x) {
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 false)
  br label %coro.ret

coro.ret:
  ret void
}

this is not correct (SimplifyCFG only hoists llvm.coro.save):

define void @f(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 0
  %0 = call token @llvm.coro.save(ptr null)
  br i1 %cmp, label %await.suspend, label %final.suspend

await.suspend:
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %coro.ret

final.suspend:
  %3 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %coro.ret

coro.ret:
  ret void
}

but this is correct (SimplifyCFG hoists llvm.coro.suspend too):

define void @f(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 0
  %0 = call token @llvm.coro.save(ptr null)
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  ret void
}

I think what we need is that: only allow merging llvm.coro.save when its user llvm.coro.suspend is also merged. But I don't know a good way to express that with IR.

[ChuanqiXu]

In D129230#3634731, @ychen wrote:

so it clearly makes no sense to merge the non-final coro saves. D129025 looks good enough to me temporarily.

Agreed.

however, the only other choice would be to change llvm.coro.save in some way like:

1.add new flag parameter to llvm.coro.save
2.add llvm.coro.save.final

I'm not a fan of both.

So option2 is not valid. But I am a fan of option 1. Since it looks naturally to merge 2 @llvm.coro.save(ptr null) but it wouldn't be naturally to merge @llvm.coro.save(ptr null, i32 0) and @llvm.coro.save(ptr null, i32 1). So we solve the problem in the IR level.

it looks naturally to merge 2 @llvm.coro.save(ptr null)

Hmmm, I don't quite understand this. Do you please give an example?

I mean it might be a little bit surprising for people (who understand compiler but don't understand the coro intrinsics) that the optimizer wouldn't merge @llvm.coro.save in the example. But if the @llvm.coro.save intrinsic have another parameter, they won't be surprised.

One more complication is that for

define void @f(i32 %x) {
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 false)
  br label %coro.ret

coro.ret:
  ret void
}

this is not correct (SimplifyCFG only hoists llvm.coro.save):

define void @f(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 0
  %0 = call token @llvm.coro.save(ptr null)
  br i1 %cmp, label %await.suspend, label %final.suspend

await.suspend:
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %coro.ret

final.suspend:
  %3 = call i8 @llvm.coro.suspend(token %0, i1 false)
  br label %coro.ret

coro.ret:
  ret void
}

but this is correct (SimplifyCFG hoists llvm.coro.suspend too):

define void @f(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 0
  %0 = call token @llvm.coro.save(ptr null)
  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
  ret void
}

Oh, you're right. It won't be a problem if the @llvm.coro.suspend intrinsic get merged too.

I think what we need is that: only allow merging llvm.coro.save when its user llvm.coro.suspend is also merged. But I don't know a good way to express that with IR.

I think you're right. I don't have a good idea too. Before we get a solution, I think the suggestion from @nikic is good enough in practice.

[ychen]

In D129230#3634730, @nikic wrote:

Another possibility is to mark llvm.coro.save as IntrNoMerge, though that would prevent any merging (similar to the SimplifyCFG patch, just only for this intrinsic and not other token-returning intrinsics).

I'm not familiar with these coro intrinsics, but if these have two different ways of being lowered, then adding a flag to distinguish them generally makes sense to me.

Ah, nomerge function attribute looks very useful. Although it could not cover all the cases perfectly, I wouldn't be too worried since these cases should be relatively rare in practice.

@ChuanqiXu @nikic Thanks for the discussion. I'll create a patch to obsolete D129025.

[ychen]

Obsoleted by D129530

llvm/docs/LangRef.rst

 
 Token Type
 ^^^^^^^^^^
 
 :Overview:
 
 The token type is used when a value is associated with an instruction
 but all uses of the value must not attempt to introspect or obscure it.
-As such, it is not appropriate to have a :ref:`phi <i_phi>` or
-:ref:`select <i_select>` of type token.
+As such, it is illegal to have a :ref:`phi <i_phi>` or
+:ref:`select <i_select>` of type token. It is also illegal to merge two
+token returning instructions that belongs to different basic blocks unless their respective users could also be merged.
 
 :Syntax:
 
 ::
 
       token
 
 

llvm/lib/IR/Instruction.cpp

static bool haveSameSpecialState(const Instruction *I1, const Instruction *I2,
   if (const StoreInst *SI = dyn_cast<StoreInst>(I1))
     return SI->isVolatile() == cast<StoreInst>(I2)->isVolatile() &&
            (SI->getAlign() == cast<StoreInst>(I2)->getAlign() ||
             IgnoreAlignment) &&
            SI->getOrdering() == cast<StoreInst>(I2)->getOrdering() &&
            SI->getSyncScopeID() == cast<StoreInst>(I2)->getSyncScopeID();
   if (const CmpInst *CI = dyn_cast<CmpInst>(I1))
     return CI->getPredicate() == cast<CmpInst>(I2)->getPredicate();
-  if (const CallInst *CI = dyn_cast<CallInst>(I1))
+  if (const CallInst *CI = dyn_cast<CallInst>(I1)) {
+    // Must have the same "Final Suspend" flag.
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I1)) {
+      if (II->getIntrinsicID() == Intrinsic::coro_save && I1->getNumUses() &&
+          I2->getNumUses()) {
+        auto *CS1 = cast<IntrinsicInst>(I1->user_back());
+        auto *CS2 = cast<IntrinsicInst>(I2->user_back());
+        assert(CS1->getIntrinsicID() == Intrinsic::coro_suspend &&
+               CS2->getIntrinsicID() == Intrinsic::coro_suspend);
+        return cast<Constant>(CS1->getArgOperand(1))->isOneValue() ==
+               cast<Constant>(CS2->getArgOperand(1))->isOneValue();
+      }
+    }
     return CI->isTailCall() == cast<CallInst>(I2)->isTailCall() &&
            CI->getCallingConv() == cast<CallInst>(I2)->getCallingConv() &&
            CI->getAttributes() == cast<CallInst>(I2)->getAttributes() &&
            CI->hasIdenticalOperandBundleSchema(*cast<CallInst>(I2));
+  }
   if (const InvokeInst *CI = dyn_cast<InvokeInst>(I1))
     return CI->getCallingConv() == cast<InvokeInst>(I2)->getCallingConv() &&
            CI->getAttributes() == cast<InvokeInst>(I2)->getAttributes() &&
            CI->hasIdenticalOperandBundleSchema(*cast<InvokeInst>(I2));
   if (const CallBrInst *CI = dyn_cast<CallBrInst>(I1))
     return CI->getCallingConv() == cast<CallBrInst>(I2)->getCallingConv() &&
            CI->getAttributes() == cast<CallBrInst>(I2)->getAttributes() &&
            CI->hasIdenticalOperandBundleSchema(*cast<CallBrInst>(I2));

llvm/lib/Transforms/Utils/FunctionComparator.cpp

 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
if (int Res =
             cmpOrderings(SI->getOrdering(), cast<StoreInst>(R)->getOrdering()))
       return Res;
     return cmpNumbers(SI->getSyncScopeID(),
                       cast<StoreInst>(R)->getSyncScopeID());
   }
   if (const CmpInst *CI = dyn_cast<CmpInst>(L))
     return cmpNumbers(CI->getPredicate(), cast<CmpInst>(R)->getPredicate());
   if (auto *CBL = dyn_cast<CallBase>(L)) {
+    // Must have the same "Final Suspend" flag.
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(L)) {
+      if (II->getIntrinsicID() == Intrinsic::coro_save && L->getNumUses() &&
+          R->getNumUses()) {
+        auto *CS1 = cast<IntrinsicInst>(L->user_back());
+        auto *CS2 = cast<IntrinsicInst>(R->user_back());
+        assert(CS1->getIntrinsicID() == Intrinsic::coro_suspend &&
+               CS2->getIntrinsicID() == Intrinsic::coro_suspend);
+        bool IsFinal1 = cast<Constant>(CS1->getArgOperand(1))->isOneValue();
+        bool IsFinal2 = cast<Constant>(CS2->getArgOperand(1))->isOneValue();
+        if (int Res = cmpNumbers(IsFinal1, IsFinal2))
+          return Res;
+      }
+    }
+
     auto *CBR = cast<CallBase>(R);
     if (int Res = cmpNumbers(CBL->getCallingConv(), CBR->getCallingConv()))
       return Res;
     if (int Res = cmpAttrs(CBL->getAttributes(), CBR->getAttributes()))
       return Res;
     if (int Res = cmpOperandBundlesSchema(*CBL, *CBR))
       return Res;
     if (const CallInst *CI = dyn_cast<CallInst>(L))

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

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)