[VectorUtils] Teach findScalarElement to return splat value.

[sdesmalen]

Summary

If the vector is a splat of some scalar value, findScalarElement()
can simply return the scalar value if it knows the requested lane
is in the vector.

This is only needed for scalable vectors, because the InsertElement/ShuffleVector
case is already handled explicitly for the fixed-width case.

This helps to recognize an InstCombine fold like:

extractelt(bitcast(splat(%v))) -> bitcast(%v)

[sdesmalen]

llvm/lib/Analysis/VectorUtils.cpp

329

@spatel, I wasn't entirely sure if similar arguments apply as mentioned here: https://reviews.llvm.org/D104867#2840537

[frasercrmck]

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

253

Are these undefs supposed to be poison nowadays? I think I've seen both recently so it's as much for my own knowledge :)

[sdesmalen]

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

253

I think for this specific example it doesn't really matter, but the commit message of D93586 shows an example where it does. I'll change it to poison anyway to match what's done in other tests.

[sdesmalen]

Changed undef -> poison in test.

[frasercrmck]

LGTM but I think official approval would best come from the other reviewers.

[spatel]

Sorry for the delay - LGTM

[spatel]

llvm/lib/Analysis/VectorUtils.cpp

329

Oops - missed this comment. If it's a splat, then do we care what EltNo is?

Ie, if we have this:
define i64 @ext_lane_from_splat(i64 %v) {

%in = insertelement <vscale x 4 x i64> poison, i64 %v, i32 0
%splat = shufflevector <vscale x 4 x i64> %in, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
%r = extractelement <vscale x 4 x i64> %splat, i32 4 ; what does it mean if there are only 4 elts at runtime?
ret i64 %r

}

Should it simplify to ret i64 %v?

[sdesmalen]

llvm/lib/Analysis/VectorUtils.cpp

329

The LangRef states that if there are only 4 elements at runtime, that extracting element 4 results in a poison value. I'm tempted to argue that if executed on a machine with a wider vector width, returning %v would be safe, so we can equally assume returning %v should not lead to undefined behaviour on a machine that has shorter vectors. But my knowledge of poison is a bit lacking here, so I'm not sure if that assumption is safe to make.

llvm/lib/Analysis/VectorUtils.cpp

if (SVI && isa<FixedVectorType>(SVI->getType())) {
     if (InEl < (int)LHSWidth)
       return findScalarElement(SVI->getOperand(0), InEl);
     return findScalarElement(SVI->getOperand(1), InEl - LHSWidth);
   }
 
   // Extract a value from a vector add operation with a constant zero.
   // TODO: Use getBinOpIdentity() to generalize this.
   Value *Val; Constant *C;
   if (match(V, m_Add(m_Value(Val), m_Constant(C))))
     if (Constant *Elt = C->getAggregateElement(EltNo))
       if (Elt->isNullValue())
         return findScalarElement(Val, EltNo);
 
+  // If the vector is a splat then we can trivially find the scalar element.
+  if (isa<ScalableVectorType>(VTy))
+    if (Value *Splat = getSplatValue(V))
+      if (EltNo < VTy->getElementCount().getKnownMinValue())
+        return Splat;
+
   // Otherwise, we don't know.
   return nullptr;
 }
 
 int llvm::getSplatIndex(ArrayRef<int> Mask) {
   int SplatIndex = -1;
   for (int M : Mask) {
     // Ignore invalid (undefined) mask elements.

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

 
 ; Check that we can extract more complex cases where the stepvector is
 ; involved in a binary operation prior to the lane being extracted.
 
 define i64 @ext_lane0_from_add_with_stepvec(i64 %i) {
 ; CHECK-LABEL: @ext_lane0_from_add_with_stepvec(
 ; CHECK-NEXT:    ret i64 [[I:%.*]]
 ;
   %tmp = insertelement <vscale x 2 x i64> poison, i64 %i, i32 0
   %splatofi = shufflevector <vscale x 2 x i64> %tmp, <vscale x 2 x i64> poison, <vscale x  2 x i32> zeroinitializer
   %stepvec = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
   %add = add <vscale x 2 x i64> %splatofi, %stepvec
   %res = extractelement <vscale x 2 x i64> %add, i32 0
   ret i64 %res
 }
 
 define i1 @ext_lane1_from_cmp_with_stepvec(i64 %i) {
 ; CHECK-LABEL: @ext_lane1_from_cmp_with_stepvec(
 ; CHECK-NEXT:    [[RES:%.*]] = icmp eq i64 [[I:%.*]], 1
 ; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %tmp = insertelement <vscale x 2 x i64> poison, i64 %i, i32 0
   %splatofi = shufflevector <vscale x 2 x i64> %tmp, <vscale x 2 x i64> poison, <vscale x  2 x i32> zeroinitializer
   %stepvec = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
   %cmp = icmp eq <vscale x 2 x i64> %splatofi, %stepvec
   %res = extractelement <vscale x 2 x i1> %cmp, i32 1
   ret i1 %res
 }
 
+define i64* @ext_lane_from_bitcast_of_splat(i32* %v) {
+; CHECK-LABEL: @ext_lane_from_bitcast_of_splat(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[R:%.*]] = bitcast i32* [[V:%.*]] to i64*
+; CHECK-NEXT:    ret i64* [[R]]
+;
+entry:
+  %in = insertelement <vscale x 4 x i32*> poison, i32* %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32*> %in, <vscale x 4 x i32*> poison, <vscale x 4 x i32> zeroinitializer
+  %bc = bitcast <vscale x 4 x i32*> %splat to <vscale x 4 x i64*>
+  %r = extractelement <vscale x 4 x i64*> %bc, i32 3
+  ret i64* %r
+}
+
 declare <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
 declare <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
 declare <vscale x 4 x i32> @llvm.experimental.stepvector.nxv4i32()
 declare <vscale x 512 x i8> @llvm.experimental.stepvector.nxv512i8()