AMDGPU: Promote array alloca if used by memmove/memcpy

[arsenm]

Shouldn't these have been expanded into loads and stores already?

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

508

Extra ;

[ruiling]

In D140599#4014540, @arsenm wrote:

Shouldn't these have been expanded into loads and stores already?

We only split mem transfer intrinsics using a loop if its length is > 1024.

[arsenm]

In D140599#4014542, @ruiling wrote:

In D140599#4014540, @arsenm wrote:

Shouldn't these have been expanded into loads and stores already?

We only split mem transfer intrinsics using a loop if its length is > 1024.

I think InstCombine and/or SROA do some too. Also that threshold is arbitrary, and we could start expanding the small ones in IR too

[ruiling]

In D140599#4015680, @arsenm wrote:

In D140599#4014542, @ruiling wrote:

In D140599#4014540, @arsenm wrote:

Shouldn't these have been expanded into loads and stores already?

We only split mem transfer intrinsics using a loop if its length is > 1024.

I think InstCombine and/or SROA do some too. Also that threshold is arbitrary, and we could start expanding the small ones in IR too

If I understand correctly, InstCombine focus on simplifying mem transfer intrinsic instead of expanding them and only work in very limited cases.

In fact, the change here shares similar idea with SROA, it will analyze whether the Alloca involved in memmove/memcpy can be promoted to vector/scalar operation. If it is promotable, then we will expand the mem transfer intrinsic as part of the promotion process. I would say the important thing here is to get optimal code generation for small alloca, expanding smaller mem transfer is not quite important. Does this sound reasonable to you? Or could you share how would you like we expand small mem transfer intrinsic?

In the long term, I would hope we can enhance the SROA to cover the optimization we have done in tryPromoteAllocaToVector(). But that need much more time/effort.

[ruiling]

ping

[ruiling]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

508

I guess you might misread here? I think we always put a semicolon at the end of lambda expression.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

502

[=]?

593–595

probably should dyn_cast to MemIntrinsicInst here to defend agains these becoming out of sync

611

No reason to care about type pointers anymore

llvm/test/CodeGen/AMDGPU/promote-alloca-array-aggregate.ll

219

Needs tests for the different address space cases you mentioned in the todo

219

also memcpy_inline

[ruiling]

rebase and ping.

[ruiling]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

611

Our graphics compiler has not fully switched to opaque pointers yet. @foad am I right?

[ruiling]

address review comments.

[ruiling]

Thanks for the careful review @arsenm, I have fixed most of them. Please help take a second look. Thanks!

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

502

We are capturing GEPVectorIdx and Alloca, so I think capture by reference is preferred over by copy here. Please correct me if I am wrong.

593–595

Agree, I have changed to cast to catch the possible broken case.

llvm/test/CodeGen/AMDGPU/promote-alloca-array-aggregate.ll

219

Good idea, I have added all of them.

[foad]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

544

This can be cast instead of dyn_cast.

611

Yes.

[ruiling]

use cast instead of dyn_cast.

[ruiling]

Use alignment from alloca.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

593–595

This still has the switch over specific IDs instead of the matching cast

611

Well, that needs to be fixed soon since code is going to start getting ripped out very shortly

llvm/test/CodeGen/AMDGPU/promote-alloca-array-aggregate.ll

307

Another edge case to maybe test is the identity copy of the full alloca to itself

[foad]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

611

We have top engineers working on it! https://github.com/GPUOpen-Drivers/llpc/pull/2210

[ruiling]

address review comments.

[ruiling]

llvm/test/CodeGen/AMDGPU/promote-alloca-array-aggregate.ll

307

I have added a case for it, but would not bother to optimize for it as such code pattern are trivially optimized away early.

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

static Value *GEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
   uint64_t Rem;
   APInt::udivrem(ConstOffset, VecElemSize, Quot, Rem);
   if (Rem != 0)
     return nullptr;
 
   return ConstantInt::get(GEP->getContext(), Quot);
 }
 
