[Loads] Check context instruction for context-sensitive derefability (#109277)

If a dereferenceability fact is provided through `!dereferenceable` (or
similar), it may only hold on the given control flow path. When we use
`isSafeToSpeculativelyExecute()` to check multiple instructions, we
might make use of `!dereferenceable` information that does not hold at
the speculation target. This doesn't happen when speculating
instructions one by one, because `!dereferenceable` will be dropped
while speculating.

Fix this by checking whether the instruction with `!dereferenceable`
dominates the context instruction. If this is not the case, it means we
are speculating, and cannot guarantee that it holds at the speculation
target.

Fixes #108854.

Author: nikic

clang/test/CodeGenOpenCL/builtins-amdgcn.cl

@@ -638,11 +638,7 @@ void test_get_workgroup_size(int d, global int *out)
 
 // CHECK-LABEL: @test_get_grid_size(
 // CHECK: {{.*}}call align 4 dereferenceable(64){{.*}} ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
-// CHECK: getelementptr inbounds i8, ptr addrspace(4) %{{.*}}, i64 12
-// CHECK: load i32, ptr addrspace(4) %{{.*}}, align 4, !invariant.load
-// CHECK: getelementptr inbounds i8, ptr addrspace(4) %{{.*}}, i64 16
-// CHECK: load i32, ptr addrspace(4) %{{.*}}, align 4, !invariant.load
-// CHECK: getelementptr inbounds i8, ptr addrspace(4) %{{.*}}, i64 20
+// CHECK: getelementptr inbounds i8, ptr addrspace(4) %{{.*}}, i64 %.sink
 // CHECK: load i32, ptr addrspace(4) %{{.*}}, align 4, !invariant.load
 void test_get_grid_size(int d, global int *out)
 {

llvm/include/llvm/Analysis/ValueTracking.h

@@ -805,7 +805,9 @@ bool onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V);
 ///
 /// If the CtxI is specified this method performs context-sensitive analysis
 /// and returns true if it is safe to execute the instruction immediately
-/// before the CtxI.
+/// before the CtxI. If the instruction has (transitive) operands that don't
+/// dominate CtxI, the analysis is performed under the assumption that these
+/// operands will also be speculated to a point before CxtI.
 ///
 /// If the CtxI is NOT specified this method only looks at the instruction
 /// itself and its operands, so if this method returns true, it is safe to

llvm/lib/Analysis/Loads.cpp

@@ -104,6 +104,17 @@ static bool isDereferenceableAndAlignedPointer(
     if (CheckForNonNull &&
         !isKnownNonZero(V, SimplifyQuery(DL, DT, AC, CtxI)))
       return false;
+    // When using something like !dereferenceable on a load, the
+    // dereferenceability may only be valid on a specific control-flow path.
+    // If the instruction doesn't dominate the context instruction, we're
+    // asking about dereferenceability under the assumption that the
+    // instruction has been speculated to the point of the context instruction,
+    // in which case we don't know if the dereferenceability info still holds.
+    // We don't bother handling allocas here, as they aren't speculatable
+    // anyway.
+    auto *I = dyn_cast<Instruction>(V);
+    if (I && !isa<AllocaInst>(I))
+      return CtxI && isValidAssumeForContext(I, CtxI, DT);
     return true;
   };
   if (IsKnownDeref()) {

llvm/lib/Analysis/MemDerefPrinter.cpp

@@ -30,10 +30,10 @@ PreservedAnalyses MemDerefPrinterPass::run(Function &F,
   for (auto &I : instructions(F)) {
     if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
       Value *PO = LI->getPointerOperand();
-      if (isDereferenceablePointer(PO, LI->getType(), DL))
+      if (isDereferenceablePointer(PO, LI->getType(), DL, LI))
         Deref.push_back(PO);
       if (isDereferenceableAndAlignedPointer(PO, LI->getType(), LI->getAlign(),
-                                             DL))
+                                             DL, LI))
         DerefAndAligned.insert(PO);
     }
   }

llvm/lib/CodeGen/MachineOperand.cpp

@@ -1047,7 +1047,8 @@ bool MachinePointerInfo::isDereferenceable(unsigned Size, LLVMContext &C,
     return false;
 
   return isDereferenceableAndAlignedPointer(
-      BasePtr, Align(1), APInt(DL.getPointerSizeInBits(), Offset + Size), DL);
+      BasePtr, Align(1), APInt(DL.getPointerSizeInBits(), Offset + Size), DL,
+      dyn_cast<Instruction>(BasePtr));
 }
 
 /// getConstantPool - Return a MachinePointerInfo record that refers to the

llvm/test/Transforms/SimplifyCFG/speculate-derefable-load.ll

@@ -77,14 +77,17 @@ exit:
   ret i64 %res
 }
 
-; FIXME: This is a miscompile.
 define i64 @deref_no_hoist(i1 %c, ptr align 8 dereferenceable(8) %p1) {
 ; CHECK-LABEL: define i64 @deref_no_hoist(
 ; CHECK-SAME: i1 [[C:%.*]], ptr align 8 dereferenceable(8) [[P1:%.*]]) {
-; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P1]], align 8, !align [[META0:![0-9]+]]
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br i1 [[C]], label %[[IF:.*]], label %[[EXIT:.*]]
+; CHECK:       [[IF]]:
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P1]], align 8, !dereferenceable [[META0:![0-9]+]], !align [[META0]]
 ; CHECK-NEXT:    [[V:%.*]] = load i64, ptr [[P2]], align 8
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[C]], i64 [[V]], i64 0
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[RES:%.*]] = phi i64 [ [[V]], %[[IF]] ], [ 0, %[[ENTRY]] ]
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
 entry: