[InstSimplify] Add constant fold for extractelement + splat for scalable vectors

[CarolineConcatto]

Summary

This patch allows that scalable vector can fold extractelement and splat of a constant
only when the lane index is lower than the minimum number of elements of the vector.

[CarolineConcatto]

The implementation in this patch was previously in:
https://reviews.llvm.org/D101916.

[foad]

Thanks for splitting it up. I like small patches.

llvm/lib/IR/ConstantFold.cpp

910–911

You should be able to remove the two lines above handling the CAZ case, because your new code will handle it.

[spatel]

In D103180#2782861, @foad wrote:

Thanks for splitting it up. I like small patches.

Definitely agree. :)

For the test where the index is presumed invalid (i32 -1), is that enforced somehow? Is there a hard limit somewhere that says i32 0xffffffff must be invalid?

[CarolineConcatto]

Hi @foad,
I don't know if I understood correct your suggestion.
I did not change any line.
I have tried some combinations of IF removals, but they all broke some tests

llvm/lib/IR/ConstantFold.cpp

910–911

a

910–911

Hi @foad
I believe we should keep these 2 tests:
!CIdx->uge(ValSVTy->getMinNumElements())
because it limits the index to be lower than the minimum width of the scalable vector.

and it should be only for scalable vector because of this

return CAZ->getElementValue(CIdx->getZExtValue());

transformation.
If we change to be vector type many other tests fail. For instance this one:
CodeGen/AMDGPU/amdgpu-codegenprepare-idiv.ll
Transforms/LoopVectorize/X86/gather_scatter.ll

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

910–911

Hi @CarolineConcatto, I think @foad meant removing the lines:

if (auto *CAZ = dyn_cast<ConstantAggregateZero>(Val))
  return CAZ->getElementValue(CIdx->getZExtValue());

If Val is a constant aggregate zero (<=> all zeroes), then this is a splat value that's recognized by getSplatValue, so your new lines added in this patch will handle the same case making the above lines redundant.

[CarolineConcatto]

• Remove vector of zeros check by a splat vector check

[CarolineConcatto]

Hi @foad,
Sorry for my miss interpretation.
I believe now it is as you suggested.

llvm/lib/IR/ConstantFold.cpp

910–911

Thank you @sdesmalen, that makes sense.
Change made.

[CarolineConcatto]

In D103180#2782981, @spatel wrote:

In D103180#2782861, @foad wrote:

Thanks for splitting it up. I like small patches.

Definitely agree. :)

For the test where the index is presumed invalid (i32 -1), is that enforced somehow? Is there a hard limit somewhere that says i32 0xffffffff must be invalid?

Sorry, @spatel
I've missed your comment above
I think the reason it says is invalid is because of the test:
uge(ValSVTy->getMinNumElements() -> unsigned greater or equal

[spatel]

In D103180#2793718, @CarolineConcatto wrote:

In D103180#2782981, @spatel wrote:

For the test where the index is presumed invalid (i32 -1), is that enforced somehow? Is there a hard limit somewhere that says i32 0xffffffff must be invalid?

I think the reason it says is invalid is because of the test:
uge(ValSVTy->getMinNumElements() -> unsigned greater or equal

Hmm...but those are uint64_t types/compares, so I don't see how it would automatically be considered invalid to have a "i32 -1". I don't think it really matters to this patch. I was just curious if there's some implicit semantic difference being verified between the tests "extractconstant_shuffle_maybe_out_of_range" and "extractconstant_shuffle_invalid_index".

So no objection from me, but someone more directly involved with scalable vectors should probably give final approval.

[sdesmalen]

The patch looks fine to me with the nit addressed.

llvm/lib/IR/ConstantFold.cpp

909–910

nit: s/!CIdx->uge/CIdx->ult/

[CarolineConcatto]

llvm/lib/IR/ConstantFold.cpp

909–910

Hey @sdesmalen
I may be completely wrong, but I believe your suggestion is not possible
https://llvm.org/doxygen/classllvm_1_1ConstantInt.html
If I looked the documentation correctly llvm::constant only has UGE

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

909–910

Ah I hadn't realised that. In that case, could you write it more explicitly as

CIdx->getZExtValue() < ValSVTy->getMinNumElements()

[dmgreen]

llvm/lib/IR/ConstantFold.cpp

909–910

Do we know the type of CIdx is always less that 64bits?

uge() is just a convenience wrapper around APInt::uge(). You can use CIdx->getValue().ult(..) if you feel the ult is better than the !uge

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

909–910

Agreed, that's better, my main concern was the !uge.

[CarolineConcatto]

• replace the use of !CIdx->ult() by !CIdx->getValue().uge(..)

[foad]

The patch looks OK to me. I added some inline comments but you don't need to do anything about them, they are just observations.

llvm/lib/IR/ConstantFold.cpp

909–910

Or you could implement ConstantInt::ult.

909–910

I don't really understand why this special case for scalable types exists here. Why not just fall through to the call to Val->getAggregateElement below? It looks like getAggregateElement already has some support for scalable types.

[CarolineConcatto]

• update test vector_splat_indices_nxv2i64_ext0 in gep-vector-indices.ll

[Matt]

llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll

15

(Tiny nit:) Perhaps "extracting" would fit a little bit better here?

[CarolineConcatto]

• s/extract/extracting in test commemts

[CarolineConcatto]

• remove scalable vector check because getAggregateElement has safeguards for scalable vectors

llvm/lib/IR/ConstantFold.cpp

if (auto *GEP = dyn_cast<GEPOperator>(CE)) {
                                 APSInt(CIdx->getValue()))) {
           return CE->getOperand(1);
         } else {
           return ConstantExpr::getExtractElement(CE->getOperand(0), CIdx);
         }
       }
     }
   }
 
