[InstSimplify] fold extractelement of splat even with variable index

[spatel]

Summary

We unnecessarily restricted a fold of a splat to a constant vector.

We overlooked this fold in D102404 and earlier patches, but the fixed vector variant is shown in:
https://llvm.org/PR50817

Alive2 agrees on that:
https://alive2.llvm.org/ce/z/HpijPC

I don't think Alive2 understands scalable vectors, but the same logic applies IIUC.

[frasercrmck]

I agree that fixed- and scalable-vectors should be treated identically, at least according to my reading.

As for the legality, so even though %idx may exceed the length of the vector and return a poison value, it's acceptable for us to "relax" that into undef and then immediately into the splat value? Is that what's going on?

[spatel]

In D104867#2840170, @frasercrmck wrote:

I agree that fixed- and scalable-vectors should be treated identically, at least according to my reading.

As for the legality, so even though %idx may exceed the length of the vector and return a poison value, it's acceptable for us to "relax" that into undef and then immediately into the splat value? Is that what's going on?

Yes, exactly - in the post-undef world, I think we would say: since the index is out-of-bounds, the result of the operation is poison (arbitrary), so we are choosing the splat value for convenience in making this fold.

If we wanted to, we could try to detect if the index is always out-of-bounds (via known bits) and return poison instead of the splat value (similar to the existing code that handles a constant index value). But that would have potentially large compile-time cost and no practical value IMO.

[spatel]

Ping.

[RKSimon]

LGTM - thanks @spatel

llvm/lib/Analysis/InstructionSimplify.cpp

 /// If not, this returns null.
 static Value *SimplifyExtractElementInst(Value *Vec, Value *Idx,
                                          const SimplifyQuery &Q, unsigned) {
   auto *VecVTy = cast<VectorType>(Vec->getType());
   if (auto *CVec = dyn_cast<Constant>(Vec)) {
     if (auto *CIdx = dyn_cast<Constant>(Idx))
       return ConstantExpr::getExtractElement(CVec, CIdx);
 
-    // The index is not relevant if our vector is a splat.
-    if (auto *Splat = CVec->getSplatValue())
-      return Splat;
-
     if (Q.isUndefValue(Vec))
       return UndefValue::get(VecVTy->getElementType());
   }
 
   // An undef extract index can be arbitrarily chosen to be an out-of-range
   // index value, which would result in the instruction being poison.
   if (Q.isUndefValue(Idx))
     return PoisonValue::get(VecVTy->getElementType());
 
   // If extracting a specified index from the vector, see if we can recursively
   // find a previously computed scalar that was inserted into the vector.
   if (auto *IdxC = dyn_cast<ConstantInt>(Idx)) {
     // For fixed-length vector, fold into undef if index is out of bounds.
     unsigned MinNumElts = VecVTy->getElementCount().getKnownMinValue();
     if (isa<FixedVectorType>(VecVTy) && IdxC->getValue().uge(MinNumElts))
       return PoisonValue::get(VecVTy->getElementType());
     // Handle case where an element is extracted from a splat.
     if (IdxC->getValue().ult(MinNumElts))
       if (auto *Splat = getSplatValue(Vec))
         return Splat;
     if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
       return Elt;
+  } else {
+    // The index is not relevant if our vector is a splat.
+    if (Value *Splat = getSplatValue(Vec))
+      return Splat;
   }
   return nullptr;
 }
 
 Value *llvm::SimplifyExtractElementInst(Value *Vec, Value *Idx,
                                         const SimplifyQuery &Q) {
   return ::SimplifyExtractElementInst(Vec, Idx, Q, RecursionLimit);
 }

llvm/test/Transforms/InstSimplify/extract-element.ll

 ; CHECK-NEXT:    ret float 2.000000e+00
 ;
   %r = extractelement <4 x float> <float 2.0, float 2.0, float 2.0, float 2.0>, i32 %y
   ret float %r
 }
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %index) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <3 x i32> poison, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <3 x i32> [[IN]], <3 x i32> poison, <3 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <3 x i32> [[SPLAT]], i32 [[INDEX:%.*]]
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 [[V:%.*]]
 ;
   %in = insertelement <3 x i32> poison, i32 %v, i32 0
   %splat = shufflevector <3 x i32> %in, <3 x i32> poison, <3 x i32> zeroinitializer
   %r = extractelement <3 x i32> %splat, i32 %index
   ret i32 %r
 }

llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll

;
   %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
   %r = extractelement <vscale x 4 x i32> %splat, i32 1
   ret i32 %r
 }
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %idx) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 [[IDX:%.*]]
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 [[V:%.*]]
 ;
   %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> undef, <vscale x 4 x i32> zeroinitializer
   %r = extractelement <vscale x 4 x i32> %splat, i32 %idx
   ret i32 %r
 }

llvm/test/Transforms/InstSimplify/vscale.ll

;
   %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
   %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %r = extractelement <vscale x 4 x i32> %splat, i32 1
   ret i32 %r
 }
 
 define i32 @extractelement_splat_variable_index(i32 %v, i32 %idx) {
 ; CHECK-LABEL: @extractelement_splat_variable_index(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 [[IDX:%.*]]
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 [[V:%.*]]
 ;
   %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
   %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
   %r = extractelement <vscale x 4 x i32> %splat, i32 %idx
   ret i32 %r
 }