diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp --- a/llvm/lib/IR/ConstantFold.cpp +++ b/llvm/lib/IR/ConstantFold.cpp @@ -907,14 +907,10 @@ } } - // CAZ of type ScalableVectorType and n < CAZ->getMinNumElements() => - // extractelt CAZ, n -> 0 - if (auto *ValSVTy = dyn_cast(Val->getType())) { - if (!CIdx->uge(ValSVTy->getMinNumElements())) { - if (auto *CAZ = dyn_cast(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); diff --git a/llvm/test/Transforms/InstCombine/gep-vector-indices.ll b/llvm/test/Transforms/InstCombine/gep-vector-indices.ll --- a/llvm/test/Transforms/InstCombine/gep-vector-indices.ll +++ b/llvm/test/Transforms/InstCombine/gep-vector-indices.ll @@ -13,7 +13,7 @@ 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 ( shufflevector ( insertelement ( poison, i64 4, i32 0), poison, zeroinitializer), i32 0) +; CHECK-NEXT: [[RES:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 4 ; CHECK-NEXT: ret i32* [[RES]] ; %tmp = insertelement poison, i64 4, i32 0 diff --git a/llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll b/llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll --- a/llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll +++ b/llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll @@ -1,4 +1,4 @@ -; RUN: opt -instcombine -S < %s | FileCheck %s +; RUN: opt -S -instsimplify < %s | FileCheck %s ; CHECK-LABEL: definitely_in_bounds ; CHECK: ret i8 0 @@ -11,3 +11,35 @@ define i8 @maybe_in_bounds() { ret i8 extractelement ( 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 undef, i32 1024, i32 0 + %splat = shufflevector %in, undef, zeroinitializer + %r = extractelement %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 ( shufflevector ( insertelement ( undef, i32 1024, i32 0), undef, zeroinitializer), i32 4) +; + %in = insertelement undef, i32 1024, i32 0 + %splat = shufflevector %in, undef, zeroinitializer + %r = extractelement %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 ( shufflevector ( insertelement ( undef, i32 1024, i32 0), undef, zeroinitializer), i32 -1) +; + %in = insertelement undef, i32 1024, i32 0 + %splat = shufflevector %in, undef, zeroinitializer + %r = extractelement %splat, i32 -1 + ret i32 %r +}