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

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/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 https://github.com/llvm/llvm-project/issues/108854.

>From d9ccc859090e4e51add89939b63fe52b0895bd92 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Tue, 17 Sep 2024 15:48:42 +0200
Subject: [PATCH] [Loads] Check context instruction for context-sensitive
 derefability

If a dereferenceability fact is provided through `!dereferenceable`,
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 https://github.com/llvm/llvm-project/issues/108854.
---
 llvm/lib/Analysis/Loads.cpp                           | 11 +++++++++++
 llvm/lib/Analysis/MemDerefPrinter.cpp                 |  4 ++--
 llvm/lib/CodeGen/MachineOperand.cpp                   |  3 ++-
 .../SimplifyCFG/speculate-derefable-load.ll           | 11 +++++++----
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Analysis/Loads.cpp b/llvm/lib/Analysis/Loads.cpp
index 957ac883490c45..11f3807ffacf6e 100644
--- a/llvm/lib/Analysis/Loads.cpp
+++ b/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()) {
diff --git a/llvm/lib/Analysis/MemDerefPrinter.cpp b/llvm/lib/Analysis/MemDerefPrinter.cpp
index e858d941435441..68cb8859488f70 100644
--- a/llvm/lib/Analysis/MemDerefPrinter.cpp
+++ b/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);
     }
   }
diff --git a/llvm/lib/CodeGen/MachineOperand.cpp b/llvm/lib/CodeGen/MachineOperand.cpp
index 6ee47624f31c54..89d32c3f005e00 100644
--- a/llvm/lib/CodeGen/MachineOperand.cpp
+++ b/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
diff --git a/llvm/test/Transforms/SimplifyCFG/speculate-derefable-load.ll b/llvm/test/Transforms/SimplifyCFG/speculate-derefable-load.ll
index 8c7afa4598bd4b..0138433312ed84 100644
--- a/llvm/test/Transforms/SimplifyCFG/speculate-derefable-load.ll
+++ b/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: