[llvm] Range assume bundles (PR #112758)

Andreas Jonson via llvm-commits llvm-commits at lists.llvm.org
Thu Oct 17 11:21:49 PDT 2024


https://github.com/andjo403 created https://github.com/llvm/llvm-project/pull/112758

Adding support for Range assume bundles and Handling in value tracking.

This was shortly discussed in end of https://discourse.llvm.org/t/rfc-metadata-attachments-for-function-arguments/76420 but maybe a separat RFC is needed.

Looked in to adding to the getKnowledgeFromBundle but did not find a good way to do it so added a separat handling for range bundles.

Did not find a good test for the computeConstantRange as the AssumptionCache is not added in most calls to computeConstantRange.

Can split the PR if needed but wanted to have at least one use of the new bundle.

>From e872852412fc225134833899b38520324e0ce4c2 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Wed, 16 Oct 2024 20:04:00 +0200
Subject: [PATCH 1/5] [IR] Allow range assume bundles

---
 llvm/lib/Analysis/AssumeBundleQueries.cpp      | 18 +++++++++++++++---
 llvm/lib/IR/Verifier.cpp                       | 17 +++++++++++++++++
 llvm/test/Verifier/assume-bundles.ll           |  8 ++++++++
 .../Analysis/AssumeBundleQueriesTest.cpp       |  8 +++++---
 4 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Analysis/AssumeBundleQueries.cpp b/llvm/lib/Analysis/AssumeBundleQueries.cpp
index 21530693c5f180..229993a49e98ed 100644
--- a/llvm/lib/Analysis/AssumeBundleQueries.cpp
+++ b/llvm/lib/Analysis/AssumeBundleQueries.cpp
@@ -69,8 +69,14 @@ bool llvm::hasAttributeInAssume(AssumeInst &Assume, Value *IsOn,
 
 void llvm::fillMapFromAssume(AssumeInst &Assume, RetainedKnowledgeMap &Result) {
   for (auto &Bundles : Assume.bundle_op_infos()) {
-    std::pair<Value *, Attribute::AttrKind> Key{
-        nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())};
+    Attribute::AttrKind AttrKind =
+        Attribute::getAttrKindFromName(Bundles.Tag->getKey());
+
+    if (!Attribute::isEnumAttrKind(AttrKind) &&
+        !Attribute::isIntAttrKind(AttrKind))
+      continue;
+
+    std::pair<Value *, Attribute::AttrKind> Key{nullptr, AttrKind};
     if (bundleHasArgument(Bundles, ABA_WasOn))
       Key.first = getValueFromBundleOpInfo(Assume, Bundles, ABA_WasOn);
 
@@ -101,8 +107,14 @@ llvm::getKnowledgeFromBundle(AssumeInst &Assume,
   RetainedKnowledge Result;
   if (!DebugCounter::shouldExecute(AssumeQueryCounter))
     return Result;
+  Attribute::AttrKind AttrKind =
+      Attribute::getAttrKindFromName(BOI.Tag->getKey());
+
+  if (!Attribute::isEnumAttrKind(AttrKind) &&
+      !Attribute::isIntAttrKind(AttrKind))
+    return Result;
 
-  Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
+  Result.AttrKind = AttrKind;
   if (bundleHasArgument(BOI, ABA_WasOn))
     Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, ABA_WasOn);
   auto GetArgOr1 = [&](unsigned Idx) -> uint64_t {
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index f34fe7594c8602..6388c7775e4856 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -5416,6 +5416,23 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
                 "third argument should be an integer if present", Call);
         return;
       }
+      if (Kind == Attribute::Range) {
+        Check(ArgCount == 3, "range assumptions should have 3 arguments", Call);
+        Type *FirstType = Call.getOperand(Elem.Begin)->getType();
+        Check(FirstType->isIntOrIntVectorTy(),
+              "first argument should be an integer or vector of integers",
+              Call);
+        Type *ST = FirstType->getScalarType();
+        Check(Call.getOperand(Elem.Begin + 1)->getType() == ST,
+              "second argument should be an integer with same bit width as the "
+              "first argument",
+              Call);
+        Check(Call.getOperand(Elem.Begin + 2)->getType() == ST,
+              "third argument should be an integer with same bit width as the "
+              "first argument",
+              Call);
+        return;
+      }
       Check(ArgCount <= 2, "too many arguments", Call);
       if (Kind == Attribute::None)
         break;
diff --git a/llvm/test/Verifier/assume-bundles.ll b/llvm/test/Verifier/assume-bundles.ll
index afe6cc0ab4c674..717de735cf6538 100644
--- a/llvm/test/Verifier/assume-bundles.ll
+++ b/llvm/test/Verifier/assume-bundles.ll
@@ -28,5 +28,13 @@ define void @func(ptr %P, i32 %P1, ptr %P2, ptr %P3) {
   call void @llvm.assume(i1 true) ["separate_storage"(ptr %P)]
 ; CHECK: arguments to separate_storage assumptions should be pointers
   call void @llvm.assume(i1 true) ["separate_storage"(ptr %P, i32 123)]
+; CHECK: range assumptions should have 3 arguments
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, i32 0)]
+; CHECK: first argument should be an integer or vector of integers
+  call void @llvm.assume(i1 true) ["range"(ptr %P, i32 %P1, i32 4)]
+; CHECK: second argument should be an integer with same bit width as the first argument
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, ptr %P, i32 4)]
+; CHECK: third argument should be an integer with same bit width as the first argument
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, i32 4, i8 10)]
   ret void
 }
