[llvm] [IPSCCP] Infer nonnull return attribute (PR #106553)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/106553

Similarly to the existing range attribute inference, also infer the nonnull attribute on function return values.

I think in practice FunctionAttrs will handle nearly all cases, the main one I think it doesn't is cases involving branch conditions. But as we already have the information here, we may as well materialize it.

>From 813356917612ad40a7f4a9308de3fdf1470119c2 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 29 Aug 2024 11:49:15 +0200
Subject: [PATCH] [IPSCCP] Infer nonnull return attribute

Similarly to the existing range attribute inference, also infer
the nonnull attribute on function return values.

I think in practice FunctionAttrs will handle nearly all cases,
the main one I think it doesn't is cases involving branch
conditions. But as we already have the information here, we
may as well materialize it.
---
 llvm/lib/Transforms/IPO/SCCP.cpp             |  8 ++++++++
 llvm/test/Transforms/PGOProfile/memprof.ll   |  4 ++--
 llvm/test/Transforms/SCCP/pointer-nonnull.ll | 10 +++++++---
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/SCCP.cpp b/llvm/lib/Transforms/IPO/SCCP.cpp
index 6fa2724e0116ec..9eabc8f7ee021b 100644
--- a/llvm/lib/Transforms/IPO/SCCP.cpp
+++ b/llvm/lib/Transforms/IPO/SCCP.cpp
@@ -296,6 +296,14 @@ static bool runIPSCCP(
       F->addRangeRetAttr(CR);
       continue;
     }
+    // Infer nonnull return attribute.
+    if (F->getReturnType()->isPointerTy() && ReturnValue.isNotConstant() &&
+        ReturnValue.getNotConstant()->isNullValue() &&
+        !F->hasRetAttribute(Attribute::NonNull) &&
+        !F->hasRetAttribute(Attribute::Dereferenceable)) {
+      F->addRetAttr(Attribute::NonNull);
+      continue;
+    }
     if (F->getReturnType()->isVoidTy())
       continue;
     if (SCCPSolver::isConstant(ReturnValue) || ReturnValue.isUnknownOrUndef())
diff --git a/llvm/test/Transforms/PGOProfile/memprof.ll b/llvm/test/Transforms/PGOProfile/memprof.ll
index b6407f7b123d3f..e1457ca7251ed8 100644
--- a/llvm/test/Transforms/PGOProfile/memprof.ll
+++ b/llvm/test/Transforms/PGOProfile/memprof.ll
@@ -81,7 +81,7 @@ target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: mustprogress noinline optnone uwtable
-; ALL-LABEL: define dso_local noundef ptr @_Z3foov()
+; ALL-LABEL: define dso_local noundef{{.*}}ptr @_Z3foov()
 ; There should be some PGO metadata
 ; PGO: !prof
 define dso_local noundef ptr @_Z3foov() #0 !dbg !10 {
@@ -96,7 +96,7 @@ entry:
 declare noundef nonnull ptr @_Znam(i64 noundef) #1
 
 ; Function Attrs: mustprogress noinline optnone uwtable
-; ALL-LABEL: define dso_local noundef ptr @_Z4foo2v()
+; ALL-LABEL: define dso_local noundef{{.*}}ptr @_Z4foo2v()
 define dso_local noundef ptr @_Z4foo2v() #0 !dbg !15 {
 entry:
   ; MEMPROF: call {{.*}} @_Z3foov{{.*}} !callsite ![[C2:[0-9]+]]
diff --git a/llvm/test/Transforms/SCCP/pointer-nonnull.ll b/llvm/test/Transforms/SCCP/pointer-nonnull.ll
index 0c4425468638c3..08d4a76345bb63 100644
--- a/llvm/test/Transforms/SCCP/pointer-nonnull.ll
+++ b/llvm/test/Transforms/SCCP/pointer-nonnull.ll
@@ -232,9 +232,13 @@ define i1 @ip_test_nonnull_caller(ptr %p) {
 }
 
 define ptr @ret_nonnull_pointer(ptr nonnull %p) {
-; CHECK-LABEL: define ptr @ret_nonnull_pointer(
-; CHECK-SAME: ptr nonnull [[P:%.*]]) {
-; CHECK-NEXT:    ret ptr [[P]]
+; SCCP-LABEL: define ptr @ret_nonnull_pointer(
+; SCCP-SAME: ptr nonnull [[P:%.*]]) {
+; SCCP-NEXT:    ret ptr [[P]]
+;
+; IPSCCP-LABEL: define nonnull ptr @ret_nonnull_pointer(
+; IPSCCP-SAME: ptr nonnull [[P:%.*]]) {
+; IPSCCP-NEXT:    ret ptr [[P]]
 ;
   ret ptr %p
 }