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

[Alexandros Lamprineas 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");
----------------
labrinea wrote:

I am just thinking visit(phi1) would have returned C, likewise visit(phi2). So the incoming values of phi0 would be in KnownConstants already. Anything below this line handles TransitivePhis which isn't the case in the above example.

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