[AA] Cache (optionally) estimated PartialAlias offsets.

[dfukalov]

Summary

For the cases of two clobbering loads and one loaded object is fully contained
in the second BasicAAResult::aliasGEP returns just PartialAlias that
is actually more common case of partial overlap, it doesn't say anything about
actual overlapping sizes.

AA users such as GVN and DSE have no functionality to estimate aliasing of GEPs
with non-constant offsets. The change stores estimated relative offsets so they
can be used further.

[nikic]

llvm/lib/Analysis/AliasAnalysis.cpp

971

Why is it okay to treat the offset symmetrically? I'd expect the offset either be negative if swapped, or be not defined at all, depending on the contract of this function is supposed to be (probably needs a doc comment).

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1435

Why do you use a separate VisitedPhi flag, rather than checking !VisitedPhiBBs.empty(), as we do elsewhere?

This kind of flag breaks BatchAA, because it will leak across queries.

llvm/lib/Analysis/CaptureTracking.cpp

417 ↗

(On Diff #312734)

This looks unrelated, land separately?

[dfukalov]

Rebased, change reworked as requested.

[dfukalov]

llvm/lib/Analysis/AliasAnalysis.cpp

971

My thought was to mark "undefined" offset as -1 and store pointers in pairs in "positive offset" order. It was based on GVN ability to use positive offset only for clobbering.
But you're right, it is more effective to store pointers pairs in cache in the same way as it is implemented for AliasCache.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1435

Thanks, good catch, missed this leak.

[dfukalov]

Ping...

[dfukalov]

Ping...

[dfukalov]

Ping. I'm going to submit the change on Friday if no objections are observed.

[lebedev.ri]

This should probably be split up into BasicAA (+unit test), MemoryDependenceAnalysis(?), and GVN patches.

[fhahn]

IIUC this is similar to the handling of partial overwrites in DSE (https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp#L548)?

If so, could we just generalize the helper there, rather than threading this through the different caches from BasicAA to GVN? Yes, we have to do some extra work in GVN, but there would be no need to maintain a cached value not used by other clients, it will work with any AA implementations returning partial alias & there will be no extra work when moving away from MemDepAnalysis.

[dfukalov]

In D93529#2504303, @fhahn wrote:

IIUC this is similar to the handling of partial overwrites in DSE (https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp#L548)?

If so, could we just generalize the helper there, rather than threading this through the different caches from BasicAA to GVN? Yes, we have to do some extra work in GVN, but there would be no need to maintain a cached value not used by other clients, it will work with any AA implementations returning partial alias & there will be no extra work when moving away from MemDepAnalysis.

GVN has similar partial aliased memory objects functionality in https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/Utils/VNCoercion.cpp but it has the same issue: they use GetPointerBaseWithConstantOffset() and give up in a case when GEP offset is not constant (please refer the trivial test case in the patch). But BasicAA has an ability to go further and prove partial alias plus estimate overlapping offsets. It seems all these operations (including PHIs processing) should not be re-applied in an AA user.

The issue was masked for different targets because usually vector types are splitted before GVN so it's just see MustAlias parts of vectors and successfully eliminate them.
Actually current DSE fails with the same pattern of stores, I planned to fix it after the patch.

I'm going to split the change by Roman request and update the description to clarify these circumstances.

[dfukalov]

AA part of the splitted patch.

[asbirlea]

Adding this only seems to make sense when using BatchAA.
Could we restrict populating this cache? e.g. extend AAQueryInfo constructor to take bool BatchMode = false, and set it to true for BatchAA. Then update the get/set methods to first check this (check in the setter and assert in the getter).

llvm/include/llvm/Analysis/AliasAnalysis.h

370

Document (add comments) what this cache stores, and its purpose.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1173

referenced

[dfukalov]

Comments addressed.

[dfukalov]

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1173

Thanks!

[asbirlea]

Thanks! Diff looks reasonable to me, but I'd like to see the use case changes as well. Could you (re-)upload a concrete use case as dependent patch?
One reason I have in mind is we may want to concretely enable this caching solely for a particular BatchAA instance, i.e. pass true to BatchAA and change the meaning in AAQI from BatchMode to CacheOffsets (and, same as now, only allow BatchAA to set this bool). Then only pass true at the BatchAA callsites where it will be used.

llvm/include/llvm/Analysis/AliasAnalysis.h

350

complicated

389

only

[nikic]

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1175

I'm not sure this code is the right place to perform these access size checks. Consider the case where BatchAA is queried for the same pointers with different location sizes in a row. The first query is such that they are PartialAlias and the LocationSize fits into the larger access. The second one is also PartialAlias, but does not fit. In this case, the getClobberOffset() API will still return a result though, and a caller may incorrectly assume that the access is nested based on this check, even though it isn't.

[dfukalov]

Added memory sizes to map key, BatchAA can be created without caching.

[dfukalov]

In D93529#2517400, @asbirlea wrote:

... Could you (re-)upload a concrete use case as dependent patch?

I've just created D95543 with the test for upcoming changes in GVN. It is actually the same code as one used in unit test in this patch.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1175

Thanks for good point! Added sizes to cache key, updated test to check that offset is returned only for exact the same set of two pointers and two sizes.

[dfukalov]

Gentle ping...

[dfukalov]

Ping...

[dfukalov]

Ping...

[dfukalov]

@asbirlea, @nikic, would you please take a look at the updated patch?

[nikic]

I think the proper solution here would be to change AliasResult into a class that carries both the type of the result and an offset for PartialAlias (probably of reduced width), rather than making this information available via a cache side channel.

That would be a larger change though, and I'm okay with doing this in the meantime to make some progress.

llvm/include/llvm/Analysis/AliasAnalysis.h

399

nit: Unnecessary braces.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

1153

No need to use pointer for LocationSize, the type is okay to copy.

1155

nit: Off.isNegative()

1168

nit: auto -> uint64_t. Type is not obvious from context.

llvm/unittests/Analysis/AliasAnalysisTest.cpp

323

Wrong test :)

[dfukalov]

Rebased, addressed comments.

llvm/include/llvm/Analysis/AliasAnalysis.h

 /// query. By default, the root query is stateless and starts with a freshly
 /// constructed info object. Specific alias analyses can use this query info to
 /// store per-query state that is important for recursive or nested queries to
 /// avoid recomputing. To enable preserving this state across multiple queries
 /// where safe (due to the IR not changing), use a `BatchAAResults` wrapper.
 /// The information stored in an `AAQueryInfo` is currently limitted to the
 /// caches used by BasicAA, but can further be extended to fit other AA needs.
 class AAQueryInfo {
+  /// Storage for estimated relative offsets between two partially aliased
+  /// values. Used to optimize out redundant parts of loads/stores (in GVN/DSE).
+  /// These users cannot process quite complicated addresses (e.g. GEPs with
+  /// non-constant offsets). Used by BatchAAResults only.
+  bool CacheOffsets = false;
+  SmallDenseMap<std::pair<std::pair<const Value *, const Value *>,
+                          std::pair<uint64_t, uint64_t>>,
+                int64_t, 4>
+      ClobberOffsets;
+
 public:
   using LocPair = std::pair<MemoryLocation, MemoryLocation>;
   struct CacheEntry {
     AliasResult Result;
     /// Number of times a NoAlias assumption has been used.
     /// 0 for assumptions that have not been used, -1 for definitive results.
     int NumAssumptionUses;
     /// Whether this is a definitive (non-assumption) result.
     bool isDefinitive() const { return NumAssumptionUses < 0; }
   };
   using AliasCacheT = SmallDenseMap<LocPair, CacheEntry, 8>;
   AliasCacheT AliasCache;
 
   using IsCapturedCacheT = SmallDenseMap<const Value *, bool, 8>;
   IsCapturedCacheT IsCapturedCache;
 
   /// Query depth used to distinguish recursive queries.
   unsigned Depth = 0;
 
   /// How many active NoAlias assumption uses there are.
   int NumAssumptionUses = 0;
 
   /// Location pairs for which an assumption based result is currently stored.
   /// Used to remove all potentially incorrect results from the cache if an
   /// assumption is disproven.
   SmallVector<AAQueryInfo::LocPair, 4> AssumptionBasedResults;
 
-  AAQueryInfo() : AliasCache(), IsCapturedCache() {}
+  AAQueryInfo(bool CacheOffsets = false)
+      : CacheOffsets(CacheOffsets), ClobberOffsets(), AliasCache(),
+        IsCapturedCache() {}
 
   /// Create a new AAQueryInfo based on this one, but with the cache cleared.
   /// This is used for recursive queries across phis, where cache results may
   /// not be valid.
   AAQueryInfo withEmptyCache() {
     AAQueryInfo NewAAQI;
     NewAAQI.Depth = Depth;
     return NewAAQI;
   }
+
+  Optional<int64_t> getClobberOffset(const Value *Ptr1, const Value *Ptr2,
+                                     uint64_t Size1, uint64_t Size2) const {
+    assert(CacheOffsets && "Clobber offset cached in batch mode only!");
+    const bool Swapped = Ptr1 > Ptr2;
+    if (Swapped) {
+      std::swap(Ptr1, Ptr2);
+      std::swap(Size1, Size2);
+    }
+    const auto IOff = ClobberOffsets.find({{Ptr1, Ptr2}, {Size1, Size2}});
+    if (IOff != ClobberOffsets.end())
+      return Swapped ? -IOff->second : IOff->second;
+    return None;
+  }
+
+  void setClobberOffset(const Value *Ptr1, const Value *Ptr2, uint64_t Size1,
+                        uint64_t Size2, int64_t Offset) {
+    // Cache offset for batch mode only.
+    if (!CacheOffsets)
+      return;
+    if (Ptr1 > Ptr2) {
+      std::swap(Ptr1, Ptr2);
+      std::swap(Size1, Size2);
+      Offset = -Offset;
+    }
+    ClobberOffsets[{{Ptr1, Ptr2}, {Size1, Size2}}] = Offset;
+  }
 };
 
 class BatchAAResults;
 
 class AAResults {
 public:
   // Make these results default constructable and movable. We have to spell
   // these out because MSVC won't synthesize them.

 /// esentially making AA work in "batch mode". The internal state cannot be
 /// cleared, so to go "out-of-batch-mode", the user must either use AAResults,
 /// or create a new BatchAAResults.
 class BatchAAResults {
   AAResults &AA;
   AAQueryInfo AAQI;
 
 public:
-  BatchAAResults(AAResults &AAR) : AA(AAR), AAQI() {}
+  BatchAAResults(AAResults &AAR, bool CacheOffsets = false)
+      : AA(AAR), AAQI(CacheOffsets) {}
   AliasResult alias(const MemoryLocation &LocA, const MemoryLocation &LocB) {
     return AA.alias(LocA, LocB, AAQI);
   }
   bool pointsToConstantMemory(const MemoryLocation &Loc, bool OrLocal = false) {
     return AA.pointsToConstantMemory(Loc, AAQI, OrLocal);
   }
   ModRefInfo getModRefInfo(const CallBase *Call, const MemoryLocation &Loc) {
     return AA.getModRefInfo(Call, Loc, AAQI);
public:
   }
   bool isMustAlias(const MemoryLocation &LocA, const MemoryLocation &LocB) {
     return alias(LocA, LocB) == MustAlias;
   }
   bool isMustAlias(const Value *V1, const Value *V2) {
     return alias(MemoryLocation(V1, LocationSize::precise(1)),
                  MemoryLocation(V2, LocationSize::precise(1))) == MustAlias;
   }
+  Optional<int64_t> getClobberOffset(const MemoryLocation &LocA,
+                                     const MemoryLocation &LocB) const;
 };
 
 /// Temporary typedef for legacy code that uses a generic \c AliasAnalysis
 /// pointer or reference.
 using AliasAnalysis = AAResults;
 
 /// A private abstract base class describing the concept of an individual alias
 /// analysis implementation.

llvm/lib/Analysis/AliasAnalysis.cpp

if (auto *WrapperPass = P.getAnalysisIfAvailable<CFLSteensAAWrapperPass>())
     AAR.addAAResult(WrapperPass->getResult());
   if (auto *WrapperPass = P.getAnalysisIfAvailable<ExternalAAWrapperPass>())
     if (WrapperPass->CB)
       WrapperPass->CB(P, F, AAR);
 
   return AAR;
 }
 
+Optional<int64_t>
+BatchAAResults::getClobberOffset(const MemoryLocation &LocA,
+                                 const MemoryLocation &LocB) const {
+  if (!LocA.Size.hasValue() || !LocB.Size.hasValue())
+    return None;
+  const Value *V1 = LocA.Ptr->stripPointerCastsForAliasAnalysis();
+  const Value *V2 = LocB.Ptr->stripPointerCastsForAliasAnalysis();
+  return AAQI.getClobberOffset(V1, V2, LocA.Size.getValue(),
+                               LocB.Size.getValue());
+}
+
 bool llvm::isNoAliasCall(const Value *V) {
   if (const auto *Call = dyn_cast<CallBase>(V))
     return Call->hasRetAttr(Attribute::NoAlias);
   return false;
 }
 
 static bool isNoAliasOrByValArgument(const Value *V) {
   if (const Argument *A = dyn_cast<Argument>(V))

llvm/lib/Analysis/BasicAliasAnalysis.cpp

if (BaseAlias != MustAlias) {
     return BaseAlias;
   }
 
   // If there is a constant difference between the pointers, but the difference
   // is less than the size of the associated memory object, then we know
   // that the objects are partially overlapping.  If the difference is
   // greater, we know they do not overlap.
   if (DecompGEP1.Offset != 0 && DecompGEP1.VarIndices.empty()) {
-    if (DecompGEP1.Offset.sge(0)) {
-      if (V2Size.hasValue()) {
-        if (DecompGEP1.Offset.ult(V2Size.getValue()))
-          return PartialAlias;
-        return NoAlias;
-      }
-    } else {
-      // We have the situation where:
+    APInt &Off = DecompGEP1.Offset;
+
+    // Initialize for Off >= 0 (V2 <= GEP1) case.
+    const Value *LeftPtr = V2;
+    const Value *RightPtr = GEP1;
+    LocationSize VLeftSize = V2Size;
+    LocationSize VRightSize = V1Size;
+
+    if (Off.isNegative()) {
+      // Swap if we have the situation where:
       // +                +
       // | BaseOffset     |
       // ---------------->|
       // |-->V1Size       |-------> V2Size
       // GEP1             V2
-      if (V1Size.hasValue()) {
-        if ((-DecompGEP1.Offset).ult(V1Size.getValue()))
-          return PartialAlias;
-        return NoAlias;
+      std::swap(LeftPtr, RightPtr);
+      std::swap(VLeftSize, VRightSize);
+      Off = -Off;
+    }
+
+    if (VLeftSize.hasValue()) {
+      const uint64_t LSize = VLeftSize.getValue();
+      if (Off.ult(LSize)) {
+        // Conservatively drop processing if a phi was visited and/or offset is
+        // too big.
+        if (VisitedPhiBBs.empty() && VRightSize.hasValue() &&
+            Off.ule(INT64_MAX)) {
+          // Memory referenced by right pointer is nested. Save the offset in
+          // cache.
+          const uint64_t RSize = VRightSize.getValue();
+          if ((Off + RSize).ule(LSize))
+            AAQI.setClobberOffset(LeftPtr, RightPtr, LSize, RSize,
+                                  Off.getSExtValue());
+        }
       }
+      return NoAlias;
     }
   }
 
   if (!DecompGEP1.VarIndices.empty()) {
     APInt GCD;
     bool AllNonNegative = DecompGEP1.Offset.isNonNegative();
     bool AllNonPositive = DecompGEP1.Offset.isNonPositive();
     for (unsigned i = 0, e = DecompGEP1.VarIndices.size(); i != e; ++i) {
if (V1Srcs.empty())
     return MayAlias;
 
   // If this PHI node is recursive, indicate that the pointer may be moved
   // across iterations. We can only prove NoAlias if different underlying
   // objects are involved.
   if (isRecursive)
     PNSize = LocationSize::beforeOrAfterPointer();
 
   // In the recursive alias queries below, we may compare values from two
   // different loop iterations. Keep track of visited phi blocks, which will
   // be used when determining value equivalence.
   bool BlockInserted = VisitedPhiBBs.insert(PN->getParent()).second;
   auto _ = make_scope_exit([&]() {
     if (BlockInserted)
       VisitedPhiBBs.erase(PN->getParent());
   });
 

llvm/unittests/Analysis/AliasAnalysisTest.cpp

TEST_F(AliasAnalysisTest, BatchAAPhiAssumption) {
   EXPECT_EQ(MayAlias, AA.alias(ALoc, BLoc));
   EXPECT_EQ(MayAlias, AA.alias(ANextLoc, BNextLoc));
 
   BatchAAResults BatchAA(AA);
   EXPECT_EQ(MayAlias, BatchAA.alias(ALoc, BLoc));
   EXPECT_EQ(MayAlias, BatchAA.alias(ANextLoc, BNextLoc));
 }
 
+// Check that two aliased GEPs with non-constant offsets are correctly
+// analyzed and their relative offset can be requested from AA.
+TEST_F(AliasAnalysisTest, PartialAliasOffset) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(R"(
+    define void @foo(float* %arg, i32 %i) {
+    bb:
+      %i2 = zext i32 %i to i64
+      %i3 = getelementptr inbounds float, float* %arg, i64 %i2
+      %i4 = bitcast float* %i3 to <2 x float>*
+      %L2 = load <2 x float>, <2 x float>* %i4, align 16
+      %i7 = add nuw nsw i32 %i, 1
+      %i8 = zext i32 %i7 to i64
+      %i9 = getelementptr inbounds float, float* %arg, i64 %i8
+      %L1 = load float, float* %i9, align 4
+      ret void
+    }
+  )",
+                                                  Err, C);
+
+  if (!M)
+    Err.print("PartialAliasOffset", errs());
+
+  Function *F = M->getFunction("foo");
+  const auto Loc1 = MemoryLocation::get(getInstructionByName(*F, "L1"));
+  auto Loc2 = MemoryLocation::get(getInstructionByName(*F, "L2"));
+
+  auto &AA = getAAResults(*F);
+
+  BatchAAResults BatchAA(AA, /*CacheOffsets =*/true);
+  EXPECT_EQ(PartialAlias, BatchAA.alias(Loc1, Loc2));
+  EXPECT_EQ(-4, BatchAA.getClobberOffset(Loc1, Loc2).getValueOr(0));
+  EXPECT_EQ(4, BatchAA.getClobberOffset(Loc2, Loc1).getValueOr(0));
+
+  // Check that no offset returned for different size.
+  Loc2.Size = LocationSize::precise(42);
+  EXPECT_EQ(0, BatchAA.getClobberOffset(Loc1, Loc2).getValueOr(0));
+}
+
 class AAPassInfraTest : public testing::Test {
 protected:
   LLVMContext C;
   SMDiagnostic Err;
   std::unique_ptr<Module> M;
 
 public:
   AAPassInfraTest()