[llvm] [DAG] computeKnownFPClass - add ISD::EXTRACT_VECTOR_ELT handling (PR #190307)

[Xinlong Chen via llvm-commits]

https://github.com/Xinlong-Chen updated https://github.com/llvm/llvm-project/pull/190307

>From 2ee1c2d50803c337ca109ebd467e1522c4d2c42a Mon Sep 17 00:00:00 2001
From: xinlongchen <xinlongchen at tencent.com>
Date: Fri, 3 Apr 2026 11:49:18 +0800
Subject: [PATCH] [DAG] computeKnownFPClass - add ISD::EXTRACT_VECTOR_ELT
 handling

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 20 +++++++
 llvm/test/CodeGen/RISCV/combine-is_fpclass.ll | 55 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 6eb8853550a19..a572ce62ab56d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6076,6 +6076,26 @@ KnownFPClass SelectionDAG::computeKnownFPClass(SDValue Op,
     }
     break;
   }
+  case ISD::EXTRACT_VECTOR_ELT: {
+    SDValue Src = Op.getOperand(0);
+    auto *CIdx = dyn_cast<ConstantSDNode>(Op.getOperand(1));
+    EVT SrcVT = Src.getValueType();
+    if (SrcVT.isFixedLengthVector() && CIdx) {
+      if (CIdx->getAPIntValue().ult(SrcVT.getVectorNumElements())) {
+        APInt DemandedSrcElts = APInt::getOneBitSet(
+            SrcVT.getVectorNumElements(), CIdx->getZExtValue());
+        Known = computeKnownFPClass(Src, DemandedSrcElts, InterestedClasses,
+                                    Depth + 1);
+      } else {
+        // Out of bounds index is poison.
+        Known.KnownFPClasses = fcNone;
+        Known.SignBit = false;
+      }
+    } else {
+      Known = computeKnownFPClass(Src, InterestedClasses, Depth + 1);
+    }
+    break;
+  }
   case ISD::BITCAST: {
     // FIXME: It should not be necessary to check for an elementwise bitcast.
     // If a bitcast is not elementwise between vector / scalar types,
diff --git a/llvm/test/CodeGen/RISCV/combine-is_fpclass.ll b/llvm/test/CodeGen/RISCV/combine-is_fpclass.ll
index 0292df415a655..90cacfe3b1a91 100644
--- a/llvm/test/CodeGen/RISCV/combine-is_fpclass.ll
+++ b/llvm/test/CodeGen/RISCV/combine-is_fpclass.ll
@@ -21,3 +21,58 @@ define i8 @iszero_constant_v4f32() nounwind {
   %r = bitcast <8 x i1> %f to i8
   ret i8 %r
 }
+
+define i1 @extract_nonconst_idx_const_vec_isnan(i32 %idx) nounwind {
+; CHECK-LABEL: extract_nonconst_idx_const_vec_isnan:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ret
+  %elt = extractelement <4 x float> <float 1.0, float 2.0, float 3.0, float 4.0>, i32 %idx
+  %res = call i1 @llvm.is.fpclass.f32(float %elt, i32 3)  ; 0x3 = nan
+  ret i1 %res
+}
+
+define i1 @extract_oob_idx_const_vec_isnan() nounwind {
+; CHECK-LABEL: extract_oob_idx_const_vec_isnan:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    fclass.s a0, fa5
+; CHECK-NEXT:    andi a0, a0, 768
+; CHECK-NEXT:    snez a0, a0
+; CHECK-NEXT:    ret
+  %elt = extractelement <4 x float> <float 1.0, float 2.0, float 3.0, float 4.0>, i32 6
+  %res = call i1 @llvm.is.fpclass.f32(float %elt, i32 3)  ; 0x3 = nan
+  ret i1 %res
+}
+
+define i1 @extract_fabs_vec_isneg(<4 x float> %a0) nounwind {
+; CHECK-LABEL: extract_fabs_vec_isneg:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ret
+  %abs = call <4 x float> @llvm.fabs.v4f32(<4 x float> %a0)
+  %elt = extractelement <4 x float> %abs, i32 2
+  %res = call i1 @llvm.is.fpclass.f32(float %elt, i32 60)  ; 0x3C = negative
+  ret i1 %res
+}
+
+define i1 @extract_fabs_unknown_idx_isneg(<4 x float> %a0, i32 %idx) nounwind {
+; CHECK-LABEL: extract_fabs_unknown_idx_isneg:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ret
+  %abs = call <4 x float> @llvm.fabs.v4f32(<4 x float> %a0)
+  %elt = extractelement <4 x float> %abs, i32 %idx
+  %res = call i1 @llvm.is.fpclass.f32(float %elt, i32 60)  ; 0x3C = negative
+  ret i1 %res
+}
+
+define i1 @extract_scalable_fabs_isneg(<vscale x 4 x float> %a0) nounwind {
+; CHECK-LABEL: extract_scalable_fabs_isneg:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ret
+  %abs = call <vscale x 4 x float> @llvm.fabs.nxv4f32(<vscale x 4 x float> %a0)
+  %elt = extractelement <vscale x 4 x float> %abs, i32 0
+  %res = call i1 @llvm.is.fpclass.f32(float %elt, i32 60)  ; 0x3C = negative
+  ret i1 %res
+}