From 67896f494e8a310d561bc88e6890cf5c981ddc73 Mon Sep 17 00:00:00 2001
From: Philip Reames <listmail@philipreames.com>
Date: Wed, 6 Oct 2021 11:50:51 -0700
Subject: [PATCH] Returning poison from a function w/ noundef return attribute
 is UB

This does for readability of returns within said function as what we do for the caller side when reasoning about what might be poison.

Differential Revision: https://reviews.llvm.org/D111180
---
 llvm/lib/Analysis/ValueTracking.cpp                |  5 ++++-
 .../Analysis/ScalarEvolution/flags-from-poison.ll  | 22 +++++-----------------
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index a21a7a8..b385810 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5418,7 +5418,10 @@ void llvm::getGuaranteedWellDefinedOps(
       }
       break;
     }
-
+    case Instruction::Ret:
+      if (I->getFunction()->hasRetAttribute(Attribute::NoUndef))
+        Operands.insert(I->getOperand(0));
+      break;
     default:
       break;
   }
diff --git a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
index 0423854..79e76c1 100644
--- a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
+++ b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
@@ -1621,21 +1621,15 @@ cont6:                                            ; preds = %cont1, %if.then
   br label %for.cond
 }
 
-; TODO: once D111180 lands, remove the udiv from these *-basic tests.
-; noundef really should be enough
-
 define noundef i32 @add-basic(i32 %a, i32 %b) {
 ; CHECK-LABEL: 'add-basic'
 ; CHECK-NEXT:  Classifying expressions for: @add-basic
 ; CHECK-NEXT:    %res = add nuw nsw i32 %a, %b
 ; CHECK-NEXT:    --> (%a + %b)<nuw><nsw> U: full-set S: full-set
-; CHECK-NEXT:    %res2 = udiv i32 255, %res
-; CHECK-NEXT:    --> (255 /u (%a + %b)<nuw><nsw>) U: [0,256) S: [0,256)
 ; CHECK-NEXT:  Determining loop execution counts for: @add-basic
 ;
   %res = add nuw nsw i32 %a, %b
-  %res2 = udiv i32 255, %res
-  ret i32 %res2
+  ret i32 %res
 }
 
 define noundef i32 @sub-basic(i32 %a, i32 %b) {
@@ -1643,13 +1637,10 @@ define noundef i32 @sub-basic(i32 %a, i32 %b) {
 ; CHECK-NEXT:  Classifying expressions for: @sub-basic
 ; CHECK-NEXT:    %res = sub nuw nsw i32 %a, %b
 ; CHECK-NEXT:    --> ((-1 * %b) + %a) U: full-set S: full-set
-; CHECK-NEXT:    %res2 = udiv i32 255, %res
-; CHECK-NEXT:    --> (255 /u ((-1 * %b) + %a)) U: [0,256) S: [0,256)
 ; CHECK-NEXT:  Determining loop execution counts for: @sub-basic
 ;
   %res = sub nuw nsw i32 %a, %b
-  %res2 = udiv i32 255, %res
-  ret i32 %res2
+  ret i32 %res
 }
 
 define noundef i32 @mul-basic(i32 %a, i32 %b) {
@@ -1657,13 +1648,10 @@ define noundef i32 @mul-basic(i32 %a, i32 %b) {
 ; CHECK-NEXT:  Classifying expressions for: @mul-basic
 ; CHECK-NEXT:    %res = mul nuw nsw i32 %a, %b
 ; CHECK-NEXT:    --> (%a * %b)<nuw><nsw> U: full-set S: full-set
-; CHECK-NEXT:    %res2 = udiv i32 255, %res
-; CHECK-NEXT:    --> (255 /u (%a * %b)<nuw><nsw>) U: [0,256) S: [0,256)
 ; CHECK-NEXT:  Determining loop execution counts for: @mul-basic
 ;
   %res = mul nuw nsw i32 %a, %b
-  %res2 = udiv i32 255, %res
-  ret i32 %res2
+  ret i32 %res
 }
 
 @gA = external global i32
@@ -1746,9 +1734,9 @@ define noundef i32 @mul-recurse() {
 ; CHECK-NEXT:    %d = load i32, i32* @gD, align 4
 ; CHECK-NEXT:    --> %d U: full-set S: full-set
 ; CHECK-NEXT:    %x = mul nuw i32 %a, %b
-; CHECK-NEXT:    --> (%a * %b) U: full-set S: full-set
+; CHECK-NEXT:    --> (%a * %b)<nuw> U: full-set S: full-set
 ; CHECK-NEXT:    %y = mul nuw i32 %c, %d
-; CHECK-NEXT:    --> (%c * %d) U: full-set S: full-set
+; CHECK-NEXT:    --> (%c * %d)<nuw> U: full-set S: full-set
 ; CHECK-NEXT:    %res = mul nuw i32 %x, %y
 ; CHECK-NEXT:    --> (%a * %b * %c * %d) U: full-set S: full-set
 ; CHECK-NEXT:  Determining loop execution counts for: @mul-recurse
-- 
2.7.4