Skip bitcasts while looking for GEP in LoadStoreVectorizer

[arsenm]

I don't think this should be needed? Is this representative of the case where you saw this as a problem? If I run this testcase through-separate-const-offset-from-gep first it handles it. This and other cases are expected to be cleaned up by that first

[rampitec]

In D32101#727681, @arsenm wrote:

I don't think this should be needed? Is this representative of the case where you saw this as a problem? If I run this testcase through-separate-const-offset-from-gep first it handles it. This and other cases are expected to be cleaned up by that first

It is quite representative. Here is the real piece of code from app, right as it comes to the vectorizer:

%97 = zext i32 %96 to i64
%98 = getelementptr inbounds float, float addrspace(1)* %1, i64 %97
%99 = bitcast float addrspace(1)* %98 to i32 addrspace(1)*
%100 = load i32, i32 addrspace(1)* %99, align 4, !tbaa !10
%101 = add i32 %96, 1
%102 = zext i32 %101 to i64
%103 = getelementptr inbounds float, float addrspace(1)* %1, i64 %102
%104 = bitcast float addrspace(1)* %103 to i32 addrspace(1)*
%105 = load i32, i32 addrspace(1)* %104, align 4, !tbaa !10
%106 = add i32 %96, 2
%107 = zext i32 %106 to i64
%108 = getelementptr inbounds float, float addrspace(1)* %1, i64 %107
%109 = bitcast float addrspace(1)* %108 to i32 addrspace(1)*
%110 = load i32, i32 addrspace(1)* %109, align 4, !tbaa !10
%111 = add i32 %96, 3
%112 = zext i32 %111 to i64
%113 = getelementptr inbounds float, float addrspace(1)* %1, i64 %112
%114 = bitcast float addrspace(1)* %113 to i32 addrspace(1)*
%115 = load i32, i32 addrspace(1)* %114, align 4, !tbaa !10

[arsenm]

In D32101#727682, @rampitec wrote:

In D32101#727681, @arsenm wrote:

I don't think this should be needed? Is this representative of the case where you saw this as a problem? If I run this testcase through-separate-const-offset-from-gep first it handles it. This and other cases are expected to be cleaned up by that first

It is quite representative. Here is the real piece of code from app, right as it comes to the vectorizer:

%97 = zext i32 %96 to i64
%98 = getelementptr inbounds float, float addrspace(1)* %1, i64 %97
%99 = bitcast float addrspace(1)* %98 to i32 addrspace(1)*
%100 = load i32, i32 addrspace(1)* %99, align 4, !tbaa !10
%101 = add i32 %96, 1
%102 = zext i32 %101 to i64
%103 = getelementptr inbounds float, float addrspace(1)* %1, i64 %102
%104 = bitcast float addrspace(1)* %103 to i32 addrspace(1)*
%105 = load i32, i32 addrspace(1)* %104, align 4, !tbaa !10
%106 = add i32 %96, 2
%107 = zext i32 %106 to i64
%108 = getelementptr inbounds float, float addrspace(1)* %1, i64 %107
%109 = bitcast float addrspace(1)* %108 to i32 addrspace(1)*
%110 = load i32, i32 addrspace(1)* %109, align 4, !tbaa !10
%111 = add i32 %96, 3
%112 = zext i32 %111 to i64
%113 = getelementptr inbounds float, float addrspace(1)* %1, i64 %112
%114 = bitcast float addrspace(1)* %113 to i32 addrspace(1)*
%115 = load i32, i32 addrspace(1)* %114, align 4, !tbaa !10

What does it look like immediately after SeparateConstOffsetFromGEP? Does one of the other passes break this somehow?

[rampitec]

In D32101#727691, @arsenm wrote:

What does it look like immediately after SeparateConstOffsetFromGEP? Does one of the other passes break this somehow?

This is after SeparateConstOffsetFromGEP:

