[llvm] Attributor: Do not assume function context in AANoCapture (PR #91462)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm updated https://github.com/llvm/llvm-project/pull/91462

>From e07e5e6cd6cf31c90d5ba0040cba47166ba979d8 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Wed, 8 May 2024 11:31:58 +0200
Subject: [PATCH] Attributor: Do not assume function context in AANoCapture

If the calling function has the null_pointer_is_valid attribute,
somehow a null constant reaches here. I'm not sure why exactly,
it doesn't happen for other types of constants.

Fixes #87856
---
 .../Transforms/IPO/AttributorAttributes.cpp   | 10 ++-
 llvm/test/Transforms/Attributor/issue87856.ll | 61 +++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/Attributor/issue87856.ll

diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 41b66aafe7d34..1b3bf3c732ed0 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -5690,6 +5690,9 @@ bool AANoCapture::isImpliedByIR(Attributor &A, const IRPosition &IRP,
     return V.use_empty();
 
   // You cannot "capture" null in the default address space.
+  //
+  // FIXME: This should use NullPointerIsDefined to account for the function
+  // attribute.
   if (isa<UndefValue>(V) || (isa<ConstantPointerNull>(V) &&
                              V.getType()->getPointerAddressSpace() == 0)) {
     return true;
@@ -5899,10 +5902,13 @@ ChangeStatus AANoCaptureImpl::updateImpl(Attributor &A) {
 
   const Function *F =
       isArgumentPosition() ? IRP.getAssociatedFunction() : IRP.getAnchorScope();
-  assert(F && "Expected a function!");
-  const IRPosition &FnPos = IRPosition::function(*F);
+
+  // TODO: Is the checkForAllUses below useful for constants?
+  if (!F)
+    return indicatePessimisticFixpoint();
 
   AANoCapture::StateType T;
+  const IRPosition &FnPos = IRPosition::function(*F);
 
   // Readonly means we cannot capture through memory.
   bool IsKnown;
diff --git a/llvm/test/Transforms/Attributor/issue87856.ll b/llvm/test/Transforms/Attributor/issue87856.ll
new file mode 100644
index 0000000000000..4da29cc4448d4
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/issue87856.ll
@@ -0,0 +1,61 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 4
+; RUN: opt -S -passes=attributor < %s | FileCheck %s
+
+define void @null_ptr_is_valid_call_with_null() #0 {
+; CHECK-LABEL: define void @null_ptr_is_valid_call_with_null(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    call void @store_as0(ptr nofree noundef writeonly align 4294967296 null) #[[ATTR4:[0-9]+]]
+; CHECK-NEXT:    ret void
+;
+  call void @store_as0(ptr null)
+  ret void
+}
+
+define void @null_ptr_is_valid_call_with_undef() #0 {
+; CHECK-LABEL: define void @null_ptr_is_valid_call_with_undef(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:    call void @store_as0(ptr undef) #[[ATTR4]]
+; CHECK-NEXT:    ret void
+;
+  call void @store_as0(ptr undef)
+  ret void
+}
+
+define void @store_as0(ptr %0) {
+; CHECK-LABEL: define void @store_as0(
+; CHECK-SAME: ptr nocapture nofree noundef nonnull writeonly align 2 dereferenceable(2) [[TMP0:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT:    store i16 0, ptr [[TMP0]], align 2
+; CHECK-NEXT:    ret void
+;
+  store i16 0, ptr %0, align 2
+  ret void
+}
+
+define void @call_store_as1() {
+; CHECK-LABEL: define void @call_store_as1(
+; CHECK-SAME: ) #[[ATTR3:[0-9]+]] {
+; CHECK-NEXT:    call void @store_as1(ptr addrspace(1) nocapture nofree noundef writeonly align 4294967296 null) #[[ATTR4]]
+; CHECK-NEXT:    ret void
+;
+  call void @store_as1(ptr addrspace(1) null)
+  ret void
+}
+
+define void @store_as1(ptr addrspace(1) %arg) {
+; CHECK-LABEL: define void @store_as1(
+; CHECK-SAME: ptr addrspace(1) nocapture nofree noundef writeonly align 2 dereferenceable_or_null(2) [[ARG:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    store i16 0, ptr addrspace(1) [[ARG]], align 2
+; CHECK-NEXT:    ret void
+;
+  store i16 0, ptr addrspace(1) %arg, align 2
+  ret void
+}
+
+attributes #0 = { null_pointer_is_valid }
+;.
+; CHECK: attributes #[[ATTR0]] = { mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(write) }
+; CHECK: attributes #[[ATTR1]] = { mustprogress nofree norecurse nosync nounwind null_pointer_is_valid willreturn memory(none) }
+; CHECK: attributes #[[ATTR2]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(argmem: write) }
+; CHECK: attributes #[[ATTR3]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) }
+; CHECK: attributes #[[ATTR4]] = { nofree nosync nounwind willreturn memory(write) }
+;.