[LV] Extend trunc optimization to all IVs with constant integer steps

This patch extends the optimization of truncations whose operand is an
induction variable with a constant integer step. Previously we were only
applying this optimization to the primary induction variable. However, the cost
model assumes the optimization is applied to the truncation of all integer
induction variables (even regardless of step type). The transformation is now
applied to the other induction variables, and I've updated the cost model to
ensure it is better in sync with the transformation we actually perform.

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

llvm-svn: 294967

1 files changed, 20 insertions, 12 deletions

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 401b942d307..5513c7fd62a 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4879,12 +4879,15 @@ void InnerLoopVectorizer::vectorizeBlockInLoop(BasicBlock *BB, PhiVector *PV) {
       // induction variable. Notice that we can only optimize the 'trunc' case
       // because (a) FP conversions lose precision, (b) sext/zext may wrap, and
       // (c) other casts depend on pointer size.
-      auto ID = Legal->getInductionVars()->lookup(OldInduction);
-      if (isa<TruncInst>(CI) && CI->getOperand(0) == OldInduction &&
-          ID.getConstIntStepValue()) {
-        widenIntInduction(OldInduction, cast<TruncInst>(CI));
-        break;
-      }
+      if (auto *Trunc = dyn_cast<TruncInst>(CI))
+        if (auto *Phi = dyn_cast<PHINode>(Trunc->getOperand(0))) {
+          auto II = Legal->getInductionVars()->find(Phi);
+          if (II != Legal->getInductionVars()->end())
+            if (II->second.getConstIntStepValue()) {
+              widenIntInduction(Phi, Trunc);
+              break;
+            }
+        }
 
       /// Vectorize casts.
       Type *DestTy =
@@ -7224,12 +7227,17 @@ unsigned LoopVectorizationCostModel::getInstructionCost(Instruction *I,
   case Instruction::Trunc:
   case Instruction::FPTrunc:
   case Instruction::BitCast: {
-    // We optimize the truncation of induction variable.
-    // The cost of these is the same as the scalar operation.
-    if (I->getOpcode() == Instruction::Trunc &&
-        Legal->isInductionVariable(I->getOperand(0)))
-      return TTI.getCastInstrCost(I->getOpcode(), I->getType(),
-                                  I->getOperand(0)->getType());
+    // We optimize the truncation of induction variables having constant
+    // integer steps. The cost of these truncations is the same as the scalar
+    // operation.
+    if (auto *Trunc = dyn_cast<TruncInst>(I))
+      if (auto *Phi = dyn_cast<PHINode>(Trunc->getOperand(0))) {
+        auto II = Legal->getInductionVars()->find(Phi);
+        if (II != Legal->getInductionVars()->end())
+          if (II->second.getConstIntStepValue())
+            return TTI.getCastInstrCost(Instruction::Trunc, Trunc->getDestTy(),
+                                        Trunc->getSrcTy());
+      }
 
     Type *SrcScalarTy = I->getOperand(0)->getType();
     Type *SrcVecTy = ToVectorTy(SrcScalarTy, VF);