diff --git a/llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp b/llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp
index 8ad15ca41510f2..47c4305c4c7739 100644
--- a/llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp
+++ b/llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp
@@ -411,8 +411,10 @@ static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
   std::mt19937 Rng(Seed);
   std::uniform_int_distribution<int> DistCount(MinCount, MaxCount);
   std::uniform_int_distribution<unsigned> DistValue(0, MaxValue);
-  std::uniform_int_distribution<unsigned> DistAttr(0,
-                                                   Attribute::EndAttrKinds - 1);
+  std::uniform_int_distribution<unsigned> DistEnumAttr(Attribute::FirstEnumAttr,
+                                                       Attribute::LastEnumAttr);
+  std::uniform_int_distribution<unsigned> DistIntAttr(Attribute::FirstIntAttr,
+                                                      Attribute::LastIntAttr);
 
   std::unique_ptr<Module> Mod = std::make_unique<Module>("AssumeQueryAPI", C);
   if (!Mod)
@@ -449,7 +451,7 @@ static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
   for (int i = 0; i < Size; i++) {
     int count = DistCount(Rng);
     int value = DistValue(Rng);
-    int attr = DistAttr(Rng);
+    int attr = count > 1 ? DistIntAttr(Rng) : DistEnumAttr(Rng);
     std::string str;
     raw_string_ostream ss(str);
     ss << Attribute::getNameFromAttrKind(

>From df32fd97538f20c6ab4dd9c1c9921e60f2e377b4 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Wed, 16 Oct 2024 21:11:39 +0200
Subject: [PATCH 2/5] [ValueTracking] Tests with range assume bundles (NFC)

---
 llvm/test/Analysis/BasicAA/range.ll           | 12 +++++
 .../Analysis/ValueTracking/known-non-zero.ll  | 52 +++++++++++++++++++
 .../InstSimplify/shift-knownbits.ll           | 45 ++++++++++++++++
 3 files changed, 109 insertions(+)

diff --git a/llvm/test/Analysis/BasicAA/range.ll b/llvm/test/Analysis/BasicAA/range.ll
index e5dfb60c8b8784..85a97c44cffba9 100644
--- a/llvm/test/Analysis/BasicAA/range.ll
+++ b/llvm/test/Analysis/BasicAA/range.ll
@@ -271,6 +271,18 @@ entry:
   ret i32 %load_
 }
 
+; CHECK-LABEL: Function: range_assume
+; CHECK: MayAlias: i32* %gep1, i32* %gep2
+define void @range_assume(ptr %s, ptr %q) {
+  %in_array = load i32, ptr %q
+  call void @llvm.assume(i1 true) ["range"(i32 %in_array, i32 0, i32 2)]
+  %gep1 = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1, i32 %in_array
+  %gep2 = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 2
+  load i32, ptr %gep1
+  load i32, ptr %gep2
+  ret void
+}
+
 declare void @llvm.assume(i1)
 
 !0 = !{ i32 0, i32 2 }
diff --git a/llvm/test/Analysis/ValueTracking/known-non-zero.ll b/llvm/test/Analysis/ValueTracking/known-non-zero.ll
index db2c4f3a1ed650..4699fdc0c58268 100644
--- a/llvm/test/Analysis/ValueTracking/known-non-zero.ll
+++ b/llvm/test/Analysis/ValueTracking/known-non-zero.ll
@@ -1546,4 +1546,56 @@ define i1 @vec_reverse_non_zero_demanded_fail(<4 x i8> %xx) {
   ret i1 %r
 }
 
+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]]
+;
+  call void @llvm.assume(i1 true) ["range"(i8 %x, i8 1, i8 0)]
+  %or = or i8 %y, %x
+  %cmp = icmp eq i8 %or, 0
+  ret i1 %cmp
+}
+
+define i1 @neg_range_assum(i8 %x, i8 %y) {
+; CHECK-LABEL: @neg_range_assum(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i8 [[X:%.*]], i8 -1, i8 1) ]
+; CHECK-NEXT:    [[OR:%.*]] = or i8 [[Y:%.*]], [[X]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[OR]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  call void @llvm.assume(i1 true) ["range"(i8 %x, i8 -1, i8 1)]
+  %or = or i8 %y, %x
+  %cmp = icmp eq i8 %or, 0
+  ret i1 %cmp
+}
+
+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]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i8> %x, i8 1, i8 0)]
+  %or = or <2 x i8> %y, %x
+  %cmp = icmp ne <2 x i8> %or, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
+define <2 x i1> @neg_range_assum_vec(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @neg_range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i8> [[X:%.*]], i8 -1, i8 1) ]
+; 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]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i8> %x, i8 -1, i8 1)]
+  %or = or <2 x i8> %y, %x
+  %cmp = icmp ne <2 x i8> %or, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
 declare i32 @llvm.experimental.get.vector.length.i32(i32, i32, i1)
