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 @@ -398,6 +398,20 @@ SQ.getWithInstruction(&EI))) return replaceInstUsesWith(EI, V); + // extractelt (select %x, %vec1, %vec2), %const -> + // select %x, %vec1[%const], %vec2[%const] + // TODO: Support constant folding of multiple select operands: + // extractelt (select %x, %vec1, %vec2), (select %x, %c1, %c2) + // If the extractelement will for instance try to do out of bounds accesses + // because of the values of %c1 and/or %c2, the sequence could be optimized + // early. This is currently not possible because constant folding will reach + // an unreachable assertion if it doesn't find a constant operand. + if (SelectInst *SI = dyn_cast(EI.getVectorOperand())) + if (SI->getCondition()->getType()->isIntegerTy() && + isa(EI.getIndexOperand())) + if (Instruction *R = FoldOpIntoSelect(EI, SI)) + return R; + // If extracting a specified index from the vector, see if we can recursively // find a previously computed scalar that was inserted into the vector. auto *IndexC = dyn_cast(Index); diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp --- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp +++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp @@ -1069,6 +1069,9 @@ return Builder.CreateBinaryIntrinsic(IID, SO, II->getArgOperand(1)); } + if (auto *EI = dyn_cast(&I)) + return Builder.CreateExtractElement(SO, EI->getIndexOperand()); + assert(I.isBinaryOp() && "Unexpected opcode for select folding"); // Figure out if the constant is the left or the right argument. diff --git a/llvm/test/Transforms/InstCombine/extractelement.ll b/llvm/test/Transforms/InstCombine/extractelement.ll --- a/llvm/test/Transforms/InstCombine/extractelement.ll +++ b/llvm/test/Transforms/InstCombine/extractelement.ll @@ -785,8 +785,7 @@ define i32 @extelt_select_const_operand_vector(i1 %c) { ; ANY-LABEL: @extelt_select_const_operand_vector( -; ANY-NEXT: [[S:%.*]] = select i1 [[C:%.*]], <3 x i32> , <3 x i32> -; ANY-NEXT: [[R:%.*]] = extractelement <3 x i32> [[S]], i64 2 +; ANY-NEXT: [[R:%.*]] = select i1 [[C:%.*]], i32 4, i32 7 ; ANY-NEXT: ret i32 [[R]] ; %s = select i1 %c, <3 x i32> , <3 x i32> @@ -796,8 +795,7 @@ define float @extelt_select_const_operand_vector_float(i1 %c) { ; ANY-LABEL: @extelt_select_const_operand_vector_float( -; ANY-NEXT: [[S:%.*]] = select i1 [[C:%.*]], <3 x float> , <3 x float> -; ANY-NEXT: [[R:%.*]] = extractelement <3 x float> [[S]], i64 2 +; ANY-NEXT: [[R:%.*]] = select i1 [[C:%.*]], float 4.000000e+00, float 7.000000e+00 ; ANY-NEXT: ret float [[R]] ; %s = select i1 %c, <3 x float> , <3 x float> @@ -818,11 +816,10 @@ define i32 @extelt_select_const_operand_extractelt_use(i1 %c) { ; ANY-LABEL: @extelt_select_const_operand_extractelt_use( -; ANY-NEXT: [[S:%.*]] = select i1 [[C:%.*]], <3 x i32> , <3 x i32> -; ANY-NEXT: [[E:%.*]] = extractelement <3 x i32> [[S]], i64 2 -; ANY-NEXT: [[M:%.*]] = shl i32 [[E]], 1 -; ANY-NEXT: [[M_2:%.*]] = shl i32 [[E]], 2 -; ANY-NEXT: [[R:%.*]] = mul i32 [[M]], [[M_2]] +; ANY-NEXT: [[E:%.*]] = select i1 [[C:%.*]], i32 4, i32 7 +; ANY-NEXT: [[M:%.*]] = shl nuw nsw i32 [[E]], 1 +; ANY-NEXT: [[M_2:%.*]] = shl nuw nsw i32 [[E]], 2 +; ANY-NEXT: [[R:%.*]] = mul nuw nsw i32 [[M]], [[M_2]] ; ANY-NEXT: ret i32 [[R]] ; %s = select i1 %c, <3 x i32> , <3 x i32> @@ -874,8 +871,8 @@ define i32 @extelt_select_var_const_operand_vector(i1 %c, <3 x i32> %v) { ; ANY-LABEL: @extelt_select_var_const_operand_vector( -; ANY-NEXT: [[S:%.*]] = select i1 [[C:%.*]], <3 x i32> [[V:%.*]], <3 x i32> -; ANY-NEXT: [[R:%.*]] = extractelement <3 x i32> [[S]], i64 1 +; ANY-NEXT: [[TMP1:%.*]] = extractelement <3 x i32> [[V:%.*]], i64 1 +; ANY-NEXT: [[R:%.*]] = select i1 [[C:%.*]], i32 [[TMP1]], i32 6 ; ANY-NEXT: ret i32 [[R]] ; %s = select i1 %c, <3 x i32> %v, <3 x i32> @@ -885,8 +882,8 @@ define i32 @extelt_select_const_var_operand_vector(i1 %c, <3 x i32> %v) { ; ANY-LABEL: @extelt_select_const_var_operand_vector( -; ANY-NEXT: [[S:%.*]] = select i1 [[C:%.*]], <3 x i32> , <3 x i32> [[V:%.*]] -; ANY-NEXT: [[R:%.*]] = extractelement <3 x i32> [[S]], i64 0 +; ANY-NEXT: [[TMP1:%.*]] = extractelement <3 x i32> [[V:%.*]], i64 0 +; ANY-NEXT: [[R:%.*]] = select i1 [[C:%.*]], i32 5, i32 [[TMP1]] ; ANY-NEXT: ret i32 [[R]] ; %s = select i1 %c, <3 x i32> , <3 x i32> %v diff --git a/llvm/test/Transforms/InstCombine/shufflevector-div-rem-inseltpoison.ll b/llvm/test/Transforms/InstCombine/shufflevector-div-rem-inseltpoison.ll --- a/llvm/test/Transforms/InstCombine/shufflevector-div-rem-inseltpoison.ll +++ b/llvm/test/Transforms/InstCombine/shufflevector-div-rem-inseltpoison.ll @@ -8,11 +8,8 @@ ; extracting the second element in the vector). define i16 @test_srem_orig(i16 %a, i1 %cmp) { ; CHECK-LABEL: @test_srem_orig( -; CHECK-NEXT: [[SPLATINSERT:%.*]] = insertelement <2 x i16> poison, i16 [[A:%.*]], i64 0 -; CHECK-NEXT: [[TMP1:%.*]] = srem <2 x i16> [[SPLATINSERT]], -; CHECK-NEXT: [[SPLAT_OP:%.*]] = shufflevector <2 x i16> [[TMP1]], <2 x i16> poison, <2 x i32> -; CHECK-NEXT: [[T2:%.*]] = select i1 [[CMP:%.*]], <2 x i16> , <2 x i16> [[SPLAT_OP]] -; CHECK-NEXT: [[T3:%.*]] = extractelement <2 x i16> [[T2]], i64 1 +; CHECK-NEXT: [[TMP1:%.*]] = srem i16 [[A:%.*]], 2 +; CHECK-NEXT: [[T3:%.*]] = select i1 [[CMP:%.*]], i16 1, i16 [[TMP1]] ; CHECK-NEXT: ret i16 [[T3]] ; %splatinsert = insertelement <2 x i16> poison, i16 %a, i32 0 diff --git a/llvm/test/Transforms/InstCombine/shufflevector-div-rem.ll b/llvm/test/Transforms/InstCombine/shufflevector-div-rem.ll --- a/llvm/test/Transforms/InstCombine/shufflevector-div-rem.ll +++ b/llvm/test/Transforms/InstCombine/shufflevector-div-rem.ll @@ -8,11 +8,8 @@ ; extracting the second element in the vector). define i16 @test_srem_orig(i16 %a, i1 %cmp) { ; CHECK-LABEL: @test_srem_orig( -; CHECK-NEXT: [[SPLATINSERT:%.*]] = insertelement <2 x i16> undef, i16 [[A:%.*]], i64 0 -; CHECK-NEXT: [[TMP1:%.*]] = srem <2 x i16> [[SPLATINSERT]], -; CHECK-NEXT: [[SPLAT_OP:%.*]] = shufflevector <2 x i16> [[TMP1]], <2 x i16> poison, <2 x i32> -; CHECK-NEXT: [[T2:%.*]] = select i1 [[CMP:%.*]], <2 x i16> , <2 x i16> [[SPLAT_OP]] -; CHECK-NEXT: [[T3:%.*]] = extractelement <2 x i16> [[T2]], i64 1 +; CHECK-NEXT: [[TMP1:%.*]] = srem i16 [[A:%.*]], 2 +; CHECK-NEXT: [[T3:%.*]] = select i1 [[CMP:%.*]], i16 1, i16 [[TMP1]] ; CHECK-NEXT: ret i16 [[T3]] ; %splatinsert = insertelement <2 x i16> undef, i16 %a, i32 0