[LVI] Handle nonnull attributes at callsite

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

@@ -344,4 +344,72 @@ define i1 @test_store_same_block(ptr %arg) {
   ret i1 %cmp
 }
 
+
+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:    ret i1 false
+;
+entry:
+  call void @callee(ptr noundef nonnull %src)
+  %nonnull = icmp eq ptr %src, null
+  ret i1 %nonnull
+}
+
+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:    ret i1 false
+;
+entry:
+  call void @callee2(ptr nonnull %src)
+  %nonnull = icmp eq ptr %src, null
+  ret i1 %nonnull
+}
+
+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:    ret i1 false
+;
+entry:
+  call void @callee(ptr dereferenceable(1) %src)
+  %nonnull = icmp eq ptr %src, null
+  ret i1 %nonnull
+}
+
+; Negative tests
+
+define i1 @test_known_nonnull_at_callsite_without_noundef(ptr %src) {
+; CHECK-LABEL: @test_known_nonnull_at_callsite_without_noundef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @callee(ptr nonnull [[SRC:%.*]])
+; CHECK-NEXT:    [[NONNULL:%.*]] = icmp eq ptr [[SRC]], null
+; CHECK-NEXT:    ret i1 [[NONNULL]]
+;
+entry:
+  call void @callee(ptr nonnull %src)
+  %nonnull = icmp eq ptr %src, null
+  ret i1 %nonnull
+}
+
+define i1 @test_known_nonnull_at_callsite_dereferenceable_null_is_defined(ptr %src) #0 {
+; CHECK-LABEL: @test_known_nonnull_at_callsite_dereferenceable_null_is_defined(
+; 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]]
+;
+entry:
+  call void @callee(ptr dereferenceable(1) %src)
+  %nonnull = icmp eq ptr %src, null
+  ret i1 %nonnull
+}
+
+declare void @callee(ptr)
+declare void @callee2(ptr noundef)
+
 attributes #0 = { null_pointer_is_valid }