[BasicAA] Guess reasonable contexts for separate storage hints

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
+}