[BasicAA] Fix handling of indirect assumption based results

If a result is potentially based on a not yet proven assumption,
BasicAA will remember it inside AssumptionBasedResults and remove
the cache entry it an assumption higher up is later disproven.
However, we currently miss the case where another cache entry ends
up depending on such an AssumptionBased result.

Fix this by introducing an additional AssumptionBased state for
cache entries. If such a result is used, we'll still increment
AAQI.NumAssumptionUses, which means that the using entry will
also become AssumptionBased and be cleared if the assumption is
disproven.

At the end of the root query, convert remaining AssumptionBased
results into definitive results.

Fixes #98978.

llvm/include/llvm/Analysis/AliasAnalysis.h

@@ -244,12 +244,23 @@ class AAQueryInfo {
 public:
   using LocPair = std::pair<AACacheLoc, AACacheLoc>;
   struct CacheEntry {
+    /// Cache entry is neither an assumption nor does it use a (non-definitive)
+    /// assumption.
+    static constexpr int Definitive = -2;
+    /// Cache entry is not an assumption itself, may be using an assumption
+    /// from higher up the stack.
+    static constexpr int AssumptionBased = -1;
+
     AliasResult Result;
-    /// Number of times a NoAlias assumption has been used.
-    /// 0 for assumptions that have not been used, -1 for definitive results.
+    /// Number of times a NoAlias assumption has been used, 0 for assumptions
+    /// that have not been used. Can also take one of the Definitive or
+    /// AssumptionBased values documented above.
     int NumAssumptionUses;
+
     /// Whether this is a definitive (non-assumption) result.
-    bool isDefinitive() const { return NumAssumptionUses < 0; }
+    bool isDefinitive() const { return NumAssumptionUses == Definitive; }
+    /// Whether this is an assumption that has not been proven yet.
+    bool isAssumption() const { return NumAssumptionUses >= 0; }
   };
 
   // Alias analysis result aggregration using which this query is performed.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -1692,9 +1692,12 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
   if (!Pair.second) {
     auto &Entry = Pair.first->second;
     if (!Entry.isDefinitive()) {
-      // Remember that we used an assumption.
-      ++Entry.NumAssumptionUses;
+      // Remember that we used an assumption. This may either be a direct use
+      // of an assumption, or a use of an entry that may itself be based on an
+      // assumption.
       ++AAQI.NumAssumptionUses;
+      if (Entry.isAssumption())
+        ++Entry.NumAssumptionUses;
     }
     // Cache contains sorted {V1,V2} pairs but we should return original order.
     auto Result = Entry.Result;
@@ -1722,7 +1725,6 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
   Entry.Result = Result;
   // Cache contains sorted {V1,V2} pairs.
   Entry.Result.swap(Swapped);
-  Entry.NumAssumptionUses = -1;
 
   // If the assumption has been disproven, remove any results that may have
   // been based on this assumption. Do this after the Entry updates above to
@@ -1734,8 +1736,26 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
   // The result may still be based on assumptions higher up in the chain.
   // Remember it, so it can be purged from the cache later.
   if (OrigNumAssumptionUses != AAQI.NumAssumptionUses &&
-      Result != AliasResult::MayAlias)
+      Result != AliasResult::MayAlias) {
     AAQI.AssumptionBasedResults.push_back(Locs);
+    Entry.NumAssumptionUses = AAQueryInfo::CacheEntry::AssumptionBased;
+  } else {
+    Entry.NumAssumptionUses = AAQueryInfo::CacheEntry::Definitive;
+  }
+
+  // Depth is incremented before this function is called, so Depth==1 indicates
+  // a root query.
+  if (AAQI.Depth == 1) {
+    // Any remaining assumption based results must be based on proven
+    // assumptions, so convert them to definitive results.
+    for (const auto &Loc : AAQI.AssumptionBasedResults) {
+      auto It = AAQI.AliasCache.find(Loc);
+      if (It != AAQI.AliasCache.end())
+        It->second.NumAssumptionUses = AAQueryInfo::CacheEntry::Definitive;
+    }
+    AAQI.AssumptionBasedResults.clear();
+    AAQI.NumAssumptionUses = 0;
+  }
   return Result;
 }
 

llvm/test/Transforms/SLPVectorizer/X86/pr98978.ll

@@ -3,27 +3,28 @@
 
 target triple = "x86_64-redhat-linux-gnu"
 
+; Should get vectorized.
 define void @test(ptr %arg, i64 %arg1, i64 %arg2) {
 ; CHECK-LABEL: define void @test(
 ; CHECK-SAME: ptr [[ARG:%.*]], i64 [[ARG1:%.*]], i64 [[ARG2:%.*]]) {
-; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    br label %[[LOOP1:.*]]
-; CHECK:       [[LOOP1]]:
-; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[I21:%.*]], %[[BB20:.*]] ], [ null, %[[ENTRY]] ]
+; CHECK-NEXT:  [[_PREHEADER48_PREHEADER:.*]]:
+; CHECK-NEXT:    br label %[[DOTLOOPEXIT49:.*]]
+; CHECK:       [[DOTLOOPEXIT49]]:
+; CHECK-NEXT:    [[I:%.*]] = phi ptr [ [[I21:%.*]], %[[BB20:.*]] ], [ null, %[[_PREHEADER48_PREHEADER]] ]
 ; CHECK-NEXT:    br i1 false, label %[[BB22:.*]], label %[[DOTPREHEADER48_PREHEADER:.*]]