%408 = zext i32 %407 to i64
%409 = getelementptr inbounds float, float addrspace(1)* %1, i64 %408
%410 = bitcast float addrspace(1)* %409 to i32 addrspace(1)*
%411 = load i32, i32 addrspace(1)* %410, align 4, !tbaa !10
%412 = or i32 %407, 1
%413 = zext i32 %412 to i64
%414 = getelementptr inbounds float, float addrspace(1)* %1, i64 %413
%415 = bitcast float addrspace(1)* %414 to i32 addrspace(1)*
%416 = load i32, i32 addrspace(1)* %415, align 4, !tbaa !10
%417 = or i32 %407, 2
%418 = zext i32 %417 to i64
%419 = getelementptr inbounds float, float addrspace(1)* %1, i64 %418
%420 = bitcast float addrspace(1)* %419 to i32 addrspace(1)*
%421 = load i32, i32 addrspace(1)* %420, align 4, !tbaa !10
%422 = or i32 %407, 3
%423 = zext i32 %422 to i64
%424 = getelementptr inbounds float, float addrspace(1)* %1, i64 %423
%425 = bitcast float addrspace(1)* %424 to i32 addrspace(1)*
%426 = load i32, i32 addrspace(1)* %425, align 4, !tbaa !10

Bitcasts are there plus we have or's instead of adds.

[rampitec]

Replaced code with standard stripPointerCasts().

[arsenm]

Can you add a test with bitcasts between pointers with different element types? I thought the or problem was also supposed to be solved

[rampitec]

In D32101#729447, @arsenm wrote:

Can you add a test with bitcasts between pointers with different element types? I thought the or problem was also supposed to be solved

Matt, what do you exactly mean by pointers with different element types? They are different in this test, float vs i32.

[arsenm]

In D32101#729486, @rampitec wrote:

In D32101#729447, @arsenm wrote:

Can you add a test with bitcasts between pointers with different element types? I thought the or problem was also supposed to be solved

Matt, what do you exactly mean by pointers with different element types? They are different in this test, float vs i32.

I meant type sizes. float and i32 are both 4 bytes, I could see something going wrong if later code relied on this assumption if the source type were i8 for example

[rampitec]

Added check for pointee size.

[rampitec]

Ping.

[vpykhtin]

LGTM.

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

310

May be turn it into a helper function?

[rampitec]

Factored out getSourceGEP() function as per review.

[vpykhtin]

Looks good.

[rampitec]

rL301343

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

public:
   Vectorizer(Function &F, AliasAnalysis &AA, DominatorTree &DT,
              ScalarEvolution &SE, TargetTransformInfo &TTI)
       : F(F), AA(AA), DT(DT), SE(SE), TTI(TTI),
         DL(F.getParent()->getDataLayout()), Builder(SE.getContext()) {}
 
   bool run();
 
 private:
-  Value *getPointerOperand(Value *I);
+  Value *getPointerOperand(Value *I) const;
+
+  GetElementPtrInst *getSourceGEP(Value *Src) const;
 
   unsigned getPointerAddressSpace(Value *I);
 
   unsigned getAlignment(LoadInst *LI) const {
     unsigned Align = LI->getAlignment();
     if (Align != 0)
       return Align;
 
for (BasicBlock *BB : post_order(&F)) {
     std::tie(LoadRefs, StoreRefs) = collectInstructions(BB);
     Changed |= vectorizeChains(LoadRefs);
     Changed |= vectorizeChains(StoreRefs);
   }
 
   return Changed;
 }
 
