diff --git a/llvm/lib/IR/IRBuilder.cpp b/llvm/lib/IR/IRBuilder.cpp --- a/llvm/lib/IR/IRBuilder.cpp +++ b/llvm/lib/IR/IRBuilder.cpp @@ -81,8 +81,10 @@ } Value *IRBuilderBase::CreateVScale(Constant *Scaling, const Twine &Name) { - Module *M = GetInsertBlock()->getParent()->getParent(); assert(isa(Scaling) && "Expected constant integer"); + if (cast(Scaling)->isZero()) + return Scaling; + Module *M = GetInsertBlock()->getParent()->getParent(); Function *TheFn = Intrinsic::getDeclaration(M, Intrinsic::vscale, {Scaling->getType()}); CallInst *CI = createCallHelper(TheFn, {}, this, Name); 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 @@ -4773,10 +4773,9 @@ "Currently unsupported for scalable vectors"); unsigned Lanes = IsUniform ? 1 : State.VF.getFixedValue(); - Value *RuntimeVF = getRuntimeVF(Builder, PtrInd->getType(), VF); for (unsigned Part = 0; Part < UF; ++Part) { - Value *PartStart = Builder.CreateMul( - RuntimeVF, ConstantInt::get(PtrInd->getType(), Part)); + Value *PartStart = createStepForVF( + Builder, ConstantInt::get(PtrInd->getType(), Part), VF); for (unsigned Lane = 0; Lane < Lanes; ++Lane) { Value *Idx = Builder.CreateAdd( PartStart, ConstantInt::get(PtrInd->getType(), Lane)); diff --git a/llvm/unittests/IR/IRBuilderTest.cpp b/llvm/unittests/IR/IRBuilderTest.cpp --- a/llvm/unittests/IR/IRBuilderTest.cpp +++ b/llvm/unittests/IR/IRBuilderTest.cpp @@ -180,6 +180,14 @@ EXPECT_EQ(FTy->getParamType(i), ArgTys[i]->getType()); } +TEST_F(IRBuilderTest, CreateVScale) { + IRBuilder<> Builder(BB); + + Constant *Zero = Builder.getInt32(0); + Value *VScale = Builder.CreateVScale(Zero); + EXPECT_TRUE(isa(VScale) && cast(VScale)->isZero()); +} + TEST_F(IRBuilderTest, CreateStepVector) { IRBuilder<> Builder(BB);