[llvm] [ValueTracking] Support scalable vectors for ExtractElement in computeKnownFPClass. (PR #143051)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/143051

>From 1e3bc92f0752be645a696a75b6e61e32a94f90b4 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 5 Jun 2025 16:59:43 -0700
Subject: [PATCH 1/3] [ValueTracking] Support scalable vectors for
 ExtractElemenet in computeKnownFPClass.

We can support scalable vectors by setting the demanded
mask to APInt(1, 1) to demand the whole vector.
---
 llvm/lib/Analysis/ValueTracking.cpp | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 0a460786d00ea..30357c170f3d7 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5604,19 +5604,19 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
     // Look through extract element. If the index is non-constant or
     // out-of-range demand all elements, otherwise just the extracted element.
     const Value *Vec = Op->getOperand(0);
-    const Value *Idx = Op->getOperand(1);
-    auto *CIdx = dyn_cast<ConstantInt>(Idx);
 
+    APInt DemandedVecElts;
     if (auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType())) {
       unsigned NumElts = VecTy->getNumElements();
-      APInt DemandedVecElts = APInt::getAllOnes(NumElts);
+      DemandedVecElts = APInt::getAllOnes(NumElts);
+      auto *CIdx = dyn_cast<ConstantInt>(Op->getOperand(1));
       if (CIdx && CIdx->getValue().ult(NumElts))
         DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
-      return computeKnownFPClass(Vec, DemandedVecElts, InterestedClasses, Known,
-                                 Q, Depth + 1);
-    }
+    } else
+      DemandedVecElts = APInt(1, 1);
 
-    break;
+    return computeKnownFPClass(Vec, DemandedVecElts, InterestedClasses, Known,
+                               Q, Depth + 1);
   }
   case Instruction::InsertElement: {
     if (isa<ScalableVectorType>(Op->getType()))

>From c2e0519f3cd5d8240e55ce1b1109ba5afdd7382f Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 5 Jun 2025 17:10:25 -0700
Subject: [PATCH 2/3] fixup! Add missing braces

---
 llvm/lib/Analysis/ValueTracking.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 30357c170f3d7..bde30060736c1 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5612,8 +5612,9 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
       auto *CIdx = dyn_cast<ConstantInt>(Op->getOperand(1));
       if (CIdx && CIdx->getValue().ult(NumElts))
         DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
-    } else
+    } else {
       DemandedVecElts = APInt(1, 1);
+    }
 
     return computeKnownFPClass(Vec, DemandedVecElts, InterestedClasses, Known,
                                Q, Depth + 1);

>From e6439bc09103919d87e07a65a27c5bf74bd44c4b Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 5 Jun 2025 17:39:50 -0700
Subject: [PATCH 3/3] fixup! update test

---
 llvm/test/Transforms/Attributor/nofpclass.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/Attributor/nofpclass.ll b/llvm/test/Transforms/Attributor/nofpclass.ll
index cfc7877383881..8dd2aa34f94cd 100644
--- a/llvm/test/Transforms/Attributor/nofpclass.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass.ll
@@ -1474,7 +1474,7 @@ define float @returned_extractelement_index_oob(<4 x float> nofpclass(nan) %vec)
 
 define float @returned_extractelement_scalable(<vscale x 4 x float> nofpclass(nan) %vec) {
 ; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
-; CHECK-LABEL: define float @returned_extractelement_scalable
+; CHECK-LABEL: define nofpclass(nan) float @returned_extractelement_scalable
 ; CHECK-SAME: (<vscale x 4 x float> nofpclass(nan) [[VEC:%.*]]) #[[ATTR3]] {
 ; CHECK-NEXT:    [[EXTRACT:%.*]] = extractelement <vscale x 4 x float> [[VEC]], i32 0
 ; CHECK-NEXT:    ret float [[EXTRACT]]