+struct MemTransferInfo {
+  ConstantInt *SrcIndex = nullptr;
+  ConstantInt *DestIndex = nullptr;
+};
+
 static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
                                      unsigned MaxVGPRs) {
 
   if (DisablePromoteAllocaToVector) {
     LLVM_DEBUG(dbgs() << "  Promotion alloca to vector is disabled\n");
     return false;
   }
 
static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
   if (!VectorTy || VectorTy->getNumElements() > 16 ||
       VectorTy->getNumElements() < 2) {
     LLVM_DEBUG(dbgs() << "  Cannot convert type to vector\n");
     return false;
   }
 
   std::map<GetElementPtrInst*, Value*> GEPVectorIdx;
   SmallVector<Instruction *> WorkList;
+  SmallVector<Instruction *> DeferredInsts;
   SmallVector<Use *, 8> Uses;
+  DenseMap<MemTransferInst *, MemTransferInfo> TransferInfo;
+
   for (Use &U : Alloca->uses())
     Uses.push_back(&U);
 
   Type *VecEltTy = VectorTy->getElementType();
+  unsigned ElementSize = DL.getTypeSizeInBits(VecEltTy) / 8;
   while (!Uses.empty()) {
     Use *U = Uses.pop_back_val();
     Instruction *Inst = dyn_cast<Instruction>(U->getUser());
 
     if (Value *Ptr = getLoadStorePointerOperand(Inst)) {
       // This is a store of the pointer, not to the pointer.
       if (isa<StoreInst>(Inst) &&
           U->getOperandNo() != StoreInst::getPointerOperandIndex())
if (auto *GEP = dyn_cast<GetElementPtrInst>(Inst)) {
       }
 
       GEPVectorIdx[GEP] = Index;
       for (Use &U : Inst->uses())
         Uses.push_back(&U);
       continue;
     }
 
+    if (MemTransferInst *TransferInst = dyn_cast<MemTransferInst>(Inst)) {
+      if (TransferInst->isVolatile())
+        return false;
+
+      ConstantInt *Len = dyn_cast<ConstantInt>(TransferInst->getLength());
+      if (!Len || !!(Len->getZExtValue() % ElementSize))
+        return false;
+
+      if (!TransferInfo.count(TransferInst)) {
+        DeferredInsts.push_back(Inst);
+        WorkList.push_back(Inst);
+        TransferInfo[TransferInst] = MemTransferInfo();
+      }
+
+      auto getPointerIndexOfAlloca = [&](Value *Ptr) -> ConstantInt * {
+        GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Ptr);
+        if (Ptr != Alloca && !GEPVectorIdx.count(GEP))
+          return nullptr;
+
+        return dyn_cast<ConstantInt>(calculateVectorIndex(Ptr, GEPVectorIdx));
+      };
+
+      unsigned OpNum = U->getOperandNo();
+      MemTransferInfo *TI = &TransferInfo[TransferInst];
+      if (OpNum == 0) {
+        Value *Dest = TransferInst->getDest();
+        ConstantInt *Index = getPointerIndexOfAlloca(Dest);
+        if (!Index)
+          return false;
+        TI->DestIndex = Index;
+      } else {
+        assert(OpNum == 1);
+        Value *Src = TransferInst->getSource();
+        ConstantInt *Index = getPointerIndexOfAlloca(Src);
+        if (!Index)
+          return false;
+        TI->SrcIndex = Index;
+      }
+      continue;
+    }
+
     // Ignore assume-like intrinsics and comparisons used in assumes.
     if (isAssumeLikeIntrinsic(Inst))
       continue;
 
     if (isa<ICmpInst>(Inst) && all_of(Inst->users(), [](User *U) {
           return isAssumeLikeIntrinsic(cast<Instruction>(U));
         }))
       continue;
 
     // Unknown user.
     return false;
   }
 
+  while (!DeferredInsts.empty()) {
+    Instruction *Inst = DeferredInsts.pop_back_val();
+    if (MemTransferInst *TransferInst = dyn_cast<MemTransferInst>(Inst)) {
+      // TODO: Support the case if the pointers are from different alloca or
+      // from different address spaces.
+      MemTransferInfo &Info = TransferInfo[TransferInst];
+      if (!Info.SrcIndex || !Info.DestIndex)
+        return false;
+    }
+  }
+
   LLVM_DEBUG(dbgs() << "  Converting alloca to vector " << *AllocaTy << " -> "
                     << *VectorTy << '\n');
 
   for (Instruction *Inst : WorkList) {
     IRBuilder<> Builder(Inst);
     switch (Inst->getOpcode()) {
     case Instruction::Load: {
       Value *Ptr = cast<LoadInst>(Inst)->getPointerOperand();
case Instruction::Store: {
       Value *Elt = SI->getValueOperand();
       if (Elt->getType() != VecEltTy)
         Elt = Builder.CreateBitOrPointerCast(Elt, VecEltTy);
       Value *NewVecValue = Builder.CreateInsertElement(VecValue, Elt, Index);
       Builder.CreateStore(NewVecValue, BitCast);
       Inst->eraseFromParent();
       break;
     }
+    case Instruction::Call: {
+      if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(Inst)) {
+        switch (II->getIntrinsicID()) {
+          default:
+            llvm_unreachable(
+                "Unsupported intrinsic when promoting alloca to vector");
+          case Intrinsic::memcpy:
+          case Intrinsic::memcpy_inline:
+          case Intrinsic::memmove: {
+            ConstantInt *Length = cast<ConstantInt>(II->getArgOperand(2));
+            unsigned NumCopied = Length->getZExtValue() / ElementSize;
+            MemTransferInfo *TI = &TransferInfo[cast<MemTransferInst>(Inst)];
+            unsigned SrcBegin = TI->SrcIndex->getZExtValue();
+            unsigned DestBegin = TI->DestIndex->getZExtValue();
+
+            SmallVector<int> Mask;
+            for (unsigned Idx = 0; Idx < VectorTy->getNumElements(); ++Idx) {
+              if (Idx >= DestBegin && Idx < DestBegin + NumCopied) {
+                Mask.push_back(SrcBegin++);
+              } else {
+                Mask.push_back(Idx);
+              }
+            }
+            Type *VecPtrTy = VectorTy->getPointerTo(Alloca->getAddressSpace());
+            Value *BitCast = Builder.CreateBitCast(Alloca, VecPtrTy);
+            Value *VecValue = Builder.CreateLoad(VectorTy, BitCast);
+            Value *NewVecValue = Builder.CreateShuffleVector(VecValue, Mask);
+            Builder.CreateStore(NewVecValue, BitCast);
+
+            Inst->eraseFromParent();
+            break;
+          }
+        }
+      } else {
+        llvm_unreachable("Unsupported call when promoting alloca to vector");
+      }
+      break;
+    }
 
     default:
       llvm_unreachable("Inconsistency in instructions promotable to vector");
     }
   }
   return true;
 }
 

