[InstCombine] peek through bitcasted vector/array pointer GEP operand

[spatel]

Summary

The bitcast may be interfering with other combines or vectorization as shown in PR16739:
https://bugs.llvm.org/show_bug.cgi?id=16739

I haven't looked at GEPs very much, so I'm not sure if I've missed any safety checks here. I'm also not sure if this warning (still) applies:

http://llvm.org/docs/GetElementPtr.html#can-gep-index-into-vector-elements
"This hasn’t always been forcefully disallowed, though it’s not recommended. It leads to awkward special cases in the optimizers, and fundamental inconsistency in the IR. In the future, it will probably be outright disallowed."

...but it's 5 years since the bug was filed, and the IR looks the same.

[efriedma]

...but it's 5 years since the bug was filed, and the IR looks the same.

Yes, it's still a little awkward, but it would be too disruptive to change at this point.

That said, I'm surprised this matters; most pointer-related optimizations look through bitcasts.

[spatel]

In D44833#1047184, @efriedma wrote:

That said, I'm surprised this matters; most pointer-related optimizations look through bitcasts.

Yes, I think you're correct. I assumed the differing types might trip up the vectorizers, but at least SLP doesn't seem to care in the motivating example.

So not sure if we want to add code for this. It does still remove the bitcast from PR16739, so the IR shrinks even if we don't vectorize it as expected.

[spatel]

Ping.

[spatel]

Ping * 2.

[efriedma]

LGTM

lib/Transforms/InstCombine/InstructionCombining.cpp

if (auto *ASC = dyn_cast<AddrSpaceCastInst>(PtrOp)) {
     //   X = bitcast A addrspace(1)* to B addrspace(1)*
     //   Y = addrspacecast A addrspace(1)* to B addrspace(2)*
     //   Z = gep Y, <...constant indices...>
     // Into an addrspacecasted GEP of the struct.
     if (auto *BC = dyn_cast<BitCastInst>(ASC->getOperand(0)))
       PtrOp = BC;
   }
 
-  /// See if we can simplify:
-  ///   X = bitcast A* to B*
-  ///   Y = gep X, <...constant indices...>
-  /// into a gep of the original struct.  This is important for SROA and alias
-  /// analysis of unions.  If "A" is also a bitcast, wait for A/X to be merged.
   if (auto *BCI = dyn_cast<BitCastInst>(PtrOp)) {
     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) {
+      return ArrTy->getArrayElementType() == VecTy->getVectorElementType() &&
+             ArrTy->getArrayNumElements() == VecTy->getVectorNumElements();
+    };
+    if (GEP.getNumOperands() == 3 &&
+        ((GEPEltType->isArrayTy() && SrcEltType->isVectorTy() &&
+          areMatchingArrayAndVecTypes(GEPEltType, SrcEltType)) ||
+         (GEPEltType->isVectorTy() && SrcEltType->isArrayTy() &&
+          areMatchingArrayAndVecTypes(SrcEltType, GEPEltType)))) {
+      GEP.setOperand(0, SrcOp);
+      GEP.setSourceElementType(SrcEltType);
+      return &GEP;
+    }
+
+    // See if we can simplify:
+    //   X = bitcast A* to B*
+    //   Y = gep X, <...constant indices...>
+    // into a gep of the original struct. This is important for SROA and alias
+    // analysis of unions. If "A" is also a bitcast, wait for A/X to be merged.
     unsigned OffsetBits = DL.getIndexTypeSizeInBits(GEPType);
     APInt Offset(OffsetBits, 0);
     if (!isa<BitCastInst>(SrcOp) && GEP.accumulateConstantOffset(DL, Offset)) {
       // If this GEP instruction doesn't move the pointer, just replace the GEP
       // with a bitcast of the real input to the dest type.
       if (!Offset) {
         // If the bitcast is of an allocation, and the allocation will be
         // converted to match the type of the cast, don't touch this.

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
 }
 
+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:    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
+}
+
+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:    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
+}
+