[Attributor][FIX] Replace uses first, then values

Before we replaced value by registering all their uses. However, as we
replace a value old uses become stale. We now replace values explicitly
and keep track of "new values" when doing so to avoid replacing only
uses in stale/old values but not their replacements.

llvm/include/llvm/Transforms/IPO/Attributor.h

@@ -1429,12 +1429,16 @@ struct Attributor {
   /// uses should be changed too.
   bool changeValueAfterManifest(Value &V, Value &NV,
                                 bool ChangeDroppable = true) {
-    bool Changed = false;
-    for (auto &U : V.uses())
-      if (ChangeDroppable || !U.getUser()->isDroppable())
-        Changed |= changeUseAfterManifest(U, NV);
-
-    return Changed;
+    auto &Entry = ToBeChangedValues[&V];
+    Value *&CurNV = Entry.first;
+    if (CurNV && (CurNV->stripPointerCasts() == NV.stripPointerCasts() ||
+                  isa<UndefValue>(CurNV)))
+      return false;
+    assert((!CurNV || CurNV == &NV || isa<UndefValue>(NV)) &&
+           "Value replacement was registered twice with different values!");
+    CurNV = &NV;
+    Entry.second = ChangeDroppable;
+    return true;
   }
 
   /// Record that \p I is to be replaced with `unreachable` after information
@@ -1890,6 +1894,10 @@ private:
   /// then trivially dead instructions as well.
   DenseMap<Use *, Value *> ToBeChangedUses;
 
+  /// Values we replace with a new value after manifest is done. We will remove
+  /// then trivially dead instructions as well.
+  DenseMap<Value *, std::pair<Value *, bool>> ToBeChangedValues;
+
   /// Instructions we replace with `unreachable` insts after manifest is done.
   SmallDenseSet<WeakVH, 16> ToBeChangedToUnreachableInsts;
 

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -1355,32 +1355,39 @@ void Attributor::identifyDeadInternalFunctions() {
 ChangeStatus Attributor::cleanupIR() {
   TimeTraceScope TimeScope("Attributor::cleanupIR");
   // Delete stuff at the end to avoid invalid references and a nice order.
-  LLVM_DEBUG(dbgs() << "\n[Attributor] Delete at least "
+  LLVM_DEBUG(dbgs() << "\n[Attributor] Delete/replace at least "
                     << ToBeDeletedFunctions.size() << " functions and "
                     << ToBeDeletedBlocks.size() << " blocks and "
                     << ToBeDeletedInsts.size() << " instructions and "
+                    << ToBeChangedValues.size() << " values and "
                     << ToBeChangedUses.size() << " uses\n");
 
   SmallVector<WeakTrackingVH, 32> DeadInsts;
   SmallVector<Instruction *, 32> TerminatorsToFold;
 
-  for (auto &It : ToBeChangedUses) {
-    Use *U = It.first;
-    Value *NewV = It.second;
+  auto ReplaceUse = [&](Use *U, Value *NewV) {
     Value *OldV = U->get();
 
+    // If we plan to replace NewV we need to update it at this point.
+    do {
+      const auto &Entry = ToBeChangedValues.lookup(NewV);
+      if (!Entry.first)
+        break;
+      NewV = Entry.first;
+    } while (true);
+
     // Do not replace uses in returns if the value is a must-tail call we will
     // not delete.
     if (isa<ReturnInst>(U->getUser()))
       if (auto *CI = dyn_cast<CallInst>(OldV->stripPointerCasts()))
         if (CI->isMustTailCall() &&
             (!ToBeDeletedInsts.count(CI) || !isRunOn(*CI->getCaller())))
-          continue;
+          return;
 
     // Do not perform call graph altering changes outside the SCC.
     if (auto *CB = dyn_cast<CallBase>(U->getUser()))
       if (CB->isCallee(U) && !isRunOn(*CB->getCaller()))
-        continue;
+        return;
 
     LLVM_DEBUG(dbgs() << "Use " << *NewV << " in " << *U->getUser()
                       << " instead of " << *OldV << "\n");
@@ -1410,7 +1417,27 @@ ChangeStatus Attributor::cleanupIR() {
         TerminatorsToFold.push_back(UserI);
       }
     }
+  };
+
+  for (auto &It : ToBeChangedUses) {
+    Use *U = It.first;
+    Value *NewV = It.second;
+    ReplaceUse(U, NewV);
   }
+
+  SmallVector<Use *, 4> Uses;
+  for (auto &It : ToBeChangedValues) {
+    Value *OldV = It.first;
+    auto &Entry = It.second;
+    Value *NewV = Entry.first;
+    Uses.clear();
+    for (auto &U : OldV->uses())
+      if (Entry.second || !U.getUser()->isDroppable())
+        Uses.push_back(&U);
+    for (Use *U : Uses)
+      ReplaceUse(U, NewV);
+  }
+
   for (auto &V : InvokeWithDeadSuccessor)
     if (InvokeInst *II = dyn_cast_or_null<InvokeInst>(V)) {
       assert(isRunOn(*II->getFunction()) &&

llvm/test/Transforms/Attributor/potential.ll

@@ -633,7 +633,7 @@ define i32 @optimize_poison_1(i1 %c) {
 ; IS__TUNIT_NPM:       t:
 ; IS__TUNIT_NPM-NEXT:    ret i32 0
 ; IS__TUNIT_NPM:       f:
-; IS__TUNIT_NPM-NEXT:    ret i32 undef
+; IS__TUNIT_NPM-NEXT:    ret i32 0
 ;
 ; IS__CGSCC_OPM: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
 ; IS__CGSCC_OPM-LABEL: define {{[^@]+}}@optimize_poison_1
@@ -651,7 +651,7 @@ define i32 @optimize_poison_1(i1 %c) {
 ; IS__CGSCC_NPM:       t:
 ; IS__CGSCC_NPM-NEXT:    ret i32 0
 ; IS__CGSCC_NPM:       f:
-; IS__CGSCC_NPM-NEXT:    ret i32 undef
+; IS__CGSCC_NPM-NEXT:    ret i32 0
 ;
   br i1 %c, label %t, label %f
 t: