[NFC][ValueTracking] Add OffsetZero into findAllocaForValue

[vitalybuka]

Summary

For StackLifetime after finding alloca we need to check that
values ponting to the begining of alloca.

[vitalybuka]

rebase

[vitalybuka]

rebase

[eugenis]

LGTM

llvm/lib/Analysis/ValueTracking.cpp

4371

Please address the lint warning.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

4375

You can just use GEP->hasAllZeroIndices() here.

[vitalybuka]

comments

llvm/include/llvm/Analysis/ValueTracking.h

void getUnderlyingObjects(const Value *V,
                             SmallVectorImpl<const Value *> &Objects,
                             LoopInfo *LI = nullptr, unsigned MaxLookup = 6);
 
   /// This is a wrapper around getUnderlyingObjects and adds support for basic
   /// ptrtoint+arithmetic+inttoptr sequences.
   bool getUnderlyingObjectsForCodeGen(const Value *V,
                                       SmallVectorImpl<Value *> &Objects);
 
-  /// Finds alloca where the value comes from.
-  AllocaInst *findAllocaForValue(Value *V);
-  inline const AllocaInst *findAllocaForValue(const Value *V) {
-    return findAllocaForValue(const_cast<Value *>(V));
+  /// Returns unique alloca where the value comes from, or nullptr.
+  /// If OffsetZero is true check that V points to the begining of the alloca.
+  AllocaInst *findAllocaForValue(Value *V, bool OffsetZero = false);
+  inline const AllocaInst *findAllocaForValue(const Value *V,
+                                              bool OffsetZero = false) {
+    return findAllocaForValue(const_cast<Value *>(V), OffsetZero);
   }
 
   /// Return true if the only users of this pointer are lifetime markers.
   bool onlyUsedByLifetimeMarkers(const Value *V);
 
   /// Return true if the only users of this pointer are lifetime markers or
   /// droppable instructions.
   bool onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V);

llvm/lib/Analysis/ValueTracking.cpp

for (const Value *V : Objs) {
       }
       Objects.push_back(const_cast<Value *>(V));
     }
   } while (!Working.empty());
   return true;
 }
 
 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
   // given value.
   auto I = AllocaForValue.find(V);
   if (I != AllocaForValue.end())
     return I->second;
   // Store 0 while we're calculating alloca for value V to avoid
   // infinite recursion if the value references itself.
   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) {
+      for (unsigned i = 1; i < GEP->getNumOperands(); ++i) {
+        ConstantInt *OpC = dyn_cast<ConstantInt>(GEP->getOperand(i));
+        if (!OpC || !OpC->isZero())
+          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(
     const Value *V, bool AllowLifetime, bool AllowDroppable) {
   for (const User *U : V->users()) {
     const IntrinsicInst *II = dyn_cast<IntrinsicInst>(U);
     if (!II)
       return false;

llvm/unittests/Analysis/ValueTrackingTest.cpp

define i32 @test(i32 %x.1, i32 %x.2) {
     // passing Depth == MaxDepth == 6.
     ConstantRange CR2 = computeConstantRange(X2, true, &AC, I, 6);
     EXPECT_TRUE(CR2.isFullSet());
   }
 }
 
 struct FindAllocaForValueTestParams {
   const char *IR;
-  bool Result;
+  bool AnyOffsetResult;
+  bool ZeroOffsetResult;
 };
 
 class FindAllocaForValueTest
     : public ValueTrackingTest,
       public ::testing::WithParamInterface<FindAllocaForValueTestParams> {
 protected:
 };
 
 const FindAllocaForValueTestParams FindAllocaForValueTests[] = {
     {R"(
       define void @test() {
         %a = alloca i64
         %r = bitcast i64* %a to i32*
         ret void
       })",
-     true},
+     true, true},
 
     {R"(
       define void @test() {
         %a = alloca i32
         %r = getelementptr i32, i32* %a, i32 1
         ret void
       })",
-     true},
+     true, false},
+
+    {R"(
+      define void @test() {
+        %a = alloca i32
+        %r = getelementptr i32, i32* %a, i32 0
+        ret void
+      })",
+     true, true},
 
     {R"(
       define void @test(i1 %cond) {
       entry:
         %a = alloca i32
         br label %bb1
 
       bb1:
         %r = phi i32* [ %a, %entry ], [ %r, %bb1 ]
         br i1 %cond, label %bb1, label %exit
 
       exit:
         ret void
       })",
-     true},
+     true, true},
 
     {R"(
       define void @test(i1 %cond) {
         %a = alloca i32
         %r = select i1 %cond, i32* %a, i32* %a
         ret void
       })",
-     true},
+     true, true},
 
     {R"(
       define void @test(i1 %cond) {
         %a = alloca i32
         %b = alloca i32
         %r = select i1 %cond, i32* %a, i32* %b
         ret void
       })",
-     false},
+     false, false},
 
     {R"(
       define void @test(i1 %cond) {
       entry:
         %a = alloca i64
         %a32 = bitcast i64* %a to i32*
         br label %bb1
 
       bb1:
         %x = phi i32* [ %a32, %entry ], [ %x, %bb1 ]
         %r = getelementptr i32, i32* %x, i32 1
         br i1 %cond, label %bb1, label %exit
 
       exit:
         ret void
       })",
-     true},
+     true, false},
 
     {R"(
       define void @test(i1 %cond) {
       entry:
         %a = alloca i64
         %a32 = bitcast i64* %a to i32*
         br label %bb1
 
       bb1:
         %x = phi i32* [ %a32, %entry ], [ %r, %bb1 ]
         %r = getelementptr i32, i32* %x, i32 1
         br i1 %cond, label %bb1, label %exit
 
       exit:
         ret void
       })",
-     false},
+     false, false},
 
     {R"(
       define void @test(i1 %cond, i64* %a) {
       entry:
         %r = bitcast i64* %a to i32*
         ret void
       })",
-     false},
+     false, false},
 
     {R"(
       define void @test(i1 %cond) {
       entry:
         %a = alloca i32
         %b = alloca i32
         br label %bb1
 
       bb1:
         %r = phi i32* [ %a, %entry ], [ %b, %bb1 ]
         br i1 %cond, label %bb1, label %exit
 
       exit:
         ret void
       })",
-     false},
+     false, false},
 };
 
 TEST_P(FindAllocaForValueTest, findAllocaForValue) {
   auto M = parseModule(GetParam().IR);
   Function *F = M->getFunction("test");
   Instruction *I = &findInstructionByName(F, "r");
   const AllocaInst *AI = findAllocaForValue(I);
-  EXPECT_EQ(!!AI, GetParam().Result);
+  EXPECT_EQ(!!AI, GetParam().AnyOffsetResult);
+}
+
+TEST_P(FindAllocaForValueTest, findAllocaForValueZeroOffset) {
+  auto M = parseModule(GetParam().IR);
+  Function *F = M->getFunction("test");
+  Instruction *I = &findInstructionByName(F, "r");
+  const AllocaInst *AI = findAllocaForValue(I, true);
+  EXPECT_EQ(!!AI, GetParam().ZeroOffsetResult);
 }
 
 INSTANTIATE_TEST_CASE_P(FindAllocaForValueTest, FindAllocaForValueTest,
                         ::testing::ValuesIn(FindAllocaForValueTests), );