[ValueTracking] Support scalable vector splats in computeKnownFPClass

[dtcxzyw]

Address comment #169904 (comment)

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

Changes

Address comment #169904 (comment)

---

Full diff: https://github.com/llvm/llvm-project/pull/170325.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/ValueTracking.cpp (+6)
• (modified) llvm/test/Transforms/Attributor/nofpclass.ll (+6-9)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index dbceb8e557849..eb8650fd0eb60 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5877,6 +5877,12 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
     break;
   }
   case Instruction::ShuffleVector: {
+    // Handle vector splat idiom
+    if (Value *Splat = getSplatValue(V)) {
+      computeKnownFPClass(Splat, Known, InterestedClasses, Q, Depth + 1);
+      break;
+    }
+
     // For undef elements, we don't know anything about the common state of
     // the shuffle result.
     APInt DemandedLHS, DemandedRHS;
diff --git a/llvm/test/Transforms/Attributor/nofpclass.ll b/llvm/test/Transforms/Attributor/nofpclass.ll
index a9ebdaa397015..41357b2442cc3 100644
--- a/llvm/test/Transforms/Attributor/nofpclass.ll
+++ b/llvm/test/Transforms/Attributor/nofpclass.ll
@@ -2667,15 +2667,10 @@ define [4 x float] @constant_aggregate_zero() {
 }
 
 define <vscale x 4 x float> @scalable_splat_pnorm() {
-; CHECK-CV: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
-; CHECK-CV-LABEL: define noundef <vscale x 4 x float> @scalable_splat_pnorm
-; CHECK-CV-SAME: () #[[ATTR3]] {
-; CHECK-CV-NEXT:    ret <vscale x 4 x float> splat (float 1.000000e+00)
-;
-; CHECK-CI: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
-; CHECK-CI-LABEL: define noundef nofpclass(nan inf zero sub nnorm) <vscale x 4 x float> @scalable_splat_pnorm
-; CHECK-CI-SAME: () #[[ATTR3]] {
-; CHECK-CI-NEXT:    ret <vscale x 4 x float> splat (float 1.000000e+00)
+; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define noundef nofpclass(nan inf zero sub nnorm) <vscale x 4 x float> @scalable_splat_pnorm
+; CHECK-SAME: () #[[ATTR3]] {
+; CHECK-NEXT:    ret <vscale x 4 x float> splat (float 1.000000e+00)
 ;
   ret <vscale x 4 x float> splat (float 1.0)
 }
@@ -2989,5 +2984,7 @@ attributes #5 = { "denormal-fp-math"="ieee,positive-zero" }
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CGSCC-CI: {{.*}}
 ; CGSCC-CV: {{.*}}
+; CHECK-CI: {{.*}}
+; CHECK-CV: {{.*}}
 ; TUNIT-CI: {{.*}}
 ; TUNIT-CV: {{.*}}