[AMDGPU] Unroll preferences improvements

[rampitec]

Summary

Exit loop analysis early if suitable private access found.
Do not account for GEPs which are invariant to loop induction variable.
Do not account for Allocas which are too big to fit into register file anyway.
Add option for tuning: -amdgpu-unroll-threshold-private.

[arsenm]

Needs tests in test/Transforms/LoopUnroll/AMDGPU

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

32–35

I don't think we need this. The regular option I think overrides the target preference if set

58–59

Won't LICM move this out of the loop before unrolling?

100–102

I don't think we should change this away from the hardware sizes. I think this hook is only used by the vectorizers we don't use. We should define a different constant for use for the alloca heuristic

[rampitec]

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

32–35

Right, it does.

58–59

In most cases it will, but consider this:

for (int i=0; i < 16; i++)
  for (int j=0; j < 100; j++)
    arr[(j + gid) % 64] = a[i];

Here arr[j] depends on the induction variable of inner loop, so it is not hoisted. We do not want to unroll outer loop because of that though.
Actually this check is not sufficient, so I have rewritten this place. I will update review shortly.

100–102

The numbers here were just incorrect. SI to CI have 128 registers. Then it makes sense to take into consideration real register file size, which is target dependent. As a todo we need to limit it further if we have occupancy attributes.

[arsenm]

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

100–102

No, there have always been 256 VGPRs.

[rampitec]

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

100–102

Ouch. My memory is wrong. Will fix.

[rampitec]

Added test.
Reverted portion about number of registers.
Removed general unroll threshold option.
Added logic to go deep into loop nest to see if that is actually contained loop needs to be unrolled, but not outer.

[arsenm]

LGTM except for the isSized question

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

59

Can you add a test where isSized is necessary? I thought it was illegal to alloca such a type

67

!inst

[rampitec]

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

59

It happens when you have opaque type, like image. I doubt I can easily write a test like this, but when it happens getTypeAllocSize() asserts, so it is better to keep it.

[vpykhtin]

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

58–59

Please add this motivating example as a comment

[vpykhtin]

llvm/trunk/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

69 ↗

