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 @@ -94,11 +94,22 @@ } Value *IRBuilderBase::CreateStepVector(Type *DstType, const Twine &Name) { - if (isa(DstType)) - return CreateIntrinsic(Intrinsic::experimental_stepvector, {DstType}, {}, - nullptr, Name); - Type *STy = DstType->getScalarType(); + if (isa(DstType)) { + Type *StepVecType = DstType; + // TODO: We expect this special case (element type < 8 bits) to be + // temporary - once the intrinsic properly supports < 8 bits this code + // can be removed. + if (STy->getScalarSizeInBits() < 8) + StepVecType = + VectorType::get(getInt8Ty(), cast(DstType)); + Value *Res = CreateIntrinsic(Intrinsic::experimental_stepvector, + {StepVecType}, {}, nullptr, Name); + if (StepVecType != DstType) + Res = CreateTrunc(Res, DstType); + return Res; + } + unsigned NumEls = cast(DstType)->getNumElements(); // Create a vector of consecutive numbers from zero to VF. 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 @@ -214,6 +214,25 @@ EXPECT_EQ(Call->getIntrinsicID(), Intrinsic::experimental_stepvector); } +TEST_F(IRBuilderTest, CreateStepVectorI3) { + IRBuilder<> Builder(BB); + + // Scalable vectors + Type *DstVecTy = VectorType::get(IntegerType::get(Ctx, 3), 2, true); + Type *VecI8Ty = VectorType::get(Builder.getInt8Ty(), 2, true); + Value *StepVec = Builder.CreateStepVector(DstVecTy); + EXPECT_TRUE(isa(StepVec)); + TruncInst *Trunc = cast(StepVec); + EXPECT_EQ(Trunc->getDestTy(), DstVecTy); + EXPECT_EQ(Trunc->getSrcTy(), VecI8Ty); + EXPECT_TRUE(isa(Trunc->getOperand(0))); + + CallInst *Call = cast(Trunc->getOperand(0)); + FunctionType *FTy = Call->getFunctionType(); + EXPECT_EQ(FTy->getReturnType(), VecI8Ty); + EXPECT_EQ(Call->getIntrinsicID(), Intrinsic::experimental_stepvector); +} + TEST_F(IRBuilderTest, ConstrainedFP) { IRBuilder<> Builder(BB); Value *V;