[InstCombine] Fix big-endian miscompile of (bitcast (zext/trunc (bitcast)))

[bjope]

Summary

optimizeVectorResize is rewriting patterns like:

%1 = bitcast vector %src to integer
%2 = trunc/zext %1
%dst = bitcast %2 to vector

Since bitcasting between integer an vector types gives
different integer values depending on endianness, we need
to take endianness into account. As it happens the old
implementation only gave the correct result for little
endian targets.

Fixes: https://bugs.llvm.org/show_bug.cgi?id=44178

[lebedev.ri]

llvm/test/Transforms/InstCombine/cast.ll

2–4

I think this needs two runlines:

; RUN: opt < %s -instcombine -S --data-layout=<> | FileCheck %s --check-prefixes=ALL,LE
; RUN: opt < %s -instcombine -S --data-layout=<> | FileCheck %s --check-prefixes=ALL,BE

and drop hardcoded datalayout.

[bjope]

llvm/test/Transforms/InstCombine/cast.ll

2–4

Sure, I'll precommit an update to make this test case run with both big/little endian.

[bjope]

Rebase after update to run cast.ll for both big and little endian.

[lebedev.ri]

The test change looks correct to me, but i'm wondering whether there is
any cleaner way to structure those optimizeVectorResize() changes.
Or maybe that code is missing comments.

[bjope]

Minor refactoring and added some more code comments.

[bjope]

Realized it was stupid to use a temporary vector when it is possible to simply use append. So a yet another minor refactoring.

[lattner]

This LGTM but I'm not an active reviewer for this area, so I'd get a second opinion :)

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1851–1894

I personally find this to be a somewhat too convoluted approach.
What do you think about: (didn't compile-test)

SmallVector<uint32_t, 16> ShuffleMaskStorage;
ArrayRef<uint32_t> ShuffleMask;
Value *V2;

bool IsBigEndian = IC.getDataLayout().isBigEndian();
unsigned SrcElts = SrcTy->getNumElements();
unsigned DestElts = DestTy->getNumElements();

// Produce identity shuffle mask for src vector
ShuffleMaskStorage.resize(SrcElts);
std::iota(ShuffleMaskStorage.begin(), ShuffleMaskStorage.end(), 0);

if (SrcElts > DestElts) {
  // If we're shrinking the number of elements, just shuffle in the elements
  // from the input and use undef as the second shuffle input.
  V2 = UndefValue::get(SrcTy);
  ShuffleMask = ShuffleMaskStorage;
  unsigned EltDelta = SrcElts - DestElts;
  // Vector in-memory layout differs between little- and big- endian machines.
  if(IsBigEndian)
    ShuffleMask.drop_front(EltDelta)
  else
    ShuffleMask.drop_back(EltDelta)
} else {
  // If we're increasing the number of elements, shuffle in all of the
  // elements from InVal. Fill the rest of the result elements with zeros
  // from a constant zero.
  V2 = Constant::getNullValue(SrcTy);

  // Vector in-memory layout differs between little- and big- endian machines.
  decltype(ShuffleMaskStorage)::iterator InsertPos = IsBigEndian ?
                                                     ShuffleMaskStorage.begin() :
                                                     ShuffleMaskStorage.end()
  
  // And pad with zeros
  unsigned EltDelta = DestElts - SrcElts;
  ShuffleMaskStorage.insert(InsertPos, EltDelta, /*first element of second shuffle input*/SrcElts)
  
  ShuffleMask = ShuffleMaskStorage;
}

return new ShuffleVectorInst(InVal, V2,
                             ConstantDataVector::get(V2->getContext(),
                                                     ShuffleMask));

[bjope]

Some more refactoring, very much inspired by the ideas from Roman.

[lebedev.ri]

LGTM now, thank you.
Up to you if you want to wait for yet another review (@spatel ?)

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1872

maybe let's do

ShuffleMaskStorage.reserve(std::max(SrcElts, DestElts));
ShuffleMaskStorage.resize(SrcElts);

to try to avoid allocs in widening case.
Though we then waste memory in truncating case.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1866–1867

Add an assert that SrcElts != DestElts? Make sure the calling code doesn't get out-of-sync from the assumption in this function.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombiner::visitPtrToInt(PtrToIntInst &CI) {
   if (Ty->isVectorTy()) // Handle vectors of pointers.
     PtrTy = VectorType::get(PtrTy, Ty->getVectorNumElements());
 
   Value *P = Builder.CreatePtrToInt(CI.getOperand(0), PtrTy);
   return CastInst::CreateIntegerCast(P, Ty, /*isSigned=*/false);
 }
 
 /// This input value (which is known to have vector type) is being zero extended
-/// or truncated to the specified vector type.
+/// or truncated to the specified vector type. Since the zext/trunc is done
+/// using an integer type, we have a (bitcast(cast(bitcast))) pattern,
+/// endianness will impact which end of the vector that is extended or
+/// truncated.
+///
 /// Try to replace it with a shuffle (and vector/vector bitcast) if possible.
 ///
 /// The source and destination vector types may have different element types.
 static Instruction *optimizeVectorResize(Value *InVal, VectorType *DestTy,
                                          InstCombiner &IC) {
   // We can only do this optimization if the output is a multiple of the input
   // element size, or the input is a multiple of the output element size.
   // Convert the input type to have the same element type as the output.
   VectorType *SrcTy = cast<VectorType>(InVal->getType());
 
   if (SrcTy->getElementType() != DestTy->getElementType()) {
     // The input types don't need to be identical, but for now they must be the
     // same size.  There is no specific reason we couldn't handle things like
     // <4 x i16> -> <4 x i32> by bitcasting to <2 x i32> but haven't gotten
     // there yet.
     if (SrcTy->getElementType()->getPrimitiveSizeInBits() !=
         DestTy->getElementType()->getPrimitiveSizeInBits())
       return nullptr;
 
     SrcTy = VectorType::get(DestTy->getElementType(), SrcTy->getNumElements());
     InVal = IC.Builder.CreateBitCast(InVal, SrcTy);
   }
 
   // Now that the element types match, get the shuffle mask and RHS of the
   // shuffle to use, which depends on whether we're increasing or decreasing the
   // size of the input.
   SmallVector<uint32_t, 16> ShuffleMask;
   Value *V2;
 
-  if (SrcTy->getNumElements() > DestTy->getNumElements()) {
-    // If we're shrinking the number of elements, just shuffle in the low
-    // elements from the input and use undef as the second shuffle input.
+  bool IsBigEndian = IC.getDataLayout().isBigEndian();
+  unsigned SrcElts = SrcTy->getNumElements();
+  unsigned DestElts = DestTy->getNumElements();
+
+  if (SrcElts > DestElts) {
+    // If we're shrinking the number of elements, just shuffle in the elements
+    // from the input and use undef as the second shuffle input.
     V2 = UndefValue::get(SrcTy);
-    for (unsigned i = 0, e = DestTy->getNumElements(); i != e; ++i)
-      ShuffleMask.push_back(i);
+    // For big endian we should pick elements with high indices, for little
+    // endian from low indices.
+    unsigned FirstElt = IsBigEndian ? SrcElts - DestElts : 0;
+    for (unsigned i = 0, e = DestElts; i != e; ++i)
+      ShuffleMask.push_back(FirstElt + i);
 
   } else {
     // If we're increasing the number of elements, shuffle in all of the
-    // elements from InVal and fill the rest of the result elements with zeros
+    // elements from InVal. Fill the rest of the result elements with zeros
     // from a constant zero.
     V2 = Constant::getNullValue(SrcTy);
-    unsigned SrcElts = SrcTy->getNumElements();
+
+    // For big endian zeroes should be inserted at low indices. Use the first
+    // element of the zero vector.
+    if (IsBigEndian)
+      ShuffleMask.append(DestElts - SrcElts, SrcElts);
+
+    // Add the mask for picking values from the source vector <0, 1, 2 ...>.
     for (unsigned i = 0, e = SrcElts; i != e; ++i)
       ShuffleMask.push_back(i);
 
-    // The excess elements reference the first element of the zero input.
-    for (unsigned i = 0, e = DestTy->getNumElements()-SrcElts; i != e; ++i)
-      ShuffleMask.push_back(SrcElts);
+    // For little endian zeroes should be inserted at high indices. Use the
+    // first element of the zero vector.
+    if (!IsBigEndian)
+      ShuffleMask.append(DestElts - SrcElts, SrcElts);
   }
 
   return new ShuffleVectorInst(InVal, V2,
                                ConstantDataVector::get(V2->getContext(),
                                                        ShuffleMask));
 }
 
 static bool isMultipleOfTypeSize(unsigned Value, Type *Ty) {
   return Value % Ty->getPrimitiveSizeInBits() == 0;
 }
 
 static unsigned getTypeSizeIndex(unsigned Value, Type *Ty) {
   return Value / Ty->getPrimitiveSizeInBits();

llvm/test/Transforms/InstCombine/cast.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; Tests to make sure elimination of casts is working correctly
 ; RUN: opt < %s -instcombine -S -data-layout="E-p:64:64:64-p1:32:32:32-p2:64:64:64-p3:64:64:64-a0:0:8-f32:32:32-f64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-v64:64:64-v128:128:128-n8:16:32:64" | FileCheck %s --check-prefixes=ALL,BE
 ; RUN: opt < %s -instcombine -S -data-layout="e-p:64:64:64-p1:32:32:32-p2:64:64:64-p3:64:64:64-a0:0:8-f32:32:32-f64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-v64:64:64-v128:128:128-n8:16:32:64" | FileCheck %s --check-prefixes=ALL,LE
 
 @inbuf = external global [32832 x i8]
 
 define i32 @test1(i32 %A) {
 ; ALL-LABEL: @test1(
 ; ALL-NEXT:    ret i32 [[A:%.*]]
 ;
   %c1 = bitcast i32 %A to i32
;
   %F = zext i8 %B to i32
   %G = lshr i32 %F, 4
   %H = or i32 %G, %E
   %I = zext i32 %H to i64
   ret i64 %I
 }
 
 define <3 x i32> @test60(<4 x i32> %call4) {
-; ALL-LABEL: @test60(
-; ALL-NEXT:    [[P10:%.*]] = shufflevector <4 x i32> [[CALL4:%.*]], <4 x i32> undef, <3 x i32> <i32 0, i32 1, i32 2>
-; ALL-NEXT:    ret <3 x i32> [[P10]]
+; BE-LABEL: @test60(
+; BE-NEXT:    [[P10:%.*]] = shufflevector <4 x i32> [[CALL4:%.*]], <4 x i32> undef, <3 x i32> <i32 1, i32 2, i32 3>
+; BE-NEXT:    ret <3 x i32> [[P10]]
+;
+; LE-LABEL: @test60(
+; LE-NEXT:    [[P10:%.*]] = shufflevector <4 x i32> [[CALL4:%.*]], <4 x i32> undef, <3 x i32> <i32 0, i32 1, i32 2>
+; LE-NEXT:    ret <3 x i32> [[P10]]
 ;
   %p11 = bitcast <4 x i32> %call4 to i128
   %p9 = trunc i128 %p11 to i96
   %p10 = bitcast i96 %p9 to <3 x i32>
   ret <3 x i32> %p10
 
 }
 
 define <4 x i32> @test61(<3 x i32> %call4) {
-; ALL-LABEL: @test61(
-; ALL-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[CALL4:%.*]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
-; ALL-NEXT:    ret <4 x i32> [[P10]]
+; BE-LABEL: @test61(
+; BE-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[CALL4:%.*]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 3, i32 0, i32 1, i32 2>
+; BE-NEXT:    ret <4 x i32> [[P10]]
+;
+; LE-LABEL: @test61(
+; LE-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[CALL4:%.*]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; LE-NEXT:    ret <4 x i32> [[P10]]
 ;
   %p11 = bitcast <3 x i32> %call4 to i96
   %p9 = zext i96 %p11 to i128
   %p10 = bitcast i128 %p9 to <4 x i32>
   ret <4 x i32> %p10
 }
 
 define <4 x i32> @test62(<3 x float> %call4) {
-; ALL-LABEL: @test62(
-; ALL-NEXT:    [[TMP1:%.*]] = bitcast <3 x float> [[CALL4:%.*]] to <3 x i32>
-; ALL-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[TMP1]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
-; ALL-NEXT:    ret <4 x i32> [[P10]]
+; BE-LABEL: @test62(
+; BE-NEXT:    [[TMP1:%.*]] = bitcast <3 x float> [[CALL4:%.*]] to <3 x i32>
+; BE-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[TMP1]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 3, i32 0, i32 1, i32 2>
+; BE-NEXT:    ret <4 x i32> [[P10]]
+;
+; LE-LABEL: @test62(
+; LE-NEXT:    [[TMP1:%.*]] = bitcast <3 x float> [[CALL4:%.*]] to <3 x i32>
+; LE-NEXT:    [[P10:%.*]] = shufflevector <3 x i32> [[TMP1]], <3 x i32> <i32 0, i32 undef, i32 undef>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; LE-NEXT:    ret <4 x i32> [[P10]]
 ;
   %p11 = bitcast <3 x float> %call4 to i96
   %p9 = zext i96 %p11 to i128
   %p10 = bitcast i128 %p9 to <4 x i32>
   ret <4 x i32> %p10
 }
 
 ; PR7311 - Don't create invalid IR on scalar->vector cast.