ValueTracking: Handle extractelement and extractvalue in computeKnownFPClass

[foad]

Seems reasonable. It would be nice to have some tests for vector sources where only some elements are known to be non-nan, but I guess they are hard to write.

[arsenm]

In D148182#4264236, @foad wrote:

Seems reasonable. It would be nice to have some tests for vector sources where only some elements are known to be non-nan, but I guess they are hard to write.

I think they're unwritable as-is. This is testable after insertelement is handled in D148186

[arsenm]

409ef45000185d74b695a7dc94f5aeac0c4422d7

llvm/lib/Analysis/ValueTracking.cpp

if (InterestedClasses & fcInf) {
       // integer, the result of the cast must be finite.
       Type *FPTy = Op->getType()->getScalarType();
       if (ilogb(APFloat::getLargest(FPTy->getFltSemantics())) >= IntSize)
         Known.knownNot(fcInf);
     }
 
     break;
   }
+  case Instruction::ExtractElement: {
+    // 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);
+
+    if (auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType())) {
+      unsigned NumElts = VecTy->getNumElements();
+      APInt DemandedVecElts = APInt::getAllOnes(NumElts);
+      if (CIdx && CIdx->getValue().ult(NumElts))
+        DemandedVecElts = APInt::getOneBitSet(NumElts, CIdx->getZExtValue());
+      return computeKnownFPClass(Vec, DemandedVecElts, InterestedClasses, Known,
+                                 Depth + 1, Q, TLI);
+    }
+
+    break;
+  }
+  case Instruction::ExtractValue: {
+    computeKnownFPClass(Op->getOperand(0), DemandedElts, InterestedClasses,
+                        Known, Depth + 1, Q, TLI);
+    break;
+  }
   default:
     break;
   }
 }
 
 KnownFPClass llvm::computeKnownFPClass(
     const Value *V, const APInt &DemandedElts, const DataLayout &DL,
     FPClassTest InterestedClasses, unsigned Depth, const TargetLibraryInfo *TLI,

llvm/test/Transforms/Attributor/nofpclass.ll

entry:
   %class1 = call i1 @llvm.is.fpclass.f32(float %arg, i32 682)
   call void @llvm.assume(i1 %class1)
   %class2 = call i1 @llvm.is.fpclass.f32(float %arg, i32 341)
   call void @llvm.assume(i1 %class2)
   call void @extern.use(float %arg)
   ret float %arg
 }
 
+define float @returned_extractelement_dynamic_index(<4 x float> nofpclass(nan) %vec, i32 %idx) {
+; CHECK: Function Attrs: nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define nofpclass(nan) float @returned_extractelement_dynamic_index
+; CHECK-SAME: (<4 x float> nofpclass(nan) [[VEC:%.*]], i32 [[IDX:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[EXTRACT:%.*]] = extractelement <4 x float> [[VEC]], i32 [[IDX]]
+; CHECK-NEXT:    ret float [[EXTRACT]]
+;
+  %extract = extractelement <4 x float> %vec, i32 %idx
+  ret float %extract
+}
+
+define float @returned_extractelement_index0(<4 x float> nofpclass(nan) %vec) {
+; CHECK: Function Attrs: nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define nofpclass(nan) float @returned_extractelement_index0
+; CHECK-SAME: (<4 x float> nofpclass(nan) [[VEC:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[EXTRACT:%.*]] = extractelement <4 x float> [[VEC]], i32 0
+; CHECK-NEXT:    ret float [[EXTRACT]]
+;
+  %extract = extractelement <4 x float> %vec, i32 0
+  ret float %extract
+}
+
+define float @returned_extractelement_scalable(<vscale x 4 x float> nofpclass(nan) %vec) {
+; CHECK: Function Attrs: nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define float @returned_extractelement_scalable
+; CHECK-SAME: (<vscale x 4 x float> nofpclass(nan) [[VEC:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[EXTRACT:%.*]] = extractelement <vscale x 4 x float> [[VEC]], i32 0
+; CHECK-NEXT:    ret float [[EXTRACT]]
+;
+  %extract = extractelement <vscale x 4 x float> %vec, i32 0
+  ret float %extract
+}
+
+define float @returned_extractvalue([4 x float] nofpclass(nan) %array) {
+; CHECK: Function Attrs: nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define nofpclass(nan) float @returned_extractvalue
+; CHECK-SAME: ([4 x float] nofpclass(nan) [[ARRAY:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[EXTRACT:%.*]] = extractvalue [4 x float] [[ARRAY]], 0
+; CHECK-NEXT:    ret float [[EXTRACT]]
+;
+  %extract = extractvalue [4 x float] %array, 0
+  ret float %extract
+}