[Constants] Add scalable vector support to ConstantVector::getSplat. [IR support for SVE scalable vectors 2/4]

[paulwalker-arm]

Summary

Create a version of ConstantVector::getSplat that takes VT::ElementCount
and thus preserves a vector's scalable property.  The result for
non-scalable vectors remain unchanged but scalable vectors use the same
IR sequence as used for non-constant splats.  This prompted updates to
some instructions to ensure they also preserve a vector's scalable
property.

[rengolin]

Hi Paul,

I'm really not sure what this will give us in benefit. We can't represent a constant vector of unknown lanes without refactoring the IR textual notation (ex. "<n x 4 x i32> <i32 29, ...>"), so why bother replacing a sequence of instructions in SSA form for one all in the same line?

If we don't need this to proceed, and we'd like to change the representation (at least the textual one), I'd rather leave IR changing decisions for later.

cheers,
--renato

include/llvm/IR/Constants.h

489

Shouldn't you replace the current one, rather than add? Or even, make the old one a wrapper to the EC version with isScalable = false?

lib/IR/Constants.cpp

1048

This is confusing. Better to merge the two and make the other call this one with { NumElts, false }.

test/Transforms/ConstProp/splat.ll

40

I'm not sure what we gain by this being all in the same line versus separate lines.

include/llvm/IR/Constants.h

public:
   static Constant *get(ArrayRef<Constant*> V);
 
 private:
   static Constant *getImpl(ArrayRef<Constant *> V);
 
 public:
   /// Return a ConstantVector with the specified constant in each element.
   static Constant *getSplat(unsigned NumElts, Constant *Elt);
+  static Constant *getSplat(VectorType::ElementCount EC, Constant *Elt);
 
   /// Specialize the getType() method to always return a VectorType,
   /// which reduces the amount of casting needed in parts of the compiler.
   inline VectorType *getType() const {
     return cast<VectorType>(Value::getType());
   }
 
   /// If this is a splat constant, meaning that all of the elements have the

lib/IR/ConstantFold.cpp

 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
                                                      Constant *Elt,
                                                      Constant *Idx) {
   if (isa<UndefValue>(Idx))
     return UndefValue::get(Val->getType());
 
+  if (Val->getType()->getVectorIsScalable())
+    return nullptr;
+
   ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx) return nullptr;
 
   unsigned NumElts = Val->getType()->getVectorNumElements();
   if (CIdx->uge(NumElts))
     return UndefValue::get(Val->getType());
 
   SmallVector<Constant*, 16> Result;