llvm/test/CodeGen/AMDGPU/promote-alloca-array-aggregate.ll

;
   %foo9 = insertelement <4 x float> %foo8, float %foo6, i32 0
   %foo10 = insertelement <4 x float> %foo9, float %foo6, i32 1
   %foo11 = insertelement <4 x float> %foo10, float %foo6, i32 2
   %foo12 = insertelement <4 x float> %foo11, float %foo6, i32 3
   store <4 x float> %foo12, ptr addrspace(1) @pv
   ret void
 }
 
+define amdgpu_vs void @promote_memmove_aggr() #0 {
+; CHECK-LABEL: @promote_memmove_aggr(
+; CHECK-NEXT:    [[F1:%.*]] = alloca [5 x float], align 4, addrspace(5)
+; CHECK-NEXT:    store [5 x float] zeroinitializer, ptr addrspace(5) [[F1]], align 4
+; CHECK-NEXT:    [[FOO1:%.*]] = getelementptr [5 x float], ptr addrspace(5) [[F1]], i32 0, i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <5 x float> [[TMP1]], float 1.000000e+00, i64 1
+; CHECK-NEXT:    store <5 x float> [[TMP2]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[FOO2:%.*]] = getelementptr [5 x float], ptr addrspace(5) [[F1]], i32 0, i32 3
+; CHECK-NEXT:    [[TMP3:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <5 x float> [[TMP3]], float 2.000000e+00, i64 3
+; CHECK-NEXT:    store <5 x float> [[TMP4]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP5:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <5 x float> [[TMP5]], <5 x float> poison, <5 x i32> <i32 1, i32 2, i32 3, i32 4, i32 4>
+; CHECK-NEXT:    store <5 x float> [[TMP6]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP7:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <5 x float> [[TMP7]], i32 0
+; CHECK-NEXT:    store float [[TMP8]], ptr addrspace(1) @pv, align 4
+; CHECK-NEXT:    ret void
+;
+  %f1 = alloca [5 x float], addrspace(5)
+  store [5 x float] zeroinitializer, ptr addrspace(5) %f1
+  %foo1 = getelementptr [5 x float], ptr addrspace(5) %f1, i32 0, i32 1
+  store float 1.0, ptr addrspace(5) %foo1
+  %foo2 = getelementptr [5 x float], ptr addrspace(5) %f1, i32 0, i32 3
+  store float 2.0, ptr addrspace(5) %foo2
+  call void @llvm.memmove.p5i8.p5i8.i32(ptr addrspace(5) align 4 %f1, ptr addrspace(5) align 4 %foo1, i32 16, i1 false)
+  %foo3 = load float, ptr addrspace(5) %f1
+  store float %foo3, ptr addrspace(1) @pv
+  ret void
+}
+
+define amdgpu_vs void @promote_memcpy_aggr() #0 {
+; CHECK-LABEL: @promote_memcpy_aggr(
+; CHECK-NEXT:    [[F1:%.*]] = alloca [5 x float], align 4, addrspace(5)
+; CHECK-NEXT:    store [5 x float] zeroinitializer, ptr addrspace(5) [[F1]], align 4
+; CHECK-NEXT:    [[FOO1:%.*]] = getelementptr [5 x float], ptr addrspace(5) [[F1]], i32 0, i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <5 x float> [[TMP1]], float 1.000000e+00, i64 1
+; CHECK-NEXT:    store <5 x float> [[TMP2]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[FOO2:%.*]] = getelementptr [5 x float], ptr addrspace(5) [[F1]], i32 0, i32 3
+; CHECK-NEXT:    [[TMP3:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <5 x float> [[TMP3]], float 2.000000e+00, i64 3
+; CHECK-NEXT:    store <5 x float> [[TMP4]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[FOO3:%.*]] = getelementptr [[BLOCK3:%.*]], ptr addrspace(1) @block3, i32 0, i32 0
+; CHECK-NEXT:    [[FOO4:%.*]] = load i32, ptr addrspace(1) [[FOO3]], align 4
+; CHECK-NEXT:    [[FOO5:%.*]] = getelementptr [5 x float], ptr addrspace(5) [[F1]], i32 0, i32 [[FOO4]]
+; CHECK-NEXT:    [[TMP5:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <5 x float> [[TMP5]], float 3.000000e+00, i32 [[FOO4]]
+; CHECK-NEXT:    store <5 x float> [[TMP6]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP7:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP8:%.*]] = shufflevector <5 x float> [[TMP7]], <5 x float> poison, <5 x i32> <i32 3, i32 4, i32 2, i32 3, i32 4>
+; CHECK-NEXT:    store <5 x float> [[TMP8]], ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP9:%.*]] = load <5 x float>, ptr addrspace(5) [[F1]], align 32
+; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <5 x float> [[TMP9]], i32 0
+; CHECK-NEXT:    store float [[TMP10]], ptr addrspace(1) @pv, align 4
+; CHECK-NEXT:    ret void
+;
+  %f1 = alloca [5 x float], addrspace(5)
+  store [5 x float] zeroinitializer, ptr addrspace(5) %f1
+  %foo1 = getelementptr [5 x float], ptr addrspace(5) %f1, i32 0, i32 1
+  store float 1.0, ptr addrspace(5) %foo1
+  %foo2 = getelementptr [5 x float], ptr addrspace(5) %f1, i32 0, i32 3
+  store float 2.0, ptr addrspace(5) %foo2
+
+  %foo3 = getelementptr %Block3, ptr addrspace(1) @block3, i32 0, i32 0
+  %foo4 = load i32, ptr addrspace(1) %foo3
+  %foo5 = getelementptr [5 x float], ptr addrspace(5) %f1, i32 0, i32 %foo4
+  store float 3.0, ptr addrspace(5) %foo5
+
+  call void @llvm.memcpy.p5i8.p5i8.i32(ptr addrspace(5) align 4 %f1, ptr addrspace(5) align 4 %foo2, i32 8, i1 false)
+  %foo6 = load float, ptr addrspace(5) %f1
+  store float %foo6, ptr addrspace(1) @pv
+  ret void
+}
+
+declare void @llvm.memcpy.p5i8.p5i8.i32(ptr addrspace(5) nocapture writeonly, ptr addrspace(5) nocapture readonly, i32, i1 immarg)
+declare void @llvm.memmove.p5i8.p5i8.i32(ptr addrspace(5) nocapture writeonly, ptr addrspace(5) nocapture readonly, i32, i1 immarg)
+
 @tmp_g = external addrspace(1) global { [4 x double], <2 x double>, <3 x double>, <4 x double> }
 @frag_color = external addrspace(1) global <4 x float>
 
 define amdgpu_ps void @promote_double_aggr() #0 {
 ; CHECK-LABEL: @promote_double_aggr(
 ; CHECK-NEXT:    [[S:%.*]] = alloca [2 x double], align 8, addrspace(5)
 ; CHECK-NEXT:    [[FOO:%.*]] = getelementptr { [4 x double], <2 x double>, <3 x double>, <4 x double> }, ptr addrspace(1) @tmp_g, i32 0, i32 0, i32 0
 ; CHECK-NEXT:    [[FOO1:%.*]] = load double, ptr addrspace(1) [[FOO]], align 8

 ; CHECK-NEXT:    [[TMP9:%.*]] = mul nuw nsw i32 [[TMP5]], [[TMP4]]
 ; CHECK-NEXT:    [[TMP10:%.*]] = mul i32 [[TMP9]], [[TMP6]]
 ; CHECK-NEXT:    [[TMP11:%.*]] = mul nuw nsw i32 [[TMP7]], [[TMP4]]
 ; CHECK-NEXT:    [[TMP12:%.*]] = add i32 [[TMP10]], [[TMP11]]
 ; CHECK-NEXT:    [[TMP13:%.*]] = add i32 [[TMP12]], [[TMP8]]
 ; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds [1024 x [2 x %struct]], ptr addrspace(3) @alloca_struct.alloca, i32 0, i32 [[TMP13]]
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %alloca = alloca [2 x %struct], align 4, addrspace(5)
   ret void
 }