[SimplifyCFG] Don't merge invoke if this makes immarg non-constant (PR61265)

Don't merge invokes if this replaces constant operands with phis
in a place where this is not legal.

This also disallows converting operand bundles from constant to
non-constant, in line with the restriction we use in other
transforms.

Fixes https://github.com/llvm/llvm-project/issues/61265.

Differential Revision: https://reviews.llvm.org/D146723

(cherry picked from commit 11313108ff778ab6037dcf3d6057de6e1d26fb80)

3 files changed, 67 insertions, 19 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 9e0483966d3e..f662e126e378 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -2455,9 +2455,13 @@ bool CompatibleSets::shouldBelongToSameSet(ArrayRef<InvokeInst *> Invokes) {
 
   // Can we theoretically form the data operands for the merged `invoke`?
   auto IsIllegalToMergeArguments = [](auto Ops) {
-    Type *Ty = std::get<0>(Ops)->getType();
-    assert(Ty == std::get<1>(Ops)->getType() && "Incompatible types?");
-    return Ty->isTokenTy() && std::get<0>(Ops) != std::get<1>(Ops);
+    Use &U0 = std::get<0>(Ops);
+    Use &U1 = std::get<1>(Ops);
+    if (U0 == U1)
+      return false;
+    return U0->getType()->isTokenTy() ||
+           !canReplaceOperandWithVariable(cast<Instruction>(U0.getUser()),
+                                          U0.getOperandNo());
   };
   assert(Invokes.size() == 2 && "Always called with exactly two candidates.");
   if (any_of(zip(Invokes[0]->data_ops(), Invokes[1]->data_ops()),
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/merge-compatible-invokes-of-landingpad.ll b/llvm/test/Transforms/SimplifyCFG/X86/merge-compatible-invokes-of-landingpad.ll
index c29451ae2ca0..276aa83b0326 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/merge-compatible-invokes-of-landingpad.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/merge-compatible-invokes-of-landingpad.ll
@@ -1232,7 +1232,7 @@ if.end:
 }
 
 ; We need to PHI together the arguments of the operand bundles.
-define void @t21_semicompatible_operand_bundle() personality ptr @__gxx_personality_v0 {
+define void @t21_semicompatible_operand_bundle(i32 %a, i32 %b) personality ptr @__gxx_personality_v0 {
 ; CHECK-LABEL: @t21_semicompatible_operand_bundle(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C0:%.*]] = call i1 @cond()
@@ -1246,7 +1246,7 @@ define void @t21_semicompatible_operand_bundle() personality ptr @__gxx_personal
 ; CHECK-NEXT:    [[C1:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1_INVOKE]], label [[IF_END:%.*]]
 ; CHECK:       if.then1.invoke:
-; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ 0, [[IF_ELSE]] ], [ 42, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ [[B:%.*]], [[IF_ELSE]] ], [ [[A:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    invoke void @simple_throw() [ "abc"(i32 [[TMP0]]) ]
 ; CHECK-NEXT:    to label [[IF_THEN1_CONT:%.*]] unwind label [[LPAD:%.*]]
 ; CHECK:       if.then1.cont:
@@ -1260,7 +1260,7 @@ entry:
   br i1 %c0, label %if.then0, label %if.else
 
 if.then0:
-  invoke void @simple_throw() [ "abc"(i32 42) ] to label %invoke.cont0 unwind label %lpad
+  invoke void @simple_throw() [ "abc"(i32 %a) ] to label %invoke.cont0 unwind label %lpad
 
 invoke.cont0:
   unreachable
@@ -1275,7 +1275,7 @@ if.else:
   br i1 %c1, label %if.then1, label %if.end
 
 if.then1:
-  invoke void @simple_throw() [ "abc"(i32 0) ] to label %invoke.cont2 unwind label %lpad
+  invoke void @simple_throw() [ "abc"(i32 %b) ] to label %invoke.cont2 unwind label %lpad
 
 invoke.cont2:
   unreachable
@@ -2342,7 +2342,7 @@ if.end:
 }
 
 ; For indirect invokes, different operand bundle arguments are fine.
-define void @t39_different_arguments_and_operand_bundes_are_fine(ptr %callee0, ptr %callee1) personality ptr @__gxx_personality_v0 {
+define void @t39_different_arguments_and_operand_bundes_are_fine(ptr %callee0, ptr %callee1, i32 %a, i32 %b) personality ptr @__gxx_personality_v0 {
 ; CHECK-LABEL: @t39_different_arguments_and_operand_bundes_are_fine(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C0:%.*]] = call i1 @cond()
@@ -2356,7 +2356,7 @@ define void @t39_different_arguments_and_operand_bundes_are_fine(ptr %callee0, p
 ; CHECK-NEXT:    [[C1:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1_INVOKE]], label [[IF_END:%.*]]
 ; CHECK:       if.then1.invoke:
-; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ 0, [[IF_ELSE]] ], [ 42, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ [[B:%.*]], [[IF_ELSE]] ], [ [[A:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[TMP1:%.*]] = phi ptr [ [[CALLEE1:%.*]], [[IF_ELSE]] ], [ [[CALLEE0:%.*]], [[ENTRY]] ]
 ; CHECK-NEXT:    invoke void [[TMP1]]() [ "abc"(i32 [[TMP0]]) ]
 ; CHECK-NEXT:    to label [[IF_THEN1_CONT:%.*]] unwind label [[LPAD:%.*]]
@@ -2371,7 +2371,7 @@ entry:
   br i1 %c0, label %if.then0, label %if.else
 
 if.then0:
-  invoke void %callee0() [ "abc"(i32 42) ] to label %invoke.cont0 unwind label %lpad
+  invoke void %callee0() [ "abc"(i32 %a) ] to label %invoke.cont0 unwind label %lpad
 
 invoke.cont0:
   unreachable
@@ -2386,7 +2386,7 @@ if.else:
   br i1 %c1, label %if.then1, label %if.end
 
 if.then1:
-  invoke void %callee1() [ "abc"(i32 0) ] to label %invoke.cont2 unwind label %lpad
+  invoke void %callee1() [ "abc"(i32 %b) ] to label %invoke.cont2 unwind label %lpad
 
 invoke.cont2:
   unreachable
@@ -2397,7 +2397,7 @@ if.end:
 }
 
 ; For indirect invokes, both different arguments and operand bundle arguments are fine.
-define void @t40_different_arguments_and_operand_bundes_are_fine(ptr %callee0, ptr %callee1) personality ptr @__gxx_personality_v0 {
+define void @t40_different_arguments_and_operand_bundes_are_fine(ptr %callee0, ptr %callee1, i32 %a, i32 %b) personality ptr @__gxx_personality_v0 {
 ; CHECK-LABEL: @t40_different_arguments_and_operand_bundes_are_fine(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C0:%.*]] = call i1 @cond()
@@ -2412,7 +2412,7 @@ define void @t40_different_arguments_and_operand_bundes_are_fine(ptr %callee0, p
 ; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1_INVOKE]], label [[IF_END:%.*]]
 ; CHECK:       if.then1.invoke:
 ; CHECK-NEXT:    [[TMP0:%.*]] = phi i32 [ 42, [[IF_ELSE]] ], [ 0, [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[TMP1:%.*]] = phi i32 [ 0, [[IF_ELSE]] ], [ 42, [[ENTRY]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i32 [ [[B:%.*]], [[IF_ELSE]] ], [ [[A:%.*]], [[ENTRY]] ]
 ; CHECK-NEXT:    [[TMP2:%.*]] = phi ptr [ [[CALLEE1:%.*]], [[IF_ELSE]] ], [ [[CALLEE0:%.*]], [[ENTRY]] ]
 ; CHECK-NEXT:    invoke void [[TMP2]](i32 [[TMP0]]) [ "abc"(i32 [[TMP1]]) ]
 ; CHECK-NEXT:    to label [[IF_THEN1_CONT:%.*]] unwind label [[LPAD:%.*]]
@@ -2427,7 +2427,7 @@ entry:
   br i1 %c0, label %if.then0, label %if.else
 
 if.then0:
-  invoke void %callee0(i32 0) [ "abc"(i32 42) ] to label %invoke.cont0 unwind label %lpad
+  invoke void %callee0(i32 0) [ "abc"(i32 %a) ] to label %invoke.cont0 unwind label %lpad
 
 invoke.cont0:
   unreachable
@@ -2442,7 +2442,7 @@ if.else:
   br i1 %c1, label %if.then1, label %if.end
 
 if.then1:
-  invoke void %callee1(i32 42) [ "abc"(i32 0) ] to label %invoke.cont2 unwind label %lpad
+  invoke void %callee1(i32 42) [ "abc"(i32 %b) ] to label %invoke.cont2 unwind label %lpad
 
 invoke.cont2:
   unreachable
@@ -2452,6 +2452,43 @@ if.end:
   ret void
 }
 
+define void @dont_merge_different_immargs(i1 %c1) gc "statepoint-example" personality ptr null {
+; CHECK-LABEL: @dont_merge_different_immargs(
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    [[T1:%.*]] = invoke token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 1, i32 0, ptr elementtype(void (ptr addrspace(1))) null, i32 1, i32 0, ptr addrspace(1) null, i64 0, i64 0)
+; CHECK-NEXT:    to label [[UNREACHABLE:%.*]] unwind label [[LPAD:%.*]]
+; CHECK:       else:
+; CHECK-NEXT:    [[T2:%.*]] = invoke token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 0, i32 0, ptr elementtype(void (ptr addrspace(1))) null, i32 1, i32 0, ptr addrspace(1) null, i64 0, i64 0)
+; CHECK-NEXT:    to label [[UNREACHABLE]] unwind label [[LPAD]]
+; CHECK:       unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       lpad:
+; CHECK-NEXT:    [[T3:%.*]] = landingpad token
+; CHECK-NEXT:    cleanup
+; CHECK-NEXT:    ret void
+;
+  br i1 %c1, label %if, label %else
+
+if:
+  %t1 = invoke token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 1, i32 0, ptr elementtype(void (ptr addrspace(1))) null, i32 1, i32 0, ptr addrspace(1) null, i64 0, i64 0)
+  to label %unreachable unwind label %lpad
+
+else:
+  %t2 = invoke token (i64, i32, ptr, i32, i32, ...) @llvm.experimental.gc.statepoint.p0(i64 0, i32 0, ptr elementtype(void (ptr addrspace(1))) null, i32 1, i32 0, ptr addrspace(1) null, i64 0, i64 0)
+  to label %unreachable unwind label %lpad
+
+unreachable:
+  unreachable
+
+lpad:
+  %t3 = landingpad token
+  cleanup
+  ret void
+}
+
+declare token @llvm.experimental.gc.statepoint.p0(i64 immarg, i32 immarg, ptr, i32 immarg, i32 immarg, ...)
+
 declare i1 @cond()
 
 declare void @sideeffect()
diff --git a/llvm/test/Transforms/SimplifyCFG/merge-deopt-bundle-constants.ll b/llvm/test/Transforms/SimplifyCFG/merge-deopt-bundle-constants.ll
index 27c6d7e3bbf5..abdc8a325da0 100644
--- a/llvm/test/Transforms/SimplifyCFG/merge-deopt-bundle-constants.ll
+++ b/llvm/test/Transforms/SimplifyCFG/merge-deopt-bundle-constants.ll
@@ -7,11 +7,18 @@ target triple = "x86_64-unknown-linux-gnu"
 define void @test_01(i1 %cond) gc "statepoint-example" personality ptr @zot {
 ; CHECK-LABEL: @test_01(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    [[DOT:%.*]] = select i1 [[COND:%.*]], i32 0, i32 1
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BB3:%.*]], label [[BB8:%.*]]
+; CHECK:       bb3:
 ; CHECK-NEXT:    [[TMP4:%.*]] = call ptr @wibble()
-; CHECK-NEXT:    [[TMP0:%.*]] = invoke align 8 dereferenceable_or_null(8) ptr addrspace(1) [[TMP4]](ptr addrspace(1) undef) [ "deopt"(i32 [[DOT]]) ]
-; CHECK-NEXT:    to label [[BB8_CONT:%.*]] unwind label [[BB13:%.*]]
-; CHECK:       bb8.cont:
+; CHECK-NEXT:    [[TMP6:%.*]] = invoke align 8 dereferenceable_or_null(8) ptr addrspace(1) [[TMP4]](ptr addrspace(1) undef) [ "deopt"(i32 0) ]
+; CHECK-NEXT:    to label [[BB7:%.*]] unwind label [[BB13:%.*]]
+; CHECK:       bb7:
+; CHECK-NEXT:    unreachable
+; CHECK:       bb8:
+; CHECK-NEXT:    [[TMP9:%.*]] = call ptr @wibble()
+; CHECK-NEXT:    [[TMP11:%.*]] = invoke align 8 dereferenceable_or_null(8) ptr addrspace(1) [[TMP9]](ptr addrspace(1) undef) [ "deopt"(i32 1) ]
+; CHECK-NEXT:    to label [[BB12:%.*]] unwind label [[BB13]]
+; CHECK:       bb12:
 ; CHECK-NEXT:    unreachable
 ; CHECK:       bb13:
 ; CHECK-NEXT:    [[TMP14:%.*]] = landingpad { ptr, i32 }