[SSAUpdaterBulk] Add PHI simplification pass.

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

@@ -13,7 +13,6 @@
 #ifndef LLVM_TRANSFORMS_UTILS_SSAUPDATERBULK_H
 #define LLVM_TRANSFORMS_UTILS_SSAUPDATERBULK_H
 
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/PredIteratorCache.h"
 #include "llvm/Support/Compiler.h"
@@ -79,6 +78,10 @@ class SSAUpdaterBulk {
   LLVM_ABI void
   RewriteAllUses(DominatorTree *DT,
                  SmallVectorImpl<PHINode *> *InsertedPHIs = nullptr);
+
+  /// Rewrite all uses and simplify the inserted PHI nodes.
+  /// Use this method to preserve behavior when replacing SSAUpdater.
+  void RewriteAndOptimizeAllUses(DominatorTree &DT);
 };
 
 } // end namespace llvm

llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp

@@ -11,11 +11,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/SSAUpdaterBulk.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/IteratedDominanceFrontier.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
-#include "llvm/IR/Instructions.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/Value.h"
 
@@ -112,7 +112,7 @@ struct BBValueInfo {
 void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
                                     SmallVectorImpl<PHINode *> *InsertedPHIs) {
   DenseMap<BasicBlock *, BBValueInfo> BBInfos;
-  for (auto &R : Rewrites) {
+  for (RewriteInfo &R : Rewrites) {
     BBInfos.clear();
 
     // Compute locations for new phi-nodes.
@@ -145,7 +145,7 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
       BBInfos[BB].LiveOutValue = V;
 
     // We've computed IDF, now insert new phi-nodes there.
-    for (auto *FrontierBB : IDFBlocks) {
+    for (BasicBlock *FrontierBB : IDFBlocks) {
       IRBuilder<> B(FrontierBB, FrontierBB->begin());
       PHINode *PN = B.CreatePHI(R.Ty, 0, R.Name);
       BBInfos[FrontierBB].LiveInValue = PN;
@@ -156,7 +156,7 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
     // 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];
+      BBValueInfo *BBInfo = &BBInfos[BB];
 
       if (IsLiveOut && BBInfo->LiveOutValue)
         return BBInfo->LiveOutValue;
@@ -187,7 +187,7 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
       if (!V)
         V = UndefValue::get(R.Ty);
 
-      for (auto *BBInfo : Stack)
+      for (BBValueInfo *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;
@@ -196,7 +196,7 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
     };
 
     // Fill in arguments of the inserted PHIs.
-    for (auto *BB : IDFBlocks) {
+    for (BasicBlock *BB : IDFBlocks) {
       auto *PHI = cast<PHINode>(&BB->front());
       for (BasicBlock *Pred : PredCache.get(BB))
         PHI->addIncoming(ComputeValue(Pred, /*IsLiveOut=*/true), Pred);
@@ -222,3 +222,97 @@ void SSAUpdaterBulk::RewriteAllUses(DominatorTree *DT,
     }
   }
 }
+
+// Perform a single pass of simplification over the worklist of PHIs.
+// This should be called after RewriteAllUses() because simplifying PHIs
+// immediately after creation would require updating all references to those
+// PHIs in the BBValueInfo structures, which would necessitate additional
+// reference tracking overhead.
+static void simplifyPass(MutableArrayRef<PHINode *> Worklist,
+                         const DataLayout &DL) {
+  for (PHINode *&PHI : Worklist) {
+    if (Value *Simplified = simplifyInstruction(PHI, DL)) {
+      PHI->replaceAllUsesWith(Simplified);
+      PHI->eraseFromParent();
+      PHI = nullptr; // Mark as removed.
+    }
+  }
+}
+
+#ifndef NDEBUG // Should this be under EXPENSIVE_CHECKS?
+// New PHI nodes should not reference one another but they may reference
+// themselves or existing PHI nodes, and existing PHI nodes may reference new
+// PHI nodes.
+static bool
+PHIAreRefEachOther(const iterator_range<BasicBlock::phi_iterator> NewPHIs) {
+  SmallPtrSet<PHINode *, 8> NewPHISet;
+  for (PHINode &PN : NewPHIs)
+    NewPHISet.insert(&PN);
+  for (PHINode &PHI : NewPHIs) {
+    for (Value *V : PHI.incoming_values()) {
+      PHINode *IncPHI = dyn_cast<PHINode>(V);
+      if (IncPHI && IncPHI != &PHI && NewPHISet.contains(IncPHI))
+        return true;
+    }
+  }
+  return false;
+}
+#endif
+
+static bool replaceIfIdentical(PHINode &PHI, PHINode &ReplPHI) {
+  if (!PHI.isIdenticalToWhenDefined(&ReplPHI))
+    return false;
+  PHI.replaceAllUsesWith(&ReplPHI);
+  PHI.eraseFromParent();
+  return true;
+}
+
+bool EliminateNewDuplicatePHINodes(BasicBlock *BB,
+                                   BasicBlock::phi_iterator FirstExistingPN) {
+  auto NewPHIs = make_range(BB->phis().begin(), FirstExistingPN);
+  assert(!PHIAreRefEachOther(NewPHIs));
+
+  // Deduplicate new PHIs first to reduce the number of comparisons on the
+  // following new -> existing pass.
+  bool Changed = false;
+  for (auto I = BB->phis().begin(); I != FirstExistingPN; ++I) {
+    for (auto J = std::next(I); J != FirstExistingPN;) {
+      Changed |= replaceIfIdentical(*J++, *I);
+    }
+  }
+
+  // Iterate over existing PHIs and replace identical new PHIs.
+  for (PHINode &ExistingPHI : make_range(FirstExistingPN, BB->phis().end())) {
+    auto I = BB->phis().begin();
+    assert(I != FirstExistingPN); // Should be at least one new PHI.
+    do {
+      Changed |= replaceIfIdentical(*I++, ExistingPHI);
+    } while (I != FirstExistingPN);
+    if (BB->phis().begin() == FirstExistingPN)
+      return Changed;
+  }
+  return Changed;
+}
+
+static void deduplicatePass(ArrayRef<PHINode *> Worklist) {
+  SmallDenseMap<BasicBlock *, unsigned> BBs;
+  for (PHINode *PHI : Worklist) {
+    if (PHI)
+      ++BBs[PHI->getParent()];
+  }
+
+  for (auto [BB, NumNewPHIs] : BBs) {
+    auto FirstExistingPN = std::next(BB->phis().begin(), NumNewPHIs);
+    EliminateNewDuplicatePHINodes(BB, FirstExistingPN);
+  }
+}
+
+void SSAUpdaterBulk::RewriteAndOptimizeAllUses(DominatorTree &DT) {
+  SmallVector<PHINode *, 4> PHIs;
+  RewriteAllUses(&DT, &PHIs);
+  if (PHIs.empty())
+    return;
+
+  simplifyPass(PHIs, PHIs.front()->getParent()->getDataLayout());
+  deduplicatePass(PHIs);
+}

llvm/unittests/Transforms/Utils/SSAUpdaterBulkTest.cpp

@@ -308,3 +308,223 @@ TEST(SSAUpdaterBulk, TwoBBLoop) {
   EXPECT_EQ(Phi->getIncomingValueForBlock(Entry), ConstantInt::get(I32Ty, 0));
   EXPECT_EQ(Phi->getIncomingValueForBlock(Loop), I);
 }
+
+TEST(SSAUpdaterBulk, SimplifyPHIs) {
+  const char *IR = R"(
+      define void @main(i32 %val, i1 %cond) {
+      entry:
+          br i1 %cond, label %left, label %right
+      left:
+          %add = add i32 %val, 1
+          br label %exit
+      right:
+          %sub = sub i32 %val, 1
+          br label %exit
+      exit:
+          %phi = phi i32 [ %sub, %right ], [ %add, %left ]
+          %cmp = icmp slt i32 0, 42
+          ret void
+      }
+  )";
+
+  llvm::LLVMContext Context;
+  llvm::SMDiagnostic Err;
+  std::unique_ptr<llvm::Module> M = llvm::parseAssemblyString(IR, Err, Context);
+  ASSERT_NE(M, nullptr) << "Failed to parse IR: " << Err.getMessage();
+
+  Function *F = M->getFunction("main");
+  auto *Entry = &F->getEntryBlock();
+  auto *Left = Entry->getTerminator()->getSuccessor(0);
+  auto *Right = Entry->getTerminator()->getSuccessor(1);
+  auto *Exit = Left->getSingleSuccessor();
+  auto *Val = &*F->arg_begin();
+  auto *Phi = &Exit->front();
+  auto *Cmp = &*std::next(Exit->begin());
+  auto *Add = &Left->front();
+  auto *Sub = &Right->front();
+
+  SSAUpdaterBulk Updater;
+  Type *I32Ty = Type::getInt32Ty(Context);
+
+  // Use %val directly instead of creating a phi.
+  unsigned ValVar = Updater.AddVariable("Val", I32Ty);
+  Updater.AddAvailableValue(ValVar, Left, Val);
+  Updater.AddAvailableValue(ValVar, Right, Val);
+  Updater.AddUse(ValVar, &Cmp->getOperandUse(0));
+
+  // Use existing %phi for %add and %sub values.
+  unsigned AddSubVar = Updater.AddVariable("AddSub", I32Ty);
+  Updater.AddAvailableValue(AddSubVar, Left, Add);
+  Updater.AddAvailableValue(AddSubVar, Right, Sub);
+  Updater.AddUse(AddSubVar, &Cmp->getOperandUse(1));
+
+  auto ExitSizeBefore = Exit->size();
+  DominatorTree DT(*F);
+  Updater.RewriteAndOptimizeAllUses(DT);
+
+  //  Output for Exit->dump():
+  //  exit:                                             ; preds = %right, %left
+  //    %phi = phi i32 [ %sub, %right ], [ %add, %left ]
+  //    %cmp = icmp slt i32 %val, %phi
+  //    ret void
+
+  ASSERT_EQ(Exit->size(), ExitSizeBefore);
+  ASSERT_EQ(&Exit->front(), Phi);
+  EXPECT_EQ(Val, Cmp->getOperand(0));
+  EXPECT_EQ(Phi, Cmp->getOperand(1));
+}
+
+bool EliminateNewDuplicatePHINodes(BasicBlock *BB,
+                                   BasicBlock::phi_iterator FirstExistingPN);
+
+// Helper to run both versions on the same input.
+static void RunEliminateNewDuplicatePHINode(
+    const char *AsmText,
+    std::function<void(BasicBlock &,
+                       bool(BasicBlock *BB, BasicBlock::phi_iterator))>
+        Check) {
+  LLVMContext C;
+
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(AsmText, Err, C);
+  if (!M) {
+    Err.print("UtilsTests", errs());
+    return;
+  }
+
+  Function *F = M->getFunction("main");
+  auto BBIt = std::find_if(F->begin(), F->end(), [](const BasicBlock &Block) {
+    return Block.getName() == "testbb";
+  });
+  ASSERT_NE(BBIt, F->end());
+  Check(*BBIt, EliminateNewDuplicatePHINodes);
+}
+
+static BasicBlock::phi_iterator getPhiIt(BasicBlock &BB, unsigned Idx) {
+  return std::next(BB.phis().begin(), Idx);
+}
+
+static PHINode *getPhi(BasicBlock &BB, unsigned Idx) {
+  return &*getPhiIt(BB, Idx);
+}
+
+static int getNumPHIs(BasicBlock &BB) {
+  return std::distance(BB.phis().begin(), BB.phis().end());
+}
+
+TEST(SSAUpdaterBulk, EliminateNewDuplicatePHINodes_OrderExisting) {
+  RunEliminateNewDuplicatePHINode(R"(
+      define void @main() {
+      entry:
+          br label %testbb
+      testbb:
+          %np0 = phi i32 [ 1, %entry ]
+          %np1 = phi i32 [ 1, %entry ]
+          %ep0 = phi i32 [ 1, %entry ]
+          %ep1 = phi i32 [ 1, %entry ]
+          %u = add i32 %np0, %np1
+          ret void
+      }
+  )", [](BasicBlock &BB, auto *ENDPN) {
+    AssertingVH<PHINode> EP0 = getPhi(BB, 2);
+    AssertingVH<PHINode> EP1 = getPhi(BB, 3);
+    EXPECT_TRUE(ENDPN(&BB, getPhiIt(BB, 2)));
+    // Expected:
+    //   %ep0 = phi i32 [ 1, %entry ]
+    //   %ep1 = phi i32 [ 1, %entry ]
+    //   %u = add i32 %ep0, %ep0
+    EXPECT_EQ(getNumPHIs(BB), 2);
+    Instruction &Add = *BB.getFirstNonPHIIt();
+    EXPECT_EQ(Add.getOperand(0), EP0);
+    EXPECT_EQ(Add.getOperand(1), EP0);
+    (void)EP1; // Avoid "unused" warning.
+  });
+}
+
+TEST(SSAUpdaterBulk, EliminateNewDuplicatePHINodes_OrderNew) {
+  RunEliminateNewDuplicatePHINode(R"(
+      define void @main() {
+      entry:
+          br label %testbb
+      testbb:
+          %np0 = phi i32 [ 1, %entry ]
+          %np1 = phi i32 [ 1, %entry ]
+          %ep0 = phi i32 [ 2, %entry ]
+          %ep1 = phi i32 [ 2, %entry ]
+          %u = add i32 %np0, %np1
+          ret void
+      }
+  )", [](BasicBlock &BB, auto *ENDPN) {
+    AssertingVH<PHINode> NP0 = getPhi(BB, 0);
+    AssertingVH<PHINode> EP0 = getPhi(BB, 2);
+    AssertingVH<PHINode> EP1 = getPhi(BB, 3);
+    EXPECT_TRUE(ENDPN(&BB, getPhiIt(BB, 2)));
+    // Expected:
+    //   %np0 = phi i32 [ 1, %entry ]
+    //   %ep0 = phi i32 [ 2, %entry ]
+    //   %ep1 = phi i32 [ 2, %entry ]
+    //   %u = add i32 %np0, %np0
+    EXPECT_EQ(getNumPHIs(BB), 3);
+    Instruction &Add = *BB.getFirstNonPHIIt();
+    EXPECT_EQ(Add.getOperand(0), NP0);
+    EXPECT_EQ(Add.getOperand(1), NP0);
+    (void)EP0;
+    (void)EP1; // Avoid "unused" warning.
+  });
+}
+
+TEST(SSAUpdaterBulk, EliminateNewDuplicatePHINodes_NewRefExisting) {
+  RunEliminateNewDuplicatePHINode(R"(
+      define void @main() {
+      entry:
+          br label %testbb
+      testbb:
+          %np0 = phi i32 [ 1, %entry ], [ %ep0, %testbb ]
+          %np1 = phi i32 [ 1, %entry ], [ %ep1, %testbb ]
+          %ep0 = phi i32 [ 1, %entry ], [ %ep0, %testbb ]
+          %ep1 = phi i32 [ 1, %entry ], [ %ep1, %testbb ]
+          %u = add i32 %np0, %np1
+          br label %testbb
+      }
+  )", [](BasicBlock &BB, auto *ENDPN) {
+    AssertingVH<PHINode> EP0 = getPhi(BB, 2);
+    AssertingVH<PHINode> EP1 = getPhi(BB, 3);
+    EXPECT_TRUE(ENDPN(&BB, getPhiIt(BB, 2)));
+    // Expected:
+    //   %ep0 = phi i32 [ 1, %entry ], [ %ep0, %testbb ]
+    //   %ep1 = phi i32 [ 1, %entry ], [ %ep1, %testbb ]
+    //   %u = add i32 %ep0, %ep1
+    EXPECT_EQ(getNumPHIs(BB), 2);
+    Instruction &Add = *BB.getFirstNonPHIIt();
+    EXPECT_EQ(Add.getOperand(0), EP0);
+    EXPECT_EQ(Add.getOperand(1), EP1);
+  });
+}
+
+TEST(SSAUpdaterBulk, EliminateNewDuplicatePHINodes_ExistingRefNew) {
+  RunEliminateNewDuplicatePHINode(R"(
+      define void @main() {
+      entry:
+          br label %testbb
+      testbb:
+          %np0 = phi i32 [ 1, %entry ], [ %np0, %testbb ]
+          %np1 = phi i32 [ 1, %entry ], [ %np1, %testbb ]
+          %ep0 = phi i32 [ 1, %entry ], [ %np0, %testbb ]
+          %ep1 = phi i32 [ 1, %entry ], [ %np1, %testbb ]
+          %u = add i32 %np0, %np1
+          br label %testbb
+      }
+  )", [](BasicBlock &BB, auto *ENDPN) {
+    AssertingVH<PHINode> EP0 = getPhi(BB, 2);
+    AssertingVH<PHINode> EP1 = getPhi(BB, 3);
+    EXPECT_TRUE(ENDPN(&BB, getPhiIt(BB, 2)));
+    // Expected:
+    //   %ep0 = phi i32 [ 1, %entry ], [ %ep0, %testbb ]
+    //   %ep1 = phi i32 [ 1, %entry ], [ %ep1, %testbb ]
+    //   %u = add i32 %ep0, %ep1
+    EXPECT_EQ(getNumPHIs(BB), 2);
+    Instruction &Add = *BB.getFirstNonPHIIt();
+    EXPECT_EQ(Add.getOperand(0), EP0);
+    EXPECT_EQ(Add.getOperand(1), EP1);
+  });
+}