[LVI] Handle nonnull attributes at callsite

[dtcxzyw]

This patch is the followup of #124908.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

Changes

This patch is the followup of #124908.

---

Full diff: https://github.com/llvm/llvm-project/pull/125377.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/LazyValueInfo.cpp (+19-6)
• (modified) llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll (+68)

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index 20f69a0955f51c..14f07da4eefb15 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/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;
+            }
+          }
+        }
+      }
+    }
   }
 }
 
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll b/llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll
index 35d866ed2d92eb..090826e90ef568 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/non-null.ll
+++ b/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 }

[dtcxzyw]

Compile-time impact: https://llvm-compile-time-tracker.com/compare.php?from=5c3c0a8cec9205efba78583bd0b8a646270968a1&to=498ac15f6e7fd9a79649c4aecb0ab10fcf4eb00b&stat=instructions%3Au

[nikic]

LGTM

[nikic]

As a followup, it may be worth checking whether always using stripInBoundsOffset() has a positive effect on llvm-opt-benchmark. The lookup always uses stripInBoundsOffset(), so currently there is mismatch between getUnderlyingObject on one side an stripInBoundsOffset on the other and it may give better results to always use only stripInBoundsOffset.

[dtcxzyw]

It is worse than getUnderlyingObject.