[LVI] Extract information from assume operand bundles

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -14,6 +14,7 @@
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
@@ -833,13 +834,39 @@ void LazyValueInfoImpl::intersectAssumeOrGuardBlockValueConstantRange(
     // Only check assumes in the block of the context instruction. Other
     // assumes will have already been taken into account when the value was
     // propagated from predecessor blocks.
-    auto *I = cast<CallInst>(AssumeVH);
+    auto *I = cast<AssumeInst>(AssumeVH);
+
     if (I->getParent() != BB || !isValidAssumeForContext(I, BBI))
       continue;
 
-    BBLV = BBLV.intersect(*getValueFromCondition(Val, I->getArgOperand(0),
-                                                 /*IsTrueDest*/ true,
-                                                 /*UseBlockValue*/ false));
+    if (I->hasOperandBundles()) {
+      for (auto Info : I->bundle_op_infos()) {
+        if (RetainedKnowledge RK = getKnowledgeFromBundle(*I, Info)) {
+          if (RK.WasOn != Val)
+            continue;
+          switch (RK.AttrKind) {
+          case Attribute::NonNull:
+            BBLV = BBLV.intersect(ValueLatticeElement::getNot(
+                Constant::getNullValue(RK.WasOn->getType())));
+            break;
+
+          case Attribute::Dereferenceable:
+            if (auto *CI = dyn_cast<ConstantInt>(RK.IRArgValue);
+                CI && !CI->isZero())
+              BBLV = BBLV.intersect(ValueLatticeElement::getNot(
+                  Constant::getNullValue(RK.IRArgValue->getType())));
+            break;
+
+          default:
+            break;
+          }
+        }
+      }
+    } else {
+      BBLV = BBLV.intersect(*getValueFromCondition(Val, I->getArgOperand(0),
+                                                   /*IsTrueDest*/ true,
+                                                   /*UseBlockValue*/ false));
+    }
   }
 
   // If guards are not used in the module, don't spend time looking for them

llvm/test/Transforms/JumpThreading/assume-operand-bundles.ll

@@ -0,0 +1,60 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=jump-threading,dce %s | FileCheck %s
+
+define i32 @nonull(ptr %ptr) {
+; CHECK-LABEL: define i32 @nonull(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:  [[TRUE:.*:]]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[PTR]]) ]
+; CHECK-NEXT:    ret i32 1
+;
+  call void @llvm.assume(i1 true) [ "nonnull"(ptr %ptr) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  br i1 %is_nonnull, label %true, label %false
+
+true:
+  ret i32 1
+
+false:
+  ret i32 2
+}
+
+define i32 @nonull_on_other_ptr(ptr %ptr, ptr %other_ptr) {
+; CHECK-LABEL: define i32 @nonull_on_other_ptr(
+; CHECK-SAME: ptr [[PTR:%.*]], ptr [[OTHER_PTR:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[OTHER_PTR]]) ]
+; CHECK-NEXT:    [[IS_NONNULL:%.*]] = icmp ne ptr [[PTR]], null
+; CHECK-NEXT:    br i1 [[IS_NONNULL]], label %[[TRUE:.*]], label %[[FALSE:.*]]
+; CHECK:       [[TRUE]]:
+; CHECK-NEXT:    ret i32 1
+; CHECK:       [[FALSE]]:
+; CHECK-NEXT:    ret i32 2
+;
+  call void @llvm.assume(i1 true) [ "nonnull"(ptr %other_ptr) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  br i1 %is_nonnull, label %true, label %false
+
+true:
+  ret i32 1
+
+false:
+  ret i32 2
+}
+
+define i32 @dereferenceable(ptr %ptr) {
+; CHECK-LABEL: define i32 @dereferenceable(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:  [[TRUE:.*:]]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[PTR]], i32 1) ]
+; CHECK-NEXT:    ret i32 1
+;
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %ptr, i32 1) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  br i1 %is_nonnull, label %true, label %false
+
+true:
+  ret i32 1
+
+false:
+  ret i32 2
+}