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 @@ -872,6 +872,11 @@ if (!CIdx) return nullptr; + if (Constant *SplatVal = Val->getSplatValue()) { + unsigned NumElts = ValVTy->getElementCount().getKnownMinValue(); + if (CIdx->getZExtValue() < NumElts) + return SplatVal; + } if (auto *ValFVTy = dyn_cast(Val->getType())) { // ee({w,x,y,z}, wrong_value) -> poison if (CIdx->uge(ValFVTy->getNumElements())) diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp --- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp +++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp @@ -347,6 +347,12 @@ ElementCount EC = EI.getVectorOperandType()->getElementCount(); unsigned NumElts = EC.getKnownMinValue(); + if (auto *I = dyn_cast(SrcVec)) + if (auto *SE = dyn_cast(I)) + if (isSplatValue(SE) && IndexC->getZExtValue() < NumElts) + if (auto *IE = dyn_cast(I->getOperand(0))) + return replaceInstUsesWith(EI, IE->getOperand(1)); + // InstSimplify should handle cases where the index is invalid. // For fixed-length vector, it's invalid to extract out-of-range element. if (!EC.isScalable() && IndexC->getValue().uge(NumElts)) diff --git a/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll --- a/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll +++ b/llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll @@ -55,13 +55,9 @@ 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 poison, i32 [[V:%.*]], i32 0 -; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector [[IN]], poison, zeroinitializer -; CHECK-NEXT: [[R:%.*]] = extractelement [[SPLAT]], i32 1 -; CHECK-NEXT: ret i32 [[R]] +; CHECK-NEXT: ret i32 %v ; %in = insertelement poison, i32 %v, i32 0 %splat = shufflevector %in, poison, zeroinitializer @@ -69,6 +65,16 @@ ret i32 %r } +define i32 @extractconstant_shuffle_in_range() { +; CHECK-LABEL: @extractconstant_shuffle_in_range( +; CHECK-NEXT: ret i32 1024 +; + %in = insertelement poison, i32 1024, i32 0 + %splat = shufflevector %in, poison, zeroinitializer + %r = extractelement %splat, i32 1 + ret i32 %r +} + define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) { ; CHECK-LABEL: @extractelement_shuffle_maybe_out_of_range( ; CHECK-NEXT: [[IN:%.*]] = insertelement poison, i32 [[V:%.*]], i32 0 @@ -82,6 +88,16 @@ ret i32 %r } +define i32 @extractconstant_shuffle_maybe_out_of_range() { +; CHECK-LABEL: @extractconstant_shuffle_maybe_out_of_range( +; CHECK-NEXT: ret i32 extractelement ( shufflevector ( insertelement ( poison, i32 1024, i32 0), poison, zeroinitializer), i32 4) +; + %in = insertelement poison, i32 1024, i32 0 + %splat = shufflevector %in, poison, zeroinitializer + %r = extractelement %splat, i32 4 + ret i32 %r +} + define i32 @extractelement_shuffle_invalid_index(i32 %v) { ; CHECK-LABEL: @extractelement_shuffle_invalid_index( ; CHECK-NEXT: [[IN:%.*]] = insertelement poison, i32 [[V:%.*]], i32 0 @@ -95,6 +111,15 @@ ret i32 %r } +define i32 @extractconstant_shuffle_invalid_index() { +; CHECK-LABEL: @extractconstant_shuffle_invalid_index( +; CHECK-NEXT: ret i32 extractelement ( shufflevector ( insertelement ( poison, i32 1024, i32 0), poison, zeroinitializer), i32 -1) +; + %in = insertelement poison, i32 1024, i32 0 + %splat = shufflevector %in, poison, zeroinitializer + %r = extractelement %splat, i32 -1 + ret i32 %r +} define i32 @extractelement_shuffle_symbolic_index(i32 %v, i32 %idx) { ; CHECK-LABEL: @extractelement_shuffle_symbolic_index( @@ -109,6 +134,18 @@ ret i32 %r } +; SimplifyExtractElementInst assumes that any index in a splat +; should return the value +define i32 @extractconstant_shuffle_symbolic_index(i32 %idx) { +; CHECK-LABEL: @extractconstant_shuffle_symbolic_index( +; CHECK-NEXT: ret i32 1024 +; + %in = insertelement poison, i32 1024, i32 0 + %splat = shufflevector %in, poison, zeroinitializer + %r = extractelement %splat, i32 %idx + ret i32 %r +} + define @extractelement_insertelement_same_positions( %vec) { ; CHECK-LABEL: @extractelement_insertelement_same_positions( ; CHECK-NEXT: ret [[VEC:%.*]] diff --git a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll --- a/llvm/test/Transforms/InstCombine/vscale_extractelement.ll +++ b/llvm/test/Transforms/InstCombine/vscale_extractelement.ll @@ -55,13 +55,9 @@ 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 undef, i32 [[V:%.*]], i32 0 -; CHECK-NEXT: [[SPLAT:%.*]] = shufflevector [[IN]], undef, zeroinitializer -; CHECK-NEXT: [[R:%.*]] = extractelement [[SPLAT]], i32 1 -; CHECK-NEXT: ret i32 [[R]] +; CHECK-NEXT: ret i32 %v ; %in = insertelement undef, i32 %v, i32 0 %splat = shufflevector %in, undef, zeroinitializer @@ -69,6 +65,16 @@ ret i32 %r } +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 @extractelement_shuffle_maybe_out_of_range(i32 %v) { ; CHECK-LABEL: @extractelement_shuffle_maybe_out_of_range( ; CHECK-NEXT: [[IN:%.*]] = insertelement undef, i32 [[V:%.*]], i32 0 @@ -82,6 +88,16 @@ 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 @extractelement_shuffle_invalid_index(i32 %v) { ; CHECK-LABEL: @extractelement_shuffle_invalid_index( ; CHECK-NEXT: [[IN:%.*]] = insertelement undef, i32 [[V:%.*]], i32 0 @@ -95,6 +111,15 @@ 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 +} define i32 @extractelement_shuffle_symbolic_index(i32 %v, i32 %idx) { ; CHECK-LABEL: @extractelement_shuffle_symbolic_index( @@ -109,6 +134,18 @@ ret i32 %r } +; SimplifyExtractElementInst assumes that any index in a splat +; should return the value +define i32 @extractconstant_shuffle_symbolic_index(i32 %v, i32 %idx) { +; CHECK-LABEL: @extractconstant_shuffle_symbolic_index( +; CHECK-NEXT: ret i32 1024 +; + %in = insertelement undef, i32 1024, i32 0 + %splat = shufflevector %in, undef, zeroinitializer + %r = extractelement %splat, i32 %idx + ret i32 %r +} + define @extractelement_insertelement_same_positions( %vec) { ; CHECK-LABEL: @extractelement_insertelement_same_positions( ; CHECK-NEXT: ret [[VEC:%.*]]