(On Diff #86931)

Why is this check nessesary? Is this an early exit when GEP is dependent on more than 1 induction variable?

[• tstellarAMD]

llvm/trunk/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

88 ↗

(On Diff #86931)

Do you also want to set PartialThreshold here?

[vpykhtin]

llvm/trunk/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

88 ↗

(On Diff #86931)

I thought partialy unrolled loops won't make it possible to SROA private arrays. What are the benefits of partial unrolling on AMDGPU btw? What comes in mind: mem ops clustering/widening, less branches? What else?

[• tstellarAMD]

llvm/trunk/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

88 ↗

(On Diff #86931)

I had a test case where bumping the PartialThreshold helped more non-partial loops be unrolled, but I looked at the case again and increasing the normal Threshold has the same affect, so I don't think this is needed.

[rampitec]

llvm/trunk/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

69 ↗

(On Diff #86931)

This is just a check that real dependency is in the inner loop, which really needs to be unrolled. When you iterate through blocks a a loop you will get those belonging to inner loops as well. I'm just checking we are about to unroll a right one.

88 ↗

(On Diff #86931)

Actually I agree, partial unroll does not help to SROA an array. There can be other motivation, but not this.

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

 #include "llvm/IR/Intrinsics.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Target/CostTable.h"
 #include "llvm/Target/TargetLowering.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "AMDGPUtti"
 
+static cl::opt<unsigned> UnrollThresholdPrivate(
+  "amdgpu-unroll-threshold-private",
+  cl::desc("Unroll threshold for AMDGPU if private memory used in a loop"),
+  cl::init(800), cl::Hidden);
 
 void AMDGPUTTIImpl::getUnrollingPreferences(Loop *L,
                                             TTI::UnrollingPreferences &UP) {
   UP.Threshold = 300; // Twice the default.
   UP.MaxCount = UINT_MAX;
   UP.Partial = true;
 
   // TODO: Do we want runtime unrolling?
 
+  // Maximum alloca size than can fit registers. Reserve 16 registers.
+  const unsigned MaxAlloca = (256 - 16) * 4;
   for (const BasicBlock *BB : L->getBlocks()) {
     const DataLayout &DL = BB->getModule()->getDataLayout();
     for (const Instruction &I : *BB) {
       const GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(&I);
       if (!GEP || GEP->getAddressSpace() != AMDGPUAS::PRIVATE_ADDRESS)
         continue;
 
       const Value *Ptr = GEP->getPointerOperand();
       const AllocaInst *Alloca =
           dyn_cast<AllocaInst>(GetUnderlyingObject(Ptr, DL));
       if (Alloca) {
+        Type *Ty = Alloca->getAllocatedType();
+        unsigned AllocaSize = Ty->isSized() ? DL.getTypeAllocSize(Ty) : 0;
+        if (AllocaSize > MaxAlloca)
+          continue;
+
+        // Check if GEP depends on a value defined by this loop itself.
+        bool HasLoopDef = false;
+        for (const Value *Op : GEP->operands()) {
+          const Instruction *Inst = dyn_cast<Instruction>(Op);
+          if (Inst ==nullptr || L->isLoopInvariant(Op))
+            continue;
+          if (any_of(L->getSubLoops(), [Inst](const Loop* SubLoop) {
+               return SubLoop->contains(Inst); }))
+            continue;
+          HasLoopDef = true;
+          break;
+        }
+        if (!HasLoopDef)
+          continue;
+
         // We want to do whatever we can to limit the number of alloca
         // instructions that make it through to the code generator.  allocas
         // require us to use indirect addressing, which is slow and prone to
         // compiler bugs.  If this loop does an address calculation on an
         // alloca ptr, then we want to use a higher than normal loop unroll
         // threshold. This will give SROA a better chance to eliminate these
         // allocas.
         //
         // Don't use the maximum allowed value here as it will make some
         // programs way too big.
-        UP.Threshold = 800;
+        UP.Threshold = UnrollThresholdPrivate;
+        return;
       }
     }
   }
 }
 
 unsigned AMDGPUTTIImpl::getNumberOfRegisters(bool Vec) {
   if (Vec)
     return 0;
 
   // Number of VGPRs on SI.
   if (ST->getGeneration() >= AMDGPUSubtarget::SOUTHERN_ISLANDS)
     return 256;
 
   return 4 * 128; // XXX - 4 channels. Should these count as vector instead?
 }
 
 unsigned AMDGPUTTIImpl::getRegisterBitWidth(bool Vector) {
   return Vector ? 0 : 32;
 }
 
 unsigned AMDGPUTTIImpl::getLoadStoreVecRegBitWidth(unsigned AddrSpace) const {

test/Transforms/LoopUnroll/AMDGPU/unroll-for-private.ll

+; RUN: opt -mtriple=amdgcn-unknown-amdhsa -loop-unroll -S -amdgpu-unroll-threshold-private=20000 %s | FileCheck %s
+
+; Check that we full unroll loop to be able to eliminate alloca
+; CHECK-LABEL: @non_invariant_ind
+; CHECK:       for.body:
+; CHECK-NOT:   br
+; CHECK:       store i32 %tmp15, i32 addrspace(1)* %arrayidx7, align 4
+; CHECK:       ret void
+
+define void @non_invariant_ind(i32 addrspace(1)* nocapture %a, i32 %x) {
+entry:
+  %arr = alloca [64 x i32], align 4
+  %tmp1 = tail call i32 @llvm.amdgcn.workitem.id.x() #1
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  %arrayidx5 = getelementptr inbounds [64 x i32], [64 x i32]* %arr, i32 0, i32 %x
+  %tmp15 = load i32, i32* %arrayidx5, align 4
+  %arrayidx7 = getelementptr inbounds i32, i32 addrspace(1)* %a, i32 %tmp1
+  store i32 %tmp15, i32 addrspace(1)* %arrayidx7, align 4
+  ret void
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.015 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %idxprom = sext i32 %i.015 to i64
+  %arrayidx = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %idxprom
+  %tmp16 = load i32, i32 addrspace(1)* %arrayidx, align 4
+  %add = add nsw i32 %i.015, %tmp1
+  %rem = srem i32 %add, 64
+  %arrayidx3 = getelementptr inbounds [64 x i32], [64 x i32]* %arr, i32 0, i32 %rem
+  store i32 %tmp16, i32* %arrayidx3, align 4
+  %inc = add nuw nsw i32 %i.015, 1
+  %exitcond = icmp eq i32 %inc, 100
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+}
+
+; Check that we unroll inner loop but not outer
+; CHECK-LABEL: @invariant_ind
+; CHECK:       %[[exitcond:[^ ]+]] = icmp eq i32 %{{.*}}, 32
+; CHECK:       br i1 %[[exitcond]]
+; CHECK-NOT:   icmp eq i32 %{{.*}}, 100
+
+define void @invariant_ind(i32 addrspace(1)* nocapture %a, i32 %x) {
+entry:
+  %arr = alloca [64 x i32], align 4
+  %tmp1 = tail call i32 @llvm.amdgcn.workitem.id.x() #1
+  br label %for.cond2.preheader
+
+for.cond2.preheader:                              ; preds = %for.cond.cleanup5, %entry
+  %i.026 = phi i32 [ 0, %entry ], [ %inc10, %for.cond.cleanup5 ]
+  %idxprom = sext i32 %i.026 to i64
+  %arrayidx = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %idxprom
+  %tmp15 = load i32, i32 addrspace(1)* %arrayidx, align 4
+  br label %for.body6
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup5
+  %arrayidx13 = getelementptr inbounds [64 x i32], [64 x i32]* %arr, i32 0, i32 %x
+  %tmp16 = load i32, i32* %arrayidx13, align 4
+  %arrayidx15 = getelementptr inbounds i32, i32 addrspace(1)* %a, i32 %tmp1
+  store i32 %tmp16, i32 addrspace(1)* %arrayidx15, align 4
+  ret void
+
+for.cond.cleanup5:                                ; preds = %for.body6
+  %inc10 = add nuw nsw i32 %i.026, 1
+  %exitcond27 = icmp eq i32 %inc10, 32
+  br i1 %exitcond27, label %for.cond.cleanup, label %for.cond2.preheader
+
+for.body6:                                        ; preds = %for.body6, %for.cond2.preheader
+  %j.025 = phi i32 [ 0, %for.cond2.preheader ], [ %inc, %for.body6 ]
+  %add = add nsw i32 %j.025, %tmp1
+  %rem = srem i32 %add, 64
+  %arrayidx8 = getelementptr inbounds [64 x i32], [64 x i32]* %arr, i32 0, i32 %rem
+  store i32 %tmp15, i32* %arrayidx8, align 4
+  %inc = add nuw nsw i32 %j.025, 1
+  %exitcond = icmp eq i32 %inc, 100
+  br i1 %exitcond, label %for.cond.cleanup5, label %for.body6
+}
+
+; Check we do not enforce unroll if alloca is too big
+; CHECK-LABEL: @too_big
+; CHECK:       for.body:
+; CHECK:       icmp eq i32 %{{.*}}, 100
+; CHECK:       br
+
+define void @too_big(i32 addrspace(1)* nocapture %a, i32 %x) {
+entry:
+  %arr = alloca [256 x i32], align 4
+  %tmp1 = tail call i32 @llvm.amdgcn.workitem.id.x() #1
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  %arrayidx5 = getelementptr inbounds [256 x i32], [256 x i32]* %arr, i32 0, i32 %x
+  %tmp15 = load i32, i32* %arrayidx5, align 4
+  %arrayidx7 = getelementptr inbounds i32, i32 addrspace(1)* %a, i32 %tmp1
+  store i32 %tmp15, i32 addrspace(1)* %arrayidx7, align 4
+  ret void
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.015 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %idxprom = sext i32 %i.015 to i64
+  %arrayidx = getelementptr inbounds i32, i32 addrspace(1)* %a, i64 %idxprom
+  %tmp16 = load i32, i32 addrspace(1)* %arrayidx, align 4
+  %add = add nsw i32 %i.015, %tmp1
+  %rem = srem i32 %add, 64
+  %arrayidx3 = getelementptr inbounds [256 x i32], [256 x i32]* %arr, i32 0, i32 %rem
+  store i32 %tmp16, i32* %arrayidx3, align 4
+  %inc = add nuw nsw i32 %i.015, 1
+  %exitcond = icmp eq i32 %inc, 100
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+}
+
+declare i8 addrspace(2)* @llvm.amdgcn.dispatch.ptr() #1
+
+declare i32 @llvm.amdgcn.workitem.id.x() #1
+
+declare i32 @llvm.amdgcn.workgroup.id.x() #1
+
+declare i8 addrspace(2)* @llvm.amdgcn.implicitarg.ptr() #1
+
+attributes #1 = { nounwind readnone }