[ValueTracking] Handle range assume bundles in isKnownNonZero

Author: andjo403

llvm/include/llvm/Analysis/AssumeBundleQueries.h

@@ -16,11 +16,13 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include <optional>
 
 namespace llvm {
 class AssumptionCache;
 class DominatorTree;
 class Instruction;
+class ConstantRange;
 
 /// Index of elements in the operand bundle.
 /// If the element exist it is guaranteed to be what is specified in this enum
@@ -170,6 +172,9 @@ RetainedKnowledge getKnowledgeValidInContext(
 RetainedKnowledge getKnowledgeFromBundle(AssumeInst &Assume,
                                          const CallBase::BundleOpInfo &BOI);
 
+std::optional<ConstantRange>
+getRangeFromBundle(AssumeInst &Assume, const CallBase::BundleOpInfo &BOI);
+
 } // namespace llvm
 
 #endif

llvm/lib/Analysis/AssumeBundleQueries.cpp

@@ -10,6 +10,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -214,3 +215,21 @@ RetainedKnowledge llvm::getKnowledgeValidInContext(
                                 return isValidAssumeForContext(I, CtxI, DT);
                               });
 }
+
+std::optional<ConstantRange>
+llvm::getRangeFromBundle(AssumeInst &Assume,
+                         const CallBase::BundleOpInfo &BOI) {
+  if (Attribute::getAttrKindFromName(BOI.Tag->getKey()) != Attribute::Range)
+    return std::nullopt;
+
+  assert(BOI.End - BOI.Begin > ABA_Argument + 1 &&
+         "range assumptions should have 3 arguments");
+
+  if (auto *Lower = dyn_cast<ConstantInt>(
+          getValueFromBundleOpInfo(Assume, BOI, ABA_Argument)))
+    if (auto *Upper = dyn_cast<ConstantInt>(
+            getValueFromBundleOpInfo(Assume, BOI, ABA_Argument + 1)))
+      if (Lower->getValue() != Upper->getValue())
+        return ConstantRange(Lower->getValue(), Upper->getValue());
+  return std::nullopt;
+}

llvm/lib/Analysis/ValueTracking.cpp

@@ -595,18 +595,25 @@ static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
            "Got assumption for the wrong function!");
 
     if (Elem.Index != AssumptionCache::ExprResultIdx) {
-      if (!V->getType()->isPointerTy())
-        continue;
+      if (V->getType()->isPointerTy()) {
       if (RetainedKnowledge RK = getKnowledgeFromBundle(
               *I, I->bundle_op_info_begin()[Elem.Index])) {
         if (RK.WasOn == V &&
             (RK.AttrKind == Attribute::NonNull ||
              (RK.AttrKind == Attribute::Dereferenceable &&
-              !NullPointerIsDefined(Q.CxtI->getFunction(),
+                !NullPointerIsDefined(
+                    Q.CxtI->getFunction(),
                                     V->getType()->getPointerAddressSpace()))) &&
             isValidAssumeForContext(I, Q.CxtI, Q.DT))
           return true;
       }
+      } else if (V->getType()->isIntOrIntVectorTy()) {
+        if (std::optional<ConstantRange> Range =
+                getRangeFromBundle(*I, I->bundle_op_info_begin()[Elem.Index]))
+          if (!Range->contains(APInt::getZero(Range->getBitWidth())) &&
+              isValidAssumeForContext(I, Q.CxtI, Q.DT))
+            return true;
+      }
       continue;
     }
 

llvm/test/Analysis/ValueTracking/known-non-zero.ll

@@ -1549,9 +1549,7 @@ define i1 @vec_reverse_non_zero_demanded_fail(<4 x i8> %xx) {
 define i1 @range_assume(i8 %x, i8 %y) {
 ; CHECK-LABEL: @range_assume(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i8 [[X:%.*]], i8 1, i8 0) ]
-; CHECK-NEXT:    [[OR:%.*]] = or i8 [[Y:%.*]], [[X]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[OR]], 0
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   call void @llvm.assume(i1 true) ["range"(i8 %x, i8 1, i8 0)]
   %or = or i8 %y, %x
@@ -1575,9 +1573,7 @@ define i1 @neg_range_assum(i8 %x, i8 %y) {
 define <2 x i1> @range_assum_vec(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @range_assum_vec(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i8> [[X:%.*]], i8 1, i8 0) ]
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[Y:%.*]], [[X]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i8> [[OR]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+; CHECK-NEXT:    ret <2 x i1> <i1 true, i1 true>
 ;
   call void @llvm.assume(i1 true) ["range"(<2 x i8> %x, i8 1, i8 0)]
   %or = or <2 x i8> %y, %x