[LVI] Handle nonnull at callsite

Author: dtcxzyw

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -780,12 +780,25 @@ void LazyValueInfoImpl::intersectAssumeOrGuardBlockValueConstantRange(
   }
 
   if (BBLV.isOverdefined()) {
-    // Check whether we're checking at the terminator, and the pointer has
-    // been dereferenced in this block.
-    PointerType *PTy = dyn_cast<PointerType>(Val->getType());
-    if (PTy && BB->getTerminator() == BBI &&
-        isNonNullAtEndOfBlock(Val, BB))
-      BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
+    if (PointerType *PTy = dyn_cast<PointerType>(Val->getType())) {
+      // Check whether we're checking at the terminator, and the pointer has
+      // been dereferenced in this block.
+      if (BB->getTerminator() == BBI && isNonNullAtEndOfBlock(Val, BB))
+        BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
+      else {
+        for (Use &U : Val->uses()) {
+          if (auto *CB = dyn_cast<CallBase>(U.getUser())) {
+            if (CB->isArgOperand(&U) &&
+                CB->paramHasNonNullAttr(CB->getArgOperandNo(&U),
+                                        /*AllowUndefOrPoison=*/false) &&
+                isValidAssumeForContext(CB, BBI)) {
+              BBLV = ValueLatticeElement::getNot(ConstantPointerNull::get(PTy));
+              break;
+            }
+          }
+        }
+      }
+    }
   }
 }
 

llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll

@@ -349,8 +349,7 @@ define i1 @test_known_nonnull_at_callsite(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_at_callsite(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee(ptr noundef nonnull [[SRC:%.*]])
-; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
-; CHECK-NEXT:    ret i1 [[NONNULL]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   call void @callee(ptr noundef nonnull %src)
@@ -362,8 +361,7 @@ define i1 @test_known_nonnull_mixed(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_mixed(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee2(ptr nonnull [[SRC:%.*]])
-; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
-; CHECK-NEXT:    ret i1 [[NONNULL]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   call void @callee2(ptr nonnull %src)
@@ -375,8 +373,7 @@ define i1 @test_known_nonnull_at_callsite_dereferenceable(ptr %src) {
 ; CHECK-LABEL: @test_known_nonnull_at_callsite_dereferenceable(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @callee(ptr dereferenceable(1) [[SRC:%.*]])
-; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
-; CHECK-NEXT:    ret i1 [[NONNULL]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   call void @callee(ptr dereferenceable(1) %src)