[InstCombine] remove useless insertelement

extractelement (bitcast (insertelement (Vec, b)), a) ->
extractelement (bitcast (Vec), a)

Reviewed By: RKSimon

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

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -228,8 +228,9 @@ Instruction *InstCombinerImpl::foldBitcastExtElt(ExtractElementInst &Ext) {
   // truncate a subset of scalar bits of an insert op.
   if (NumSrcElts.getKnownMinValue() < NumElts.getKnownMinValue()) {
     Value *Scalar;
+    Value *Vec;
     uint64_t InsIndexC;
-    if (!match(X, m_InsertElt(m_Value(), m_Value(Scalar),
+    if (!match(X, m_InsertElt(m_Value(Vec), m_Value(Scalar),
                               m_ConstantInt(InsIndexC))))
       return nullptr;
 
@@ -239,8 +240,19 @@ Instruction *InstCombinerImpl::foldBitcastExtElt(ExtractElementInst &Ext) {
     // of elements 4-7 of the bitcasted vector.
     unsigned NarrowingRatio =
         NumElts.getKnownMinValue() / NumSrcElts.getKnownMinValue();
-    if (ExtIndexC / NarrowingRatio != InsIndexC)
+
+    if (ExtIndexC / NarrowingRatio != InsIndexC) {
+      // Remove insertelement, if we don't use the inserted element.
+      // extractelement (bitcast (insertelement (Vec, b)), a) ->
+      // extractelement (bitcast (Vec), a)
+      // FIXME: this should be removed to SimplifyDemandedVectorElts,
+      // once scale vectors are supported.
+      if (X->hasOneUse() && Ext.getVectorOperand()->hasOneUse()) {
+        Value *NewBC = Builder.CreateBitCast(Vec, Ext.getVectorOperandType());
+        return ExtractElementInst::Create(NewBC, Ext.getIndexOperand());
+      }
       return nullptr;
+    }
 
     // We are extracting part of the original scalar. How that scalar is
     // inserted into the vector depends on the endian-ness. Example:

llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll

@@ -1,6 +1,10 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
+
+declare void @use_vscale_2_i32(<vscale x 2 x i32>)
+declare void @use_vscale_8_i8(<vscale x 8 x i8>)
+
 define i32 @extractelement_in_range(<vscale x 4 x i32> %a) {
 ; CHECK-LABEL: @extractelement_in_range(
 ; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i64 1
@@ -41,16 +45,46 @@ define i8 @extractelement_bitcast_to_trunc(<vscale x 2 x i32> %a, i32 %x) {
   ret i8 %r
 }
 
-; TODO: Instcombine could remove the insert.
-define i8 @extractelement_bitcast_wrong_insert(<vscale x 2 x i32> %a, i32 %x) {
-; CHECK-LABEL: @extractelement_bitcast_wrong_insert(
+define i8 @extractelement_bitcast_useless_insert(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_useless_insert(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <vscale x 2 x i32> [[A:%.*]] to <vscale x 8 x i8>
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[TMP1]], i64 2
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1
+  %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
+  %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
+  ret i8 %r
+}
+
+; extra use tests
+
+define i8 @extractelement_bitcast_insert_extra_use_insert(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_insert_extra_use_insert(
 ; CHECK-NEXT:    [[VEC:%.*]] = insertelement <vscale x 2 x i32> [[A:%.*]], i32 [[X:%.*]], i64 1
+; CHECK-NEXT:    call void @use_vscale_2_i32(<vscale x 2 x i32> [[VEC]])
 ; CHECK-NEXT:    [[VEC_CAST:%.*]] = bitcast <vscale x 2 x i32> [[VEC]] to <vscale x 8 x i8>
 ; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[VEC_CAST]], i64 2
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
-  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1 ; <- This insert could be removed.
+  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1
+  call void @use_vscale_2_i32(<vscale x 2 x i32> %vec)
+  %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
+  %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
+  ret i8 %r
+}
+
+define i8 @extractelement_bitcast_insert_extra_use_bitcast(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_insert_extra_use_bitcast(
+; CHECK-NEXT:    [[VEC:%.*]] = insertelement <vscale x 2 x i32> [[A:%.*]], i32 [[X:%.*]], i64 1
+; CHECK-NEXT:    [[VEC_CAST:%.*]] = bitcast <vscale x 2 x i32> [[VEC]] to <vscale x 8 x i8>
+; CHECK-NEXT:    call void @use_vscale_8_i8(<vscale x 8 x i8> [[VEC_CAST]])
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[VEC_CAST]], i64 2
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1
   %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
+  call void @use_vscale_8_i8(<vscale x 8 x i8> %vec_cast)
   %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
   ret i8 %r
 }

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

@@ -41,12 +41,10 @@ define i8 @extractelement_bitcast_to_trunc(<vscale x 2 x i32> %a, i32 %x) {
   ret i8 %r
 }
 
-; TODO: Instcombine could remove the insert.
-define i8 @extractelement_bitcast_wrong_insert(<vscale x 2 x i32> %a, i32 %x) {
-; CHECK-LABEL: @extractelement_bitcast_wrong_insert(
-; CHECK-NEXT:    [[VEC:%.*]] = insertelement <vscale x 2 x i32> [[A:%.*]], i32 [[X:%.*]], i64 1
-; CHECK-NEXT:    [[VEC_CAST:%.*]] = bitcast <vscale x 2 x i32> [[VEC]] to <vscale x 8 x i8>
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[VEC_CAST]], i64 2
+define i8 @extractelement_bitcast_useless_insert(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_useless_insert(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <vscale x 2 x i32> [[A:%.*]] to <vscale x 8 x i8>
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[TMP1]], i64 2
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1 ; <- This insert could be removed.