AMDGPU: Promote array alloca if used by memmove/memcpy

Reviewed by: arsenm

Differential Revision: https://reviews.llvm.org/D140599

Change-Id: If4627299a1d12195cf1ccd6f39b77b21b571f54d

Author: ruiling

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

@@ -379,6 +379,11 @@ static Value *GEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
   return ConstantInt::get(GEP->getContext(), Quot);
 }
 
+struct MemTransferInfo {
+  ConstantInt *SrcIndex = nullptr;
+  ConstantInt *DestIndex = nullptr;
+};
+
 static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
                                      unsigned MaxVGPRs) {
 
@@ -419,11 +424,15 @@ static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
 
   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());
@@ -476,6 +485,47 @@ static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
       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;
@@ -489,6 +539,16 @@ static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
     return false;
   }
 
+  while (!DeferredInsts.empty()) {
+    Instruction *Inst = DeferredInsts.pop_back_val();
+    MemTransferInst *TransferInst = 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');
 
@@ -523,6 +583,35 @@ static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
       Inst->eraseFromParent();
       break;
     }
+    case Instruction::Call: {
+      if (const MemTransferInst *MTI = dyn_cast<MemTransferInst>(Inst)) {
+        ConstantInt *Length = cast<ConstantInt>(MTI->getLength());
+        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.CreateAlignedLoad(VectorTy, BitCast, Alloca->getAlign());
+        Value *NewVecValue = Builder.CreateShuffleVector(VecValue, Mask);
+        Builder.CreateAlignedStore(NewVecValue, BitCast, Alloca->getAlign());
+
+        Inst->eraseFromParent();
+      } else {
+        llvm_unreachable("Unsupported call when promoting alloca to vector");
+      }
+      break;
+    }
 
     default:
       llvm_unreachable("Inconsistency in instructions promotable to vector");