AMDGPU/SI: Handle BitCast of GEP in promoting alloca to vector

[cfang]

Summary

This patch handles the case that the pointer is a BITCAST of GEP to change the data type. Before this patch, the check stops
at the BITCAST (pointer) and the load/store are not actually handled at all.

We also remove an invalid LIT test because we haven't handled the case yet, and check always pass with or without the promotion because
the index is an immediate and the load/store could be easily optimized away in other optimizations.

[cfang]

ping

[cfang]

Add reviewers and ping.

Thanks;

[arsenm]

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

330

I think there's another bug here (which should be fixed in a separate patch). This probably also needs to check if it's atomic

340–341

This shouldn't need to special case GEP users. You really want something like stripPointerCasts, but not one that looks through addrspacecast.

Could use a test with 0 index GEPs which should show the same problem.

348

Typo othereise

354

This looks broken. addrspacecast should forbid the vector promotion. It should only be allowed for the LDS promotion. Separate patch?

357

Same as with the load

433–436

Use dyn_cast instead of separate isa and cast

452

Same as the load case

test/CodeGen/AMDGPU/vector-alloca.ll

154

needs some more tests with multiple uses of the bitcast source, with accesses through different types as well as non-bitcast users

[cfang]

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

340–341

I could not understand this comment.

Here we are handling the case of load(bitcast(gep))!
We are checking the bitcase and collect the load. We will transfer to vector load later, which depends on the gep for the index.
So we have to have the gep, and collect the load/store.

[cfang]

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

330

Done in a separate patch under review.

354

Done in a separate patch under review.

[cfang]

Update based on Matt's comments.

[arsenm]

Is this still needed?

[rampitec]

I have run these new tests with ToT opt and there were no allocas remaining. I assume this change is not needed anymore.

[cfang]

This is no longer needed.

lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

static Value* GEPToVectorIndex(GetElementPtrInst *GEP) {
 
   ConstantInt *I0 = dyn_cast<ConstantInt>(GEP->getOperand(1));
   if (!I0 || !I0->isZero())
     return nullptr;
 
   return GEP->getOperand(2);
 }
 
-// Not an instruction handled below to turn into a vector.
+// NOTE: We mainly check whether a load or a store is vectorizable here.
+// A special case here is BITCAST of a GEP, in which case we check
+// whether all users of the BITCAST is vectorizable.
 //
 // TODO: Check isTriviallyVectorizable for calls and handle other
 // instructions.
