[IR] Change CreateStepVector to work with element types smaller than i8

[david-arm]

Summary

Currently the stepvector intrinsic only supports element types that
are integers of size 8 bits or more. This patch adds support for the
creation of stepvectors with smaller element types by creating
the intrinsic with i8 elements that we then truncate to the requested
size.

It's not currently possible to write a vectoriser test to exercise
this code path so I have added a unit test here:

llvm/unittests/IR/IRBuilderTest.cpp

[efriedma]

Where does the restriction on the element type of stepvector come from? I mean, I see the code in the verifier, but I'm not sure why it's there...

[paulwalker-arm]

When first introduced, the stepvector intrinsic defined wrapping as undefined behaviour. At the time we pointed out this effectively made its usage for i1 vectors practically useless because only the first two lanes of a scalable vector could ever have a defined value. This was discussed during an SVE sync call and there wasn't agreement to define the wrapping behaviour so as a compromised the supported element type was restricted to i8 or bigger.

Since then, ISD::STEP_VECTOR has changed so its wrapping behaviour is now defined and thus I see little reason for the intrinsic's definition to be different from the ISD node, however. We cannot just change the definition because there's code generation work required to ensure the new combinations work. This is something we'll fix but in the mean time there exists critical clang failures that result from (currently) invalid stepvector intrinsics being created and so we'd like to fix those issues first as a priority by modelling the necessary functionality whilst maintaining the current definition.

[efriedma]

Sure, makes sense.

Maybe leave a comment in the code noting that you expect this special case to be temporary.

[david-arm]

HI all, thanks for the review comments! I've added a TODO to CreateStepVector explaining this is temporary.

llvm/lib/IR/IRBuilder.cpp

Function *TheFn =
       Intrinsic::getDeclaration(M, Intrinsic::vscale, {Scaling->getType()});
   CallInst *CI = createCallHelper(TheFn, {}, this, Name);
   return cast<ConstantInt>(Scaling)->getSExtValue() == 1
              ? CI
              : CreateMul(CI, Scaling);
 }
 
 Value *IRBuilderBase::CreateStepVector(Type *DstType, const Twine &Name) {
-  if (isa<ScalableVectorType>(DstType))
-    return CreateIntrinsic(Intrinsic::experimental_stepvector, {DstType}, {},
-                           nullptr, Name);
-
   Type *STy = DstType->getScalarType();
+  if (isa<ScalableVectorType>(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<ScalableVectorType>(DstType));
+    Value *Res = CreateIntrinsic(Intrinsic::experimental_stepvector,
+                                 {StepVecType}, {}, nullptr, Name);
+    if (StepVecType != DstType)
+      Res = CreateTrunc(Res, DstType);
+    return Res;
+  }
+
   unsigned NumEls = cast<FixedVectorType>(DstType)->getNumElements();
 
   // Create a vector of consecutive numbers from zero to VF.
   SmallVector<Constant *, 8> Indices;
   for (unsigned i = 0; i < NumEls; ++i)
     Indices.push_back(ConstantInt::get(STy, i));
 
   // Add the consecutive indices to the vector value.

llvm/unittests/IR/IRBuilderTest.cpp

TEST_F(IRBuilderTest, CreateStepVector) {
   StepVec = Builder.CreateStepVector(DstVecTy);
   EXPECT_TRUE(isa<CallInst>(StepVec));
   CallInst *Call = cast<CallInst>(StepVec);
   FunctionType *FTy = Call->getFunctionType();
   EXPECT_EQ(FTy->getReturnType(), DstVecTy);
   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<TruncInst>(StepVec));
+  TruncInst *Trunc = cast<TruncInst>(StepVec);
+  EXPECT_EQ(Trunc->getDestTy(), DstVecTy);
+  EXPECT_EQ(Trunc->getSrcTy(), VecI8Ty);
+  EXPECT_TRUE(isa<CallInst>(Trunc->getOperand(0)));
+
+  CallInst *Call = cast<CallInst>(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;
   Value *VDouble;
   Value *VInt;
   CallInst *Call;
   IntrinsicInst *II;
   GlobalVariable *GVDouble = new GlobalVariable(*M, Type::getDoubleTy(Ctx),