Ensure that we don't update Uses after MarkLive

[Aaron1011]

Summary

Fixes bug #47444

Previously, we would call MarkValue for each RetOrArg, either calling
MarkLive or inserting dependencies into Uses. However, this can cause us
to end up calling MarkLive for a RetOrArg, and then later inserting
additional dependencies into Uses. As a result, the liveness propagation
during MarkLive will miss some dependencies, causing us to incorrectly
see a RetOrArg as dead.

I've fixed this by replacing MarkValue with a new function MarkValues,
which takes in a vector of RetOrArgs to inspect. It works in two passes
- we first insert all dependencies into Uses, and then call MarkLive for
any live values. This ensures that liveness is correctly propagated to
all dependencies.

To ensure that this type of issue doesn't occur again, I've added an
assertion to verify that we don't try to insert dependencies into Uses
for a value that has been already marked live.

I've included a minimized version of the IR from bug #47444 as a test
case. Without this pass, return value #1 is incorrectly removed from
the function 'g'.

[Aaron1011]

The pre-merge check seems to want markValues instead of MarkValues, even though the method is curretly 'MarkValue'. Should I change this?

[rnk]

Arthur, do you mind taking a look?

[aeubanks]

Yep I'll take a look.

[aeubanks]

https://reviews.llvm.org/D88529 seems like a nicer solution, wdyt?

llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h

public:
   /// This set contains all values that have been determined to be live.
   LiveSet LiveValues;
 
   /// This set contains all values that are cannot be changed in any way.
   LiveFuncSet LiveFunctions;
 
   using UseVector = SmallVector<RetOrArg, 5>;
 
+  using Livenesses = SmallVector<Liveness, 5>;
+  using RetOrArgUses = SmallVector<UseVector, 5>;
+
   /// This allows this pass to do double-duty as the dead arg hacking pass
   /// (used only by bugpoint).
   bool ShouldHackArguments = false;
 
 private:
   Liveness MarkIfNotLive(RetOrArg Use, UseVector &MaybeLiveUses);
   Liveness SurveyUse(const Use *U, UseVector &MaybeLiveUses,
                      unsigned RetValNum = -1U);
   Liveness SurveyUses(const Value *V, UseVector &MaybeLiveUses);
 
   void SurveyFunction(const Function &F);
-  void MarkValue(const RetOrArg &RA, Liveness L,
-                 const UseVector &MaybeLiveUses);
+  void MarkValues(const UseVector &RetOrArgs, Livenesses &L,
+                  const RetOrArgUses &RetOrArgUses);
   void MarkLive(const RetOrArg &RA);
   void MarkLive(const Function &F);
   void PropagateLiveness(const RetOrArg &RA);
   bool RemoveDeadStuffFromFunction(Function *F);
   bool DeleteDeadVarargs(Function &Fn);
   bool RemoveDeadArgumentsFromCallers(Function &Fn);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_DEADARGUMENTELIMINATION_H

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

void DeadArgumentEliminationPass::SurveyFunction(const Function &F) {
   if (F.hasFnAttribute(Attribute::Naked)) {
     MarkLive(F);
     return;
   }
 
   unsigned RetCount = NumRetVals(&F);
 
   // Assume all return values are dead
-  using RetVals = SmallVector<Liveness, 5>;
-
-  RetVals RetValLiveness(RetCount, MaybeLive);
-
-  using RetUses = SmallVector<UseVector, 5>;
+  Livenesses RetValLiveness(RetCount, MaybeLive);
+  UseVector RetValues;
 
   // These vectors map each return value to the uses that make it MaybeLive, so
   // we can add those to the Uses map if the return value really turns out to be
   // MaybeLive. Initialized to a list of RetCount empty lists.
-  RetUses MaybeLiveRetUses(RetCount);
+  RetOrArgUses MaybeLiveRetUses(RetCount);
 
   bool HasMustTailCalls = false;
 
   for (Function::const_iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
     if (const ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
       if (RI->getNumOperands() != 0 && RI->getOperand(0)->getType()
           != F.getFunctionType()->getReturnType()) {
         // We don't support old style multiple return values.
void DeadArgumentEliminationPass::SurveyFunction(const Function &F) {
   }
 
   if (HasMustTailCallers) {
     LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - " << F.getName()
                       << " has musttail callers\n");
   }
 
   // Now we've inspected all callers, record the liveness of our return values.
-  for (unsigned Ri = 0; Ri != RetCount; ++Ri)
-    MarkValue(CreateRet(&F, Ri), RetValLiveness[Ri], MaybeLiveRetUses[Ri]);
+  for (unsigned Ri = 0; Ri != RetCount; ++Ri) {
+    RetValues.push_back(CreateRet(&F, Ri));
+  }
+  MarkValues(RetValues, RetValLiveness, MaybeLiveRetUses);
 
   LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Inspecting args for fn: "
                     << F.getName() << "\n");
 
   // Now, check all of our arguments.
   unsigned ArgI = 0;
-  UseVector MaybeLiveArgUses;
+  UseVector ArgValues;
+  RetOrArgUses MaybeLiveArgUses;
+  Livenesses ArgLivenesses;
   for (Function::const_arg_iterator AI = F.arg_begin(), E = F.arg_end();
        AI != E; ++AI, ++ArgI) {
     Liveness Result;
+    UseVector ArgUses;
     if (F.getFunctionType()->isVarArg() || HasMustTailCallers ||
         HasMustTailCalls) {
       // Variadic functions will already have a va_arg function expanded inside
       // them, making them potentially very sensitive to ABI changes resulting
       // from removing arguments entirely, so don't. For example AArch64 handles
       // register and stack HFAs very differently, and this is reflected in the
       // IR which has already been generated.
       //
       // `musttail` calls to this function restrict argument removal attempts.
       // The signature of the caller must match the signature of the function.
       //
       // `musttail` calls in this function prevents us from changing its
       // signature
       Result = Live;
     } else {
       // See what the effect of this use is (recording any uses that cause
       // MaybeLive in MaybeLiveArgUses).
-      Result = SurveyUses(&*AI, MaybeLiveArgUses);
+      Result = SurveyUses(&*AI, ArgUses);
     }
 
-    // Mark the result.
-    MarkValue(CreateArg(&F, ArgI), Result, MaybeLiveArgUses);
-    // Clear the vector again for the next iteration.
-    MaybeLiveArgUses.clear();
+    ArgValues.push_back(std::move(CreateArg(&F, ArgI)));
+    ArgLivenesses.push_back(Result);
+    MaybeLiveArgUses.push_back(ArgUses);
   }
+  MarkValues(ArgValues, ArgLivenesses, MaybeLiveArgUses);
 }
 
-/// MarkValue - This function marks the liveness of RA depending on L. If L is
-/// MaybeLive, it also takes all uses in MaybeLiveUses and records them in Uses,
-/// such that RA will be marked live if any use in MaybeLiveUses gets marked
-/// live later on.
-void DeadArgumentEliminationPass::MarkValue(const RetOrArg &RA, Liveness L,
-                                            const UseVector &MaybeLiveUses) {
+/// MarkValues - This function marks the liveness of each value in RetOrArgs
+/// depending on their corresponding livness L in Lives.
+/// This is done in two passes:
+//
+/// First, all RA values that have an L of MaybeLive have their corresponding
+/// uses in RetOrArgUses recorded in Uses, such that RA will be marked live if
+/// any use in MaybeLiveUses gets marked live later on.
+//
+// Second, we call MarkLive for all RA values that have an L of Live.
+// This is done as a separate step so that any MaybeLive dependencies have
+// been recorded in Uses, to be read by MarkLive.
+void DeadArgumentEliminationPass::MarkValues(const UseVector &RetOrArgs,
+                                             Livenesses &Lives,
+                                             const RetOrArgUses &RetOrArgUses) {
+  for (unsigned i = 0; i < RetOrArgs.size(); i++) {
+    const auto RA = RetOrArgs[i];
+    const auto L = Lives[i];
+    const auto MaybeLiveUses = RetOrArgUses[i];
+
     switch (L) {
+    // Don't do anything yet - we will mark these as live in the
+    // second pass, after we have recorded all MaybeLive dependencies
+    // in Uses
     case Live:
-      MarkLive(RA);
       break;
     case MaybeLive:
       // Note any uses of this value, so this return value can be
       // marked live whenever one of the uses becomes live.
-      for (const auto &MaybeLiveUse : MaybeLiveUses)
-        Uses.insert(std::make_pair(MaybeLiveUse, RA));
+      for (const auto &Use : MaybeLiveUses) {
+        assert(!(LiveFunctions.count(Use.F) || LiveValues.count(Use)) &&
+               "Use is already live!");
+        Uses.insert(std::make_pair(Use, RA));
+      }
       break;
     }
   }
 
+  for (unsigned i = 0; i < RetOrArgs.size(); i++) {
+    const auto RA = RetOrArgs[i];
+    const auto L = Lives[i];
+
+    switch (L) {
+    case Live:
+      MarkLive(RA);
+      break;
+    // Handled in the first pass
+    case MaybeLive:
+      break;
+    }
+  }
+}
+
 /// MarkLive - Mark the given Function as alive, meaning that it cannot be
 /// changed in any way. Additionally,
 /// mark any values that are used as this function's parameters or by its return
 /// values (according to Uses) live as well.
 void DeadArgumentEliminationPass::MarkLive(const Function &F) {
   LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Intrinsically live fn: "
                     << F.getName() << "\n");
   // Mark the function as live.

llvm/test/Transforms/DeadArgElim/preserve-used-ret.ll

+; RUN: opt -S -deadargelim %s | FileCheck %s
+
+define internal { i64, i64 } @f(i64 %a, i64 %b) {
+start:
+  %0 = insertvalue { i64, i64 } undef, i64 %a, 0
+  %1 = insertvalue { i64, i64 } %0, i64 %b, 1
+  ret { i64, i64 } %1
+}
+
+; Check that we don't delete either of g's return values
+
+; CHECK-LABEL: define internal { i64, i64 } @g(i64 %a, i64 %b)
+define internal { i64, i64 } @g(i64 %a, i64 %b) {
+start:
+  %0 = call { i64, i64 } @f(i64 %a, i64 %b)
+  ret { i64, i64 } %0
+}
+
+declare dso_local i32 @test(i64, i64)
+
+define i32 @main(i32 %argc, i8** %argv) {
+start:
+  %x = call { i64, i64 } @g(i64 13, i64 42)
+  %x.0 = extractvalue { i64, i64 } %x, 0
+  %x.1 = extractvalue { i64, i64 } %x, 1
+  %z = bitcast i64 %x.0 to i64
+  %y = call { i64, i64 } @f(i64 %x.0, i64 %x.1)
+  %y.1 = extractvalue { i64, i64 } %y, 1
+  %0 = call i32 @test(i64 %x.0, i64 %y.1)
+  ret i32 %0
+}