[llvm] 6fd1604 - [InstCombine] Add instcombine fold for extractelement + splat for scalable vectors

[Caroline Concatto via llvm-commits]

Author: Caroline Concatto
Date: 2021-06-08T10:43:38+01:00
New Revision: 6fd1604d14335a31268bbc477de27e81310f39ef

URL: https://github.com/llvm/llvm-project/commit/6fd1604d14335a31268bbc477de27e81310f39ef
DIFF: https://github.com/llvm/llvm-project/commit/6fd1604d14335a31268bbc477de27e81310f39ef.diff

LOG: [InstCombine] Add instcombine fold for extractelement + splat for scalable vectors

This patch allows that scalable vector can also use the fold that already
exists for fixed vector, only when the lane index is lower than the minimum
number of elements of the vector.

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

Added: 
    

Modified: 
    llvm/lib/Analysis/InstructionSimplify.cpp
    llvm/test/Transforms/InstCombine/gep-vector-indices.ll
    llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
    llvm/test/Transforms/InstCombine/vscale_extractelement.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index c6dabf8870fd..61be2dcd0ad6 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -4495,9 +4495,13 @@ static Value *SimplifyExtractElementInst(Value *Vec, Value *Idx,
   // find a previously computed scalar that was inserted into the vector.
   if (auto *IdxC = dyn_cast<ConstantInt>(Idx)) {
     // For fixed-length vector, fold into undef if index is out of bounds.
-    if (isa<FixedVectorType>(VecVTy) &&
-        IdxC->getValue().uge(cast<FixedVectorType>(VecVTy)->getNumElements()))
+    unsigned MinNumElts = VecVTy->getElementCount().getKnownMinValue();
+    if (isa<FixedVectorType>(VecVTy) && IdxC->getValue().uge(MinNumElts))
       return PoisonValue::get(VecVTy->getElementType());
+    // Handle case where an element is extracted from a splat.
+    if (IdxC->getValue().ult(MinNumElts))
+      if (auto *Splat = getSplatValue(Vec))
+        return Splat;
     if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
       return Elt;
   }

diff  --git a/llvm/test/Transforms/InstCombine/gep-vector-indices.ll b/llvm/test/Transforms/InstCombine/gep-vector-indices.ll
index e0830f24fd95..709c315b5068 100644
--- a/llvm/test/Transforms/InstCombine/gep-vector-indices.ll
+++ b/llvm/test/Transforms/InstCombine/gep-vector-indices.ll
@@ -61,10 +61,7 @@ define i32* @vector_splat_ptrs_v2i64_ext0(i32* %a, i64 %index) {
 
 define i32* @vector_splat_ptrs_nxv2i64_ext0(i32* %a, i64 %index) {
 ; CHECK-LABEL: @vector_splat_ptrs_nxv2i64_ext0(
-; CHECK-NEXT:    [[TMP:%.*]] = insertelement <vscale x 2 x i32*> poison, i32* [[A:%.*]], i32 0
-; CHECK-NEXT:    [[SPLATOFA:%.*]] = shufflevector <vscale x 2 x i32*> [[TMP]], <vscale x 2 x i32*> poison, <vscale x 2 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP0:%.*]] = extractelement <vscale x 2 x i32*> [[SPLATOFA]], i32 0
-; CHECK-NEXT:    [[RES:%.*]] = getelementptr i32, i32* [[TMP0]], i64 [[INDEX:%.*]]
+; CHECK-NEXT:    [[RES:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret i32* [[RES]]
 ;
   %tmp = insertelement <vscale x 2 x i32*> poison, i32* %a, i32 0

diff  --git a/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
index de04e4562d1a..07ca96d44de2 100644
--- a/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
+++ b/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
@@ -55,13 +55,9 @@ define i8 @extractelement_bitcast_wrong_insert(<vscale x 2 x i32> %a, i32 %x) {
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 %v
 ;
   %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

diff  --git a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
index 8d49ae712a19..a5d2d9443c86 100644
--- a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll
@@ -55,13 +55,9 @@ define i8 @extractelement_bitcast_wrong_insert(<vscale x 2 x i32> %a, i32 %x) {
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 %v
 ;
   %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
   %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer