[InstCombine] Fold extractelt with select of constants

An extractelt with a constant index which extracts an element from the
two vector operands of a select can be directly folded into a select.

extractelt (select %x, %vec1, %vec2), %const ->
select %x, %vec1[%const], %vec2[%const]

Note: the implementation currently only works for constant vector operands.

Reviewed By: foad, spatel

Differential Revision: https://reviews.llvm.org/D137934

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 5964c96..782af54 100644 (file)
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -398,6 +398,20 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
                                             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<SelectInst>(EI.getVectorOperand()))
+    if (SI->getCondition()->getType()->isIntegerTy() &&
+        isa<Constant>(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<ConstantInt>(Index);

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 81467cf..ee20520 100644 (file)
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1069,6 +1069,9 @@ static Value *foldOperationIntoSelectOperand(Instruction &I, Value *SO,
     return Builder.CreateBinaryIntrinsic(IID, SO, II->getArgOperand(1));
   }
 
+  if (auto *EI = dyn_cast<ExtractElementInst>(&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
index 63277d1..16a2483 100644 (file)
--- a/llvm/test/Transforms/InstCombine/extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/extractelement.ll
@@ -785,8 +785,7 @@ define i1 @bit_extract_cmp(i64 %x) {
 
 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> <i32 poison, i32 poison, i32 4>, <3 x i32> <i32 poison, i32 poison, i32 7>
-; 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> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
@@ -796,8 +795,7 @@ define i32 @extelt_select_const_operand_vector(i1 %c) {
 
 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> <float poison, float poison, float 4.000000e+00>, <3 x float> <float poison, float poison, float 7.000000e+00>
-; 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> <float 2.0, float 3.0, float 4.0>, <3 x float> <float 5.0, float 6.0, float 7.0>
@@ -818,11 +816,10 @@ define i32 @extelt_vecselect_const_operand_vector(<3 x i1> %c) {
 
 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> <i32 poison, i32 poison, i32 4>, <3 x i32> <i32 poison, i32 poison, i32 7>
-; 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> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
@@ -874,8 +871,8 @@ define i32 @extelt_select_const_operand_vector_var_index(i1 %c, i32 %e) {
 
 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> <i32 poison, i32 6, i32 poison>
-; 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> <i32 5, i32 6, i32 7>
@@ -885,8 +882,8 @@ define i32 @extelt_select_var_const_operand_vector(i1 %c, <3 x i32> %v) {
 
 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> <i32 5, i32 poison, i32 poison>, <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> <i32 5, i32 6, i32 7>, <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
index 4e6e492..d7684aa 100644 (file)
--- 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]], <i16 2, i16 1>
-; CHECK-NEXT:    [[SPLAT_OP:%.*]] = shufflevector <2 x i16> [[TMP1]], <2 x i16> poison, <2 x i32> <i32 undef, i32 0>
-; CHECK-NEXT:    [[T2:%.*]] = select i1 [[CMP:%.*]], <2 x i16> <i16 poison, i16 1>, <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
index 9c07f99..6345a58 100644 (file)
--- 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]], <i16 2, i16 1>
-; CHECK-NEXT:    [[SPLAT_OP:%.*]] = shufflevector <2 x i16> [[TMP1]], <2 x i16> poison, <2 x i32> <i32 undef, i32 0>
-; CHECK-NEXT:    [[T2:%.*]] = select i1 [[CMP:%.*]], <2 x i16> <i16 poison, i16 1>, <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