ValueTracking] Check across single predecessors in willNotFreeBetween…

llvm/lib/Analysis/ValueTracking.cpp

@@ -557,25 +557,44 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv,
 
 bool llvm::willNotFreeBetween(const Instruction *Assume,
                               const Instruction *CtxI) {
-  if (CtxI->getParent() != Assume->getParent() || !Assume->comesBefore(CtxI))
-    return false;
+  // Helper to check if there are any calls in the range that may free memory.
+  auto hasNoFreeCalls = [](auto Range) {
+    for (const auto &[Idx, I] : enumerate(Range)) {
+      if (Idx > MaxInstrsToCheckForFree)
+        return false;
+      if (const auto *CB = dyn_cast<CallBase>(&I))
+        if (!CB->hasFnAttr(Attribute::NoFree))
+          return false;
+    }
+    return true;
+  };
+
   // Make sure the current function cannot arrange for another thread to free on
   // its behalf.
   if (!CtxI->getFunction()->hasNoSync())
     return false;
 
-  // Check if there are any calls between the assume and CtxI that may
-  // free memory.
-  for (const auto &[Idx, I] :
-       enumerate(make_range(Assume->getIterator(), CtxI->getIterator()))) {
-    // Limit number of instructions to walk.
-    if (Idx > MaxInstrsToCheckForFree)
+  // Handle cross-block case: CtxI in a successor of Assume's block.
+  const BasicBlock *CtxBB = CtxI->getParent();
+  const BasicBlock *AssumeBB = Assume->getParent();
+  BasicBlock::const_iterator CtxIter = CtxI->getIterator();
+  if (CtxBB != AssumeBB) {
+    if (CtxBB->getSinglePredecessor() != AssumeBB)
+      return false;
+
+    if (!hasNoFreeCalls(make_range(CtxBB->begin(), CtxBB->end())))
+      return false;
+
+    CtxIter = AssumeBB->end();
+  } else {
+    // Same block case: check that Assume comes before CtxI.
+    if (!Assume->comesBefore(CtxI))
       return false;
-    if (const auto *CB = dyn_cast<CallBase>(&I))
-      if (!CB->hasFnAttr(Attribute::NoFree))
-        return false;
   }
-  return true;
+
+  // Check if there are any calls between Assume and CtxIter that may free
+  // memory.
+  return hasNoFreeCalls(make_range(Assume->getIterator(), CtxIter));
 }
 
 // TODO: cmpExcludesZero misses many cases where `RHS` is non-constant but