[CallSiteSplitting] Simplify isPredicateOnPHI & continue checking PHIs.

As pointed out by @thakis, currently CallSiteSplitting bails out after
checking the first PHI node. We should check all PHI nodes, until we
find one where call site splitting is beneficial.

This patch also slightly simplifies the code using BasicBlock::phis().

Reviewers: davidxl, junbuml, thakis

Reviewed By: davidxl

Differential Revision: https://reviews.llvm.org/D77089
diff --git a/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp b/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
index e34c011..286e237 100644
--- a/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
+++ b/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
@@ -424,22 +424,20 @@
   if (Instr != Parent->getFirstNonPHIOrDbg())
     return false;
 
-  for (auto &BI : *Parent) {
-    if (PHINode *PN = dyn_cast<PHINode>(&BI)) {
-      for (auto &I : CS.args())
-        if (&*I == PN) {
-          assert(PN->getNumIncomingValues() == 2 &&
-                 "Unexpected number of incoming values");
-          if (PN->getIncomingBlock(0) == PN->getIncomingBlock(1))
-            return false;
-          if (PN->getIncomingValue(0) == PN->getIncomingValue(1))
-            continue;
-          if (isa<Constant>(PN->getIncomingValue(0)) &&
-              isa<Constant>(PN->getIncomingValue(1)))
-            return true;
-        }
+  for (auto &PN : Parent->phis()) {
+    for (auto &Arg : CS.args()) {
+      if (&*Arg != &PN)
+        continue;
+      assert(PN.getNumIncomingValues() == 2 &&
+             "Unexpected number of incoming values");
+      if (PN.getIncomingBlock(0) == PN.getIncomingBlock(1))
+        return false;
+      if (PN.getIncomingValue(0) == PN.getIncomingValue(1))
+        continue;
+      if (isa<Constant>(PN.getIncomingValue(0)) &&
+          isa<Constant>(PN.getIncomingValue(1)))
+        return true;
     }
-    break;
   }
   return false;
 }
diff --git a/llvm/test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll b/llvm/test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll
index cd2edb1..83b265b 100644
--- a/llvm/test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll
+++ b/llvm/test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll
@@ -586,3 +586,37 @@
 End:
   ret void
 }
+
+; CHECK-LABEL: i32 @test_multiple_phis(
+; CHECK:      Header.split:
+; CHECK-NEXT:   %r2 = call i32 @callee(i32* null, i32 1, i32 5)
+; CHECK-NEXT:   br label %Tail
+
+; CHECK:      TBB.split:
+; CHECK-NEXT:   %r1 = call i32 @callee(i32* null, i32 2, i32 10)
+; CHECK-NEXT:   br label %Tail
+
+; CHECK:        Tail:
+; CHECK-NEXT:   %phi.call = phi i32 [ %r1, %TBB.split ], [ %r2, %Header.split ]
+; CHECK-NEXT:   %p.0 = phi i32 [ 0, %Header.split ], [ 99, %TBB.split ]
+; CHECK-NEXT:   %res = add i32 %phi.call, %p.0
+; CHECK-NEXT:   ret i32 %phi.call
+;
+define i32 @test_multiple_phis(i1 %c.1) {
+Header:
+  br i1 %c.1, label %Tail, label %TBB
+
+TBB:
+  br label %Tail
+
+Tail:
+  %p.0 = phi i32 [0, %Header], [99, %TBB]
+  %p.1 = phi i32[1, %Header], [2, %TBB]
+  %p.2 = phi i32 [5, %Header], [10, %TBB]
+  %r = call i32 @callee(i32* null, i32 %p.1, i32 %p.2)
+  %res = add i32 %r, %p.0
+  ret i32 %r
+
+End:
+  ret i32 10
+}