1 files changed, 21 insertions, 15 deletions
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5be566b..bcb182a 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -4331,7 +4331,8 @@ bool llvm::getUnderlyingObjectsForCodeGen(const Value *V,
 }
 
 static AllocaInst *
-findAllocaForValue(Value *V, DenseMap<Value *, AllocaInst *> &AllocaForValue) {
+findAllocaForValue(Value *V, DenseMap<Value *, AllocaInst *> &AllocaForValue,
+                   bool OffsetZero) {
   if (AllocaInst *AI = dyn_cast<AllocaInst>(V))
     return AI;
   // See if we've already calculated (or started to calculate) alloca for a
@@ -4344,36 +4345,41 @@ findAllocaForValue(Value *V, DenseMap<Value *, AllocaInst *> &AllocaForValue) {
   AllocaForValue[V] = nullptr;
   AllocaInst *Res = nullptr;
   if (CastInst *CI = dyn_cast<CastInst>(V))
-    Res = findAllocaForValue(CI->getOperand(0), AllocaForValue);
-  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)) {
+    Res = findAllocaForValue(CI->getOperand(0), AllocaForValue, OffsetZero);
+  else if (PHINode *PN = dyn_cast<PHINode>(V)) {
     for (Value *IncValue : PN->incoming_values()) {
       // Allow self-referencing phi-nodes.
       if (IncValue == PN)
         continue;
-      AllocaInst *IncValueAI = findAllocaForValue(IncValue, AllocaForValue);
+      AllocaInst *IncValueAI =
+          findAllocaForValue(IncValue, AllocaForValue, OffsetZero);
       // AI for incoming values should exist and should all be equal.
       if (IncValueAI == nullptr || (Res != nullptr && IncValueAI != Res))
         return nullptr;
       Res = IncValueAI;
     }
-  } else if (GetElementPtrInst *EP = dyn_cast<GetElementPtrInst>(V)) {
-    Res = findAllocaForValue(EP->getPointerOperand(), AllocaForValue);
+  } else if (auto *SI = dyn_cast<SelectInst>(V)) {
+    Res = findAllocaForValue(SI->getTrueValue(), AllocaForValue, OffsetZero);
+    if (!Res)
+      return nullptr;
+    AllocaInst *F =
+        findAllocaForValue(SI->getFalseValue(), AllocaForValue, OffsetZero);
+    if (F != Res)
+      return nullptr;
+  } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(V)) {
+    if (OffsetZero && !GEP->hasAllZeroIndices())
+      return nullptr;
+    Res = findAllocaForValue(GEP->getPointerOperand(), AllocaForValue,
+                             OffsetZero);
   }
   if (Res)
     AllocaForValue[V] = Res;
   return Res;
 }
 
-AllocaInst *llvm::findAllocaForValue(Value *V) {
+AllocaInst *llvm::findAllocaForValue(Value *V, bool OffsetZero) {
   DenseMap<Value *, AllocaInst *> AllocaForValue;
-  return ::findAllocaForValue(V, AllocaForValue);
+  return ::findAllocaForValue(V, AllocaForValue, OffsetZero);
 }
 
 static bool onlyUsedByLifetimeMarkersOrDroppableInstsHelper(