diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -610,7 +610,8 @@ /// represented as. void truncateToMinimalBitwidths(VPTransformState &State); - /// This function adds (StartIdx, StartIdx + Step, StartIdx + 2*Step, ...) + /// This function adds + /// (StartIdx * Step, (StartIdx + 1) * Step, (StartIdx + 2) * Step, ...) /// to each vector element of Val. The sequence starts at StartIndex. /// \p Opcode is relevant for FP induction variable. virtual Value *getStepVector(Value *Val, int StartIdx, Value *Step, @@ -2452,8 +2453,10 @@ Value *InnerLoopVectorizer::getStepVector(Value *Val, int StartIdx, Value *Step, Instruction::BinaryOps BinOp) { // Create and check the types. - auto *ValVTy = cast(Val->getType()); - int VLen = ValVTy->getNumElements(); + assert(isa(Val->getType()) && + "Creation of scalable step vector not yet supported"); + auto *ValVTy = cast(Val->getType()); + ElementCount VLen = ValVTy->getElementCount(); Type *STy = Val->getType()->getScalarType(); assert((STy->isIntegerTy() || STy->isFloatingPointTy()) && @@ -2462,34 +2465,36 @@ SmallVector Indices; - if (STy->isIntegerTy()) { - // Create a vector of consecutive numbers from zero to VF. - for (int i = 0; i < VLen; ++i) - Indices.push_back(ConstantInt::get(STy, StartIdx + i)); + // Create a vector of consecutive numbers from zero to VF. + VectorType *InitVecValVTy = ValVTy; + Type *InitVecValSTy = STy; + if (STy->isFloatingPointTy()) { + InitVecValSTy = + IntegerType::get(STy->getContext(), STy->getScalarSizeInBits()); + InitVecValVTy = VectorType::get(InitVecValSTy, VLen); + } + Value *InitVec = Builder.CreateStepVector(InitVecValVTy); + + // Add on StartIdx + Value *StartIdxSplat = Builder.CreateVectorSplat( + VLen, ConstantInt::get(InitVecValSTy, StartIdx)); + InitVec = Builder.CreateAdd(InitVec, StartIdxSplat); - // Add the consecutive indices to the vector value. - Constant *Cv = ConstantVector::get(Indices); - assert(Cv->getType() == Val->getType() && "Invalid consecutive vec"); + if (STy->isIntegerTy()) { Step = Builder.CreateVectorSplat(VLen, Step); assert(Step->getType() == Val->getType() && "Invalid step vec"); // FIXME: The newly created binary instructions should contain nsw/nuw flags, // which can be found from the original scalar operations. - Step = Builder.CreateMul(Cv, Step); + Step = Builder.CreateMul(InitVec, Step); return Builder.CreateAdd(Val, Step, "induction"); } // Floating point induction. assert((BinOp == Instruction::FAdd || BinOp == Instruction::FSub) && "Binary Opcode should be specified for FP induction"); - // Create a vector of consecutive numbers from zero to VF. - for (int i = 0; i < VLen; ++i) - Indices.push_back(ConstantFP::get(STy, (double)(StartIdx + i))); - - // Add the consecutive indices to the vector value. - // Floating-point operations inherit FMF via the builder's flags. - Constant *Cv = ConstantVector::get(Indices); + InitVec = Builder.CreateUIToFP(InitVec, ValVTy); Step = Builder.CreateVectorSplat(VLen, Step); - Value *MulOp = Builder.CreateFMul(Cv, Step); + Value *MulOp = Builder.CreateFMul(InitVec, Step); return Builder.CreateBinOp(BinOp, Val, MulOp, "induction"); } @@ -4766,12 +4771,12 @@ // phi as base and a vectorized version of the step value // () as offset. for (unsigned Part = 0; Part < State.UF; ++Part) { - SmallVector Indices; + Type *VecPhiType = VectorType::get(PhiType, State.VF); + Value *StartOffset = + ConstantInt::get(VecPhiType, Part * State.VF.getKnownMinValue()); // Create a vector of consecutive numbers from zero to VF. - for (unsigned i = 0; i < State.VF.getKnownMinValue(); ++i) - Indices.push_back( - ConstantInt::get(PhiType, i + Part * State.VF.getKnownMinValue())); - Constant *StartOffset = ConstantVector::get(Indices); + StartOffset = + Builder.CreateAdd(StartOffset, Builder.CreateStepVector(VecPhiType)); Value *GEP = Builder.CreateGEP( ScStValueType->getPointerElementType(), NewPointerPhi,