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

The previous implementation incorrectly calculated incoming values from
loop backedges, as demonstrated by the tests. The issue was that it did
not distinguish between live-in and live-out values for blocks. This
patch addresses the problem and fixes
llvm#131761.

To avoid bloating storage in `R.Defines`, processing data has been moved
to a temporary map `BBInfos`. This change helps manage heap allocation
more efficiently and likely improves caching.

Author: vpykhtin

llvm/include/llvm/Transforms/Utils/SSAUpdaterBulk.h

@@ -38,7 +38,7 @@ class DominatorTree;
 /// is used).
 class SSAUpdaterBulk {
   struct RewriteInfo {
-    DenseMap<BasicBlock *, Value *> Defines;
+    SmallVector<std::pair<BasicBlock *, Value *>, 4> Defines;
     SmallVector<Use *, 4> Uses;
     StringRef Name;
     Type *Ty;
@@ -49,8 +49,6 @@ class SSAUpdaterBulk {
 
   PredIteratorCache PredCache;
 
-  Value *computeValueAt(BasicBlock *BB, RewriteInfo &R, DominatorTree *DT);
-
 public:
   explicit SSAUpdaterBulk() = default;
   SSAUpdaterBulk(const SSAUpdaterBulk &) = delete;

llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp

@@ -53,7 +53,7 @@ void SSAUpdaterBulk::AddAvailableValue(unsigned Var, BasicBlock *BB, Value *V) {
   LLVM_DEBUG(dbgs() << "SSAUpdater: Var=" << Var
                     << ": added new available value " << *V << " in "
                     << BB->getName() << "\n");
-  Rewrites[Var].Defines[BB] = V;
+  Rewrites[Var].Defines.emplace_back(BB, V);
 }
 
 /// Record a use of the symbolic value. This use will be updated with a
@@ -65,21 +65,6 @@ void SSAUpdaterBulk::AddUse(unsigned Var, Use *U) {
   Rewrites[Var].Uses.push_back(U);
 }
 
-// Compute value at the given block BB. We either should already know it, or we
-// should be able to recursively reach it going up dominator tree.
-Value *SSAUpdaterBulk::computeValueAt(BasicBlock *BB, RewriteInfo &R,
-                                      DominatorTree *DT) {
-  if (!R.Defines.count(BB)) {
-    if (DT->isReachableFromEntry(BB) && PredCache.get(BB).size()) {
-      BasicBlock *IDom = DT->getNode(BB)->getIDom()->getBlock();
-      Value *V = computeValueAt(IDom, R, DT);
-      R.Defines[BB] = V;
-    } else
-      R.Defines[BB] = UndefValue::get(R.Ty);
-  }
-  return R.Defines[BB];
-}
-
 /// Given sets of UsingBlocks and DefBlocks, compute the set of LiveInBlocks.
 /// This is basically a subgraph limited by DefBlocks and UsingBlocks.
 static void
@@ -117,11 +102,19 @@ ComputeLiveInBlocks(const SmallPtrSetImpl<BasicBlock *> &UsingBlocks,
   }
 }
 
+struct BBValueInfo {
+  Value *LiveInValue = nullptr;
+  Value *LiveOutValue = nullptr;
+};
+
 /// Perform all the necessary updates, including new PHI-nodes insertion and the
 /// requested uses update.
 void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
                                     SmallVectorImpl<PHINode *> *InsertedPHIs) {
+  DenseMap<BasicBlock *, BBValueInfo> BBInfos;
   for (auto &R : Rewrites) {
+    BBInfos.clear();
+
     // Compute locations for new phi-nodes.
     // For that we need to initialize DefBlocks from definitions in R.Defines,
     // UsingBlocks from uses in R.Uses, then compute LiveInBlocks, and then use
@@ -132,8 +125,8 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
                       << " use(s)\n");
 
     SmallPtrSet<BasicBlock *, 2> DefBlocks;
-    for (auto &Def : R.Defines)
-      DefBlocks.insert(Def.first);
+    for (auto [BB, V] : R.Defines)
+      DefBlocks.insert(BB);
     IDF.setDefiningBlocks(DefBlocks);
 
     SmallPtrSet<BasicBlock *, 2> UsingBlocks;
@@ -143,34 +136,82 @@ 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 sufficient buckets to prevent map growth. [1]
+    BBInfos.reserve(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).
+    auto ComputeValue = [&](BasicBlock *BB, bool IsLiveOut) -> Value * {
+      auto *BBInfo = &BBInfos[BB];
+
+      if (IsLiveOut && BBInfo->LiveOutValue)
+        return BBInfo->LiveOutValue;
+
+      if (BBInfo->LiveInValue)
+        return BBInfo->LiveInValue;
+
+      SmallVector<BBValueInfo *, 4> Stack = {BBInfo};
+      Value *V = nullptr;
+
+      while (DT->isReachableFromEntry(BB) && !PredCache.get(BB).empty() &&
+             (BB = DT->getNode(BB)->getIDom()->getBlock())) {
+        BBInfo = &BBInfos[BB];
+
+        if (BBInfo->LiveOutValue) {
+          V = BBInfo->LiveOutValue;
+          break;
+        }
+
+        if (BBInfo->LiveInValue) {
+          V = BBInfo->LiveInValue;
+          break;
+        }
+
+        Stack.emplace_back(BBInfo);
+      }
+
+      if (!V)
+        V = UndefValue::get(R.Ty);
+
+      for (auto *BBInfo : Stack)
+        // Loop above can insert new entries into the BBInfos map: assume the
+        // map shouldn't grow due to [1] and BBInfo references are valid.
+        BBInfo->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());
       Value *OldVal = U->get();
       assert(OldVal && "Invalid use!");
       // Notify that users of the existing value that it is being replaced.

llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp

@@ -15,7 +15,6 @@
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/SourceMgr.h"
-#include "gtest/gtest-spi.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -236,7 +235,7 @@ TEST(SSAUpdaterBulk, SingleBBLoop) {
   Loop->dump();
   // Output:
   // loop: ; preds = %loop, %entry
-  //   %i.prev = phi i32 [ %i.prev, %loop ], [ 0, %entry ]
+  //   %i.prev = phi i32 [ %i, %loop ], [ 0, %entry ]
   //   %i = add i32 %i.prev, 1
   //   %cmp = icmp slt i32 %i, 42
   //   br i1 %cmp, label %loop, label %exit
@@ -246,8 +245,7 @@ TEST(SSAUpdaterBulk, SingleBBLoop) {
   PHINode *Phi = Inserted[0];
   EXPECT_EQ(Phi, dyn_cast<PHINode>(I->getOperand(0)));
   EXPECT_EQ(Phi->getIncomingValueForBlock(Entry), ConstantInt::get(I32Ty, 0));
-  EXPECT_NONFATAL_FAILURE(EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), I),
-                          "Expected equality of these values");
+  EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), I);
 }
 
 TEST(SSAUpdaterBulk, TwoBBLoop) {
@@ -298,15 +296,14 @@ TEST(SSAUpdaterBulk, TwoBBLoop) {
   //   %i.prev = phi i32 [ %i, %loop ], [ 0, %entry ]
   //   br label %loop
   // loop:                                             ; preds = %loop_header
-  //   %i = add i32 %i, 1
+  //   %i = add i32 %i.prev, 1
   //   %cmp = icmp slt i32 %i, 42
   //   br i1 %cmp, label %loop_header, label %exit
 #endif
 
   ASSERT_EQ(Inserted.size(), 1u);
   PHINode *Phi = Inserted[0];
-  EXPECT_NONFATAL_FAILURE(EXPECT_EQ(Phi, dyn_cast<PHINode>(I->getOperand(0))),
-                          "Expected equality of these values");
+  EXPECT_EQ(Phi, dyn_cast<PHINode>(I->getOperand(0)));
   EXPECT_EQ(Phi->getParent(), LoopHdr);
   EXPECT_EQ(Phi->getIncomingValueForBlock(Entry), ConstantInt::get(I32Ty, 0));
   EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), I);