[llvm] [SDAG] Lower range attribute to AssertZext (PR #95450)

[Andreas Jonson via llvm-commits]

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

split out as requested in https://github.com/llvm/llvm-project/pull/94847#discussion_r1638231377
CC @nikic 

>From 931c16b984379905932edb2c9768bb55454806dd Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Thu, 13 Jun 2024 20:15:22 +0200
Subject: [PATCH 1/2] [SDAG] Test lower range attribute to AssertZext

---
 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..5fa1649b15c70 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_range_attr(ptr %x) nounwind {
+; CHECK-LABEL: widen_assertzext_range_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 range(i64 0, 2) <7 x i64> @test2()
+  %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 f6ffb4ab04b0b2adbc3253a12f31306a30df4f60 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Thu, 13 Jun 2024 20:36:05 +0200
Subject: [PATCH 2/2] [SDAG] Lower range attribute to AssertZext

---
 .../SelectionDAG/SelectionDAGBuilder.cpp      | 23 +++++++++++++------
 .../CodeGen/X86/legalize-vec-assertzext.ll    |  3 +++
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index be5e0f6ef058b..5d21a905c316c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -4487,6 +4487,18 @@ static const MDNode *getRangeMetadata(const Instruction &I) {
   return I.getMetadata(LLVMContext::MD_range);
 }
 
+static std::optional<ConstantRange> getRange(const Instruction &I) {
+  if (const auto *CB = dyn_cast<CallBase>(&I)) {
+    // see comment in getRangeMetadata about this check
+    if (CB->hasRetAttr(Attribute::NoUndef))
+      return CB->getRange();
+  }
+  if (const MDNode *Range = getRangeMetadata(I)) {
+    return getConstantRangeFromMetadata(*Range);
+  }
+  return std::nullopt;
+}
+
 void SelectionDAGBuilder::visitLoad(const LoadInst &I) {
   if (I.isAtomic())
     return visitAtomicLoad(I);
@@ -10229,19 +10241,16 @@ void SelectionDAGBuilder::visitVACopy(const CallInst &I) {
 SDValue SelectionDAGBuilder::lowerRangeToAssertZExt(SelectionDAG &DAG,
                                                     const Instruction &I,
                                                     SDValue Op) {
-  const MDNode *Range = getRangeMetadata(I);
-  if (!Range)
-    return Op;
+  std::optional<ConstantRange> CR = getRange(I);
 
-  ConstantRange CR = getConstantRangeFromMetadata(*Range);
-  if (CR.isFullSet() || CR.isEmptySet() || CR.isUpperWrapped())
+  if (!CR || CR->isFullSet() || CR->isEmptySet() || CR->isUpperWrapped())
     return Op;
 
-  APInt Lo = CR.getUnsignedMin();
+  APInt Lo = CR->getUnsignedMin();
   if (!Lo.isMinValue())
     return Op;
 
-  APInt Hi = CR.getUnsignedMax();
+  APInt Hi = CR->getUnsignedMax();
   unsigned Bits = std::max(Hi.getActiveBits(),
                            static_cast<unsigned>(IntegerType::MIN_INT_BITS));
 
diff --git a/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll b/llvm/test/CodeGen/X86/legalize-vec-assertzext.ll
index 5fa1649b15c70..2cf37c68b8b40 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_range_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