[llvm] [ValueTracking][WIP] Check non-null through ops (PR #126009)

[via llvm-commits]

https://github.com/goldsteinn created https://github.com/llvm/llvm-project/pull/126009

None

>From e29643a30f739d31574d4c6b33ba6ceeceb21edc Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Thu, 6 Feb 2025 00:38:15 -0500
Subject: [PATCH] [ValueTracking][WIP] Check non-null through ops

---
 llvm/lib/Analysis/ValueTracking.cpp | 67 +++++++++++++++++++++++++----
 1 file changed, 59 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index f4958162accfe33..2e6084f666ba73f 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2625,6 +2625,47 @@ static bool isGEPKnownNonNull(const GEPOperator *GEP, unsigned Depth,
   return false;
 }
 
+static bool transferNonZero(const Value *Op, const Value *V, unsigned Depth) {
+  if (Op == V)
+    return true;
+  if (Depth++ > MaxAnalysisRecursionDepth / 2)
+    return false;
+  const Operator *I = dyn_cast<Operator>(V);
+  if (!I)
+    return false;
+  switch (I->getOpcode()) {
+  case Instruction::And:
+    return transferNonZero(I->getOperand(0), V, Depth) ||
+           transferNonZero(I->getOperand(1), V, Depth);
+  case Instruction::Select: {
+    const Value *Unused0, *Unused1;
+    return matchSelectPattern(I, Unused0, Unused1).Flavor == SPF_UMIN;
+  }
+  case Instruction::Call:
+  case Instruction::Invoke:
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+      switch (II->getIntrinsicID()) {
+      default:
+        break;
+      case Intrinsic::ctpop:
+      case Intrinsic::bitreverse:
+      case Intrinsic::bswap:
+      case Intrinsic::abs:
+        return transferNonZero(II->getArgOperand(0), V, Depth);
+      case Intrinsic::umin:
+        return transferNonZero(II->getArgOperand(0), V, Depth) ||
+               transferNonZero(II->getArgOperand(1), V, Depth);
+      case Intrinsic::ptrmask:
+        return transferNonZero(II->getArgOperand(0), V, Depth);
+      }
+    }
+    break;
+  default:
+    break;
+  }
+  return false;
+}
+
 static bool isKnownNonNullFromDominatingCondition(const Value *V,
                                                   const Instruction *CtxI,
                                                   const DominatorTree *DT) {
@@ -2654,22 +2695,32 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
     }
 
     // If the value is used as a load/store, then the pointer must be non null.
-    if (V == getLoadStorePointerOperand(UI)) {
-      if (!NullPointerIsDefined(UI->getFunction(),
-                                V->getType()->getPointerAddressSpace()) &&
-          DT->dominates(UI, CtxI))
-        return true;
+    if (auto *LV = getLoadStorePointerOperand(UI)) {
+      if (transferNonZero(LV, V, /*Depth=*/0)) {
+        if (!NullPointerIsDefined(UI->getFunction(),
+                                  V->getType()->getPointerAddressSpace()) &&
+            DT->dominates(UI, CtxI))
+          return true;
+      }
     }
 
-    if ((match(UI, m_IDiv(m_Value(), m_Specific(V))) ||
-         match(UI, m_IRem(m_Value(), m_Specific(V)))) &&
+    Value *X;
+    if ((match(UI, m_IDiv(m_Value(), m_Value(X))) ||
+         match(UI, m_IRem(m_Value(), m_Value(X)))) &&
+        transferNonZero(X, V, /*Depth=*/0) &&
         isValidAssumeForContext(UI, CtxI, DT))
       return true;
 
     // Consider only compare instructions uniquely controlling a branch
     Value *RHS;
     CmpPredicate Pred;
-    if (!match(UI, m_c_ICmp(Pred, m_Specific(V), m_Value(RHS))))
+    if (!match(UI, m_ICmp(Pred, m_Value(X), m_Value(RHS))))
+      continue;
+
+    if (transferNonZero(RHS, V, /*Depth=*/0)) {
+      Pred = CmpInst::getSwappedPredicate(Pred);
+      std::swap(RHS, X);
+    } else if (!transferNonZero(X, V, /*Depth=*/0))
       continue;
 
     bool NonNullIfTrue;