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

[mssimpso]

Summary

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 transform 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.

[mkuper]

LGTM

[mssimpso]

Thanks Michael!

[mssimpso]

I'm reopening this review since I reverted the original patch. It caused execution time slowdowns in a few test-suite programs that were caught by the clang-cmake-aarch64-quick bot. It looks like the cause is related to the fact that we're not cost-modeling this optimization very well. Some truncations are free - if we replace those truncations with an induction variable, we'll end up with an extra instruction (the IV update) on every iteration of the loop. I'll post an updated patch shortly.

[mssimpso]

• I updated patch so that we extend the truncate optimization to non-primary induction variables only if the truncate isn't already free.
• I factored out the logic into a helper so we can keep things in sync easier.
• I added a test case for AArch64 showing that we keep a free truncate.

[mkuper]

Heh, interesting.
Anyway, LGTM.

lib/Transforms/Vectorize/LoopVectorize.cpp

2037

Unrelated, but after this lands, I think I'll change the name to getPrimaryInduction() or getPrimaryIV(). This is a really bad name at this point.

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

2037

I agree. Thanks Michael.

lib/Transforms/Vectorize/LoopVectorize.cpp

private:
   /// vector widths. The cost that is returned is *not* normalized by
   /// the factor width.
   VectorizationCostTy expectedCost(unsigned VF);
 
   /// Returns the execution time cost of an instruction for a given vector
   /// width. Vector width of one means scalar.
   VectorizationCostTy getInstructionCost(Instruction *I, unsigned VF);
 
   /// The cost-computation logic from getInstructionCost which provides
   /// the vector type as an output parameter.
   unsigned getInstructionCost(Instruction *I, unsigned VF, Type *&VectorTy);
 
   /// Calculate vectorization cost of memory instruction \p I.
   unsigned getMemoryInstructionCost(Instruction *I, unsigned VF);
 
   /// The cost computation for scalarized memory instruction.
   unsigned getMemInstScalarizationCost(Instruction *I, unsigned VF);
for (Instruction &I : *BB) {
     case Instruction::BitCast: {
       auto *CI = dyn_cast<CastInst>(&I);
       setDebugLocFromInst(Builder, CI);
 
       // Optimize the special case where the source is a constant integer
       // 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));
+      if (isa<TruncInst>(CI))
+        if (auto *Phi = dyn_cast<PHINode>(CI->getOperand(0))) {
+          auto II = Legal->getInductionVars()->find(Phi);
+          if (II != Legal->getInductionVars()->end())
+            if (II->second.getConstIntStepValue()) {
+              widenIntInduction(Phi, cast<TruncInst>(CI));
               break;
             }
+        }
 
       /// Vectorize casts.
       Type *DestTy =
           (VF == 1) ? CI->getType() : VectorType::get(CI->getType(), VF);
 
       const VectorParts &A = getVectorValue(CI->getOperand(0));
       VectorParts Entry(UF);
       for (unsigned Part = 0; Part < UF; ++Part)
unsigned LoopVectorizationCostModel::getInstructionCost(Instruction *I,
   case Instruction::FPExt:
   case Instruction::PtrToInt:
   case Instruction::IntToPtr:
   case Instruction::SIToFP:
   case Instruction::UIToFP:
   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)))
+    // We optimize the truncation of induction variables having constant
+    // integer steps. The cost of these truncations is the same as the scalar
+    // operation.
+    if (isa<TruncInst>(I))
+      if (auto *Phi = dyn_cast<PHINode>(I->getOperand(0))) {
+        auto II = Legal->getInductionVars()->find(Phi);
+        if (II != Legal->getInductionVars()->end())
+          if (II->second.getConstIntStepValue())
             return TTI.getCastInstrCost(I->getOpcode(), I->getType(),
                                         I->getOperand(0)->getType());
+      }
 
     Type *SrcScalarTy = I->getOperand(0)->getType();
     Type *SrcVecTy = ToVectorTy(SrcScalarTy, VF);
     if (canTruncateToMinimalBitwidth(I, VF)) {
       // This cast is going to be shrunk. This may remove the cast or it might
       // turn it into slightly different cast. For example, if MinBW == 16,
       // "zext i8 %1 to i32" becomes "zext i8 %1 to i16".
       //

test/Transforms/LoopVectorize/induction.ll

for.body:
   store i32 %indvars.iv, i32* %arrayidx, align 4
   %indvars.iv.next = add nuw nsw i32 %indvars.iv, 1
   %exitcond = icmp eq i32 %indvars.iv.next, %k
   br i1 %exitcond, label %exit, label %for.body
 
 exit:
   ret void
 }
+
+; CHECK-LABEL: @non_primary_iv_trunc(
+; CHECK:       vector.body:
+; CHECK-NEXT:    %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+; CHECK:         [[VEC_IND:%.*]] = phi <2 x i32> [ <i32 0, i32 2>, %vector.ph ], [ [[VEC_IND_NEXT:%.*]], %vector.body ]
+; CHECK:         [[TMP3:%.*]] = add i64 %index, 0
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i32, i32* %a, i64 [[TMP3]]
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i32, i32* [[TMP4]], i32 0
+; CHECK-NEXT:    [[TMP6:%.*]] = bitcast i32* [[TMP5]] to <2 x i32>*
+; CHECK-NEXT:    store <2 x i32> [[VEC_IND]], <2 x i32>* [[TMP6]], align 4
+; CHECK-NEXT:    %index.next = add i64 %index, 2
+; CHECK:         [[VEC_IND_NEXT]] = add <2 x i32> [[VEC_IND]], <i32 4, i32 4>
+; CHECK:         br i1 {{.*}}, label %middle.block, label %vector.body
+define void @non_primary_iv_trunc(i32* %a, i64 %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
+  %j = phi i64 [ %j.next, %for.body ], [ 0, %entry ]
+  %tmp0 = getelementptr inbounds i32, i32* %a, i64 %i
+  %tmp1 = trunc i64 %j to i32
+  store i32 %tmp1, i32* %tmp0, align 4
+  %i.next = add nuw nsw i64 %i, 1
+  %j.next = add nuw nsw i64 %j, 2
+  %cond = icmp slt i64 %i.next, %n
+  br i1 %cond, label %for.body, label %for.end
+
+for.end:
+  ret void
+}

test/Transforms/LoopVectorize/reverse_iter.ll

 ; RUN: opt < %s  -loop-vectorize -force-vector-interleave=1 -force-vector-width=4 -dce -instcombine -S | FileCheck %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
-; Make sure that the reverse iterators are calculated using 64bit arithmetic, not 32.
+; PR15882: This test ensures that we do not produce wrapping arithmetic when
+; creating constant reverse step vectors.
 ;
 ; int foo(int n, int *A) {
 ;   int sum;
 ;   for (int i=n; i > 0; i--)
 ;     sum += A[i*2];
 ;   return sum;
 ; }
 ;
 
 ;CHECK-LABEL: @foo(
-;CHECK:  <i64 0, i64 -1, i64 -2, i64 -3>
+;CHECK:  <i32 0, i32 -1, i32 -2, i32 -3>
 ;CHECK: ret
 define i32 @foo(i32 %n, i32* nocapture %A) {
   %1 = icmp sgt i32 %n, 0
   br i1 %1, label %.lr.ph, label %._crit_edge
 
 .lr.ph:                                           ; preds = %0
   %2 = sext i32 %n to i64
   br label %3