[llvm] 96f6785 - [VectorUtils] Teach findScalarElement to return splat value.

[Sander de Smalen via llvm-commits]

Author: Sander de Smalen
Date: 2021-09-06T10:56:06+01:00
New Revision: 96f6785bc9fe3219e9486ff09b22b31fa0c73b34

URL: https://github.com/llvm/llvm-project/commit/96f6785bc9fe3219e9486ff09b22b31fa0c73b34
DIFF: https://github.com/llvm/llvm-project/commit/96f6785bc9fe3219e9486ff09b22b31fa0c73b34.diff

LOG: [VectorUtils] Teach findScalarElement to return splat value.

If the vector is a splat of some scalar value, findScalarElement()
can simply return the scalar value if it knows the requested lane
is in the vector.

This is only needed for scalable vectors, because the InsertElement/ShuffleVector
case is already handled explicitly for the fixed-width case.

This helps to recognize an InstCombine fold like:
  extractelt(bitcast(splat(%v))) -> bitcast(%v)

Reviewed By: spatel

Differential Revision: https://reviews.llvm.org/D107254

Added: 
    

Modified: 
    llvm/lib/Analysis/VectorUtils.cpp
    llvm/test/Transforms/InstCombine/vscale_extractelement.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index 543f3790fee0a..78ea00b1f46f4 100644
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -331,6 +331,12 @@ Value *llvm::findScalarElement(Value *V, unsigned EltNo) {
       if (Elt->isNullValue())
         return findScalarElement(Val, EltNo);
 
+  // If the vector is a splat then we can trivially find the scalar element.
+  if (isa<ScalableVectorType>(VTy))
+    if (Value *Splat = getSplatValue(V))
+      if (EltNo < VTy->getElementCount().getKnownMinValue())
+        return Splat;
+
   // Otherwise, we don't know.
   return nullptr;
 }

diff  --git a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
index e946b4e238819..aadadc912a524 100644
--- a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
@@ -271,6 +271,20 @@ define i1 @ext_lane1_from_cmp_with_stepvec(i64 %i) {
   ret i1 %res
 }
 
+define i64* @ext_lane_from_bitcast_of_splat(i32* %v) {
+; CHECK-LABEL: @ext_lane_from_bitcast_of_splat(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[R:%.*]] = bitcast i32* [[V:%.*]] to i64*
+; CHECK-NEXT:    ret i64* [[R]]
+;
+entry:
+  %in = insertelement <vscale x 4 x i32*> poison, i32* %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32*> %in, <vscale x 4 x i32*> poison, <vscale x 4 x i32> zeroinitializer
+  %bc = bitcast <vscale x 4 x i32*> %splat to <vscale x 4 x i64*>
+  %r = extractelement <vscale x 4 x i64*> %bc, i32 3
+  ret i64* %r
+}
+
 declare <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
 declare <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
 declare <vscale x 4 x i32> @llvm.experimental.stepvector.nxv4i32()