[llvm] [Inline][Cloning] Defer constant propagation after phi-nodes resolution (PR #87963)

[Antonio Frighetto via llvm-commits]

================
@@ -823,52 +804,46 @@ void llvm::CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
     }
   }
 
-  // Make a second pass over the PHINodes now that all of them have been
-  // remapped into the new function, simplifying the PHINode and performing any
-  // recursive simplifications exposed. This will transparently update the
-  // WeakTrackingVH in the VMap. Notably, we rely on that so that if we coalesce
-  // two PHINodes, the iteration over the old PHIs remains valid, and the
-  // mapping will just map us to the new node (which may not even be a PHI
-  // node).
-  const DataLayout &DL = NewFunc->getParent()->getDataLayout();
-  SmallSetVector<const Value *, 8> Worklist;
-  for (unsigned Idx = 0, Size = PHIToResolve.size(); Idx != Size; ++Idx)
-    if (isa<PHINode>(VMap[PHIToResolve[Idx]]))
-      Worklist.insert(PHIToResolve[Idx]);
-
-  // Note that we must test the size on each iteration, the worklist can grow.
-  for (unsigned Idx = 0; Idx != Worklist.size(); ++Idx) {
-    const Value *OrigV = Worklist[Idx];
-    auto *I = dyn_cast_or_null<Instruction>(VMap.lookup(OrigV));
-    if (!I)
-      continue;
-
-    // Skip over non-intrinsic callsites, we don't want to remove any nodes from
-    // the CGSCC.
-    CallBase *CB = dyn_cast<CallBase>(I);
-    if (CB && CB->getCalledFunction() &&
-        !CB->getCalledFunction()->isIntrinsic())
-      continue;
-
-    // See if this instruction simplifies.
-    Value *SimpleV = simplifyInstruction(I, DL);
-    if (!SimpleV)
-      continue;
-
-    // Stash away all the uses of the old instruction so we can check them for
-    // recursive simplifications after a RAUW. This is cheaper than checking all
-    // uses of To on the recursive step in most cases.
-    for (const User *U : OrigV->users())
-      Worklist.insert(cast<Instruction>(U));
-
-    // Replace the instruction with its simplified value.
-    I->replaceAllUsesWith(SimpleV);
+  auto GetOriginalValueInVMap = [&](const auto &I) -> Value * {
+    for (const auto &[OrigI, NewI] : VMap) {
+      if (NewI == I)
+        return const_cast<Value *>(OrigI);
+    }
+    return nullptr;
+  };
 
-    // If the original instruction had no side effects, remove it.
-    if (isInstructionTriviallyDead(I))
-      I->eraseFromParent();
-    else
-      VMap[OrigV] = I;
+  // As phi-nodes have been now remapped, allow incremental simplification of
+  // newly-cloned instructions.
+  const DataLayout &DL = NewFunc->getParent()->getDataLayout();
+  auto *NewEntryBlock = cast<BasicBlock>(VMap[StartingBB]);
+  ReversePostOrderTraversal<BasicBlock *> RPOT(NewEntryBlock);
+  for (BasicBlock *BB : RPOT) {
+    for (Instruction &NewI : llvm::make_early_inc_range(*BB)) {
+      if (isa<DbgInfoIntrinsic>(&NewI))
+        continue;
+
+      // Skip over non-intrinsic callsites, we don't want to remove any nodes
+      // from the CGSCC.
+      CallBase *CB = dyn_cast<CallBase>(&NewI);
+      if (CB && CB->getCalledFunction() &&
+          !CB->getCalledFunction()->isIntrinsic())
+        continue;
+
+      if (Value *V = simplifyInstruction(&NewI, DL)) {
+        if (!NewI.use_empty())
+          NewI.replaceAllUsesWith(V);
+
+        if (isInstructionTriviallyDead(&NewI)) {
+          NewI.eraseFromParent();
+        } else {
+          // Did not erase it? Restore the new instruction into VMap.
+          auto *OriginalI = GetOriginalValueInVMap(&NewI);
+          assert(OriginalI != nullptr &&
+                 "Previously remapped value was not found?");
+          VMap[OriginalI] = &NewI;
----------------
antoniofrighetto wrote:

`replaceAllUsesWith` will drop the old value in `VMap`, and update it with the newly-simplified one via `ValueHandleBase::ValueIsRAUWd`. If we do not erase the instruction, we have to manually restore `NewI`. I must have missed something while refactoring the code, since clearly we have to retrieve `OriginalI` through the just-replaced value `V`, not via `NewI`, as `VMap` has been changed. No need to do anything when the instruction gets erased. Fixed it now, thanks.

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