[InstCombine] check alloc size in bitcast of geps fold (PR44321)

[spatel]

Summary

We missed a constraint in D44833 when folding a bitcast into a GEP with vector/array types.
If the alloc sizes specified by the datalayout don't match, this could miscompile as shown in:
https://bugs.llvm.org/show_bug.cgi?id=44321

[aqjune]

LGTM!

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (auto *BCI = dyn_cast<BitCastInst>(ASCStrippedPtrOp)) {
     Value *SrcOp = BCI->getOperand(0);
     PointerType *SrcType = cast<PointerType>(BCI->getSrcTy());
     Type *SrcEltType = SrcType->getElementType();
 
     // GEP directly using the source operand if this GEP is accessing an element
     // of a bitcasted pointer to vector or array of the same dimensions:
     // gep (bitcast <c x ty>* X to [c x ty]*), Y, Z --> gep X, Y, Z
     // gep (bitcast [c x ty]* X to <c x ty>*), Y, Z --> gep X, Y, Z
-    auto areMatchingArrayAndVecTypes = [](Type *ArrTy, Type *VecTy) {
+    auto areMatchingArrayAndVecTypes = [](Type *ArrTy, Type *VecTy,
+                                          const DataLayout &DL) {
       return ArrTy->getArrayElementType() == VecTy->getVectorElementType() &&
-             ArrTy->getArrayNumElements() == VecTy->getVectorNumElements();
+             ArrTy->getArrayNumElements() == VecTy->getVectorNumElements() &&
+             DL.getTypeAllocSize(ArrTy) == DL.getTypeAllocSize(VecTy);
     };
     if (GEP.getNumOperands() == 3 &&
         ((GEPEltType->isArrayTy() && SrcEltType->isVectorTy() &&
-          areMatchingArrayAndVecTypes(GEPEltType, SrcEltType)) ||
+          areMatchingArrayAndVecTypes(GEPEltType, SrcEltType, DL)) ||
          (GEPEltType->isVectorTy() && SrcEltType->isArrayTy() &&
-          areMatchingArrayAndVecTypes(SrcEltType, GEPEltType)))) {
+          areMatchingArrayAndVecTypes(SrcEltType, GEPEltType, DL)))) {
 
       // Create a new GEP here, as using `setOperand()` followed by
       // `setSourceElementType()` won't actually update the type of the
       // existing GEP Value. Causing issues if this Value is accessed when
       // constructing an AddrSpaceCastInst
       Value *NGEP =
           GEP.isInBounds()
               ? Builder.CreateInBoundsGEP(SrcEltType, SrcOp, {Ops[1], Ops[2]})

llvm/test/Transforms/InstCombine/gep-vector.ll

 define <2 x i8*> @vectorindex3() {
 ; CHECK-LABEL: @vectorindex3(
 ; CHECK-NEXT:    ret <2 x i8*> getelementptr inbounds ([64 x [8192 x i8]], [64 x [8192 x i8]]* @block, <2 x i64> zeroinitializer, <2 x i64> <i64 0, i64 2>, <2 x i64> <i64 8191, i64 1>)
 ;
   %1 = getelementptr inbounds [64 x [8192 x i8]], [64 x [8192 x i8]]* @block, i64 0, <2 x i64> <i64 0, i64 1>, <2 x i64> <i64 8191, i64 8193>
   ret <2 x i8*> %1
 }
 
+; Negative test - datalayout's alloc size for the 2 types must match.
+
 define i32* @bitcast_vec_to_array_gep(<7 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_gep(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <7 x i32>, <7 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[ARR_PTR:%.*]] = bitcast <7 x i32>* [[X:%.*]] to [7 x i32]*
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [7 x i32], [7 x i32]* [[ARR_PTR]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
   %gep = getelementptr [7 x i32], [7 x i32]* %arr_ptr, i64 %y, i64 %z
   ret i32* %gep
 }
 
+; Negative test - datalayout's alloc size for the 2 types must match.
+
 define i32* @bitcast_array_to_vec_gep([3 x i32]* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_array_to_vec_gep(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x i32], [3 x i32]* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[VEC_PTR:%.*]] = bitcast [3 x i32]* [[X:%.*]] to <3 x i32>*
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <3 x i32>, <3 x i32>* [[VEC_PTR]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %vec_ptr = bitcast [3 x i32]* %x to <3 x i32>*
   %gep = getelementptr inbounds <3 x i32>, <3 x i32>* %vec_ptr, i64 %y, i64 %z
   ret i32* %gep
 }
 
+; Sizes and types match - safe to remove bitcast.
+
 define i32* @bitcast_vec_to_array_gep_matching_alloc_size(<4 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_gep_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <4 x i32>, <4 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %arr_ptr = bitcast <4 x i32>* %x to [4 x i32]*
   %gep = getelementptr [4 x i32], [4 x i32]* %arr_ptr, i64 %y, i64 %z
   ret i32* %gep
 }
 
+; Sizes and types match - safe to remove bitcast.
+
 define i32* @bitcast_array_to_vec_gep_matching_alloc_size([4 x i32]* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_array_to_vec_gep_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [4 x i32], [4 x i32]* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %vec_ptr = bitcast [4 x i32]* %x to <4 x i32>*
   %gep = getelementptr inbounds <4 x i32>, <4 x i32>* %vec_ptr, i64 %y, i64 %z
   ret i32* %gep
 }
 
+; Negative test - datalayout's alloc size for the 2 types must match.
+
 define i32 addrspace(3)* @bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_addrspace(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <7 x i32>, <7 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
-; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+; CHECK-NEXT:    [[ARR_PTR:%.*]] = bitcast <7 x i32>* [[X:%.*]] to [7 x i32]*
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast [7 x i32]* [[ARR_PTR]] to [7 x i32] addrspace(3)*
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [7 x i32], [7 x i32] addrspace(3)* [[ASC]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    ret i32 addrspace(3)* [[GEP]]
 ;
   %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
   %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
   %gep = getelementptr [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
+; Negative test - datalayout's alloc size for the 2 types must match.
+
 define i32 addrspace(3)* @inbounds_bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @inbounds_bitcast_vec_to_array_addrspace(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <7 x i32>, <7 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
-; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+; CHECK-NEXT:    [[ARR_PTR:%.*]] = bitcast <7 x i32>* [[X:%.*]] to [7 x i32]*
+; CHECK-NEXT:    [[ASC:%.*]] = addrspacecast [7 x i32]* [[ARR_PTR]] to [7 x i32] addrspace(3)*
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [7 x i32], [7 x i32] addrspace(3)* [[ASC]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    ret i32 addrspace(3)* [[GEP]]
 ;
   %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
   %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
   %gep = getelementptr inbounds [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
+; Sizes and types match - safe to remove bitcast.
+
 define i32 addrspace(3)* @bitcast_vec_to_array_addrspace_matching_alloc_size(<4 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @bitcast_vec_to_array_addrspace_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <4 x i32>, <4 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
 ; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
 ;
   %arr_ptr = bitcast <4 x i32>* %x to [4 x i32]*
   %asc = addrspacecast [4 x i32]* %arr_ptr to [4 x i32] addrspace(3)*
   %gep = getelementptr [4 x i32], [4 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }
 
+; Sizes and types match - safe to remove bitcast.
+
 define i32 addrspace(3)* @inbounds_bitcast_vec_to_array_addrspace_matching_alloc_size(<4 x i32>* %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: @inbounds_bitcast_vec_to_array_addrspace_matching_alloc_size(
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
 ; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
 ;
   %arr_ptr = bitcast <4 x i32>* %x to [4 x i32]*
   %asc = addrspacecast [4 x i32]* %arr_ptr to [4 x i32] addrspace(3)*
   %gep = getelementptr inbounds [4 x i32], [4 x i32] addrspace(3)* %asc, i64 %y, i64 %z
   ret i32 addrspace(3)* %gep
 }