[llvm] [SSAUpdaterBulk] Fix incorrect live-in values for a block. (PR #131762)

[Matt Arsenault via llvm-commits]

================
@@ -143,34 +136,66 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
     SmallVector<BasicBlock *, 32> IDFBlocks;
     SmallPtrSet<BasicBlock *, 32> LiveInBlocks;
     ComputeLiveInBlocks(UsingBlocks, DefBlocks, LiveInBlocks, PredCache);
-    IDF.resetLiveInBlocks();
     IDF.setLiveInBlocks(LiveInBlocks);
     IDF.calculate(IDFBlocks);
 
+    // Reserve map large enough to reduce growth.
+    BBInfos.init(LiveInBlocks.size() + DefBlocks.size());
+
+    for (auto [BB, V] : R.Defines)
+      BBInfos[BB].LiveOutValue = V;
+
     // We've computed IDF, now insert new phi-nodes there.
-    SmallVector<PHINode *, 4> InsertedPHIsForVar;
     for (auto *FrontierBB : IDFBlocks) {
       IRBuilder<> B(FrontierBB, FrontierBB->begin());
       PHINode *PN = B.CreatePHI(R.Ty, 0, R.Name);
-      R.Defines[FrontierBB] = PN;
-      InsertedPHIsForVar.push_back(PN);
+      BBInfos[FrontierBB].LiveInValue = PN;
       if (InsertedPHIs)
         InsertedPHIs->push_back(PN);
     }
 
+    // IsLiveOut indicates whether we are computing live-out values (true) or
+    // live-in values (false).
+    std::function<Value *(BasicBlock *, bool)> computeValue =
+        [&](BasicBlock *BB, bool IsLiveOut) -> Value * {
+      auto &BBInfo = BBInfos[BB];
+
+      if (IsLiveOut && BBInfo.LiveOutValue)
+        return BBInfo.LiveOutValue;
+
+      if (BBInfo.LiveInValue)
+        return BBInfo.LiveInValue;
+
+      Value *V = DT->isReachableFromEntry(BB) && PredCache.get(BB).size()
+                     ? computeValue(DT->getNode(BB)->getIDom()->getBlock(),
+                                    /* IsLiveOut = */ true)
+                     : UndefValue::get(R.Ty);
+
+      // The call to computeValue for the dominator block can insert another
+      // entry into the map, potentially causing the map to grow and
+      // invalidating the BBInfo reference. Therefore, we need to perform
+      // another map lookup. Simply reserving map size may not be sufficient
+      // as the map could grow further.
+      BBInfos[BB].LiveInValue = V;
+      return V;
+    };
+
     // Fill in arguments of the inserted PHIs.
-    for (auto *PN : InsertedPHIsForVar) {
-      BasicBlock *PBB = PN->getParent();
-      for (BasicBlock *Pred : PredCache.get(PBB))
-        PN->addIncoming(computeValueAt(Pred, R, DT), Pred);
+    for (auto *BB : IDFBlocks) {
+      auto *PHI = cast<PHINode>(&BB->front());
+      for (BasicBlock *Pred : PredCache.get(BB))
+        PHI->addIncoming(computeValue(Pred, /* IsLiveOut = */ true), Pred);
     }
 
     // Rewrite actual uses with the inserted definitions.
     SmallPtrSet<Use *, 4> ProcessedUses;
     for (Use *U : R.Uses) {
       if (!ProcessedUses.insert(U).second)
         continue;
-      Value *V = computeValueAt(getUserBB(U), R, DT);
+
+      auto *User = cast<Instruction>(U->getUser());
+      BasicBlock *BB = getUserBB(U);
+      Value *V = computeValue(BB, /* IsLiveOut = */ BB != User->getParent());
----------------
arsenm wrote:

```suggestion
      Value *V = computeValue(BB, /*IsLiveOut=*/ BB != User->getParent());
```

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