[llvm] 3d3c0ed - [InstSimplify] fold extractelement of splat with variable extract index

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2021-07-05T08:19:40-04:00
New Revision: 3d3c0ed9323ba387dd6030337c672d2c9c4816f0

URL: https://github.com/llvm/llvm-project/commit/3d3c0ed9323ba387dd6030337c672d2c9c4816f0
DIFF: https://github.com/llvm/llvm-project/commit/3d3c0ed9323ba387dd6030337c672d2c9c4816f0.diff

LOG: [InstSimplify] fold extractelement of splat with variable extract index

We already have a fold for variable index with constant vector,
but if we can determine a scalar splat value, then it does not
matter whether that value is constant or not.

We overlooked this fold in D102404 and earlier patches,
but the fixed vector variant is shown in:
https://llvm.org/PR50817

Alive2 agrees on that:
https://alive2.llvm.org/ce/z/HpijPC

The same logic applies to scalable vectors.

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

Added: 
    

Modified: 
    llvm/lib/Analysis/InstructionSimplify.cpp
    llvm/test/Transforms/InstSimplify/extract-element.ll
    llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll
    llvm/test/Transforms/InstSimplify/vscale.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index f713d5317b8c..35d1e34188d9 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -4531,10 +4531,6 @@ static Value *SimplifyExtractElementInst(Value *Vec, Value *Idx,
     if (auto *CIdx = dyn_cast<Constant>(Idx))
       return ConstantExpr::getExtractElement(CVec, CIdx);
 
-    // The index is not relevant if our vector is a splat.
-    if (auto *Splat = CVec->getSplatValue())
-      return Splat;
-
     if (Q.isUndefValue(Vec))
       return UndefValue::get(VecVTy->getElementType());
   }
@@ -4557,6 +4553,10 @@ static Value *SimplifyExtractElementInst(Value *Vec, Value *Idx,
         return Splat;
     if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
       return Elt;
+  } else {
+    // The index is not relevant if our vector is a splat.
+    if (Value *Splat = getSplatValue(Vec))
+      return Splat;
   }
   return nullptr;
 }

diff  --git a/llvm/test/Transforms/InstSimplify/extract-element.ll b/llvm/test/Transforms/InstSimplify/extract-element.ll
index 8259598a3ea5..fef09528bbf2 100644
--- a/llvm/test/Transforms/InstSimplify/extract-element.ll
+++ b/llvm/test/Transforms/InstSimplify/extract-element.ll
@@ -62,10 +62,7 @@ define float @extract_element_splat_constant_vector_variable_index(i32 %y) {
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %index) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <3 x i32> poison, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <3 x i32> [[IN]], <3 x i32> poison, <3 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <3 x i32> [[SPLAT]], i32 [[INDEX:%.*]]
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 [[V:%.*]]
 ;
   %in = insertelement <3 x i32> poison, i32 %v, i32 0
   %splat = shufflevector <3 x i32> %in, <3 x i32> poison, <3 x i32> zeroinitializer

diff  --git a/llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll b/llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll
index a7739490ce7a..11fae86c830c 100644
--- a/llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll
+++ b/llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll
@@ -210,10 +210,7 @@ define i32 @extractelement_splat_constant_index(i32 %v) {
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %idx) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; 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> undef, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 [[IDX:%.*]]
-; 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> undef, <vscale x 4 x i32> zeroinitializer

diff  --git a/llvm/test/Transforms/InstSimplify/vscale.ll b/llvm/test/Transforms/InstSimplify/vscale.ll
index 108a9bb7726a..79ada2af95b7 100644
--- a/llvm/test/Transforms/InstSimplify/vscale.ll
+++ b/llvm/test/Transforms/InstSimplify/vscale.ll
@@ -210,10 +210,7 @@ define i32 @extractelement_splat_constant_index(i32 %v) {
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %idx) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; 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 [[IDX:%.*]]
-; 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