[BasicAA] Guess reasonable contexts for separate storage hints (#76770)

The definition of the pointer of the memory location being queried is
always one such context. Even this conservative guess can be better than
no guess at all in some cases.

Fixes #64666

Co-authored-by: David Goldblatt <[email protected]>

llvm/include/llvm/Analysis/ValueTracking.h

@@ -810,9 +810,14 @@ bool isAssumeLikeIntrinsic(const Instruction *I);
 
 /// Return true if it is valid to use the assumptions provided by an
 /// assume intrinsic, I, at the point in the control-flow identified by the
-/// context instruction, CxtI.
+/// context instruction, CxtI. By default, ephemeral values of the assumption
+/// are treated as an invalid context, to prevent the assumption from being used
+/// to optimize away its argument. If the caller can ensure that this won't
+/// happen, it can call with AllowEphemerals set to true to get more valid
+/// assumptions.
 bool isValidAssumeForContext(const Instruction *I, const Instruction *CxtI,
-                             const DominatorTree *DT = nullptr);
+                             const DominatorTree *DT = nullptr,
+                             bool AllowEphemerals = false);
 
 enum class OverflowResult {
   /// Always overflows in the direction of signed/unsigned min value.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -1543,7 +1543,7 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
           TLI, NullIsValidLocation)))
     return AliasResult::NoAlias;
 
-  if (CtxI && EnableSeparateStorageAnalysis) {
+  if (EnableSeparateStorageAnalysis) {
     for (AssumptionCache::ResultElem &Elem : AC.assumptionsFor(O1)) {
       if (!Elem || Elem.Index == AssumptionCache::ExprResultIdx)
         continue;
@@ -1559,10 +1559,30 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
         const Value *HintO1 = getUnderlyingObject(Hint1);
         const Value *HintO2 = getUnderlyingObject(Hint2);
 
-        if (((O1 == HintO1 && O2 == HintO2) ||
-             (O1 == HintO2 && O2 == HintO1)) &&
-            isValidAssumeForContext(Assume, CtxI, DT))
-          return AliasResult::NoAlias;
+        auto ValidAssumeForPtrContext = [&](const Value *Ptr) {
+          if (const Instruction *PtrI = dyn_cast<Instruction>(Ptr)) {
+            return isValidAssumeForContext(Assume, PtrI, DT,
+                                           /* AllowEphemerals */ true);
+          }
+          if (const Argument *PtrA = dyn_cast<Argument>(Ptr)) {
+            const Instruction *FirstI =
+                &*PtrA->getParent()->getEntryBlock().begin();
+            return isValidAssumeForContext(Assume, FirstI, DT,
+                                           /* AllowEphemerals */ true);
+          }
+          return false;
+        };
+
+        if ((O1 == HintO1 && O2 == HintO2) || (O1 == HintO2 && O2 == HintO1)) {
+          // Note that we go back to V1 and V2 for the
+          // ValidAssumeForPtrContext checks; they're dominated by O1 and O2,
+          // so strictly more assumptions are valid for them.
+          if ((CtxI && isValidAssumeForContext(Assume, CtxI, DT,
+                                               /* AllowEphemerals */ true)) ||
+              ValidAssumeForPtrContext(V1) || ValidAssumeForPtrContext(V2)) {
+            return AliasResult::NoAlias;
+          }
+        }
       }
     }
   }

llvm/lib/Analysis/ValueTracking.cpp

@@ -485,7 +485,8 @@ bool llvm::isAssumeLikeIntrinsic(const Instruction *I) {
 
 bool llvm::isValidAssumeForContext(const Instruction *Inv,
                                    const Instruction *CxtI,
-                                   const DominatorTree *DT) {
+                                   const DominatorTree *DT,
+                                   bool AllowEphemerals) {
   // There are two restrictions on the use of an assume:
   //  1. The assume must dominate the context (or the control flow must
   //     reach the assume whenever it reaches the context).
@@ -503,7 +504,7 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv,
     // Don't let an assume affect itself - this would cause the problems
     // `isEphemeralValueOf` is trying to prevent, and it would also make
     // the loop below go out of bounds.
-    if (Inv == CxtI)
+    if (!AllowEphemerals && Inv == CxtI)
       return false;
 
     // The context comes first, but they're both in the same block.
@@ -516,7 +517,7 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv,
     if (!isGuaranteedToTransferExecutionToSuccessor(Range, 15))
       return false;
 
-    return !isEphemeralValueOf(Inv, CxtI);
+    return AllowEphemerals || !isEphemeralValueOf(Inv, CxtI);
   }
 
   // Inv and CxtI are in different blocks.

llvm/test/Analysis/BasicAA/separate_storage-alias-sets.ll

@@ -0,0 +1,43 @@
+; We want BasicAA to make a reasonable conservative guess as to context for
+; separate storage hints. This lets alias analysis users (such as the alias set
+; tracker) who can't be context-sensitive still get the benefits of hints.
+
+; RUN: opt < %s -basic-aa-separate-storage -S -passes=print-alias-sets 2>&1 | FileCheck %s
+
+declare void @llvm.assume(i1)
+
+; CHECK-LABEL: Alias sets for function 'arg_arg'
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %a1, LocationSize::precise(1))
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %a2, LocationSize::precise(1))
+define void @arg_arg(ptr %a1, ptr %a2) {
+entry:
+  call void @llvm.assume(i1 true) [ "separate_storage"(ptr %a1, ptr %a2) ]
+  %0 = load i8, ptr %a1
+  %1 = load i8, ptr %a2
+  ret void
+}
+
+; CHECK-LABEL: Alias sets for function 'arg_inst'
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %a1, LocationSize::precise(1))
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %0, LocationSize::precise(1))
+define void @arg_inst(ptr %a1, ptr %a2) {
+entry:
+  %0 = getelementptr inbounds i8, ptr %a2, i64 20
+  call void @llvm.assume(i1 true) [ "separate_storage"(ptr %a1, ptr %0) ]
+  %1 = load i8, ptr %a1
+  %2 = load i8, ptr %0
+  ret void
+}
+
+; CHECK-LABEL: Alias sets for function 'inst_inst'
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %0, LocationSize::precise(1))
+; CHECK: AliasSet[{{.*}}, 1] must alias, Ref Pointers: (ptr %1, LocationSize::precise(1))
+define void @inst_inst(ptr %a1, ptr %a2) {
+entry:
+  %0 = getelementptr inbounds i8, ptr %a1, i64 20
+  %1 = getelementptr inbounds i8, ptr %a2, i64 20
+  call void @llvm.assume(i1 true) [ "separate_storage"(ptr %0, ptr %1) ]
+  %2 = load i8, ptr %0
+  %3 = load i8, ptr %1
+  ret void
+}