From 469debe0275b5800a2231ae670d235a83e6fbb27 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Thu, 27 Aug 2020 02:01:26 -0700
Subject: [PATCH] [ValueTracking] Support select in findAllocaForValue

---
 llvm/lib/Analysis/ValueTracking.cpp           |  9 ++++++++-
 llvm/unittests/Analysis/ValueTrackingTest.cpp | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 36998cd90697..5be566b693f4 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -4345,7 +4345,14 @@ findAllocaForValue(Value *V, DenseMap<Value *, AllocaInst *> &AllocaForValue) {
   AllocaInst *Res = nullptr;
   if (CastInst *CI = dyn_cast<CastInst>(V))
     Res = findAllocaForValue(CI->getOperand(0), AllocaForValue);
-  else if (PHINode *PN = dyn_cast<PHINode>(V)) {
+  else if (auto *SI = dyn_cast<SelectInst>(V)) {
+    Res = findAllocaForValue(SI->getTrueValue(), AllocaForValue);
+    if (!Res)
+      return nullptr;
+    AllocaInst *F = findAllocaForValue(SI->getFalseValue(), AllocaForValue);
+    if (F != Res)
+      return nullptr;
+  } else if (PHINode *PN = dyn_cast<PHINode>(V)) {
     for (Value *IncValue : PN->incoming_values()) {
       // Allow self-referencing phi-nodes.
       if (IncValue == PN)
diff --git a/llvm/unittests/Analysis/ValueTrackingTest.cpp b/llvm/unittests/Analysis/ValueTrackingTest.cpp
index 732f71e4de90..7951d6dec1f8 100644
--- a/llvm/unittests/Analysis/ValueTrackingTest.cpp
+++ b/llvm/unittests/Analysis/ValueTrackingTest.cpp
@@ -1469,6 +1469,23 @@ const std::pair<const char *, bool> FindAllocaForValueTests[] = {
       })",
      true},
 
+    {R"(
+      define void @test(i1 %cond) {
+        %a = alloca i32
+        %r = select i1 %cond, i32* %a, i32* %a
+        ret void
+      })",
+     true},
+
+    {R"(
+      define void @test(i1 %cond) {
+        %a = alloca i32
+        %b = alloca i32
+        %r = select i1 %cond, i32* %a, i32* %b
+        ret void
+      })",
+     false},
+
     {R"(
       define void @test(i1 %cond) {
       entry: