Index: llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp =================================================================== --- llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp +++ llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -539,13 +539,12 @@ /// can be a truncate instruction). void buildScalarSteps(Value *ScalarIV, Value *Step, Value *EntryVal); - /// Create a vector induction phi node based on an existing scalar one. This - /// currently only works for integer induction variables with a constant - /// step. \p EntryVal is the value from the original loop that maps to the - /// vector phi node. If \p EntryVal is a truncate instruction, instead of - /// widening the original IV, we widen a version of the IV truncated to \p - /// EntryVal's type. - void createVectorIntInductionPHI(const InductionDescriptor &II, + /// Create a vector induction phi node based on an existing scalar one. \p + /// EntryVal is the value from the original loop that maps to the vector phi + /// node, and \p Step is the loop-invariant step. If \p EntryVal is a + /// truncate instruction, instead of widening the original IV, we widen a + /// version of the IV truncated to \p EntryVal's type. + void createVectorIntInductionPHI(const InductionDescriptor &II, Value *Step, Instruction *EntryVal); /// Widen an integer induction variable \p IV. If \p Trunc is provided, the @@ -2038,16 +2037,7 @@ return false; // If the truncated value is not an induction variable, return false. - if (!Legal->isInductionVariable(Op)) - return false; - - // Lastly, we only consider an induction variable truncate to be - // optimizable if it has a constant step. - // - // TODO: Expand optimizable truncates to include truncations of induction - // variables having loop-invariant steps. - auto ID = Legal->getInductionVars()->lookup(cast(Op)); - return ID.getConstIntStepValue(); + return Legal->isInductionVariable(Op); } private: @@ -2366,26 +2356,34 @@ } void InnerLoopVectorizer::createVectorIntInductionPHI( - const InductionDescriptor &II, Instruction *EntryVal) { + const InductionDescriptor &II, Value *Step, Instruction *EntryVal) { Value *Start = II.getStartValue(); - ConstantInt *Step = II.getConstIntStepValue(); - assert(Step && "Can not widen an IV with a non-constant step"); + assert(Step->getType()->isIntegerTy() && + "Cannot widen an IV having a step with a non-integer type"); // Construct the initial value of the vector IV in the vector loop preheader auto CurrIP = Builder.saveIP(); Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator()); if (isa(EntryVal)) { auto *TruncType = cast(EntryVal->getType()); - Step = ConstantInt::getSigned(TruncType, Step->getSExtValue()); + Step = Builder.CreateTrunc(Step, TruncType); Start = Builder.CreateCast(Instruction::Trunc, Start, TruncType); } Value *SplatStart = Builder.CreateVectorSplat(VF, Start); Value *SteppedStart = getStepVector(SplatStart, 0, Step); + + // Create a vector splat to use in the induction update. + // + // FIXME: If the step is non-constant, we create the vector splat with + // IRBuilder. IRBuilder can constant-fold the multiply, but it doesn't + // handle a constant vector splat. + auto *ConstVF = ConstantInt::getSigned(Step->getType(), VF); + auto *Mul = Builder.CreateMul(Step, ConstVF); + Value *SplatVF = isa(Mul) + ? ConstantVector::getSplat(VF, cast(Mul)) + : Builder.CreateVectorSplat(VF, Mul); Builder.restoreIP(CurrIP); - Value *SplatVF = - ConstantVector::getSplat(VF, ConstantInt::getSigned(Start->getType(), - VF * Step->getSExtValue())); // We may need to add the step a number of times, depending on the unroll // factor. The last of those goes into the PHI. PHINode *VecInd = PHINode::Create(SteppedStart->getType(), 2, "vec.ind", @@ -2440,9 +2438,6 @@ // induction variable. Value *ScalarIV = nullptr; - // The step of the induction. - Value *Step = nullptr; - // The value from the original loop to which we are mapping the new induction // variable. Instruction *EntryVal = Trunc ? cast(Trunc) : IV; @@ -2455,44 +2450,42 @@ // least one user in the loop that is not widened. auto NeedsScalarIV = VF > 1 && needsScalarInduction(EntryVal); - // If the induction variable has a constant integer step value, go ahead and - // get it now. - if (ID.getConstIntStepValue()) - Step = ID.getConstIntStepValue(); + // Generate code for the induction step. Note that induction steps are + // required to be loop-invariant + assert(PSE.getSE()->isLoopInvariant(ID.getStep(), OrigLoop) && + "Induction step should be loop invariant"); + auto &DL = OrigLoop->getHeader()->getModule()->getDataLayout(); + SCEVExpander Exp(*PSE.getSE(), DL, "induction"); + Value *Step = Exp.expandCodeFor(ID.getStep(), ID.getStep()->getType(), + LoopVectorPreHeader->getTerminator()); // Try to create a new independent vector induction variable. If we can't // create the phi node, we will splat the scalar induction variable in each // loop iteration. - if (VF > 1 && Step && !shouldScalarizeInstruction(EntryVal)) { - createVectorIntInductionPHI(ID, EntryVal); + if (VF > 1 && !shouldScalarizeInstruction(EntryVal)) { + createVectorIntInductionPHI(ID, Step, EntryVal); VectorizedIV = true; } // If we haven't yet vectorized the induction variable, or if we will create // a scalar one, we need to define the scalar induction variable and step // values. If we were given a truncation type, truncate the canonical - // induction variable and constant step. Otherwise, derive these values from - // the induction descriptor. + // induction variable and step. Otherwise, derive these values from the + // induction descriptor. if (!VectorizedIV || NeedsScalarIV) { if (Trunc) { auto *TruncType = cast(Trunc->getType()); - assert(Step && "Truncation requires constant integer step"); - auto StepInt = cast(Step)->getSExtValue(); + assert(Step->getType()->isIntegerTy() && + "Truncation requires an integer step"); ScalarIV = Builder.CreateCast(Instruction::Trunc, Induction, TruncType); - Step = ConstantInt::getSigned(TruncType, StepInt); + Step = Builder.CreateTrunc(Step, TruncType); } else { ScalarIV = Induction; - auto &DL = OrigLoop->getHeader()->getModule()->getDataLayout(); if (IV != OldInduction) { ScalarIV = Builder.CreateSExtOrTrunc(ScalarIV, IV->getType()); ScalarIV = ID.transform(Builder, ScalarIV, PSE.getSE(), DL); ScalarIV->setName("offset.idx"); } - if (!Step) { - SCEVExpander Exp(*PSE.getSE(), DL, "induction"); - Step = Exp.expandCodeFor(ID.getStep(), ID.getStep()->getType(), - &*Builder.GetInsertPoint()); - } } } Index: llvm/trunk/test/Transforms/LoopVectorize/induction-step.ll =================================================================== --- llvm/trunk/test/Transforms/LoopVectorize/induction-step.ll +++ llvm/trunk/test/Transforms/LoopVectorize/induction-step.ll @@ -12,11 +12,30 @@ ;} ; CHECK-LABEL: @induction_with_global( -; CHECK: %[[INT_INC:.*]] = load i32, i32* @int_inc, align 4 -; CHECK: vector.body: -; CHECK: %[[VAR1:.*]] = insertelement <8 x i32> undef, i32 %[[INT_INC]], i32 0 -; CHECK: %[[VAR2:.*]] = shufflevector <8 x i32> %[[VAR1]], <8 x i32> undef, <8 x i32> zeroinitializer -; CHECK: mul <8 x i32> , %[[VAR2]] +; CHECK: for.body.lr.ph: +; CHECK-NEXT: [[TMP0:%.*]] = load i32, i32* @int_inc, align 4 +; CHECK: vector.ph: +; CHECK-NEXT: [[DOTSPLATINSERT:%.*]] = insertelement <8 x i32> undef, i32 %init, i32 0 +; CHECK-NEXT: [[DOTSPLAT:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: [[DOTSPLATINSERT2:%.*]] = insertelement <8 x i32> undef, i32 [[TMP0]], i32 0 +; CHECK-NEXT: [[DOTSPLAT3:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT2]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: [[TMP6:%.*]] = mul <8 x i32> , [[DOTSPLAT3]] +; CHECK-NEXT: [[INDUCTION4:%.*]] = add <8 x i32> [[DOTSPLAT]], [[TMP6]] +; CHECK-NEXT: [[TMP7:%.*]] = mul i32 [[TMP0]], 8 +; CHECK-NEXT: [[DOTSPLATINSERT5:%.*]] = insertelement <8 x i32> undef, i32 [[TMP7]], i32 0 +; CHECK-NEXT: [[DOTSPLAT6:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT5]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: br label %vector.body +; CHECK: vector.body: +; CHECK-NEXT: %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ] +; CHECK-NEXT: %vec.ind = phi <8 x i32> [ [[INDUCTION4]], %vector.ph ], [ %vec.ind.next, %vector.body ] +; CHECK: [[TMP8:%.*]] = add i64 %index, 0 +; CHECK-NEXT: [[TMP9:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[TMP8]] +; CHECK-NEXT: [[TMP10:%.*]] = getelementptr i32, i32* [[TMP9]], i32 0 +; CHECK-NEXT: [[TMP11:%.*]] = bitcast i32* [[TMP10]] to <8 x i32>* +; CHECK-NEXT: store <8 x i32> %vec.ind, <8 x i32>* [[TMP11]], align 4 +; CHECK: %index.next = add i64 %index, 8 +; CHECK-NEXT: %vec.ind.next = add <8 x i32> %vec.ind, [[DOTSPLAT6]] +; CHECK: br i1 {{.*}}, label %middle.block, label %vector.body target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128" @@ -66,13 +85,28 @@ ;} ; CHECK-LABEL: @induction_with_loop_inv( -; CHECK: for.cond1.preheader: -; CHECK: %[[INDVAR0:.*]] = phi i32 [ 0, -; CHECK: %[[INDVAR1:.*]] = phi i32 [ 0, -; CHECK: vector.body: -; CHECK: %[[VAR1:.*]] = insertelement <8 x i32> undef, i32 %[[INDVAR1]], i32 0 -; CHECK: %[[VAR2:.*]] = shufflevector <8 x i32> %[[VAR1]], <8 x i32> undef, <8 x i32> zeroinitializer -; CHECK: mul <8 x i32> , %[[VAR2]] +; CHECK: vector.ph: +; CHECK-NEXT: [[DOTSPLATINSERT:%.*]] = insertelement <8 x i32> undef, i32 %x.011, i32 0 +; CHECK-NEXT: [[DOTSPLAT:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: [[DOTSPLATINSERT2:%.*]] = insertelement <8 x i32> undef, i32 %j.012, i32 0 +; CHECK-NEXT: [[DOTSPLAT3:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT2]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: [[TMP4:%.*]] = mul <8 x i32> , [[DOTSPLAT3]] +; CHECK-NEXT: [[INDUCTION4:%.*]] = add <8 x i32> [[DOTSPLAT]], [[TMP4]] +; CHECK-NEXT: [[TMP5:%.*]] = mul i32 %j.012, 8 +; CHECK-NEXT: [[DOTSPLATINSERT5:%.*]] = insertelement <8 x i32> undef, i32 [[TMP5]], i32 0 +; CHECK-NEXT: [[DOTSPLAT6:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT5]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: br label %vector.body +; CHECK: vector.body: +; CHECK-NEXT: %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ] +; CHECK-NEXT: %vec.ind = phi <8 x i32> [ [[INDUCTION4]], %vector.ph ], [ %vec.ind.next, %vector.body ] +; CHECK: [[TMP6:%.*]] = add i64 %index, 0 +; CHECK-NEXT: [[TMP7:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[TMP6]] +; CHECK-NEXT: [[TMP8:%.*]] = getelementptr i32, i32* [[TMP7]], i32 0 +; CHECK-NEXT: [[TMP9:%.*]] = bitcast i32* [[TMP8]] to <8 x i32>* +; CHECK-NEXT: store <8 x i32> %vec.ind, <8 x i32>* [[TMP9]], align 4 +; CHECK: %index.next = add i64 %index, 8 +; CHECK-NEXT: %vec.ind.next = add <8 x i32> %vec.ind, [[DOTSPLAT6]] +; CHECK: br i1 {{.*}}, label %middle.block, label %vector.body define i32 @induction_with_loop_inv(i32 %init, i32* noalias nocapture %A, i32 %N, i32 %M) { entry: @@ -122,3 +156,46 @@ %x.0.lcssa = phi i32 [ %init, %entry ], [ %x.1.lcssa.lcssa, %for.end6.loopexit ] ret i32 %x.0.lcssa } + + +; CHECK-LABEL: @non_primary_iv_loop_inv_trunc( +; CHECK: vector.ph: +; CHECK: [[TMP3:%.*]] = trunc i64 %step to i32 +; CHECK-NEXT: [[DOTSPLATINSERT5:%.*]] = insertelement <8 x i32> undef, i32 [[TMP3]], i32 0 +; CHECK-NEXT: [[DOTSPLAT6:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT5]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: [[TMP4:%.*]] = mul <8 x i32> , [[DOTSPLAT6]] +; CHECK-NEXT: [[INDUCTION7:%.*]] = add <8 x i32> zeroinitializer, [[TMP4]] +; CHECK-NEXT: [[TMP5:%.*]] = mul i32 [[TMP3]], 8 +; CHECK-NEXT: [[DOTSPLATINSERT8:%.*]] = insertelement <8 x i32> undef, i32 [[TMP5]], i32 0 +; CHECK-NEXT: [[DOTSPLAT9:%.*]] = shufflevector <8 x i32> [[DOTSPLATINSERT8]], <8 x i32> undef, <8 x i32> zeroinitializer +; CHECK-NEXT: br label %vector.body +; CHECK: vector.body: +; CHECK-NEXT: %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ] +; CHECK: [[VEC_IND10:%.*]] = phi <8 x i32> [ [[INDUCTION7]], %vector.ph ], [ [[VEC_IND_NEXT11:%.*]], %vector.body ] +; CHECK: [[TMP6:%.*]] = add i64 %index, 0 +; CHECK-NEXT: [[TMP7:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[TMP6]] +; CHECK-NEXT: [[TMP8:%.*]] = getelementptr i32, i32* [[TMP7]], i32 0 +; CHECK-NEXT: [[TMP9:%.*]] = bitcast i32* [[TMP8]] to <8 x i32>* +; CHECK-NEXT: store <8 x i32> [[VEC_IND10]], <8 x i32>* [[TMP9]], align 4 +; CHECK-NEXT: %index.next = add i64 %index, 8 +; CHECK: [[VEC_IND_NEXT11]] = add <8 x i32> [[VEC_IND10]], [[DOTSPLAT9]] +; CHECK: br i1 {{.*}}, label %middle.block, label %vector.body + +define void @non_primary_iv_loop_inv_trunc(i32* %a, i64 %n, i64 %step) { +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, %step + %cond = icmp slt i64 %i.next, %n + br i1 %cond, label %for.body, label %for.end + +for.end: + ret void +}