Constant *llvm::ConstantFoldShuffleVectorInstruction(Constant *V1,
 
   // Undefined shuffle mask -> undefined value.
   if (isa<UndefValue>(Mask))
     return UndefValue::get(VectorType::get(EltTy, MaskNumElts));
 
   // Don't break the bitcode reader hack.
   if (isa<ConstantExpr>(Mask)) return nullptr;
 
+  if (V1->getType()->getVectorIsScalable())
+    return nullptr;
+
   unsigned SrcNumElts = V1->getType()->getVectorNumElements();
 
   // Loop over the shuffle mask, evaluating each element.
   SmallVector<Constant*, 32> Result;
   for (unsigned i = 0; i != MaskNumElts; ++i) {
     int Elt = ShuffleVectorInst::getMaskValue(Mask, i);
     if (Elt == -1) {
       Result.push_back(UndefValue::get(EltTy));

lib/IR/Constants.cpp

Constant *Constant::getIntegerValue(Type *Ty, const APInt &V) {
   Constant *C = ConstantInt::get(Ty->getContext(), V);
 
   // Convert an integer to a pointer, if necessary.
   if (PointerType *PTy = dyn_cast<PointerType>(ScalarTy))
     C = ConstantExpr::getIntToPtr(C, PTy);
 
   // Broadcast a scalar to a vector, if necessary.
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    C = ConstantVector::getSplat(VTy->getNumElements(), C);
+    C = ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 Constant *Constant::getAllOnesValue(Type *Ty) {
   if (IntegerType *ITy = dyn_cast<IntegerType>(Ty))
     return ConstantInt::get(Ty->getContext(),
                             APInt::getAllOnesValue(ITy->getBitWidth()));
 
   if (Ty->isFloatingPointTy()) {
     APFloat FL = APFloat::getAllOnesValue(Ty->getPrimitiveSizeInBits(),
                                           !Ty->isPPC_FP128Ty());
     return ConstantFP::get(Ty->getContext(), FL);
   }
 
   VectorType *VTy = cast<VectorType>(Ty);
-  return ConstantVector::getSplat(VTy->getNumElements(),
+  return ConstantVector::getSplat(VTy->getElementCount(),
                                   getAllOnesValue(VTy->getElementType()));
 }
 
 Constant *Constant::getAggregateElement(unsigned Elt) const {
   if (const ConstantAggregate *CC = dyn_cast<ConstantAggregate>(this))
     return Elt < CC->getNumOperands() ? CC->getOperand(Elt) : nullptr;
 
   if (const ConstantAggregateZero *CAZ = dyn_cast<ConstantAggregateZero>(this))

 Constant *ConstantInt::getTrue(Type *Ty) {
   VectorType *VTy = dyn_cast<VectorType>(Ty);
   if (!VTy) {
     assert(Ty->isIntegerTy(1) && "True must be i1 or vector of i1.");
     return ConstantInt::getTrue(Ty->getContext());
   }
   assert(VTy->getElementType()->isIntegerTy(1) &&
          "True must be vector of i1 or i1.");
-  return ConstantVector::getSplat(VTy->getNumElements(),
+  return ConstantVector::getSplat(VTy->getElementCount(),
                                   ConstantInt::getTrue(Ty->getContext()));
 }
 
 Constant *ConstantInt::getFalse(Type *Ty) {
   VectorType *VTy = dyn_cast<VectorType>(Ty);
   if (!VTy) {
     assert(Ty->isIntegerTy(1) && "False must be i1 or vector of i1.");
     return ConstantInt::getFalse(Ty->getContext());
   }
   assert(VTy->getElementType()->isIntegerTy(1) &&
          "False must be vector of i1 or i1.");
-  return ConstantVector::getSplat(VTy->getNumElements(),
+  return ConstantVector::getSplat(VTy->getElementCount(),
                                   ConstantInt::getFalse(Ty->getContext()));
 }
 
 // Get a ConstantInt from an APInt.
 ConstantInt *ConstantInt::get(LLVMContext &Context, const APInt &V) {
   // get an existing value or the insertion position
   LLVMContextImpl *pImpl = Context.pImpl;
   std::unique_ptr<ConstantInt> &Slot = pImpl->IntConstants[V];
   if (!Slot) {
     // Get the corresponding integer type for the bit width of the value.
     IntegerType *ITy = IntegerType::get(Context, V.getBitWidth());
     Slot.reset(new ConstantInt(ITy, V));
   }
   assert(Slot->getType() == IntegerType::get(Context, V.getBitWidth()));
   return Slot.get();
 }
 
 Constant *ConstantInt::get(Type *Ty, uint64_t V, bool isSigned) {
   Constant *C = get(cast<IntegerType>(Ty->getScalarType()), V, isSigned);
 
   // For vectors, broadcast the value.
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 ConstantInt *ConstantInt::get(IntegerType *Ty, uint64_t V, bool isSigned) {
   return get(Ty->getContext(), APInt(Ty->getBitWidth(), V, isSigned));
 }
 
 ConstantInt *ConstantInt::getSigned(IntegerType *Ty, int64_t V) {
   return get(Ty, V, true);
 }
 
 Constant *ConstantInt::getSigned(Type *Ty, int64_t V) {
   return get(Ty, V, true);
 }
 
 Constant *ConstantInt::get(Type *Ty, const APInt& V) {
   ConstantInt *C = get(Ty->getContext(), V);
   assert(C->getType() == Ty->getScalarType() &&
          "ConstantInt type doesn't match the type implied by its value!");
 
   // For vectors, broadcast the value.
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 ConstantInt *ConstantInt::get(IntegerType* Ty, StringRef Str, uint8_t radix) {
   return get(Ty->getContext(), APInt(Ty->getBitWidth(), Str, radix));
 }
 
Constant *ConstantFP::get(Type *Ty, double V) {
   APFloat FV(V);
   bool ignored;
   FV.convert(*TypeToFloatSemantics(Ty->getScalarType()),
              APFloat::rmNearestTiesToEven, &ignored);
   Constant *C = get(Context, FV);
 
   // For vectors, broadcast the value.
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 
 Constant *ConstantFP::get(Type *Ty, StringRef Str) {
   LLVMContext &Context = Ty->getContext();
 
   APFloat FV(*TypeToFloatSemantics(Ty->getScalarType()), Str);
   Constant *C = get(Context, FV);
 
   // For vectors, broadcast the value.
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C; 
 }
 
 Constant *ConstantFP::getNaN(Type *Ty, bool Negative, unsigned Type) {
   const fltSemantics &Semantics = *TypeToFloatSemantics(Ty->getScalarType());
   APFloat NaN = APFloat::getNaN(Semantics, Negative, Type);
   Constant *C = get(Ty->getContext(), NaN);
 
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 Constant *ConstantFP::getNegativeZero(Type *Ty) {
   const fltSemantics &Semantics = *TypeToFloatSemantics(Ty->getScalarType());
   APFloat NegZero = APFloat::getZero(Semantics, /*Negative=*/true);
   Constant *C = get(Ty->getContext(), NegZero);
 
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 
 Constant *ConstantFP::getZeroValueForNegation(Type *Ty) {
   if (Ty->isFPOrFPVectorTy())
     return getNegativeZero(Ty);
ConstantFP* ConstantFP::get(LLVMContext &Context, const APFloat& V) {
   return Slot.get();
 }
 
 Constant *ConstantFP::getInfinity(Type *Ty, bool Negative) {
   const fltSemantics &Semantics = *TypeToFloatSemantics(Ty->getScalarType());
   Constant *C = get(Ty->getContext(), APFloat::getInf(Semantics, Negative));
 
   if (VectorType *VTy = dyn_cast<VectorType>(Ty))
-    return ConstantVector::getSplat(VTy->getNumElements(), C);
+    return ConstantVector::getSplat(VTy->getElementCount(), C);
 
   return C;
 }
 
 ConstantFP::ConstantFP(Type *Ty, const APFloat &V)
     : ConstantData(Ty, ConstantFPVal), Val(V) {
   assert(&V.getSemantics() == TypeToFloatSemantics(Ty) &&
          "FP type Mismatch");
Constant *ConstantVector::getSplat(unsigned NumElts, Constant *V) {
   if ((isa<ConstantFP>(V) || isa<ConstantInt>(V)) &&
       ConstantDataSequential::isElementTypeCompatible(V->getType()))
     return ConstantDataVector::getSplat(NumElts, V);
 
   SmallVector<Constant*, 32> Elts(NumElts, V);
   return get(Elts);
 }
 
+Constant *ConstantVector::getSplat(VectorType::ElementCount EC, Constant *V) {
+  if (!EC.isScalable())
+    return getSplat(EC.getNumElements(), V);
+
+  Type *Ty = VectorType::get(V->getType(), EC);
+
+  if (V->isNullValue())
+    return ConstantAggregateZero::get(Ty);
+  else if (isa<UndefValue>(V))
+    return UndefValue::get(Ty);
+
+  // Move splat value into vector[0].
+  Type *Int32Ty = Type::getInt32Ty(Ty->getContext());
+  Constant *Zero = ConstantInt::get(Int32Ty, 0);
+  V = ConstantExpr::getInsertElement(UndefValue::get(Ty), V, Zero);
+
+  // Broadcast element zero to all lanes.
+  Constant *Zeros = ConstantInt::get(VectorType::get(Int32Ty, EC), 0);
+  return ConstantExpr::getShuffleVector(V, UndefValue::get(Ty), Zeros);
+}
+
 ConstantTokenNone *ConstantTokenNone::get(LLVMContext &Context) {
   LLVMContextImpl *pImpl = Context.pImpl;
   if (!pImpl->TheNoneToken)
     pImpl->TheNoneToken.reset(new ConstantTokenNone(Context));
   return pImpl->TheNoneToken.get();
 }
 
 /// Remove the constant from the constant table.

 Constant *ConstantExpr::getShuffleVector(Constant *V1, Constant *V2,
                                          Constant *Mask, Type *OnlyIfReducedTy) {
   assert(ShuffleVectorInst::isValidOperands(V1, V2, Mask) &&
          "Invalid shuffle vector constant expr operands!");
 
   if (Constant *FC = ConstantFoldShuffleVectorInstruction(V1, V2, Mask))
     return FC;          // Fold a few common cases.
 
-  unsigned NElts = Mask->getType()->getVectorNumElements();
+  auto NElts = cast<VectorType>(Mask->getType())->getElementCount();
   Type *EltTy = V1->getType()->getVectorElementType();
   Type *ShufTy = VectorType::get(EltTy, NElts);
 
   if (OnlyIfReducedTy == ShufTy)
     return nullptr;
 
   // Look up the constant in the table first to ensure uniqueness
   Constant *ArgVec[] = { V1, V2, Mask };

lib/IR/ConstantsContext.h

 public:
   // allocate space for exactly three operands
   void *operator new(size_t s) {
     return User::operator new(s, 3);
   }
   ShuffleVectorConstantExpr(Constant *C1, Constant *C2, Constant *C3)
   : ConstantExpr(VectorType::get(
                    cast<VectorType>(C1->getType())->getElementType(),
-                   cast<VectorType>(C3->getType())->getNumElements()),
+                   cast<VectorType>(C3->getType())->getElementCount()),
                  Instruction::ShuffleVector,
                  &Op<0>(), 3) {
     Op<0>() = C1;
     Op<1>() = C2;
     Op<2>() = C3;
   }
   /// Transparently provide more efficient getOperand methods.
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);

lib/IR/Instructions.cpp

 //===----------------------------------------------------------------------===//
 //                      ShuffleVectorInst Implementation
 //===----------------------------------------------------------------------===//
 
 ShuffleVectorInst::ShuffleVectorInst(Value *V1, Value *V2, Value *Mask,
                                      const Twine &Name,
                                      Instruction *InsertBefore)
 : Instruction(VectorType::get(cast<VectorType>(V1->getType())->getElementType(),
-                cast<VectorType>(Mask->getType())->getNumElements()),
+                cast<VectorType>(Mask->getType())->getElementCount()),
               ShuffleVector,
               OperandTraits<ShuffleVectorInst>::op_begin(this),
               OperandTraits<ShuffleVectorInst>::operands(this),
               InsertBefore) {
   assert(isValidOperands(V1, V2, Mask) &&
          "Invalid shuffle vector instruction operands!");
   Op<0>() = V1;
   Op<1>() = V2;
   Op<2>() = Mask;
   setName(Name);
 }
 
 ShuffleVectorInst::ShuffleVectorInst(Value *V1, Value *V2, Value *Mask,
                                      const Twine &Name,
                                      BasicBlock *InsertAtEnd)
 : Instruction(VectorType::get(cast<VectorType>(V1->getType())->getElementType(),
-                cast<VectorType>(Mask->getType())->getNumElements()),
+                cast<VectorType>(Mask->getType())->getElementCount()),
               ShuffleVector,
               OperandTraits<ShuffleVectorInst>::op_begin(this),
               OperandTraits<ShuffleVectorInst>::operands(this),
               InsertAtEnd) {
   assert(isValidOperands(V1, V2, Mask) &&
          "Invalid shuffle vector instruction operands!");
 
   Op<0>() = V1;

test/Transforms/ConstProp/splat.ll

+; RUN: opt < %s -constprop -S | FileCheck %s
+
+define <4 x i32> @test1() {
+; CHECK-LABEL: @test1
+; CHECK: ret <4 x i32> undef
+  ret <4 x i32> undef
+}
+
+define <n x 4 x i32> @test2() {
+; CHECK-LABEL: @test2
+; CHECK: ret <n x 4 x i32> undef
+  ret <n x 4 x i32> undef
+}
+
+define <4 x i32> @test3() {
+; CHECK-LABEL: @test3
+; CHECK: ret <4 x i32> zeroinitializer
+  ret <4 x i32> zeroinitializer
+}
+
+define <n x 4 x i32> @test4() {
+; CHECK-LABEL: @test4
+; CHECK: ret <n x 4 x i32> zeroinitializer
+  ret <n x 4 x i32> zeroinitializer
+}
+
+; Test const splat sequences become const vectors for non-scalable types.
+define <4 x i32> @test5() {
+; CHECK-LABEL: @test5
+; CHECK: ret <4 x i32> <i32 29, i32 29, i32 29, i32 29>
+  %a = insertelement <4 x i32> undef, i32 29, i32 0
+  %b = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> zeroinitializer
+  ret <4 x i32> %b
+}
+
+; Test const splat sequences are maintained for scalable vectors with the
+; exception that all nodes are constant.
+define <n x 4 x i32> @test6() {
+; CHECK-LABEL: @test6
+; CHECK: ret <n x 4 x i32> shufflevector (<n x 4 x i32>
+; CHECK-SAME:  insertelement (<n x 4 x i32> undef, i32 29, i32 0),
+; CHECK-SAME:  <n x 4 x i32> undef,
+; CHECK-SAME:  <n x 4 x i32> zeroinitializer)
+  %a = insertelement <n x 4 x i32> undef, i32 29, i32 0
+  %b = shufflevector <n x 4 x i32> %a, <n x 4 x i32> undef, <n x 4 x i32> zeroinitializer
+  ret <n x 4 x i32> %b
+}