[SVE] Allow scalable vector splats to be gep indices

[ctetreau]

Summary

getSplatValue(): Make getSplatValue work for scalable vector splats
of the form returned by ConstantVector::getSplat().

getUniqueInteger(): The value returned by getSplatValue() is the same
value that is returned by getAggregateElement(0). No need to have the
second call, and this allows this function to work for scalable splats.

getIndexType() and getCastOpcode(): These can call getElementCount() and
VectorType::get(Type *, VectorType *)

Added test for gep using scalable vector splat

[efriedma]

I'll assume you didn't mean to post this for review

[ctetreau]

fixed issues with previous implementation, used new functionality to enable scalable vector indices to GEP, and added test

[ctetreau]

rebase

[efriedma]

Posted https://reviews.llvm.org/D82061 to try to make the GEP-related bits more sane.

llvm/lib/IR/Constants.cpp

1572

Constant::getUniqueInteger() is a terrible API. I'd prefer to fix the callers so it doesn't need to handle scalable vectors.

[ctetreau]

Split getSplatVal stuff into its own commit, rebase

[ctetreau]

@efriedma I think so long as the language reference states that gep into a struct accepts any vector splat index, then this patch should be accepted. (assuming it has no technical issues) It's a trivial change, that is probably an improvement regardless of if we consider the use case to be valid.

[ctetreau]

llvm/lib/IR/Constants.cpp

1572

That may be true, but my version is better than what existed previously.

llvm/lib/IR/Constants.cpp

if (isa<UndefValue>(Elt))
       Elt = OpC;
 
     if (OpC != Elt)
       return nullptr;
   }
   return Elt;
 }
 
 const APInt &Constant::getUniqueInteger() const {
   if (const ConstantInt *CI = dyn_cast<ConstantInt>(this))
     return CI->getValue();
-  assert(this->getSplatValue() && "Doesn't contain a unique integer!");
-  const Constant *C = this->getAggregateElement(0U);
-  assert(C && isa<ConstantInt>(C) && "Not a vector of numbers!");
+  Constant *C = this->getSplatValue();
+  assert(C && isa<ConstantInt>(C) && "Doesn't contain a unique integer!");
   return cast<ConstantInt>(C)->getValue();
 }
 
 //---- ConstantPointerNull::get() implementation.
 //
 
 ConstantPointerNull *ConstantPointerNull::get(PointerType *Ty) {
   std::unique_ptr<ConstantPointerNull> &Entry =

llvm/lib/IR/Instructions.cpp

assert(SrcTy->isFirstClassType() && DestTy->isFirstClassType() &&
          "Only first class types are castable!");
 
   if (SrcTy == DestTy)
     return BitCast;
 
   // FIXME: Check address space sizes here
   if (VectorType *SrcVecTy = dyn_cast<VectorType>(SrcTy))
     if (VectorType *DestVecTy = dyn_cast<VectorType>(DestTy))
-      if (SrcVecTy->getNumElements() == DestVecTy->getNumElements()) {
+      if (SrcVecTy->getElementCount() == DestVecTy->getElementCount()) {
         // An element by element cast.  Find the appropriate opcode based on the
         // element types.
         SrcTy = SrcVecTy->getElementType();
         DestTy = DestVecTy->getElementType();
       }
 
   // Get the bit sizes, we'll need these
   unsigned SrcBits = SrcTy->getPrimitiveSizeInBits();   // 0 for ptr

llvm/test/Transforms/InstSimplify/gep.ll

 define <vscale x 4 x float*> @scalable_vector_idx_mix_scalar_vector() {
 ; CHECK-LABEL: @scalable_vector_idx_mix_scalar_vector(
 ; CHECK-NEXT:    ret <vscale x 4 x float*> getelementptr (%struct, <vscale x 4 x %struct*> zeroinitializer, <vscale x 4 x i64> zeroinitializer, i32 1)
 ;
   %gep = getelementptr %struct, <vscale x 4 x %struct*> zeroinitializer, i32 0, i32 1
   ret <vscale x 4 x float*> %gep
 }
 
+define <vscale x 4 x float*> @scalable_vector_idx_vector() {
+; CHECK-LABEL: @scalable_vector_idx_mix_scalar_vector(
+; CHECK-NEXT:    ret <vscale x 4 x float*> getelementptr (%struct, <vscale x 4 x %struct*> zeroinitializer, <vscale x 4 x i64> zeroinitializer, <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer))
+;
+  %gep = getelementptr %struct, <vscale x 4 x %struct*> zeroinitializer, i32 0, <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+  ret <vscale x 4 x float*> %gep
+}
+
 ; Check ConstantExpr::getGetElementPtr() using ElementCount for size queries - end.