[InstCombine] canonicalize trunc + insert as bitcast + shuffle, part 2

[spatel]

Summary

This enhances the base fold from part 1 to allow mapping a right-shift to an insert index.

Example of translating a middle chunk of the scalar to vector for either endian:
https://alive2.llvm.org/ce/z/fRXCOZ

This only allows creating an identity shuffle (with optional shortening/lengthening) because that is considered the safe baseline for any target (can be inverted if needed). If we tried this fold with target-specific costs/legality, then we could do the transform more generally.

[spatel]

Rebased on top of part 1 ( a4c466766db7 ).

[RKSimon]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1553

Make it clear that we only allow an 'identity' shuffle mask here

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

if (ScalarWidth % VecEltWidth != 0)
     return nullptr;
 
   unsigned NumEltsInScalar = ScalarWidth / VecEltWidth;
   Value *X = T;
   if (IndexC == (IsBigEndian ? NumEltsInScalar - 1 : 0)) {
     // The insert is to the LSB end of the vector (depends on endian).
     // That's all we need.
   } else {
-    // TODO: Look through a shift-right and translate the insert index.
+    // If not, we must match a right-shift to translate the insert index.
+    uint64_t ShiftC;
+    if (!match(T, m_OneUse(m_LShr(m_Value(X), m_ConstantInt(ShiftC)))))
+      return nullptr;
+
+    // Check the shift amount to see if this can be folded to an identity
+    // shuffle (assuming we are shuffling with an undef base vector).
+    // Big endian has MSB at vector index 0, so the insert index is flipped.
+    if (ShiftC != (IsBigEndian ? (NumEltsInScalar - 1 - IndexC) * VecEltWidth
+                               : IndexC * VecEltWidth))
       return nullptr;
   }
 
   // Bitcast the scalar to a vector type with the destination element type.
   Type *CastTy = FixedVectorType::get(VTy->getElementType(), NumEltsInScalar);
   Value *VecX = Builder.CreateBitCast(X, CastTy, "vec." + X->getName());
 
   unsigned NumElts = VTy->getNumElements();
   if (NumElts > NumEltsInScalar) {

llvm/test/Transforms/InstCombine/insert-trunc.ll

 ;
   %t = trunc i64 %x to i8
   call void @use(i8 %t)
   %r = insertelement <4 x i8> poison, i8 %t, i64 0
   ret <4 x i8> %r
 }
 
 define <4 x i16> @lshr_same_length_poison_basevec_le(i64 %x) {
-; ALL-LABEL: @lshr_same_length_poison_basevec_le(
-; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
-; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; ALL-NEXT:    [[R:%.*]] = insertelement <4 x i16> poison, i16 [[T]], i64 2
-; ALL-NEXT:    ret <4 x i16> [[R]]
+; BE-LABEL: @lshr_same_length_poison_basevec_le(
+; BE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
+; BE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
+; BE-NEXT:    [[R:%.*]] = insertelement <4 x i16> poison, i16 [[T]], i64 2
+; BE-NEXT:    ret <4 x i16> [[R]]
+;
+; LE-LABEL: @lshr_same_length_poison_basevec_le(
+; LE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; LE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <4 x i32> <i32 undef, i32 undef, i32 2, i32 undef>
+; LE-NEXT:    ret <4 x i16> [[R]]
 ;
   %s = lshr i64 %x, 32
   %t = trunc i64 %s to i16
   %r = insertelement <4 x i16> poison, i16 %t, i64 2
   ret <4 x i16> %r
 }
 
 define <4 x i16> @lshr_same_length_poison_basevec_be(i64 %x) {
-; ALL-LABEL: @lshr_same_length_poison_basevec_be(
-; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
-; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; ALL-NEXT:    [[R:%.*]] = insertelement <4 x i16> poison, i16 [[T]], i64 1
-; ALL-NEXT:    ret <4 x i16> [[R]]
+; BE-LABEL: @lshr_same_length_poison_basevec_be(
+; BE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; BE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <4 x i32> <i32 undef, i32 1, i32 undef, i32 undef>
+; BE-NEXT:    ret <4 x i16> [[R]]
+;
+; LE-LABEL: @lshr_same_length_poison_basevec_be(
+; LE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
+; LE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
+; LE-NEXT:    [[R:%.*]] = insertelement <4 x i16> poison, i16 [[T]], i64 1
+; LE-NEXT:    ret <4 x i16> [[R]]
 ;
   %s = lshr i64 %x, 32
   %t = trunc i64 %s to i16
   %r = insertelement <4 x i16> poison, i16 %t, i64 1
   ret <4 x i16> %r
 }
 
 define <4 x i16> @lshr_same_length_poison_basevec_both_endian(i64 %x) {
 ; BE-LABEL: @lshr_same_length_poison_basevec_both_endian(
-; BE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 48
-; BE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; BE-NEXT:    [[R:%.*]] = insertelement <4 x i16> poison, i16 [[T]], i64 0
+; BE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; BE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
 ; BE-NEXT:    ret <4 x i16> [[R]]
 ;
 ; LE-LABEL: @lshr_same_length_poison_basevec_both_endian(
 ; LE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 48
 ; LE-NEXT:    [[VEC_S:%.*]] = bitcast i64 [[S]] to <4 x i16>
 ; LE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_S]], <4 x i16> poison, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
 ; LE-NEXT:    ret <4 x i16> [[R]]
 ;

 define <8 x i16> @lshr_longer_length_poison_basevec_le(i64 %x) {
 ; BE-LABEL: @lshr_longer_length_poison_basevec_le(
 ; BE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 48
 ; BE-NEXT:    [[VEC_S:%.*]] = bitcast i64 [[S]] to <4 x i16>
 ; BE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_S]], <4 x i16> poison, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
 ; BE-NEXT:    ret <8 x i16> [[R]]
 ;
 ; LE-LABEL: @lshr_longer_length_poison_basevec_le(
-; LE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 48
-; LE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; LE-NEXT:    [[R:%.*]] = insertelement <8 x i16> poison, i16 [[T]], i64 3
+; LE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; LE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
 ; LE-NEXT:    ret <8 x i16> [[R]]
 ;
   %s = lshr i64 %x, 48
   %t = trunc i64 %s to i16
   %r = insertelement <8 x i16> poison, i16 %t, i64 3
   ret <8 x i16> %r
 }
 
 define <8 x i16> @lshr_longer_length_poison_basevec_be(i64 %x) {
-; ALL-LABEL: @lshr_longer_length_poison_basevec_be(
-; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
-; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; ALL-NEXT:    [[R:%.*]] = insertelement <8 x i16> poison, i16 [[T]], i64 1
-; ALL-NEXT:    ret <8 x i16> [[R]]
+; BE-LABEL: @lshr_longer_length_poison_basevec_be(
+; BE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; BE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <8 x i32> <i32 undef, i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; BE-NEXT:    ret <8 x i16> [[R]]
+;
+; LE-LABEL: @lshr_longer_length_poison_basevec_be(
+; LE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 32
+; LE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
+; LE-NEXT:    [[R:%.*]] = insertelement <8 x i16> poison, i16 [[T]], i64 1
+; LE-NEXT:    ret <8 x i16> [[R]]
 ;
   %s = lshr i64 %x, 32
   %t = trunc i64 %s to i16
   %r = insertelement <8 x i16> poison, i16 %t, i64 1
   ret <8 x i16> %r
 }
 
 define <8 x i16> @lshr_wrong_index_longer_length_poison_basevec(i64 %x) {
 ; ALL-LABEL: @lshr_wrong_index_longer_length_poison_basevec(
 ; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 16
 ; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
 ; ALL-NEXT:    [[R:%.*]] = insertelement <8 x i16> poison, i16 [[T]], i64 6
 ; ALL-NEXT:    ret <8 x i16> [[R]]
 ;
   %s = lshr i64 %x, 16
   %t = trunc i64 %s to i16
   %r = insertelement <8 x i16> poison, i16 %t, i64 6
   ret <8 x i16> %r
 }
 
 define <2 x i16> @lshr_shorter_length_poison_basevec_le(i64 %x) {
-; ALL-LABEL: @lshr_shorter_length_poison_basevec_le(
-; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 16
-; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
-; ALL-NEXT:    [[R:%.*]] = insertelement <2 x i16> poison, i16 [[T]], i64 1
-; ALL-NEXT:    ret <2 x i16> [[R]]
+; BE-LABEL: @lshr_shorter_length_poison_basevec_le(
+; BE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 16
+; BE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i16
+; BE-NEXT:    [[R:%.*]] = insertelement <2 x i16> poison, i16 [[T]], i64 1
+; BE-NEXT:    ret <2 x i16> [[R]]
+;
+; LE-LABEL: @lshr_shorter_length_poison_basevec_le(
+; LE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <4 x i16>
+; LE-NEXT:    [[R:%.*]] = shufflevector <4 x i16> [[VEC_X]], <4 x i16> poison, <2 x i32> <i32 undef, i32 1>
+; LE-NEXT:    ret <2 x i16> [[R]]
 ;
   %s = lshr i64 %x, 16
   %t = trunc i64 %s to i16
   %r = insertelement <2 x i16> poison, i16 %t, i64 1
   ret <2 x i16> %r
 }
 
 define <4 x i8> @lshr_shorter_length_poison_basevec_be(i64 %x) {
-; ALL-LABEL: @lshr_shorter_length_poison_basevec_be(
-; ALL-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 40
-; ALL-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i8
-; ALL-NEXT:    [[R:%.*]] = insertelement <4 x i8> poison, i8 [[T]], i64 2
-; ALL-NEXT:    ret <4 x i8> [[R]]
+; BE-LABEL: @lshr_shorter_length_poison_basevec_be(
+; BE-NEXT:    [[VEC_X:%.*]] = bitcast i64 [[X:%.*]] to <8 x i8>
+; BE-NEXT:    [[R:%.*]] = shufflevector <8 x i8> [[VEC_X]], <8 x i8> poison, <4 x i32> <i32 undef, i32 undef, i32 2, i32 undef>
+; BE-NEXT:    ret <4 x i8> [[R]]
+;
+; LE-LABEL: @lshr_shorter_length_poison_basevec_be(
+; LE-NEXT:    [[S:%.*]] = lshr i64 [[X:%.*]], 40
+; LE-NEXT:    [[T:%.*]] = trunc i64 [[S]] to i8
+; LE-NEXT:    [[R:%.*]] = insertelement <4 x i8> poison, i8 [[T]], i64 2
+; LE-NEXT:    ret <4 x i8> [[R]]
 ;
   %s = lshr i64 %x, 40
   %t = trunc i64 %s to i8
   %r = insertelement <4 x i8> poison, i8 %t, i64 2
   ret <4 x i8> %r
 }
 
 define <4 x i8> @lshr_wrong_index_shorter_length_poison_basevec(i64 %x) {