[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,38 @@ 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 (AssumeVH.Index != AssumptionCache::ExprResultIdx) {
+      if (RetainedKnowledge RK = getKnowledgeFromBundle(
+              *I, I->bundle_op_info_begin()[AssumeVH.Index])) {
+        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/CorrelatedValuePropagation/assume-operand-bundles.ll

@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=correlated-propagation %s | FileCheck %s
+
+define i1 @nonnull(ptr %ptr) {
+; CHECK-LABEL: define i1 @nonnull(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[PTR]]) ]
+; CHECK-NEXT:    ret i1 true
+;
+  call void @llvm.assume(i1 true) [ "nonnull"(ptr %ptr) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  ret i1 %is_nonnull
+}
+
+define i1 @nonnull_on_other_ptr(ptr %ptr, ptr %other_ptr) {
+; CHECK-LABEL: define i1 @nonnull_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:    ret i1 [[IS_NONNULL]]
+;
+  call void @llvm.assume(i1 true) [ "nonnull"(ptr %other_ptr) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  ret i1 %is_nonnull
+}
+
+define i1 @dereferenceable(ptr %ptr) {
+; CHECK-LABEL: define i1 @dereferenceable(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[PTR]], i32 1) ]
+; CHECK-NEXT:    ret i1 true
+;
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %ptr, i32 1) ]
+  %is_nonnull = icmp ne ptr %ptr, null
+  ret i1 %is_nonnull
+}