[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

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index f713d53..35d1e34 100644 (file)
--- 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 8259598..fef0952 100644 (file)
--- 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 a773949..11fae86 100644 (file)
--- 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 108a9bb..79ada2a 100644 (file)
--- 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