SLPVectorizer: handle scalars that are extracted from vectors (using ExtractElementInst).

llvm-svn: 184325

1 files changed, 44 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Vectorize/VecUtils.cpp b/llvm/lib/Transforms/Vectorize/VecUtils.cpp
index 1e97ed400be..e79f08a56dd 100644
--- a/llvm/lib/Transforms/Vectorize/VecUtils.cpp
+++ b/llvm/lib/Transforms/Vectorize/VecUtils.cpp
@@ -331,6 +331,34 @@ int BoUpSLP::getTreeCost(ArrayRef<Value *> VL) {
   return getTreeCost_rec(VL, 0);
 }
 
+static bool CanReuseExtract(ArrayRef<Value *> VL, unsigned VF,
+                            VectorType *VecTy) {
+  // Check if all of the extracts come from the same vector and from the
+  // correct offset.
+  Value *VL0 = VL[0];
+  ExtractElementInst *E0 = cast<ExtractElementInst>(VL0);
+  Value *Vec = E0->getOperand(0);
+
+  // We have to extract from the same vector type.
+  if (Vec->getType() != VecTy)
+    return false;
+
+  // Check that all of the indices extract from the correct offset.
+  ConstantInt *CI = dyn_cast<ConstantInt>(E0->getOperand(1));
+  if (!CI || CI->getZExtValue())
+    return false;
+
+  for (unsigned i = 1, e = VF; i < e; ++i) {
+    ExtractElementInst *E = cast<ExtractElementInst>(VL[i]);
+    ConstantInt *CI = dyn_cast<ConstantInt>(E->getOperand(1));
+
+    if (!CI || CI->getZExtValue() != i || E->getOperand(0) != Vec)
+      return false;
+  }
+
+  return true;
+}
+
 void BoUpSLP::getTreeUses_rec(ArrayRef<Value *> VL, unsigned Depth) {
   if (Depth == RecursionMaxDepth) return;
 
@@ -386,6 +414,12 @@ void BoUpSLP::getTreeUses_rec(ArrayRef<Value *> VL, unsigned Depth) {
   }
 
   switch (Opcode) {
+    case Instruction::ExtractElement: {
+      VectorType *VecTy = VectorType::get(VL[0]->getType(), VL.size());
+      // No need to follow ExtractElements that are going to be optimized away.
+      if (CanReuseExtract(VL, VL.size(), VecTy)) return;
+      // Fall through.
+    }
     case Instruction::ZExt:
     case Instruction::SExt:
     case Instruction::FPToUI:
@@ -523,6 +557,11 @@ int BoUpSLP::getTreeCost_rec(ArrayRef<Value *> VL, unsigned Depth) {
         TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, i);
 
   switch (Opcode) {
+  case Instruction::ExtractElement: {
+    if (CanReuseExtract(VL, VL.size(), VecTy))
+      return 0;
+    return getScalarizationCost(VecTy);
+  }
   case Instruction::ZExt:
   case Instruction::SExt:
   case Instruction::FPToUI:
@@ -786,6 +825,11 @@ Value *BoUpSLP::vectorizeTree_rec(ArrayRef<Value *> VL, int VF) {
   }
 
   switch (Opcode) {
+  case Instruction::ExtractElement: {
+    if (CanReuseExtract(VL, VL.size(), VecTy))
+      return VL0->getOperand(0);
+    return Scalarize(VL, VecTy);
+  }
   case Instruction::ZExt:
   case Instruction::SExt:
   case Instruction::FPToUI: