[AMDGPU] Limit promote alloca to vector with VGPR budget

[rampitec]

Summary

Allow only up to 1/4 of available VGPRs for the vectorization
of any given alloca.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

439

This seems like a huge default. I thought we previously limited this to 16 VGPRs. There should also probably be a cl::opt for this too

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

439

We did not limit it to 16 VGPRs but to 16 elements. This limit is still here. The problem this limit solves is different, it is when you only have limited number of registers not to run out of them. Like wg size 1024 leaves us with 64 VGPRs and then the limit would be 16.

[rampitec]

Added cl::opt.

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

 private:
   const TargetMachine *TM;
   Module *Mod = nullptr;
   const DataLayout *DL = nullptr;
 
   // FIXME: This should be per-kernel.
   uint32_t LocalMemLimit = 0;
   uint32_t CurrentLocalMemUsage = 0;
+  unsigned MaxVGPRs;
 
   bool IsAMDGCN = false;
   bool IsAMDHSA = false;
 
   std::pair<Value *, Value *> getLocalSizeYZ(IRBuilder<> &Builder);
   Value *getWorkitemID(IRBuilder<> &Builder, unsigned N);
 
   /// BaseAlloca is the alloca root the search started from.
public:
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     FunctionPass::getAnalysisUsage(AU);
   }
 };
 
 class AMDGPUPromoteAllocaToVector : public FunctionPass {
+private:
+  unsigned MaxVGPRs;
+
 public:
   static char ID;
 
   AMDGPUPromoteAllocaToVector() : FunctionPass(ID) {}
 
   bool runOnFunction(Function &F) override;
 
   StringRef getPassName() const override {
bool AMDGPUPromoteAlloca::runOnFunction(Function &F) {
   const Triple &TT = TM->getTargetTriple();
   IsAMDGCN = TT.getArch() == Triple::amdgcn;
   IsAMDHSA = TT.getOS() == Triple::AMDHSA;
 
   const AMDGPUSubtarget &ST = AMDGPUSubtarget::get(*TM, F);
   if (!ST.isPromoteAllocaEnabled())
     return false;
 
+  if (IsAMDGCN) {
+    const GCNSubtarget &ST = TM->getSubtarget<GCNSubtarget>(F);
+    MaxVGPRs = ST.getMaxNumVGPRs(ST.getWavesPerEU(F).first);
+  } else {
+    MaxVGPRs = 128;
+  }
+
   bool SufficientLDS = hasSufficientLocalMem(F);
   bool Changed = false;
   BasicBlock &EntryBB = *F.begin();
 
   SmallVector<AllocaInst *, 16> Allocas;
   for (Instruction &I : EntryBB) {
     if (AllocaInst *AI = dyn_cast<AllocaInst>(&I))
       Allocas.push_back(AI);
return (SI->getPointerOperand() == User) &&
             UserInst->getOpcode() == Instruction::BitCast) &&
            SI->isSimple();
   }
   default:
     return false;
   }
 }
 
