From 894fd2ecc4cd1a5d4098eb1f9e5d0ff9f8a64041 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks@google.com>
Date: Tue, 29 Sep 2020 22:29:26 -0700
Subject: [PATCH] [DAE] MarkLive in MarkValue(MaybeLive) if any use is live

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.

Reviewed By: rnk

Differential Revision: https://reviews.llvm.org/D88529
---
 .../Transforms/IPO/DeadArgumentElimination.h  |  1 +
 .../IPO/DeadArgumentElimination.cpp           | 29 +++++++++++------
 .../DeadArgElim/preserve-used-ret.ll          | 32 +++++++++++++++++++
 3 files changed, 53 insertions(+), 9 deletions(-)
 create mode 100644 llvm/test/Transforms/DeadArgElim/preserve-used-ret.ll

diff --git a/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h b/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h
index 73797bc10017c..496ceea12bc96 100644
--- a/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h
+++ b/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h
@@ -128,6 +128,7 @@ class DeadArgumentEliminationPass
   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);
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 54c51b6e7161b..f2588938d9646 100644
--- a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -357,7 +357,7 @@ 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
@@ -657,10 +657,18 @@ void DeadArgumentEliminationPass::MarkValue(const RetOrArg &RA, Liveness L,
       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));
+      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);
+          break;
+        } else {
+          // Note any uses of this value, so this value can be
+          // marked live whenever one of the uses becomes live.
+          Uses.insert(std::make_pair(MaybeLiveUse, RA));
+        }
+      }
       break;
   }
 }
@@ -686,17 +694,20 @@ void DeadArgumentEliminationPass::MarkLive(const Function &F) {
 /// 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) {
diff --git a/llvm/test/Transforms/DeadArgElim/preserve-used-ret.ll b/llvm/test/Transforms/DeadArgElim/preserve-used-ret.ll
new file mode 100644
index 0000000000000..f0c2649fdb393
--- /dev/null
+++ b/llvm/test/Transforms/DeadArgElim/preserve-used-ret.ll
@@ -0,0 +1,32 @@
+; 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
+}
+