[InstCombine] Add instcombine fold for extractelement + splat for scalable vectors

[CarolineConcatto]

Summary

This patch allows that scalable vector can also use the fold that already
exists for fixed vector, only when the lane index is lower than the minimum
number of elements of the vector.

[CarolineConcatto]

• Add implementation and test to simplify constants
• Change check for splat in visitExtractElement

[foad]

llvm/lib/IR/ConstantFold.cpp

875–879 ↗

(On Diff #347632)

It looks like this handles some of the same cases as the "CAZ of ScalableVectorType" case below. Can you combine them?

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

350 ↗

(On Diff #347632)

This needs a comment explaining what case you are trying to handle, e.g. "(extractelement (shufflevector ...), ...) -> (...)"

351 ↗

(On Diff #347632)

You can dyn_cast SrcVec directly to ShuffleVectorInst, no need to go via Instruction.

[spatel]

Could this be 2 independent patches?

1. Enhancement to ConstantFold for scalable vectors
2. Add a fold to llvm::SimplifyExtractElementInst() when the base vector is known to be a splat.

If that's right, we should move the tests to test/Transforms/InstSimplify/{ConstProp}, so we're checking the minimum required analysis. I don't think we need to touch InstCombine at all.

[CarolineConcatto]

• address reviewer's comment
• reduce patch scope to only work with instcombine

[CarolineConcatto]

Thank you @foad and @spatel for the fast review.
I have splited the patch in 2 as @spatel suggested.
The constant fold is in https://reviews.llvm.org/D103180.
@spatel let me know if I understood correct what you've asked.

[spatel]

This is not what I had in mind. If we are not creating a new instruction, then we should try to get this into -instsimplify:
https://github.com/llvm/llvm-project/blob/6c92215e07f41cb566d54599e891180fe2ddbea5/llvm/lib/Analysis/InstructionSimplify.cpp#L4492

That's where we handle fixed-length vector splat+extract. Can we add to or adjust the underlying analysis there to handle scalable vectors?

[CarolineConcatto]

-move splat test from instcombine to instsimplify

[CarolineConcatto]

Hi @spatel,
Sorry for getting wrong in the first place.
I hope I get it correct now with your latest explanation.
I changed the splat test to be in SimplifyExtractElementInst (InstructionSimplify.cpp).

[sdesmalen]

llvm/lib/Analysis/InstructionSimplify.cpp

4492–4493

nit: MinNumElts

4493

This cast is redundant, because VecVTy is already a VectorType, see line 4475.

4498–4501

Why not just rely on llvm::getSplatValue here?

unsigned MinNumElts = VecVTy->getElementCount().getKnownMinValue();
if (IdxC->getValue().ult(NumElts))
  if (auto *Splat = getSplatValue(Vec))
    return Splat;

[CarolineConcatto]

• Address Sander's comments

[CarolineConcatto]

Thank you @sdesmalen for you review.
I have made the changes you've asked.

llvm/lib/Analysis/InstructionSimplify.cpp

4498–4501

Yes, it could check for Splat.
I like the other way because it explains/shows the instruction chain.
But your solution is cleaner, so I have changed.

[sdesmalen]

Thanks, LGTM with my nit addressed.

llvm/lib/Analysis/InstructionSimplify.cpp

4495

nit: Handle case where an element is extracted from a splat.

[CarolineConcatto]

• Replace comment as suggested

[CarolineConcatto]

• update test gep-vector-indices.ll

[sdesmalen]

Seems you just rebased, so it still looks fine to me.

llvm/test/Transforms/InstCombine/gep-vector-indices.ll

64 ↗

(On Diff #350539)

Nice improvement.

llvm/lib/Analysis/InstructionSimplify.cpp

if (auto *CVec = dyn_cast<Constant>(Vec)) {
     if (Q.isUndefValue(Vec))
       return UndefValue::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.
-    if (isa<FixedVectorType>(VecVTy) &&
-        IdxC->getValue().uge(cast<FixedVectorType>(VecVTy)->getNumElements()))
+    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;
   }
 
   // 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());
 
   return nullptr;
 }

llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll

 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1 ; <- This insert could be removed.
   %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
   %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; 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> poison, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
-; 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> poison, <vscale x 4 x i32> zeroinitializer
   %r = extractelement <vscale x 4 x i32> %splat, i32 1
   ret i32 %r
 }
 
 define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1 ; <- This insert could be removed.
   %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
   %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; 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 1
-; 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 1
   ret i32 %r
 }
 
 define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {