[ConstantFold] Teach how to fold a bitcast of a constant vector with some undef elements.

[andreadb]

Summary

The previous logic didn't always know how to handle bitcasts of integer vectors with some undef elements.

The problem only affected the case where the bitcast vector type had more packed elements than the input vector type.

Example:

%cast = bitcast <2 x i64> <i64 undef, i64 2> to <4 x i32>

On a little endian target, that bitcast could be folded into the following constant vector:
<4 x i32><i32 undef, i32 undef, i32 2, i32 0>

However, the input vector had some undef elements, and the constant folder didn't know how to handle them.
This patch improves the logic by teaching how to deal with undefs.

Please let me know if okay to commit.
Thanks,
Andrea

[ABataev]

lib/Analysis/ConstantFolding.cpp

234–235

Result.append(Ratio, UndefValue::get(DstEltTy));

?

[andreadb]

lib/Analysis/ConstantFolding.cpp

234–235

Thanks Alexey, I am going to change it.

[andreadb]

Patch updated.
Addressed Alexey's code review comment.

[RKSimon]

Please can you submit the current bitcast-vector-fold.ll tests to show the changes?

[andreadb]

In D24301#541053, @RKSimon wrote:

Please can you submit the current bitcast-vector-fold.ll tests to show the changes?

Sure, I am going to commit those tests now.
Then I will update the patch.

[andreadb]

Address Simon's review comment.

Patch rebased after revision 281332.

[RKSimon]

LGTM

lib/Analysis/ConstantFolding.cpp

Constant *FoldBitCast(Constant *C, Type *DestTy, const DataLayout &DL) {
   }
 
   // Handle: bitcast (<2 x i64> <i64 0, i64 1> to <4 x i32>)
   unsigned Ratio = NumDstElt/NumSrcElt;
   unsigned DstBitSize = DL.getTypeSizeInBits(DstEltTy);
 
   // Loop over each source value, expanding into multiple results.
   for (unsigned i = 0; i != NumSrcElt; ++i) {
-    auto *Src = dyn_cast_or_null<ConstantInt>(C->getAggregateElement(i));
-    if (!Src)  // Reject constantexpr elements.
+    auto *Element = C->getAggregateElement(i);
+
+    if (!Element) // Reject constantexpr elements.
+      return ConstantExpr::getBitCast(C, DestTy);
+
+    if (isa<UndefValue>(Element)) {
+      // Correctly Propagate undef values.
+      Result.append(Ratio, UndefValue::get(DstEltTy));
+      continue;
+    }
+
+    auto *Src = dyn_cast<ConstantInt>(Element);
+    if (!Src)
       return ConstantExpr::getBitCast(C, DestTy);
 
     unsigned ShiftAmt = isLittleEndian ? 0 : DstBitSize*(Ratio-1);
     for (unsigned j = 0; j != Ratio; ++j) {
       // Shift the piece of the value into the right place, depending on
       // endianness.
       Constant *Elt = ConstantExpr::getLShr(Src,
                                   ConstantInt::get(Src->getType(), ShiftAmt));

test/Transforms/InstCombine/x86-sse4a.ll

 ; CHECK-NEXT:    ret <2 x i64> <i64 65535, i64 undef>
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.extrq(<2 x i64> <i64 -1, i64 undef>, <16 x i8> <i8 16, i8 15, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>) nounwind
   ret <2 x i64> %1
 }
 
 define <2 x i64> @test_extrq_call_constexpr(<2 x i64> %x) {
 ; CHECK-LABEL: @test_extrq_call_constexpr(
-; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i64> @llvm.x86.sse4a.extrq(<2 x i64> %x, <16 x i8> bitcast (<2 x i64> <i64 0, i64 undef> to <16 x i8>))
-; CHECK-NEXT:    ret <2 x i64> [[TMP1]]
+; CHECK-NEXT:    ret <2 x i64> %x
 ;
   %1 = call <2 x i64> @llvm.x86.sse4a.extrq(<2 x i64> %x, <16 x i8> bitcast (<2 x i64> <i64 0, i64 undef> to <16 x i8>))
   ret <2 x i64> %1
 }
 
 ;
 ; EXTRQI
 ;

 ; CHECK-NEXT:    ret <2 x i64> <i64 15, i64 undef>
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.extrqi(<2 x i64> <i64 -1, i64 undef>, i8 4, i8 18)
   ret <2 x i64> %1
 }
 
 define <2 x i64> @test_extrqi_call_constexpr() {
 ; CHECK-LABEL: @test_extrqi_call_constexpr(
-; CHECK-NEXT:    ret <2 x i64> bitcast (<16 x i8> <i8 extractelement (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>), i32 2), i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef> to <2 x i64>)
+; CHECK-NEXT:    ret <2 x i64> zeroinitializer
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.extrqi(<2 x i64> bitcast (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>) to <2 x i64>), i8 8, i8 16)
   ret <2 x i64> %1
 }
 
 ;
 ; INSERTQ
 ;

 ; CHECK-NEXT:    ret <2 x i64> <i64 33, i64 undef>
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.insertq(<2 x i64> <i64 1, i64 undef>, <2 x i64> <i64 8, i64 658>) nounwind
   ret <2 x i64> %1
 }
 
 define <2 x i64> @test_insertq_call_constexpr(<2 x i64> %x) {
 ; CHECK-LABEL: @test_insertq_call_constexpr(
-; CHECK-NEXT:    [[TMP1:%.*]] = tail call <2 x i64> @llvm.x86.sse4a.insertq(<2 x i64> %x, <2 x i64> bitcast (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>) to <2 x i64>))
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i64> @llvm.x86.sse4a.insertqi(<2 x i64> %x, <2 x i64> <i64 0, i64 undef>, i8 2, i8 0)
 ; CHECK-NEXT:    ret <2 x i64> [[TMP1]]
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.insertq(<2 x i64> %x, <2 x i64> bitcast (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>) to <2 x i64>))
   ret <2 x i64> %1
 }
 
 ;
 ; INSERTQI

 ; CHECK-NEXT:    ret <2 x i64> <i64 -131055, i64 undef>
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.insertqi(<2 x i64> <i64 -1, i64 -1>, <2 x i64> <i64 8, i64 0>, i8 16, i8 1)
   ret <2 x i64> %1
 }
 
 define <2 x i64> @test_insertqi_call_constexpr(<2 x i64> %x) {
 ; CHECK-LABEL: @test_insertqi_call_constexpr(
-; CHECK-NEXT:    [[TMP1:%.*]] = tail call <2 x i64> @llvm.x86.sse4a.insertqi(<2 x i64> %x, <2 x i64> bitcast (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>) to <2 x i64>), i8 48, i8 3)
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call <2 x i64> @llvm.x86.sse4a.insertqi(<2 x i64> %x, <2 x i64> <i64 0, i64 undef>, i8 48, i8 3)
 ; CHECK-NEXT:    ret <2 x i64> [[TMP1]]
 ;
   %1 = tail call <2 x i64> @llvm.x86.sse4a.insertqi(<2 x i64> %x, <2 x i64> bitcast (<16 x i8> trunc (<16 x i16> bitcast (<4 x i64> <i64 0, i64 undef, i64 2, i64 undef> to <16 x i16>) to <16 x i8>) to <2 x i64>), i8 48, i8 3)
   ret <2 x i64> %1
 }
 
 ; The result of this insert is the second arg, since the top 64 bits of
 ; the result are undefined, and we copy the bottom 64 bits from the

test/Transforms/InstSimplify/bitcast-vector-fold.ll

 
 define <2 x double> @foo6() {
 ; CHECK-LABEL: @foo6(
 ; CHECK-NEXT:    ret <2 x double> <double 0xFFFFFFFFFFFFFFFF, double 0xFFFFFFFFFFFFFFFF>
 ;
   %cast = bitcast <4 x i32><i32 -1, i32 -1, i32 -1, i32 -1> to <2 x double>
   ret <2 x double> %cast
 }
+
+define <4 x i32> @bitcast_constexpr_4i32_2i64_u2() {
+; CHECK-LABEL: @bitcast_constexpr_4i32_2i64_u2(
+; CHECK-NEXT:    ret <4 x i32> <i32 undef, i32 undef, i32 2, i32 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 2> to <4 x i32>
+  ret <4 x i32> %cast
+}
+
+define <4 x i32> @bitcast_constexpr_4i32_2i64_1u() {
+; CHECK-LABEL: @bitcast_constexpr_4i32_2i64_1u(
+; CHECK-NEXT:    ret <4 x i32> <i32 1, i32 0, i32 undef, i32 undef>
+;
+  %cast = bitcast <2 x i64><i64 1, i64 undef> to <4 x i32>
+  ret <4 x i32> %cast
+}
+
+define <4 x i32> @bitcast_constexpr_4i32_2i64() {
+; CHECK-LABEL: @bitcast_constexpr_4i32_2i64(
+; CHECK-NEXT:    ret <4 x i32> <i32 undef, i32 undef, i32 2, i32 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 2> to <4 x i32>
+  ret <4 x i32> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_2i64_u2() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_2i64_u2(
+; CHECK-NEXT:    ret <8 x i16> <i16 undef, i16 undef, i16 undef, i16 undef, i16 2, i16 0, i16 0, i16 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 2> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_2i64_1u() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_2i64_1u(
+; CHECK-NEXT:    ret <8 x i16> <i16 1, i16 0, i16 0, i16 0, i16 undef, i16 undef, i16 undef, i16 undef>
+;
+  %cast = bitcast <2 x i64><i64 1, i64 undef> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_2i64_u65536() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_2i64_u65536(
+; CHECK-NEXT:    ret <8 x i16> <i16 undef, i16 undef, i16 undef, i16 undef, i16 0, i16 1, i16 0, i16 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 65536> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_2i64_u2() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_2i64_u2(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 2, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 2> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_2i64_256u() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_2i64_256u(
+; CHECK-NEXT:    ret <16 x i8> <i8 0, i8 1, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef>
+;
+  %cast = bitcast <2 x i64><i64 256, i64 undef> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_2i64_u256() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_2i64_u256(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 0, i8 1, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
+;
+  %cast = bitcast <2 x i64><i64 undef, i64 256> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_4i32_uu22() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_4i32_uu22(
+; CHECK-NEXT:    ret <8 x i16> <i16 undef, i16 undef, i16 undef, i16 undef, i16 2, i16 0, i16 2, i16 0>
+;
+  %cast = bitcast <4 x i32><i32 undef, i32 undef, i32 2, i32 2> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_4i32_10uu() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_4i32_10uu(
+; CHECK-NEXT:    ret <8 x i16> <i16 1, i16 0, i16 0, i16 0, i16 undef, i16 undef, i16 undef, i16 undef>
+;
+  %cast = bitcast <4 x i32><i32 1, i32 0, i32 undef, i32 undef> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <8 x i16> @bitcast_constexpr_8i16_4i32_u257u256() {
+; CHECK-LABEL: @bitcast_constexpr_8i16_4i32_u257u256(
+; CHECK-NEXT:    ret <8 x i16> <i16 undef, i16 undef, i16 0, i16 1, i16 undef, i16 undef, i16 0, i16 1>
+;
+  %cast = bitcast <4 x i32><i32 undef, i32 65536, i32 undef, i32 65536> to <8 x i16>
+  ret <8 x i16> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_4i32_u2u2() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_4i32_u2u2(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 2, i8 0, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 2, i8 0, i8 0, i8 0>
+;
+  %cast = bitcast <4 x i32><i32 undef, i32 2, i32 undef, i32 2> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_4i32_1u1u() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_4i32_1u1u(
+; CHECK-NEXT:    ret <16 x i8> <i8 1, i8 0, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 1, i8 0, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef>
+;
+  %cast = bitcast <4 x i32><i32 1, i32 undef, i32 1, i32 undef> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_4i32_u256uu() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_4i32_u256uu(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 undef, i8 undef, i8 0, i8 1, i8 0, i8 0, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef>
+;
+  %cast = bitcast <4 x i32><i32 undef, i32 256, i32 undef, i32 undef> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_8i16_u2u2u2u2() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_8i16_u2u2u2u2(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 2, i8 0, i8 undef, i8 undef, i8 2, i8 0, i8 undef, i8 undef, i8 2, i8 0, i8 undef, i8 undef, i8 2, i8 0>
+;
+  %cast = bitcast <8 x i16><i16 undef, i16 2, i16 undef, i16 2, i16 undef, i16 2, i16 undef, i16 2> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_8i16_1u1u1u1u() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_8i16_1u1u1u1u(
+; CHECK-NEXT:    ret <16 x i8> <i8 1, i8 0, i8 undef, i8 undef, i8 1, i8 0, i8 undef, i8 undef, i8 1, i8 0, i8 undef, i8 undef, i8 1, i8 0, i8 undef, i8 undef>
+;
+  %cast = bitcast <8 x i16><i16 1, i16 undef, i16 1, i16 undef, i16 1, i16 undef, i16 1, i16 undef> to <16 x i8>
+  ret <16 x i8> %cast
+}
+
+define <16 x i8> @bitcast_constexpr_16i8_8i16_u256uuu256uu() {
+; CHECK-LABEL: @bitcast_constexpr_16i8_8i16_u256uuu256uu(
+; CHECK-NEXT:    ret <16 x i8> <i8 undef, i8 undef, i8 0, i8 1, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 undef, i8 0, i8 1, i8 undef, i8 undef, i8 undef, i8 undef>
+;
+  %cast = bitcast <8 x i16><i16 undef, i16 256, i16 undef, i16 undef, i16 undef, i16 256, i16 undef, i16 undef> to <16 x i8>
+  ret <16 x i8> %cast
+}