[AMDGPU] Enable runtime unroll for LDS

[rampitec]

Summary

We want to do unroll for LDS even for runtime trip count
to combine LDS operations

[arsenm]

I’ve thought might want to always do runtime unroll

[rampitec]

In D75293#1896535, @arsenm wrote:

I’ve thought might want to always do runtime unroll

We might, but so far there is not enough justification.

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

static cl::opt<unsigned> UnrollThresholdLocal(
   cl::desc("Unroll threshold for AMDGPU if local memory used in a loop"),
   cl::init(1000), cl::Hidden);
 
 static cl::opt<unsigned> UnrollThresholdIf(
   "amdgpu-unroll-threshold-if",
   cl::desc("Unroll threshold increment for AMDGPU for each if statement inside loop"),
   cl::init(150), cl::Hidden);
 
+static cl::opt<bool> UnrollRuntimeLocal(
+  "amdgpu-unroll-runtime-local",
+  cl::desc("Allow runtime unroll for AMDGPU if local memory used in a loop"),
+  cl::init(true), cl::Hidden);
+
 static cl::opt<bool> UseLegacyDA(
   "amdgpu-use-legacy-divergence-analysis",
   cl::desc("Enable legacy divergence analysis for AMDGPU"),
   cl::init(false), cl::Hidden);
 
 static bool dependsOnLocalPhi(const Loop *L, const Value *Cond,
                               unsigned Depth = 0) {
   const Instruction *I = dyn_cast<Instruction>(Cond);
for (const Instruction &I : *BB) {
         // Inhibit unroll for local memory if we have seen addressing not to
         // a variable, most likely we will be unable to combine it.
         // Do not unroll too deep inner loops for local memory to give a chance
         // to unroll an outer loop for a more important reason.
         if (LocalGEPsSeen > 1 || L->getLoopDepth() > 2 ||
             (!isa<GlobalVariable>(GEP->getPointerOperand()) &&
              !isa<Argument>(GEP->getPointerOperand())))
           continue;
+        LLVM_DEBUG(dbgs() << "Allow unroll runtime for loop:\n"
+                          << *L << " due to LDS use.\n");
+        UP.Runtime = UnrollRuntimeLocal;
       }
 
       // 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 || L->isLoopInvariant(Op))
           continue;

llvm/test/CodeGen/AMDGPU/unroll.ll

 for.inc:                                          ; preds = %for.body, %if.then
   %inc = add nuw nsw i32 %i1, 1
   %cmp = icmp ult i32 %inc, 48
   br i1 %cmp, label %for.body, label %for.end
 
 for.end:                                          ; preds = %for.cond
   ret void
 }
+
+; Check that runtime unroll is enabled for local memory references
+
+; CHECK-LABEL: @local_memory_runtime
+; CHECK: loop.header:
+; CHECK: load i32, i32 addrspace(3)*
+; CHECK: load i32, i32 addrspace(3)*
+; CHECK: br i1
+; CHECK: loop.header.epil
+; CHECK: load i32, i32 addrspace(3)*
+; CHECK: ret
+define amdgpu_kernel void @local_memory_runtime(i32 addrspace(1)* %out, i32 addrspace(3)* %lds, i32 %n) {
+entry:
+  br label %loop.header
+
+loop.header:
+  %counter = phi i32 [0, %entry], [%inc, %loop.inc]
+  br label %loop.body
+
+loop.body:
+  %ptr_lds = getelementptr i32, i32 addrspace(3)* %lds, i32 %counter
+  %val = load i32, i32 addrspace(3)* %ptr_lds
+  %ptr_out = getelementptr i32, i32 addrspace(1)* %out, i32 %counter
+  store i32 %val, i32 addrspace(1)* %ptr_out
+  br label %loop.inc
+
+loop.inc:
+  %inc = add i32 %counter, 1
+  %cond = icmp sge i32 %counter, %n
+  br i1 %cond, label  %exit, label %loop.header
+
+exit:
+  ret void
+}