SSAUpdaterBulk: properly distinguish between the middle and end of a block

[nhaehnle]

Summary

Correct SSA update / phi insertion requires us to separately track the
value at the end of a block, i.e. after a user-provided definition, and
at the beginning, i.e. incoming values from predecessors (if any).

This change aligns the behavior of SSAUpdaterBulk with that of
SSAUpdater to avoid risk of confusion.

SSAUpdaterBulk is a separate implementation that isn't used for much at
all and seems to not have seen much testing. Its interface can sometimes
be convenient (which is why I tried to use it) and the commit of its initial
introduction claims a performance benefit to using dominator trees in
the bulk case, so it seems worth fixing this implementation.

[nhaehnle]

Ping

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

 ///
 /// This is used when code duplication or another unstructured transformation
 /// wants to rewrite a set of uses of one value with uses of a set of values.
 /// The update is done only when RewriteAllUses is called, all other methods are
 /// used for book-keeping. That helps to share some common computations between
 /// updates of different uses (which is not the case when traditional SSAUpdater
 /// is used).
 class SSAUpdaterBulk {
+  struct BlockInfo {
+    Value *Define = nullptr;
+    Value *Phi = nullptr;
+  };
   struct RewriteInfo {
-    DenseMap<BasicBlock *, Value *> Defines;
+    DenseMap<BasicBlock *, BlockInfo> Blocks;
     SmallVector<Use *, 4> Uses;
     StringRef Name;
     Type *Ty;
     RewriteInfo() = default;
     RewriteInfo(StringRef &N, Type *T) : Name(N), Ty(T){};
   };
   SmallVector<RewriteInfo, 4> Rewrites;
 
   PredIteratorCache PredCache;
 
-  Value *computeValueAt(BasicBlock *BB, RewriteInfo &R, DominatorTree *DT);
+  Value *computeValueAt(BasicBlock *BB, bool AtEnd, RewriteInfo &R,
+                        DominatorTree *DT);
 
 public:
   explicit SSAUpdaterBulk() = default;
   SSAUpdaterBulk(const SSAUpdaterBulk &) = delete;
   SSAUpdaterBulk &operator=(const SSAUpdaterBulk &) = delete;
   ~SSAUpdaterBulk() = default;
 
   /// Add a new variable to the SSA rewriter. This needs to be called before

llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp

 
 /// Indicate that a rewritten value is available in the specified block with the
 /// specified value.
 void SSAUpdaterBulk::AddAvailableValue(unsigned Var, BasicBlock *BB, Value *V) {
   assert(Var < Rewrites.size() && "Variable not found!");
   LLVM_DEBUG(dbgs() << "SSAUpdater: Var=" << Var
                     << ": added new available value" << *V << " in "
                     << BB->getName() << "\n");
-  Rewrites[Var].Defines[BB] = V;
+  Rewrites[Var].Blocks[BB].Define = V;
 }
 
 /// Record a use of the symbolic value. This use will be updated with a
 /// rewritten value when RewriteAllUses is called.
 void SSAUpdaterBulk::AddUse(unsigned Var, Use *U) {
   assert(Var < Rewrites.size() && "Variable not found!");
   LLVM_DEBUG(dbgs() << "SSAUpdater: Var=" << Var << ": added a use" << *U->get()
                     << " in " << getUserBB(U)->getName() << "\n");
   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);
+// Compute value for the given block BB. Depending on AtEnd, return either the
+// value that is valid at the end of the block (i.e., after the definition in
+// the block, if any) or at the beginning of the block (i.e., before any uses
+// added in the block).
+//
+// We either should already know it, or we should be able to recursively reach
+// it going up dominator tree.
+Value *SSAUpdaterBulk::computeValueAt(BasicBlock *BB, bool AtEnd,
+                                      RewriteInfo &R, DominatorTree *DT) {
+  auto InfoIt = R.Blocks.find(BB);
+  if (InfoIt != R.Blocks.end()) {
+    BlockInfo &Info = InfoIt->second;
+    if (AtEnd) {
+      if (Info.Define)
+        return Info.Define;
+
+      assert(Info.Phi != nullptr && "avoid creating unnecessary BlockInfos!");
+      Info.Define = Info.Phi;
+      return Info.Define;
+    }
+
+    if (Info.Phi)
+      return Info.Phi;
+  }
+
+  Value *V;
+    V = computeValueAt(IDom, true, R, DT);
+  } else {
+    V = UndefValue::get(R.Ty);
   }
-  return R.Defines[BB];
+
+  BlockInfo &Info = R.Blocks[BB];
+  Info.Phi = V;
+  if (AtEnd)
+    Info.Define = V;
+  return V;
 }
 
 /// Given sets of UsingBlocks and DefBlocks, compute the set of LiveInBlocks.
 /// This is basically a subgraph limited by DefBlocks and UsingBlocks.
 static void
 ComputeLiveInBlocks(const SmallPtrSetImpl<BasicBlock *> &UsingBlocks,
                     const SmallPtrSetImpl<BasicBlock *> &DefBlocks,
                     SmallPtrSetImpl<BasicBlock *> &LiveInBlocks,
for (auto &R : Rewrites) {
     // UsingBlocks from uses in R.Uses, then compute LiveInBlocks, and then use
     // this set for computing iterated dominance frontier (IDF).
     // The IDF blocks are the blocks where we need to insert new phi-nodes.
     ForwardIDFCalculator IDF(*DT);
     LLVM_DEBUG(dbgs() << "SSAUpdater: rewriting " << R.Uses.size()
                       << " use(s)\n");
 
     SmallPtrSet<BasicBlock *, 2> DefBlocks;
-    for (auto &Def : R.Defines)
+    for (auto &Def : R.Blocks)
       DefBlocks.insert(Def.first);
     IDF.setDefiningBlocks(DefBlocks);
 
     SmallPtrSet<BasicBlock *, 2> UsingBlocks;
     for (Use *U : R.Uses)
       UsingBlocks.insert(getUserBB(U));
 
     SmallVector<BasicBlock *, 32> IDFBlocks;
     SmallPtrSet<BasicBlock *, 32> LiveInBlocks;
     ComputeLiveInBlocks(UsingBlocks, DefBlocks, LiveInBlocks, PredCache);
     IDF.resetLiveInBlocks();
     IDF.setLiveInBlocks(LiveInBlocks);
     IDF.calculate(IDFBlocks);
 
     // 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;
+      R.Blocks[FrontierBB].Phi = PN;
       InsertedPHIsForVar.push_back(PN);
       if (InsertedPHIs)
         InsertedPHIs->push_back(PN);
     }
 
     // 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);
+        PN->addIncoming(computeValueAt(Pred, true, R, DT), 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);
+      Value *V = computeValueAt(getUserBB(U), false, R, DT);
       Value *OldVal = U->get();
       assert(OldVal && "Invalid use!");
       // Notify that users of the existing value that it is being replaced.
       if (OldVal != V && OldVal->hasValueHandle())
         ValueHandleBase::ValueIsRAUWd(OldVal, V);
       LLVM_DEBUG(dbgs() << "SSAUpdater: replacing " << *OldVal << " with " << *V
                         << "\n");
       U->set(V);
     }
   }
 }

llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp

TEST(SSAUpdaterBulk, Irreducible) {
   EXPECT_EQ(Return->getOperand(0), FirstArg);
 
   // I2 should use the new phi.
   PHINode *UpdatePhi = dyn_cast_or_null<PHINode>(I2->getOperand(0));
   EXPECT_NE(UpdatePhi, nullptr);
   EXPECT_EQ(UpdatePhi->getIncomingValueForBlock(LoopStartBB), AddOp2);
   EXPECT_EQ(UpdatePhi->getIncomingValueForBlock(IfBB), UndefValue::get(I32Ty));
 }
+
+TEST(SSAUpdaterBulk, Loop) {
+  SSAUpdaterBulk Updater;
+  LLVMContext C;
+  Module M("SSAUpdaterTest", C);
+  IRBuilder<> B(C);
+  Type *I32Ty = B.getInt32Ty();
+  auto *F = Function::Create(FunctionType::get(I32Ty, {I32Ty}, false),
+                             GlobalValue::ExternalLinkage, "F", &M);
+
+  // Generate a small program with a simple loop:
+  //     entry:
+  //       br label %loop
+  //
+  //     loop:
+  //       %1 = add i32 %0, 1
+  //       br i1 true, label %loop, label %end
+  //
+  //     end:
+  //       ret i32 %0
+  Argument *FirstArg = &*F->arg_begin();
+  BasicBlock *Entry = BasicBlock::Create(C, "entry", F);
+  BasicBlock *Loop = BasicBlock::Create(C, "loop", F);
+  BasicBlock *End = BasicBlock::Create(C, "end", F);
+
+  B.SetInsertPoint(Entry);
+  B.CreateBr(Loop);
+
+  B.SetInsertPoint(Loop);
+  auto *Add = cast<Instruction>(B.CreateAdd(FirstArg, B.getInt32(0)));
+  B.CreateCondBr(B.getTrue(), Loop, End);
+
+  B.SetInsertPoint(End);
+  ReturnInst *Return = B.CreateRet(FirstArg);
+
+  // Now rewrite uses in instructions %1 and 'ret i32 %0'. The expected result
+  // is:
+  //     entry:
+  //       br label %loop
+  //
+  //     loop:
+  //       %1 = phi i32 [ 0, %entry ], [ %2, %loop ]
+  //       %2 = add i32 %1, 1
+  //       br i1 true, label %loop, label %end
+  //
+  //     end:
+  //       ret i32 %2
+
+  // Add definitions and uses.
+  unsigned VarNum = Updater.AddVariable("c", I32Ty);
+  Updater.AddAvailableValue(VarNum, Entry, B.getInt32(0));
+  Updater.AddAvailableValue(VarNum, Loop, Add);
+
+  Updater.AddUse(VarNum, &Add->getOperandUse(0));
+  Updater.AddUse(VarNum, &Return->getOperandUse(0));
+
+  // Save all inserted phis into a vector.
+  SmallVector<PHINode *, 8> Inserted;
+  DominatorTree DT(*F);
+  Updater.RewriteAllUses(&DT, &Inserted);
+
+  // Only one phi should have been inserted, in the loop block.
+  EXPECT_EQ(Inserted.size(), 1u);
+
+  PHINode *Phi = Inserted[0];
+  EXPECT_EQ(Phi->getParent(), Loop);
+  EXPECT_EQ(Phi->getIncomingValueForBlock(Entry), B.getInt32(0));
+  EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), Add);
+
+  // Verify rewrites of uses.
+  EXPECT_EQ(Add->getOperand(0), Phi);
+  EXPECT_EQ(Return->getOperand(0), Add);
+}