-static bool canVectorizeInst(Instruction *Inst, User *User) {
+static bool canVectorizeInst(Instruction *Inst, User *Used,
+                             std::vector<Value*> &WorkList) {
   switch (Inst->getOpcode()) {
   case Instruction::Load: {
     LoadInst *LI = cast<LoadInst>(Inst);
-    // Currently only handle the case where the Pointer Operand is a GEP so check for that case.
-    return isa<GetElementPtrInst>(LI->getPointerOperand()) && !LI->isVolatile();
+    if (LI->isVolatile())
+      return false;
+    // Currently only handle the case where the Pointer Operand is a GEP
+    // or a BITCAST.
+    if (LI->getPointerOperand() != Used ||
+        (!isa<GetElementPtrInst>(Used) && !isa<BitCastInst>(Used)))
+      return false;
+    WorkList.push_back(Inst);
+    return true;
+  }
+  case Instruction::BitCast: {
+    if (isa<GetElementPtrInst>(Used)) {
+      for (User *BCUser : Inst->users()) {
+        if (!canVectorizeInst(cast<Instruction>(BCUser), Inst, WorkList))
+          return false;
+      }
+      return true;
+    }
+    // Fallthrough othereise.
+    // TODO: we do not actually have logic to handle general bitcast and
+    //       addrspacecast. We may have to be conservative here to avoid
+    //       unexpected results.
   }
-  case Instruction::BitCast:
   case Instruction::AddrSpaceCast:
     return true;
   case Instruction::Store: {
-    // Must be the stored pointer operand, not a stored value, plus
-    // since it should be canonical form, the User should be a GEP.
     StoreInst *SI = cast<StoreInst>(Inst);
-    return (SI->getPointerOperand() == User) && isa<GetElementPtrInst>(User) && !SI->isVolatile();
+    if (SI->isVolatile())
+      return false;
+    // Currently only handle the case where the Pointer Operand is a GEP
+    // or a BITCAST.
+    if (SI->getPointerOperand() != Used ||
+        (!isa<GetElementPtrInst>(Used) && !isa<BitCastInst>(Used)))
+      return false;
+    WorkList.push_back(Inst);
+    return true;
   }
   default:
     return false;
   }
 }
 
 static bool tryPromoteAllocaToVector(AllocaInst *Alloca, AMDGPUAS AS) {
 
if (!AllocaTy ||
     return false;
   }
 
   std::map<GetElementPtrInst*, Value*> GEPVectorIdx;
   std::vector<Value*> WorkList;
   for (User *AllocaUser : Alloca->users()) {
     GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(AllocaUser);
     if (!GEP) {
-      if (!canVectorizeInst(cast<Instruction>(AllocaUser), Alloca))
+      if (!canVectorizeInst(cast<Instruction>(AllocaUser), Alloca, WorkList))
         return false;
-
-      WorkList.push_back(AllocaUser);
       continue;
     }
 
     Value *Index = GEPToVectorIndex(GEP);
 
     // If we can't compute a vector index from this GEP, then we can't
     // promote this alloca to vector.
     if (!Index) {
       DEBUG(dbgs() << "  Cannot compute vector index for GEP " << *GEP << '\n');
       return false;
     }
 
     GEPVectorIdx[GEP] = Index;
     for (User *GEPUser : AllocaUser->users()) {
-      if (!canVectorizeInst(cast<Instruction>(GEPUser), AllocaUser))
+      if (!canVectorizeInst(cast<Instruction>(GEPUser), AllocaUser, WorkList))
         return false;
-
-      WorkList.push_back(GEPUser);
     }
   }
 
-  VectorType *VectorTy = arrayTypeToVecType(AllocaTy);
+  VectorType *VectorT = arrayTypeToVecType(AllocaTy);
 
   DEBUG(dbgs() << "  Converting alloca to vector "
-        << *AllocaTy << " -> " << *VectorTy << '\n');
+        << *AllocaTy << " -> " << *VectorT << '\n');
 
   for (Value *V : WorkList) {
     Instruction *Inst = cast<Instruction>(V);
     IRBuilder<> Builder(Inst);
     switch (Inst->getOpcode()) {
     case Instruction::Load: {
-      Type *VecPtrTy = VectorTy->getPointerTo(AS.PRIVATE_ADDRESS);
       Value *Ptr = cast<LoadInst>(Inst)->getPointerOperand();
+      VectorType *VectorTy = VectorT;
+      if (isa<BitCastInst>(Ptr)) {
+        VectorTy = VectorType::get(Ptr->getType()->getPointerElementType(),
+                                   AllocaTy->getNumElements());
+        Ptr = cast<BitCastInst>(Ptr)->getOperand(0);
+      }
+      Type *VecPtrTy = VectorTy->getPointerTo(AS.PRIVATE_ADDRESS);
       Value *Index = calculateVectorIndex(Ptr, GEPVectorIdx);
 
       Value *BitCast = Builder.CreateBitCast(Alloca, VecPtrTy);
       Value *VecValue = Builder.CreateLoad(BitCast);
       Value *ExtractElement = Builder.CreateExtractElement(VecValue, Index);
       Inst->replaceAllUsesWith(ExtractElement);
       Inst->eraseFromParent();
       break;
     }
     case Instruction::Store: {
-      Type *VecPtrTy = VectorTy->getPointerTo(AS.PRIVATE_ADDRESS);
-
       StoreInst *SI = cast<StoreInst>(Inst);
       Value *Ptr = SI->getPointerOperand();
+      VectorType *VectorTy = VectorT;
+      if (isa<BitCastInst>(Ptr)) {
+        VectorTy = VectorType::get(Ptr->getType()->getPointerElementType(),
+                                   AllocaTy->getNumElements());
+        Ptr = cast<BitCastInst>(Ptr)->getOperand(0);
+      }
+      Type *VecPtrTy = VectorTy->getPointerTo(AS.PRIVATE_ADDRESS);
       Value *Index = calculateVectorIndex(Ptr, GEPVectorIdx);
       Value *BitCast = Builder.CreateBitCast(Alloca, VecPtrTy);
       Value *VecValue = Builder.CreateLoad(BitCast);
       Value *NewVecValue = Builder.CreateInsertElement(VecValue,
                                                        SI->getValueOperand(),
                                                        Index);
       Builder.CreateStore(NewVecValue, BitCast);
       Inst->eraseFromParent();

test/CodeGen/AMDGPU/vector-alloca.ll

entry:
   %tmp1 = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 %w_index
   store i32 1, i32 addrspace(5)* %tmp1
   %tmp2 = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 %r_index
   %tmp3 = load i32, i32 addrspace(5)* %tmp2
   store i32 %tmp3, i32 addrspace(1)* %out
   ret void
 }
 
-; This test should be optimize to:
-; store i32 0, i32 addrspace(1)* %out
-
-; OPT-LABEL: @bitcast_gep(
-; OPT-LABEL: store i32 0, i32 addrspace(1)* %out, align 4
-
-; FUNC-LABEL: {{^}}bitcast_gep:
-; EG: STORE_RAW
-define amdgpu_kernel void @bitcast_gep(i32 addrspace(1)* %out, i32 %w_index, i32 %r_index) {
-entry:
-  %tmp = alloca [4 x i32], addrspace(5)
-  %x = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 0
-  %y = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 1
-  %z = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 2
-  %w = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 3
-  store i32 0, i32 addrspace(5)* %x
-  store i32 0, i32 addrspace(5)* %y
-  store i32 0, i32 addrspace(5)* %z
-  store i32 0, i32 addrspace(5)* %w
-  %tmp1 = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 1
-  %tmp2 = bitcast i32 addrspace(5)* %tmp1 to [4 x i32] addrspace(5)*
-  %tmp3 = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp2, i32 0, i32 0
-  %tmp4 = load i32, i32 addrspace(5)* %tmp3
-  store i32 %tmp4, i32 addrspace(1)* %out
-  ret void
-}
-
 ; OPT-LABEL: @vector_read_bitcast_gep(
 ; OPT: %0 = extractelement <4 x i32> <i32 1065353216, i32 1, i32 2, i32 3>, i32 %index
 ; OPT: store i32 %0, i32 addrspace(1)* %out, align 4
 define amdgpu_kernel void @vector_read_bitcast_gep(i32 addrspace(1)* %out, i32 %index) {
 entry:
   %tmp = alloca [4 x i32], addrspace(5)
   %x = getelementptr inbounds [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 0
   %y = getelementptr inbounds [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 1
entry:
   store i32 1, i32 addrspace(5)* %y
   store i32 2, i32 addrspace(5)* %z
   store i32 3, i32 addrspace(5)* %w
   %tmp1 = getelementptr [4 x i32], [4 x i32] addrspace(5)* %tmp, i32 0, i32 %index
   %tmp2 = load i32, i32 addrspace(5)* %tmp1
   store i32 %tmp2, i32 addrspace(1)* %out
   ret void
 }
+
+
+; OPT-LABEL: @vector_write_read_bitcast_gep(
+; OPT: %0 = insertelement <3 x float> zeroinitializer, float 1.200000e+01, i32 %w_index
+; OPT: %1 = extractelement <3 x float> %0, i32 %r_index
+; OPT: store float %1, float addrspace(1)* %out, align 4
+define amdgpu_kernel void @vector_write_read_bitcast_gep(float addrspace(1)* %out, i32 %w_index, i32 %r_index) {
+entry:
+  %scratch = alloca [3 x i32], addrspace(5)
+  %x = getelementptr [3 x i32], [3 x i32] addrspace(5)* %scratch, i32 0, i32 0
+  %y = getelementptr [3 x i32], [3 x i32] addrspace(5)* %scratch, i32 0, i32 1
+  %z = getelementptr [3 x i32], [3 x i32] addrspace(5)* %scratch, i32 0, i32 2
+  store i32 0, i32 addrspace(5)* %x
+  store i32 0, i32 addrspace(5)* %y
+  store i32 0, i32 addrspace(5)* %z
+
+  %gep_write = getelementptr inbounds [3 x i32], [3 x i32] addrspace(5)* %scratch, i32 0, i32 %w_index
+  %bc_write = bitcast i32 addrspace(5)* %gep_write to float addrspace(5)*
+  store float 12.0, float addrspace(5)* %bc_write, align 4
+
+  %gep_read = getelementptr [3 x i32], [3 x i32] addrspace(5)* %scratch, i32 0, i32 %r_index
+  %bc_read = bitcast i32 addrspace(5)* %gep_read to float addrspace(5)*
+  %result = load float, float addrspace(5)* %bc_read
+  store float %result, float addrspace(1)* %out
+
+  ret void
+}