[llvm] [SDAG] Don't transfer !range metadata without noundef to SDAG (PR #95153)

[Andreas Jonson via llvm-commits]

https://github.com/andjo403 created https://github.com/llvm/llvm-project/pull/95153

Same as https://github.com/llvm/llvm-project/commit/9deee6bffa9c331f46c68e5dd4cb4abf93dc0716 but for the noundef return attribute.

CC @nikic @arsenm

>From 345651a4e82e44f7cc95580f237c47caa5faa18a Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Tue, 11 Jun 2024 19:41:50 +0200
Subject: [PATCH 1/2] [SDAG][NFC] Test don't transfer !range metadata without
 noundef to SDAG

---
 llvm/test/CodeGen/X86/legalize-vec-assertzext.ll | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll b/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
index 1c595b7fb5e1e..bba3e530bfcbf 100644
--- a/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
+++ b/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
@@ -34,6 +34,21 @@ define i64 @widen_assertzext(ptr %x) nounwind {
   ret i64 %d
 }
 
+define i64 @widen_assertzext_noundef_ret_attr(ptr %x) nounwind {
+; CHECK-LABEL: widen_assertzext_noundef_ret_attr:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    callq test2 at PLT
+; CHECK-NEXT:    vextracti32x4 $3, %zmm0, %xmm0
+; CHECK-NEXT:    vmovq %xmm0, %rax
+; CHECK-NEXT:    popq %rcx
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+  %e = call noundef <7 x i64> @test2(), !range !0
+  %d = extractelement <7 x i64> %e, i32 6
+  ret i64 %d
+}
+
 declare  <16 x i64> @test()
 declare  <7 x i64> @test2()
 !0 = !{ i64 0, i64 2 }

>From 191e5eba203b88aa0d85620d3c3c9007bd0f4087 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Tue, 11 Jun 2024 19:52:16 +0200
Subject: [PATCH 2/2] [SDAG] Don't transfer !range metadata without noundef to
 SDAG

---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp | 7 ++++++-
 llvm/test/CodeGen/X86/legalize-vec-assertzext.ll      | 3 +++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index be5e0f6ef058b..5d03b1706358c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -4482,7 +4482,12 @@ static const MDNode *getRangeMetadata(const Instruction &I) {
   // transforms that are known not to be poison-safe, such as folding logical
   // and/or to bitwise and/or. For now, only transfer !range if !noundef is
   // also present.
-  if (!I.hasMetadata(LLVMContext::MD_noundef))
+  bool NoUndef = false;
+  if (const auto *CB = dyn_cast<CallBase>(&I)) {
+    NoUndef = CB->hasRetAttr(Attribute::NoUndef);
+  }
+  NoUndef = NoUndef || I.hasMetadata(LLVMContext::MD_noundef);
+  if (!NoUndef)
     return nullptr;
   return I.getMetadata(LLVMContext::MD_range);
 }
diff --git a/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll b/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
index bba3e530bfcbf..a4097ef49477f 100644
--- a/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
+++ b/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
@@ -39,6 +39,9 @@ define i64 @widen_assertzext_noundef_ret_attr(ptr %x) nounwind {
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushq %rax
 ; CHECK-NEXT:    callq test2 at PLT
+; CHECK-NEXT:    movb $127, %al
+; CHECK-NEXT:    kmovw %eax, %k1
+; CHECK-NEXT:    vpexpandq %zmm0, %zmm0 {%k1} {z}
 ; CHECK-NEXT:    vextracti32x4 $3, %zmm0, %xmm0
 ; CHECK-NEXT:    vmovq %xmm0, %rax
 ; CHECK-NEXT:    popq %rcx