-  // CAZ of type ScalableVectorType and n < CAZ->getMinNumElements() =>
-  //   extractelt CAZ, n -> 0
-  if (auto *ValSVTy = dyn_cast<ScalableVectorType>(Val->getType())) {
-    if (!CIdx->uge(ValSVTy->getMinNumElements())) {
-      if (auto *CAZ = dyn_cast<ConstantAggregateZero>(Val))
-        return CAZ->getElementValue(CIdx->getZExtValue());
-    }
-    return nullptr;
+  // Lane < Splat minimum vector width => extractelt Splat(x), Lane -> x
+  if (CIdx->getValue().ult(ValVTy->getElementCount().getKnownMinValue())) {
+    if (Constant *SplatVal = Val->getSplatValue())
+      return SplatVal;
   }
 
   return Val->getAggregateElement(CIdx);
 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
                                                      Constant *Elt,
                                                      Constant *Idx) {

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine %s -S | FileCheck %s
 
 define i32* @vector_splat_indices_v2i64_ext0(i32* %a) {
 ; CHECK-LABEL: @vector_splat_indices_v2i64_ext0(
 ; CHECK-NEXT:    [[RES:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 4
 ; CHECK-NEXT:    ret i32* [[RES]]
 ;
   %gep = getelementptr i32, i32* %a, <2 x i64> <i64 4, i64 4>
   %res = extractelement <2 x i32*> %gep, i32 0
   ret i32* %res
 }
 
 define i32* @vector_splat_indices_nxv2i64_ext0(i32* %a) {
 ; CHECK-LABEL: @vector_splat_indices_nxv2i64_ext0(
-; CHECK-NEXT:    [[RES:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 extractelement (<vscale x 2 x i64> shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 4, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer), i32 0)
+; CHECK-NEXT:    [[RES:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 4
 ; CHECK-NEXT:    ret i32* [[RES]]
 ;
   %tmp = insertelement <vscale x 2 x i64> poison, i64 4, i32 0
   %splatof4 = shufflevector <vscale x 2 x i64> %tmp, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %gep = getelementptr inbounds i32, i32* %a, <vscale x 2 x i64> %splatof4
   %res = extractelement <vscale x 2 x i32*> %gep, i32 0
   ret i32* %res
 }

llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll

-; RUN: opt -instcombine -S < %s | FileCheck %s
+; RUN: opt -S -instsimplify < %s | FileCheck %s
 
 ; CHECK-LABEL: definitely_in_bounds
 ; CHECK: ret i8 0
 define i8 @definitely_in_bounds() {
   ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 15)
 }
 
 ; CHECK-LABEL: maybe_in_bounds
 ; CHECK: ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 16)
 define i8 @maybe_in_bounds() {
   ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 16)
 }
+
+; Examples of extracting a lane from a splat constant
+
+define i32 @extractconstant_shuffle_in_range(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_in_range(
+; CHECK-NEXT:    ret i32 1024
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, 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 @extractconstant_shuffle_maybe_out_of_range(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_maybe_out_of_range(
+; CHECK-NEXT:    ret i32 extractelement (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1024, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer), i32 4)
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, 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 4
+  ret i32 %r
+}
+
+define i32 @extractconstant_shuffle_invalid_index(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_invalid_index(
+; CHECK-NEXT:    ret i32 extractelement (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1024, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer), i32 -1)
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, 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
+}