[DAE] MarkLive in MarkValue(MaybeLive) if any use is live

[aeubanks]

Summary

While looping through all args or all return values, we may mark a use
of a later iteration as live. Previously when we got to that later value
it would ignore that and continue adding to Uses instead of marking it
live. For example, when looping through arg#0 and arg#1,
MarkValue(arg#0, Live) may cause some use of arg#1 to be live, but
MarkValue(arg#1, MaybeLive) will not notice that and continue adding
into Uses.

Now MarkValue(RA, MaybeLive) will MarkLive(RA) if any use is live.

Fixes PR47444.

[aeubanks]

break out of loop after MarkLive()

[Aaron1011]

This looks good to me, but someone who's more familiar with this code should also take a look.

[rnk]

lgtm

I am not an expert in this pass, but I think I understand it well enough to approve the change. A use may become live after it is added to the MaybeLiveUses list, so the code has to check again later. Something like that, anyway.

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

 
 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);
+  bool IsLive(const RetOrArg &RA);
   void MarkValue(const RetOrArg &RA, Liveness L,
                  const UseVector &MaybeLiveUses);
   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

 
 /// MarkIfNotLive - This checks Use for liveness in LiveValues. If Use is not
 /// live, it adds Use to the MaybeLiveUses argument. Returns the determined
 /// liveness of Use.
 DeadArgumentEliminationPass::Liveness
 DeadArgumentEliminationPass::MarkIfNotLive(RetOrArg Use,
                                            UseVector &MaybeLiveUses) {
   // We're live if our use or its Function is already marked as live.
-  if (LiveFunctions.count(Use.F) || LiveValues.count(Use))
+  if (IsLive(Use))
     return Live;
 
   // We're maybe live otherwise, but remember that we must become live if
   // Use becomes live.
   MaybeLiveUses.push_back(Use);
   return MaybeLive;
 }
 

 /// live later on.
 void DeadArgumentEliminationPass::MarkValue(const RetOrArg &RA, Liveness L,
                                             const UseVector &MaybeLiveUses) {
   switch (L) {
     case Live:
       MarkLive(RA);
       break;
     case MaybeLive:
-      // Note any uses of this value, so this return value can be
+      assert(!IsLive(RA) && "Use is already live!");
+      for (const auto &MaybeLiveUse : MaybeLiveUses) {
+        if (IsLive(MaybeLiveUse)) {
+          // A use is live, so this value is live.
+          MarkLive(RA);
+        } else {
+          // Note any uses of this value, so this value can be
           // marked live whenever one of the uses becomes live.
-      for (const auto &MaybeLiveUse : MaybeLiveUses)
           Uses.insert(std::make_pair(MaybeLiveUse, RA));
+        }
+      }
       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) {
   for (unsigned Ri = 0, E = NumRetVals(&F); Ri != E; ++Ri)
     PropagateLiveness(CreateRet(&F, Ri));
 }
 
 /// MarkLive - Mark the given return value or argument as live. Additionally,
 /// mark any values that are used by this value (according to Uses) live as
 /// well.
 void DeadArgumentEliminationPass::MarkLive(const RetOrArg &RA) {
-  if (LiveFunctions.count(RA.F))
-    return; // Function was already marked Live.
+  if (IsLive(RA))
+    return; // Already marked Live.
 
-  if (!LiveValues.insert(RA).second)
-    return; // We were already marked Live.
+  LiveValues.insert(RA);
 
   LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Marking "
                     << RA.getDescription() << " live\n");
   PropagateLiveness(RA);
 }
 
+bool DeadArgumentEliminationPass::IsLive(const RetOrArg &RA) {
+  return LiveFunctions.count(RA.F) || LiveValues.count(RA);
+}
+
 /// PropagateLiveness - Given that RA is a live value, propagate it's liveness
 /// to any other values it uses (according to Uses).
 void DeadArgumentEliminationPass::PropagateLiveness(const RetOrArg &RA) {
   // We don't use upper_bound (or equal_range) here, because our recursive call
   // to ourselves is likely to cause the upper_bound (which is the first value
   // not belonging to RA) to become erased and the iterator invalidated.
   UseMap::iterator Begin = Uses.lower_bound(RA);
   UseMap::iterator E = Uses.end();

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
+}
+