[llvm] [LVI][ValueTracking] Merge checking whether assumes imply nonnull (PR #203523)

[Nikolas Klauser via llvm-commits]

https://github.com/philnik777 created https://github.com/llvm/llvm-project/pull/203523

None

>From 7dba520a5f54bbc62961f67b900e359a544fea16 Mon Sep 17 00:00:00 2001
From: Nikolas Klauser <nikolasklauser at berlin.de>
Date: Fri, 12 Jun 2026 15:12:13 +0200
Subject: [PATCH] [LVI][ValueTracking] Merge checking whether assumes imply
 nonnull

---
 llvm/include/llvm/IR/BundleAttributes.h |  4 ++++
 llvm/lib/Analysis/LazyValueInfo.cpp     | 20 ++------------------
 llvm/lib/Analysis/ValueTracking.cpp     | 21 +++------------------
 llvm/lib/IR/BundleAttributes.cpp        | 19 +++++++++++++++++++
 4 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/llvm/include/llvm/IR/BundleAttributes.h b/llvm/include/llvm/IR/BundleAttributes.h
index 6b2f55d16fd1a..7ca264c01f14b 100644
--- a/llvm/include/llvm/IR/BundleAttributes.h
+++ b/llvm/include/llvm/IR/BundleAttributes.h
@@ -65,6 +65,10 @@ struct AssumeSeparateStorageInfo {
 LLVM_ABI
 AssumeSeparateStorageInfo getAssumeSeparateStorageInfo(OperandBundleUse);
 
+LLVM_ABI bool assumeBundleImpliesNonNullOn(const Value *Val,
+                                           const Instruction *Context,
+                                           OperandBundleUse OBU);
+
 } // namespace llvm
 
 #endif // LLVM_IR_BUNDLE_ATTRIBUTES_H
diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index b052352b51eec..a63478128a9be 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -856,26 +856,10 @@ void LazyValueInfoImpl::intersectAssumeOrGuardBlockValueConstantRange(
       continue;
 
     if (AssumeVH.Index != AssumptionCache::ExprResultIdx) {
-      auto OBU = I->getOperandBundleAt(AssumeVH.Index);
-      switch (getBundleAttrFromOBU(OBU)) {
-      case BundleAttr::NonNull:
-        if (getAssumeNonNullInfo(OBU).Ptr != Val)
-          break;
+      if (assumeBundleImpliesNonNullOn(Val, BBI,
+                                       I->getOperandBundleAt(AssumeVH.Index)))
         BBLV = BBLV.intersect(ValueLatticeElement::getNot(
             Constant::getNullValue(Val->getType())));
-        break;
-
-      case BundleAttr::Dereferenceable: {
-        auto [Ptr, _, Count] = getAssumeDereferenceableInfo(OBU);
-        if (Ptr != Val || !Count || *Count == 0)
-          break;
-        BBLV = BBLV.intersect(ValueLatticeElement::getNot(
-            Constant::getNullValue(Val->getType())));
-      } break;
-
-      default:
-        break;
-      }
     } else {
       BBLV = BBLV.intersect(*getValueFromCondition(Val, I->getArgOperand(0),
                                                    /*IsTrueDest*/ true,
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 763620b353bb5..df6e5ae584e4e 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -834,24 +834,9 @@ static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
            "Got assumption for the wrong function!");
 
     if (Elem.Index != AssumptionCache::ExprResultIdx) {
-      bool AssumeImpliesNonNull = [&]() {
-        auto OBU = I->getOperandBundleAt(Elem.Index);
-        switch (getBundleAttrFromOBU(OBU)) {
-        case BundleAttr::Dereferenceable: {
-          auto [Ptr, _, Count] = getAssumeDereferenceableInfo(OBU);
-          return Ptr == V && Count && *Count != 0 &&
-                 !NullPointerIsDefined(Q.CxtI->getFunction(),
-                                       V->getType()->getPointerAddressSpace());
-        }
-
-        case BundleAttr::NonNull:
-          return getAssumeNonNullInfo(OBU).Ptr == V;
-
-        default:
-          return false;
-        }
-      }();
-      if (AssumeImpliesNonNull && isValidAssumeForContext(I, Q))
+      if (assumeBundleImpliesNonNullOn(V, Q.CxtI,
+                                       I->getOperandBundleAt(Elem.Index)) &&
+          isValidAssumeForContext(I, Q))
         return true;
       continue;
     }
diff --git a/llvm/lib/IR/BundleAttributes.cpp b/llvm/lib/IR/BundleAttributes.cpp
index 0ca6808d93453..08606779b13ce 100644
--- a/llvm/lib/IR/BundleAttributes.cpp
+++ b/llvm/lib/IR/BundleAttributes.cpp
@@ -72,3 +72,22 @@ llvm::getAssumeDereferenceableInfo(OperandBundleUse OBU) {
     Ret.CountVal = Size->getZExtValue();
   return Ret;
 }
+
+bool llvm::assumeBundleImpliesNonNullOn(const Value *Val,
+                                        const Instruction *Context,
+                                        OperandBundleUse OBU) {
+  switch (getBundleAttrFromOBU(OBU)) {
+  case BundleAttr::Dereferenceable: {
+    auto [Ptr, _, Count] = getAssumeDereferenceableInfo(OBU);
+    return Ptr == Val && Count && *Count != 0 &&
+           !NullPointerIsDefined(Context->getFunction(),
+                                 Val->getType()->getPointerAddressSpace());
+  }
+
+  case BundleAttr::NonNull:
+    return getAssumeNonNullInfo(OBU).Ptr == Val;
+
+  default:
+    return false;
+  }
+}