-Value *Vectorizer::getPointerOperand(Value *I) {
+Value *Vectorizer::getPointerOperand(Value *I) const {
   if (LoadInst *LI = dyn_cast<LoadInst>(I))
     return LI->getPointerOperand();
   if (StoreInst *SI = dyn_cast<StoreInst>(I))
     return SI->getPointerOperand();
   return nullptr;
 }
 
 unsigned Vectorizer::getPointerAddressSpace(Value *I) {
   if (LoadInst *L = dyn_cast<LoadInst>(I))
     return L->getPointerAddressSpace();
   if (StoreInst *S = dyn_cast<StoreInst>(I))
     return S->getPointerAddressSpace();
   return -1;
 }
 
+GetElementPtrInst *Vectorizer::getSourceGEP(Value *Src) const {
+  // First strip pointer bitcasts. Make sure pointee size is the same with
+  // and without casts.
+  // TODO: a stride set by the add instruction below can match the difference
+  // in pointee type size here. Currently it will not be vectorized.
+  Value *SrcPtr = getPointerOperand(Src);
+  Value *SrcBase = SrcPtr->stripPointerCasts();
+  if (DL.getTypeStoreSize(SrcPtr->getType()->getPointerElementType()) ==
+      DL.getTypeStoreSize(SrcBase->getType()->getPointerElementType()))
+    SrcPtr = SrcBase;
+  return dyn_cast<GetElementPtrInst>(SrcPtr);
+}
+
 // FIXME: Merge with llvm::isConsecutiveAccess
 bool Vectorizer::isConsecutiveAccess(Value *A, Value *B) {
   Value *PtrA = getPointerOperand(A);
   Value *PtrB = getPointerOperand(B);
   unsigned ASA = getPointerAddressSpace(A);
   unsigned ASB = getPointerAddressSpace(B);
 
   // Check that the address spaces match and that the pointers are valid.
if (X == PtrSCEVB)
     return true;
 
   // Sometimes even this doesn't work, because SCEV can't always see through
   // patterns that look like (gep (ext (add (shl X, C1), C2))). Try checking
   // things the hard way.
 
   // Look through GEPs after checking they're the same except for the last
   // index.
-  GetElementPtrInst *GEPA = dyn_cast<GetElementPtrInst>(getPointerOperand(A));
-  GetElementPtrInst *GEPB = dyn_cast<GetElementPtrInst>(getPointerOperand(B));
+  GetElementPtrInst *GEPA = getSourceGEP(A);
+  GetElementPtrInst *GEPB = getSourceGEP(B);
   if (!GEPA || !GEPB || GEPA->getNumOperands() != GEPB->getNumOperands())
     return false;
   unsigned FinalIndex = GEPA->getNumOperands() - 1;
   for (unsigned i = 0; i < FinalIndex; i++)
     if (GEPA->getOperand(i) != GEPB->getOperand(i))
       return false;
 
   Instruction *OpA = dyn_cast<Instruction>(GEPA->getOperand(FinalIndex));
   Instruction *OpB = dyn_cast<Instruction>(GEPB->getOperand(FinalIndex));
   if (!OpA || !OpB || OpA->getOpcode() != OpB->getOpcode() ||
       OpA->getType() != OpB->getType())
     return false;
 
   // Only look through a ZExt/SExt.
   if (!isa<SExtInst>(OpA) && !isa<ZExtInst>(OpA))
     return false;

test/Transforms/LoadStoreVectorizer/AMDGPU/gep-bitcast.ll

+; RUN: opt -S -mtriple=amdgcn--amdhsa -load-store-vectorizer < %s | FileCheck %s
+
+; Check that vectorizer can find a GEP through bitcast
+; CHECK-LABEL: @vect_zext_bitcast_f32_to_i32_idx
+; CHECK: load <4 x i32>
+define void @vect_zext_bitcast_f32_to_i32_idx(float addrspace(1)* %arg1, i32 %base) {
+  %add1 = add nuw i32 %base, 0
+  %zext1 = zext i32 %add1 to i64
+  %gep1 = getelementptr inbounds float, float addrspace(1)* %arg1, i64 %zext1
+  %f2i1 = bitcast float addrspace(1)* %gep1 to i32 addrspace(1)*
+  %load1 = load i32, i32 addrspace(1)* %f2i1, align 4
+  %add2 = add nuw i32 %base, 1
+  %zext2 = zext i32 %add2 to i64
+  %gep2 = getelementptr inbounds float, float addrspace(1)* %arg1, i64 %zext2
+  %f2i2 = bitcast float addrspace(1)* %gep2 to i32 addrspace(1)*
+  %load2 = load i32, i32 addrspace(1)* %f2i2, align 4
+  %add3 = add nuw i32 %base, 2
+  %zext3 = zext i32 %add3 to i64
+  %gep3 = getelementptr inbounds float, float addrspace(1)* %arg1, i64 %zext3
+  %f2i3 = bitcast float addrspace(1)* %gep3 to i32 addrspace(1)*
+  %load3 = load i32, i32 addrspace(1)* %f2i3, align 4
+  %add4 = add nuw i32 %base, 3
+  %zext4 = zext i32 %add4 to i64
+  %gep4 = getelementptr inbounds float, float addrspace(1)* %arg1, i64 %zext4
+  %f2i4 = bitcast float addrspace(1)* %gep4 to i32 addrspace(1)*
+  %load4 = load i32, i32 addrspace(1)* %f2i4, align 4
+  ret void
+}
+
+; CHECK-LABEL: @vect_zext_bitcast_i8_st1_to_i32_idx
+; CHECK: load i32
+; CHECK: load i32
+; CHECK: load i32
+; CHECK: load i32
+define void @vect_zext_bitcast_i8_st1_to_i32_idx(i8 addrspace(1)* %arg1, i32 %base) {
+  %add1 = add nuw i32 %base, 0
+  %zext1 = zext i32 %add1 to i64
+  %gep1 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext1
+  %f2i1 = bitcast i8 addrspace(1)* %gep1 to i32 addrspace(1)*
+  %load1 = load i32, i32 addrspace(1)* %f2i1, align 4
+  %add2 = add nuw i32 %base, 1
+  %zext2 = zext i32 %add2 to i64
+  %gep2 = getelementptr inbounds i8,i8 addrspace(1)* %arg1, i64 %zext2
+  %f2i2 = bitcast i8 addrspace(1)* %gep2 to i32 addrspace(1)*
+  %load2 = load i32, i32 addrspace(1)* %f2i2, align 4
+  %add3 = add nuw i32 %base, 2
+  %zext3 = zext i32 %add3 to i64
+  %gep3 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext3
+  %f2i3 = bitcast i8 addrspace(1)* %gep3 to i32 addrspace(1)*
+  %load3 = load i32, i32 addrspace(1)* %f2i3, align 4
+  %add4 = add nuw i32 %base, 3
+  %zext4 = zext i32 %add4 to i64
+  %gep4 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext4
+  %f2i4 = bitcast i8 addrspace(1)* %gep4 to i32 addrspace(1)*
+  %load4 = load i32, i32 addrspace(1)* %f2i4, align 4
+  ret void
+}
+
+; TODO: This can be vectorized, but currently vectorizer unable to do it.
+; CHECK-LABEL: @vect_zext_bitcast_i8_st4_to_i32_idx
+define void @vect_zext_bitcast_i8_st4_to_i32_idx(i8 addrspace(1)* %arg1, i32 %base) {
+  %add1 = add nuw i32 %base, 0
+  %zext1 = zext i32 %add1 to i64
+  %gep1 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext1
+  %f2i1 = bitcast i8 addrspace(1)* %gep1 to i32 addrspace(1)*
+  %load1 = load i32, i32 addrspace(1)* %f2i1, align 4
+  %add2 = add nuw i32 %base, 4
+  %zext2 = zext i32 %add2 to i64
+  %gep2 = getelementptr inbounds i8,i8 addrspace(1)* %arg1, i64 %zext2
+  %f2i2 = bitcast i8 addrspace(1)* %gep2 to i32 addrspace(1)*
+  %load2 = load i32, i32 addrspace(1)* %f2i2, align 4
+  %add3 = add nuw i32 %base, 8
+  %zext3 = zext i32 %add3 to i64
+  %gep3 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext3
+  %f2i3 = bitcast i8 addrspace(1)* %gep3 to i32 addrspace(1)*
+  %load3 = load i32, i32 addrspace(1)* %f2i3, align 4
+  %add4 = add nuw i32 %base, 16
+  %zext4 = zext i32 %add4 to i64
+  %gep4 = getelementptr inbounds i8, i8 addrspace(1)* %arg1, i64 %zext4
+  %f2i4 = bitcast i8 addrspace(1)* %gep4 to i32 addrspace(1)*
+  %load4 = load i32, i32 addrspace(1)* %f2i4, align 4
+  ret void
+}