diff --git a/llvm/test/Transforms/InstSimplify/shift-knownbits.ll b/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
index 6bf03779379ec7..2c15e7ae2b593c 100644
--- a/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
+++ b/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
@@ -499,3 +499,48 @@ define <1 x i64> @bitcast_noshift_vector_wrong_type(<2 x float> %v1, <1 x i64> %
   %r = shl <1 x i64> %v2, %b
   ret <1 x i64> %r
 }
+
+define i32 @shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
+; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i32 [[B:%.*]], i32 32, i32 64) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret i32 [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(i32 %b, i32 32, i32 64)]
+  %shl = shl i32 %a, %b
+  ret i32 %shl
+}
+
+define i32 @neg_shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
+; CHECK-LABEL: @neg_shl_amount_is_known_bogus_range_assum(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i32 [[B:%.*]], i32 0, i32 32) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret i32 [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(i32 %b, i32 0, i32 32)]
+  %shl = shl i32 %a, %b
+  ret i32 %shl
+}
+
+
+define <2 x i32> @shl_amount_is_known_bogus_range_assum_vec(<2 x i32> %a, <2 x i32> %b) {
+; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i32> [[B:%.*]], i32 32, i32 64) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i32> %b, i32 32, i32 64)]
+  %shl = shl <2 x i32> %a, %b
+  ret <2 x i32> %shl
+}
+
+define <2 x i32> @neg_shl_amount_is_known_bogus_range_assum_vec(<2 x i32> %a, <2 x i32> %b) {
+; CHECK-LABEL: @neg_shl_amount_is_known_bogus_range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i32> [[B:%.*]], i32 0, i32 32) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i32> %b, i32 0, i32 32)]
+  %shl = shl <2 x i32> %a, %b
+  ret <2 x i32> %shl
+}