-static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL) {
+static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
+                                     unsigned MaxVGPRs) {
 
   if (DisablePromoteAllocaToVector) {
     LLVM_DEBUG(dbgs() << "  Promotion alloca to vector is disabled\n");
     return false;
   }
 
   Type *AllocaTy = Alloca->getAllocatedType();
   auto *VectorTy = dyn_cast<FixedVectorType>(AllocaTy);
   if (auto *ArrayTy = dyn_cast<ArrayType>(AllocaTy)) {
     if (VectorType::isValidElementType(ArrayTy->getElementType()) &&
         ArrayTy->getNumElements() > 0)
       VectorTy = arrayTypeToVecType(ArrayTy);
   }
 
+  // Use up to 1/4 of available register budget for vectorization.
+  if (DL.getTypeSizeInBits(AllocaTy) * 4 > MaxVGPRs * 32) {
+    LLVM_DEBUG(dbgs() << "  Alloca too big for vectorization with "
+                      << MaxVGPRs << " registers available\n");
+    return false;
+  }
+
   LLVM_DEBUG(dbgs() << "Alloca candidate for vectorization\n");
 
   // FIXME: There is no reason why we can't support larger arrays, we
   // are just being conservative for now.
   // FIXME: We also reject alloca's of the form [ 2 x [ 2 x i32 ]] or equivalent. Potentially these
   // could also be promoted but we don't currently handle this case
   if (!VectorTy || VectorTy->getNumElements() > 16 ||
       VectorTy->getNumElements() < 2) {
bool AMDGPUPromoteAlloca::handleAlloca(AllocaInst &I, bool SufficientLDS) {
   const DataLayout &DL = Mod->getDataLayout();
   IRBuilder<> Builder(&I);
 
   // First try to replace the alloca with a vector
   Type *AllocaTy = I.getAllocatedType();
 
   LLVM_DEBUG(dbgs() << "Trying to promote " << I << '\n');
 
-  if (tryPromoteAllocaToVector(&I, DL))
+  if (tryPromoteAllocaToVector(&I, DL, MaxVGPRs))
     return true; // Promoted to vector.
 
   if (DisablePromoteAllocaToLDS)
     return false;
 
   const Function &ContainingFunction = *I.getParent()->getParent();
   CallingConv::ID CC = ContainingFunction.getCallingConv();
 
bool AMDGPUPromoteAlloca::handleAlloca(AllocaInst &I, bool SufficientLDS) {
   }
   return true;
 }
 
 bool AMDGPUPromoteAllocaToVector::runOnFunction(Function &F) {
   if (skipFunction(F) || DisablePromoteAllocaToVector)
     return false;
 
+  const TargetMachine *TM;
+  if (auto *TPC = getAnalysisIfAvailable<TargetPassConfig>())
+    TM = &TPC->getTM<TargetMachine>();
+  else
+    return false;
+
+  const AMDGPUSubtarget &ST = AMDGPUSubtarget::get(*TM, F);
+  if (!ST.isPromoteAllocaEnabled())
+    return false;
+
+  if (TM->getTargetTriple().getArch() == Triple::amdgcn) {
+    const GCNSubtarget &ST = TM->getSubtarget<GCNSubtarget>(F);
+    MaxVGPRs = ST.getMaxNumVGPRs(ST.getWavesPerEU(F).first);
+  } else {
+    MaxVGPRs = 128;
+  }
+
   bool Changed = false;
   BasicBlock &EntryBB = *F.begin();
 
   SmallVector<AllocaInst *, 16> Allocas;
   for (Instruction &I : EntryBB) {
     if (AllocaInst *AI = dyn_cast<AllocaInst>(&I))
       Allocas.push_back(AI);
   }
bool AMDGPUPromoteAllocaToVector::handleAlloca(AllocaInst &I) {
   // Array allocations are probably not worth handling, since an allocation of
   // the array type is the canonical form.
   if (!I.isStaticAlloca() || I.isArrayAllocation())
     return false;
 
   LLVM_DEBUG(dbgs() << "Trying to promote " << I << '\n');
 
   Module *Mod = I.getParent()->getParent()->getParent();
-  return tryPromoteAllocaToVector(&I, Mod->getDataLayout());
+  return tryPromoteAllocaToVector(&I, Mod->getDataLayout(), MaxVGPRs);
 }
 
 FunctionPass *llvm::createAMDGPUPromoteAlloca() {
   return new AMDGPUPromoteAlloca();
 }
 
 FunctionPass *llvm::createAMDGPUPromoteAllocaToVector() {
   return new AMDGPUPromoteAllocaToVector();
 }

llvm/test/CodeGen/AMDGPU/vector-alloca-limits.ll

+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-promote-alloca -sroa -instcombine < %s | FileCheck -check-prefix=OPT %s
+target datalayout = "A5"
+
+; OPT-LABEL: @alloca_8xi64_max1024(
+; OPT-NOT: alloca
+; OPT: <8 x i64>
+define amdgpu_kernel void @alloca_8xi64_max1024(i64 addrspace(1)* %out, i32 %index) #0 {
+entry:
+  %tmp = alloca [8 x i64], addrspace(5)
+  %x = getelementptr [8 x i64], [8 x i64] addrspace(5)* %tmp, i32 0, i32 0
+  store i64 0, i64 addrspace(5)* %x
+  %tmp1 = getelementptr [8 x i64], [8 x i64] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i64, i64 addrspace(5)* %tmp1
+  store i64 %tmp2, i64 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_9xi64_max1024(
+; OPT: alloca [9 x i64]
+; OPT-NOT: <9 x i64>
+define amdgpu_kernel void @alloca_9xi64_max1024(i64 addrspace(1)* %out, i32 %index) #0 {
+entry:
+  %tmp = alloca [9 x i64], addrspace(5)
+  %x = getelementptr [9 x i64], [9 x i64] addrspace(5)* %tmp, i32 0, i32 0
+  store i64 0, i64 addrspace(5)* %x
+  %tmp1 = getelementptr [9 x i64], [9 x i64] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i64, i64 addrspace(5)* %tmp1
+  store i64 %tmp2, i64 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_16xi64_max512(
+; OPT-NOT: alloca
+; OPT: <16 x i64>
+define amdgpu_kernel void @alloca_16xi64_max512(i64 addrspace(1)* %out, i32 %index) #1 {
+entry:
+  %tmp = alloca [16 x i64], addrspace(5)
+  %x = getelementptr [16 x i64], [16 x i64] addrspace(5)* %tmp, i32 0, i32 0
+  store i64 0, i64 addrspace(5)* %x
+  %tmp1 = getelementptr [16 x i64], [16 x i64] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i64, i64 addrspace(5)* %tmp1
+  store i64 %tmp2, i64 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_17xi64_max512(
+; OPT: alloca [17 x i64]
+; OPT-NOT: <17 x i64>
+define amdgpu_kernel void @alloca_17xi64_max512(i64 addrspace(1)* %out, i32 %index) #1 {
+entry:
+  %tmp = alloca [17 x i64], addrspace(5)
+  %x = getelementptr [17 x i64], [17 x i64] addrspace(5)* %tmp, i32 0, i32 0
+  store i64 0, i64 addrspace(5)* %x
+  %tmp1 = getelementptr [17 x i64], [17 x i64] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i64, i64 addrspace(5)* %tmp1
+  store i64 %tmp2, i64 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_9xi128_max512(
+; OPT: alloca [9 x i128]
+; OPT-NOT: <9 x i128>
+define amdgpu_kernel void @alloca_9xi128_max512(i128 addrspace(1)* %out, i32 %index) #1 {
+entry:
+  %tmp = alloca [9 x i128], addrspace(5)
+  %x = getelementptr [9 x i128], [9 x i128] addrspace(5)* %tmp, i32 0, i32 0
+  store i128 0, i128 addrspace(5)* %x
+  %tmp1 = getelementptr [9 x i128], [9 x i128] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i128, i128 addrspace(5)* %tmp1
+  store i128 %tmp2, i128 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_9xi128_max256(
+; OPT-NOT: alloca
+; OPT: <9 x i128>
+define amdgpu_kernel void @alloca_9xi128_max256(i128 addrspace(1)* %out, i32 %index) #2 {
+entry:
+  %tmp = alloca [9 x i128], addrspace(5)
+  %x = getelementptr [9 x i128], [9 x i128] addrspace(5)* %tmp, i32 0, i32 0
+  store i128 0, i128 addrspace(5)* %x
+  %tmp1 = getelementptr [9 x i128], [9 x i128] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i128, i128 addrspace(5)* %tmp1
+  store i128 %tmp2, i128 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_16xi128_max256(
+; OPT-NOT: alloca
+; OPT: <16 x i128>
+define amdgpu_kernel void @alloca_16xi128_max256(i128 addrspace(1)* %out, i32 %index) #2 {
+entry:
+  %tmp = alloca [16 x i128], addrspace(5)
+  %x = getelementptr [16 x i128], [16 x i128] addrspace(5)* %tmp, i32 0, i32 0
+  store i128 0, i128 addrspace(5)* %x
+  %tmp1 = getelementptr [16 x i128], [16 x i128] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i128, i128 addrspace(5)* %tmp1
+  store i128 %tmp2, i128 addrspace(1)* %out
+  ret void
+}
+
+; OPT-LABEL: @alloca_9xi256_max256(
+; OPT: alloca [9 x i256]
+; OPT-NOT: <9 x i256>
+define amdgpu_kernel void @alloca_9xi256_max256(i256 addrspace(1)* %out, i32 %index) #2 {
+entry:
+  %tmp = alloca [9 x i256], addrspace(5)
+  %x = getelementptr [9 x i256], [9 x i256] addrspace(5)* %tmp, i32 0, i32 0
+  store i256 0, i256 addrspace(5)* %x
+  %tmp1 = getelementptr [9 x i256], [9 x i256] addrspace(5)* %tmp, i32 0, i32 %index
+  %tmp2 = load i256, i256 addrspace(5)* %tmp1
+  store i256 %tmp2, i256 addrspace(1)* %out
+  ret void
+}
+
+attributes #0 = { "amdgpu-flat-work-group-size"="1,1024" }
+attributes #1 = { "amdgpu-flat-work-group-size"="1,512" }
+attributes #2 = { "amdgpu-flat-work-group-size"="1,256" }