-; CHECK:       [[_PREHEADER48_PREHEADER:.*:]]
-; CHECK-NEXT:    br [[DOTLOOPEXIT49:label %.*]]
-; CHECK:       [[DEAD:.*]]:
-; CHECK-NEXT:    br [[DOTLOOPEXIT49]]
-; CHECK:       [[_LOOPEXIT49:.*:]]
-; CHECK-NEXT:    [[I5:%.*]] = phi ptr [ [[I]], %[[DEAD]] ], [ [[I]], %[[DOTPREHEADER48_PREHEADER]] ]
-; CHECK-NEXT:    br label %[[DOTPREHEADER48_PREHEADER_1:.*]]
+; CHECK:       [[_PREHEADER48_PREHEADER1:.*:]]
+; CHECK-NEXT:    br [[DOTLOOPEXIT50:label %.*]]
+; CHECK:       [[_LOOPEXIT49:.*]]:
+; CHECK-NEXT:    br [[DOTLOOPEXIT50]]
 ; CHECK:       [[_PREHEADER48_PREHEADER_1:.*:]]
-; CHECK-NEXT:    br [[DOTLOOPEXIT49_1:label %.*]]
+; CHECK-NEXT:    [[I5:%.*]] = phi ptr [ [[I]], %[[_LOOPEXIT49]] ], [ [[I]], %[[DOTPREHEADER48_PREHEADER]] ]
+; CHECK-NEXT:    br label %[[DOTLOOPEXIT49_1:.*]]
 ; CHECK:       [[_LOOPEXIT42_1:.*:]]
-; CHECK-NEXT:    br i1 false, [[DOTLOOPEXIT49_1]], label %[[BB20]]
+; CHECK-NEXT:    br [[DOTLOOPEXIT49_2:label %.*]]
 ; CHECK:       [[_LOOPEXIT49_1:.*:]]
-; CHECK-NEXT:    [[I6:%.*]] = phi ptr [ [[I5]], [[DOTLOOPEXIT42_1:%.*]] ], [ [[I5]], %[[DOTPREHEADER48_PREHEADER_1]] ]
+; CHECK-NEXT:    br i1 false, [[DOTLOOPEXIT49_2]], label %[[BB20]]
+; CHECK:       [[_LOOPEXIT49_2:.*:]]
+; CHECK-NEXT:    [[I6:%.*]] = phi ptr [ [[I5]], [[DOTLOOPEXIT42_1:%.*]] ], [ [[I5]], %[[DOTLOOPEXIT49_1]] ]
 ; CHECK-NEXT:    [[I7:%.*]] = getelementptr inbounds i8, ptr [[I6]], i64 [[ARG1]]
 ; CHECK-NEXT:    br label %[[BB8:.*]]
 ; CHECK:       [[BB8]]:
@@ -36,13 +37,16 @@ define void @test(ptr %arg, i64 %arg1, i64 %arg2) {
 ; CHECK:       [[BB13]]:
 ; CHECK-NEXT:    [[I14:%.*]] = phi ptr [ [[ARG]], %[[BB13]] ], [ [[I11]], %[[BB10]] ]
 ; CHECK-NEXT:    [[I15:%.*]] = phi ptr [ null, %[[BB13]] ], [ [[I6]], %[[BB10]] ]
-; CHECK-NEXT:    [[TMP0:%.*]] = load <2 x i64>, ptr [[I14]], align 1
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <2 x i64> [[TMP0]], <2 x i64> poison, <2 x i32> <i32 1, i32 0>
-; CHECK-NEXT:    store <2 x i64> [[TMP1]], ptr [[I15]], align 1
+; CHECK-NEXT:    [[I16:%.*]] = getelementptr inbounds i8, ptr [[I14]], i64 8
+; CHECK-NEXT:    [[I17:%.*]] = load i64, ptr [[I16]], align 1
+; CHECK-NEXT:    store i64 [[I17]], ptr [[I15]], align 1
+; CHECK-NEXT:    [[I18:%.*]] = getelementptr inbounds i8, ptr [[I15]], i64 8
+; CHECK-NEXT:    [[I19:%.*]] = load i64, ptr [[I14]], align 1
+; CHECK-NEXT:    store i64 [[I19]], ptr [[I18]], align 1
 ; CHECK-NEXT:    br i1 false, label %[[BB13]], label %[[BB20]]
 ; CHECK:       [[BB20]]:
 ; CHECK-NEXT:    [[I21]] = phi ptr [ [[I5]], [[DOTLOOPEXIT42_1]] ], [ [[I6]], %[[BB13]] ]
-; CHECK-NEXT:    br label %[[LOOP1]]
+; CHECK-NEXT:    br label %[[DOTLOOPEXIT49]]
 ; CHECK:       [[BB22]]:
 ; CHECK-NEXT:    [[I23:%.*]] = getelementptr inbounds i8, ptr [[I]], i64 [[ARG2]]
 ; CHECK-NEXT:    [[I25:%.*]] = getelementptr inbounds i8, ptr [[I23]], i64 8