>From d6f1964ab18e9b1eb68f11b7cdfdc4ed9160544d Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Thu, 17 Oct 2024 18:46:17 +0200
Subject: [PATCH 3/5] [ValueTracking] Handle range assume bundles in
 isKnownNonZero

---
 .../llvm/Analysis/AssumeBundleQueries.h       |  5 +++++
 llvm/lib/Analysis/AssumeBundleQueries.cpp     | 19 +++++++++++++++++++
 llvm/lib/Analysis/ValueTracking.cpp           | 13 ++++++++++---
 .../Analysis/ValueTracking/known-non-zero.ll  |  8 ++------
 4 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/Analysis/AssumeBundleQueries.h b/llvm/include/llvm/Analysis/AssumeBundleQueries.h
index f7a893708758c5..56dc52134883f6 100644
--- a/llvm/include/llvm/Analysis/AssumeBundleQueries.h
+++ b/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
diff --git a/llvm/lib/Analysis/AssumeBundleQueries.cpp b/llvm/lib/Analysis/AssumeBundleQueries.cpp
index 229993a49e98ed..0a20871d5d59b9 100644
--- a/llvm/lib/Analysis/AssumeBundleQueries.cpp
+++ b/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;
+}
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index e9ed8b3c862b55..b7612ff4c74629 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/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;
     }
 
diff --git a/llvm/test/Analysis/ValueTracking/known-non-zero.ll b/llvm/test/Analysis/ValueTracking/known-non-zero.ll
index 4699fdc0c58268..9a53bf4c1b1a8a 100644
--- a/llvm/test/Analysis/ValueTracking/known-non-zero.ll
+++ b/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

>From 989ac2d9e884f3412ceb5adbee586cac536358ca Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Thu, 17 Oct 2024 19:15:55 +0200
Subject: [PATCH 4/5] [ValueTracking] Handle range assume bundles in
 computeConstantRange

