[llvm] [FunctionAttrs] Fix incorrect noundef inference with poison attrs (PR #89348)

[Nikita Popov via llvm-commits]

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

Currently, when inferring noundef, we only check that the return value is not undef/poison. However, we fail to account for the possibility that a poison-generating return attribute will convert the value to poison, and then violate the noundef attribute, resulting in immediate UB.

For the relevant return attributes (align, nonnull and range), check whether we can trivially re-prove the relevant property, otherwise do not infer noundef.

This fixes the FunctionAttrs side of
https://github.com/llvm/llvm-project/issues/88026.

>From 37ad6bfab69b30c1edaa7bdfe25215278b69644c Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Fri, 19 Apr 2024 14:32:16 +0900
Subject: [PATCH] [FunctionAttrs] Fix incorrect noundef inference with poison
 attrs

Currently, when inferring noundef, we only check that the return
value is not undef/poison. However, we fail to account for the
possibility that a poison-generating return attribute will convert
the value to poison, and then violate the noundef attribute,
resulting in immediate UB.

For the relevant return attributes (align, nonnull and range),
check whether we can trivially re-prove the relevant property,
otherwise do not infer noundef.

This fixes the FunctionAttrs side of
https://github.com/llvm/llvm-project/issues/88026.
---
 llvm/lib/Transforms/IPO/FunctionAttrs.cpp     | 27 ++++++++++++++++---
 llvm/test/Transforms/FunctionAttrs/noundef.ll |  9 +++----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 7ebf265e17ba1f..8e11cbf1cee469 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -1287,7 +1287,8 @@ static void addNoUndefAttrs(const SCCNodeSet &SCCNodes,
   // values.
   for (Function *F : SCCNodes) {
     // Already noundef.
-    if (F->getAttributes().hasRetAttr(Attribute::NoUndef))
+    AttributeList Attrs = F->getAttributes();
+    if (Attrs.hasRetAttr(Attribute::NoUndef))
       continue;
 
     // We can infer and propagate function attributes only when we know that the
@@ -1305,10 +1306,30 @@ static void addNoUndefAttrs(const SCCNodeSet &SCCNodes,
     if (F->getReturnType()->isVoidTy())
       continue;
 
-    if (all_of(*F, [](BasicBlock &BB) {
+    const DataLayout &DL = F->getParent()->getDataLayout();
+    if (all_of(*F, [&](BasicBlock &BB) {
           if (auto *Ret = dyn_cast<ReturnInst>(BB.getTerminator())) {
             // TODO: perform context-sensitive analysis?
-            return isGuaranteedNotToBeUndefOrPoison(Ret->getReturnValue());
+            Value *RetVal = Ret->getReturnValue();
+            if (!isGuaranteedNotToBeUndefOrPoison(RetVal))
+              return false;
+
+            // We know the original return value is not poison now, but it
+            // could still be converted to poison by another return attribute.
+            // Try to explicitly re-prove the relevant attributes.
+            if (Attrs.hasRetAttr(Attribute::NonNull) &&
+                !isKnownNonZero(RetVal, DL))
+              return false;
+
+            if (MaybeAlign Align = Attrs.getRetAlignment())
+              if (RetVal->getPointerAlignment(DL) < *Align)
+                return false;
+
+            Attribute Attr = Attrs.getRetAttr(Attribute::Range);
+            if (Attr.isValid() &&
+                !Attr.getRange().contains(
+                    computeConstantRange(RetVal, /*ForSigned=*/false)))
+              return false;
           }
           return true;
         })) {
diff --git a/llvm/test/Transforms/FunctionAttrs/noundef.ll b/llvm/test/Transforms/FunctionAttrs/noundef.ll
index 64deda546c067b..b7c583880501a2 100644
--- a/llvm/test/Transforms/FunctionAttrs/noundef.ll
+++ b/llvm/test/Transforms/FunctionAttrs/noundef.ll
@@ -167,9 +167,8 @@ define i64 @test_trunc_with_constexpr() {
   ret i64 %conv
 }
 
-; FIXME: This is a miscompile.
 define align 4 ptr @maybe_not_aligned(ptr noundef %p) {
-; CHECK-LABEL: define noundef align 4 ptr @maybe_not_aligned(
+; CHECK-LABEL: define align 4 ptr @maybe_not_aligned(
 ; CHECK-SAME: ptr noundef readnone returned [[P:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    ret ptr [[P]]
 ;
@@ -184,9 +183,8 @@ define align 4 ptr @definitely_aligned(ptr noundef align 4 %p) {
   ret ptr %p
 }
 
-; FIXME: This is a miscompile.
 define nonnull ptr @maybe_not_nonnull(ptr noundef %p) {
-; CHECK-LABEL: define noundef nonnull ptr @maybe_not_nonnull(
+; CHECK-LABEL: define nonnull ptr @maybe_not_nonnull(
 ; CHECK-SAME: ptr noundef readnone returned [[P:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    ret ptr [[P]]
 ;
@@ -201,9 +199,8 @@ define nonnull ptr @definitely_nonnull(ptr noundef nonnull %p) {
   ret ptr %p
 }
 
-; FIXME: This is a miscompile.
 define range(i8 0, 10) i8 @maybe_not_in_range(i8 noundef %v) {
-; CHECK-LABEL: define noundef range(i8 0, 10) i8 @maybe_not_in_range(
+; CHECK-LABEL: define range(i8 0, 10) i8 @maybe_not_in_range(
 ; CHECK-SAME: i8 noundef returned [[V:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    ret i8 [[V]]
 ;