[llvm] [FuncSpec] Update function specialization to handle phi-chains (PR #72903)

[Mats Petersson via llvm-commits]

================
@@ -262,29 +269,113 @@ Cost InstCostVisitor::estimateBranchInst(BranchInst &I) {
   return estimateBasicBlocks(WorkList);
 }
 
+bool InstCostVisitor::discoverTransitivelyIncomingValues(
+    Constant *Const, PHINode *Root, DenseSet<PHINode *> &TransitivePHIs,
+    SmallVectorImpl<PHINode *> &UnknownIncomingValues) {
+
+  SmallVector<PHINode *, 64> WorkList;
+  WorkList.push_back(Root);
+  unsigned Iter = 0;
+
+  while (!WorkList.empty()) {
+    PHINode *PN = WorkList.pop_back_val();
+
+    if (++Iter > MaxDiscoveryIterations ||
+        PN->getNumIncomingValues() > MaxIncomingPhiValues) {
+      // For now just collect the Phi and later we will check whether it is
+      // in the Transitive set.
+      UnknownIncomingValues.push_back(PN);
+      continue;
+      // FIXME: return false here and remove the UnknownIncomingValues entirely.
+    }
+
+    if (!TransitivePHIs.insert(PN).second)
+      continue;
+
+    for (unsigned I = 0, E = PN->getNumIncomingValues(); I != E; ++I) {
+      Value *V = PN->getIncomingValue(I);
+
+      // Disregard self-references and dead incoming values.
+      if (auto *Inst = dyn_cast<Instruction>(V))
+        if (Inst == PN || DeadBlocks.contains(PN->getIncomingBlock(I)))
+          continue;
+
+      if (Constant *C = findConstantFor(V, KnownConstants)) {
+        // Not all incoming values are the same constant. Bail immediately.
+        if (C != Const)
+          return false;
+        continue;
+      }
+
+      if (auto *Phi = dyn_cast<PHINode>(V)) {
+        WorkList.push_back(Phi);
+        continue;
+      }
+
+      // We can't reason about anything else.
+      return false;
+    }
+  }
+  return true;
+}
+
 Constant *InstCostVisitor::visitPHINode(PHINode &I) {
   if (I.getNumIncomingValues() > MaxIncomingPhiValues)
     return nullptr;
 
   bool Inserted = VisitedPHIs.insert(&I).second;
+  SmallVector<PHINode *, 8> UnknownIncomingValues;
+  DenseSet<PHINode *> TransitivePHIs;
   Constant *Const = nullptr;
+  bool HaveSeenIncomingPHI = false;
 
   for (unsigned Idx = 0, E = I.getNumIncomingValues(); Idx != E; ++Idx) {
     Value *V = I.getIncomingValue(Idx);
+
+    // Disregard self-references and dead incoming values.
     if (auto *Inst = dyn_cast<Instruction>(V))
       if (Inst == &I || DeadBlocks.contains(I.getIncomingBlock(Idx)))
         continue;
-    Constant *C = findConstantFor(V, KnownConstants);
-    if (!C) {
-      if (Inserted)
-        PendingPHIs.push_back(&I);
-      return nullptr;
+
+    if (Constant *C = findConstantFor(V, KnownConstants)) {
+      if (!Const)
+        Const = C;
+      // Not all incoming values are the same constant. Bail immediately.
+      if (C != Const)
+        return nullptr;
+      continue;
     }
-    if (!Const)
-      Const = C;
-    else if (C != Const)
+
+    if (Inserted) {
+      // First time we are seeing this phi. We will retry later, after
+      // all the constant arguments have been propagated. Bail for now.
+      PendingPHIs.push_back(&I);
       return nullptr;
+    }
+
+    if (isa<PHINode>(V)) {
+      // Perhaps it is a Transitive Phi. We will confirm later.
+      HaveSeenIncomingPHI = true;
+      continue;
+    }
+
+    // We can't reason about anything else.
+    return nullptr;
   }
+
+  assert(Const && "Should have found at least one constant incoming value");
----------------
Leporacanthicus wrote:

But isn't that covered by the `if (inserted) return nullptr;` portion of code?

I wrote this to try to make it fail:

```
define internal i64 @foo(i64 %n, i1 %c1, i1 %c2, i1 %c3, i1 %c4) {
entry:
  br label %l0
  
l1:
  %phi1 = phi i64 [ %phi0, %l0 ], [ %phi2, %l2 ]
  %add = add i64 %phi1, 1
  %div = sdiv i64 %add, 2
  br i1 %c2, label %l2, label %exit

l2:
  %phi2 = phi i64 [ %phi0, %l0 ], [ %phi1, %l1 ]
  %sub = sub i64 %phi2, 1
  %mul = mul i64 %sub, 2
  br i1 %c4, label %l1, label %exit

l0:
  %phi0 = phi i64 [ %n, %entry ]
  br i1 %c1, label %l1, label %l2

exit:
  %res = phi i64 [ %div, %l1 ], [ %mul, %l2]
  ret i64 %res
}
```

And it behaves correctly. 

https://github.com/llvm/llvm-project/pull/72903