---
 llvm/lib/Analysis/ValueTracking.cpp | 12 +++++++++---
 llvm/test/Analysis/BasicAA/range.ll |  2 +-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index b7612ff4c74629..16fca9486143ca 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -9913,11 +9913,17 @@ ConstantRange llvm::computeConstantRange(const Value *V, bool ForSigned,
     for (auto &AssumeVH : AC->assumptionsFor(V)) {
       if (!AssumeVH)
         continue;
-      CallInst *I = cast<CallInst>(AssumeVH);
+      AssumeInst *I = cast<AssumeInst>(AssumeVH);
       assert(I->getParent()->getParent() == CtxI->getParent()->getParent() &&
              "Got assumption for the wrong function!");
-      assert(I->getIntrinsicID() == Intrinsic::assume &&
-             "must be an assume intrinsic");
+
+      if (AssumeVH.Index != AssumptionCache::ExprResultIdx) {
+        if (std::optional<ConstantRange> Range = getRangeFromBundle(
+                *I, I->bundle_op_info_begin()[AssumeVH.Index]))
+          if (isValidAssumeForContext(I, CtxI, DT))
+            CR = CR.intersectWith(*Range);
+        continue;
+      }
 
       if (!isValidAssumeForContext(I, CtxI, DT))
         continue;
diff --git a/llvm/test/Analysis/BasicAA/range.ll b/llvm/test/Analysis/BasicAA/range.ll
index 85a97c44cffba9..760abd52ce6377 100644
--- a/llvm/test/Analysis/BasicAA/range.ll
+++ b/llvm/test/Analysis/BasicAA/range.ll
@@ -272,7 +272,7 @@ entry:
 }
 
 ; CHECK-LABEL: Function: range_assume
-; CHECK: MayAlias: i32* %gep1, i32* %gep2
+; CHECK: NoAlias: i32* %gep1, i32* %gep2
 define void @range_assume(ptr %s, ptr %q) {
   %in_array = load i32, ptr %q
   call void @llvm.assume(i1 true) ["range"(i32 %in_array, i32 0, i32 2)]

>From 3ef8137f37ffebeb9803fbc40a0c922404a45fd7 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Thu, 17 Oct 2024 19:22:09 +0200
Subject: [PATCH 5/5] [ValueTracking] Handle range assume bundles in
 computeKnownBits

---
 llvm/lib/Analysis/ValueTracking.cpp           | 21 ++++++++++++-------
 .../InstSimplify/shift-knownbits.ll           |  6 ++----
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 16fca9486143ca..826d44dad512b8 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -831,17 +831,22 @@ void llvm::computeKnownBitsFromContext(const Value *V, KnownBits &Known,
            "Got assumption for the wrong function!");
 
     if (Elem.Index != AssumptionCache::ExprResultIdx) {
-      if (!V->getType()->isPointerTy())
-        continue;
-      if (RetainedKnowledge RK = getKnowledgeFromBundle(
-              *I, I->bundle_op_info_begin()[Elem.Index])) {
+      if (V->getType()->isPointerTy()) {
+        if (RetainedKnowledge RK = getKnowledgeFromBundle(
+                *I, I->bundle_op_info_begin()[Elem.Index])) {
         // Allow AllowEphemerals in isValidAssumeForContext, as the CxtI might
         // be the producer of the pointer in the bundle. At the moment, align
         // assumptions aren't optimized away.
-        if (RK.WasOn == V && RK.AttrKind == Attribute::Alignment &&
-            isPowerOf2_64(RK.ArgValue) &&
-            isValidAssumeForContext(I, Q.CxtI, Q.DT, /*AllowEphemerals*/ true))
-          Known.Zero.setLowBits(Log2_64(RK.ArgValue));
+          if (RK.WasOn == V && RK.AttrKind == Attribute::Alignment &&
+              isPowerOf2_64(RK.ArgValue) &&
+              isValidAssumeForContext(I, Q.CxtI, Q.DT, /*AllowEphemerals*/ true))
+            Known.Zero.setLowBits(Log2_64(RK.ArgValue));
+        }
+      } else if (V->getType()->isIntOrIntVectorTy()) {
+        if (std::optional<ConstantRange> Range =
+                getRangeFromBundle(*I, I->bundle_op_info_begin()[Elem.Index]))
+          if (isValidAssumeForContext(I, Q.CxtI, Q.DT))
+            Known = Known.unionWith(Range->toKnownBits());
       }
       continue;
     }
diff --git a/llvm/test/Transforms/InstSimplify/shift-knownbits.ll b/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
index 2c15e7ae2b593c..f6063425d86f4a 100644
--- a/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
+++ b/llvm/test/Transforms/InstSimplify/shift-knownbits.ll
@@ -503,8 +503,7 @@ define <1 x i64> @bitcast_noshift_vector_wrong_type(<2 x float> %v1, <1 x i64> %
 define i32 @shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i32 [[B:%.*]], i32 32, i32 64) ]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret i32 [[SHL]]
+; CHECK-NEXT:    ret i32 poison
 ;
   call void @llvm.assume(i1 true) ["range"(i32 %b, i32 32, i32 64)]
   %shl = shl i32 %a, %b
@@ -526,8 +525,7 @@ define i32 @neg_shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
 define <2 x i32> @shl_amount_is_known_bogus_range_assum_vec(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum_vec(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i32> [[B:%.*]], i32 32, i32 64) ]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+; CHECK-NEXT:    ret <2 x i32> poison
 ;
   call void @llvm.assume(i1 true) ["range"(<2 x i32> %b, i32 32, i32 64)]
   %shl = shl <2 x i32> %a, %b



More information about the llvm-commits mailing list