AMDGPU: Report large stack usage for recursive calls

[arsenm]

Summary

We were previously setting an ignored bit in the kernel headers. The
current behavior is to add the large amount on top of the statically
known size of a single stack frame. I'm not sure if we should just use
the large size as the entire reported size instead.

[rampitec]

Assuming we are consistently getting norecurse attribute LGTM.

[arsenm]

c7a0c2d0f7be2f456bd72b5c3508966d5b10233b

llvm/lib/Target/AMDGPU/AMDGPUResourceUsageAnalysis.cpp

for (const MachineInstr &MI : MBB) {
         // convention does match, that would have errored earlier.
         if (Callee && AMDGPU::isEntryFunctionCC(Callee->getCallingConv()))
           report_fatal_error("invalid call to entry function");
 
         bool IsIndirect = !Callee || Callee->isDeclaration();
         if (!IsIndirect)
           I = CallGraphResourceInfo.find(Callee);
 
+        // FIXME: Call site could have norecurse on it
+        if (!Callee || !Callee->doesNotRecurse()) {
+          Info.HasRecursion = true;
+
+          // TODO: If we happen to know there is no stack usage in the
+          // callgraph, we don't need to assume an infinitely growing stack.
+          if (!MI.isReturn()) {
+            // We don't need to assume an unknown stack size for tail calls.
+
+            // FIXME: This only benefits in the case where the kernel does not
+            // directly call the tail called function. If a kernel directly
+            // calls a tail recursive function, we'll assume maximum stack size
+            // based on the regular call instruction.
+            CalleeFrameSize =
+              std::max(CalleeFrameSize,
+                       static_cast<uint64_t>(AssumedStackSizeForExternalCall));
+          }
+        }
+
         if (IsIndirect || I == CallGraphResourceInfo.end()) {
           CalleeFrameSize =
               std::max(CalleeFrameSize,
                        static_cast<uint64_t>(AssumedStackSizeForExternalCall));
 
           // Register usage of indirect calls gets handled later
           Info.UsesVCC = true;
           Info.UsesFlatScratch = ST.hasFlatAddressSpace();
           Info.HasDynamicallySizedStack = true;
           Info.HasIndirectCall = true;
         } else {
           // We force CodeGen to run in SCC order, so the callee's register
           // usage etc. should be the cumulative usage of all callees.
           MaxSGPR = std::max(I->second.NumExplicitSGPR - 1, MaxSGPR);
           MaxVGPR = std::max(I->second.NumVGPR - 1, MaxVGPR);
           MaxAGPR = std::max(I->second.NumAGPR - 1, MaxAGPR);
           CalleeFrameSize =
               std::max(I->second.PrivateSegmentSize, CalleeFrameSize);
           Info.UsesVCC |= I->second.UsesVCC;
           Info.UsesFlatScratch |= I->second.UsesFlatScratch;
           Info.HasDynamicallySizedStack |= I->second.HasDynamicallySizedStack;
           Info.HasRecursion |= I->second.HasRecursion;
           Info.HasIndirectCall |= I->second.HasIndirectCall;
         }
-
-        // FIXME: Call site could have norecurse on it
-        if (!Callee || !Callee->doesNotRecurse())
-          Info.HasRecursion = true;
       }
     }
   }
 
   Info.NumExplicitSGPR = MaxSGPR + 1;
   Info.NumVGPR = MaxVGPR + 1;
   Info.NumAGPR = MaxAGPR + 1;
   Info.PrivateSegmentSize += CalleeFrameSize;

llvm/test/CodeGen/AMDGPU/call-graph-register-usage.ll

 ; NumVgprs: 24
 ; GCN: ScratchSize: 16384
 define amdgpu_kernel void @usage_external_recurse() #0 {
   call void @external_recurse()
   ret void
 }
 
 ; GCN-LABEL: {{^}}direct_recursion_use_stack:
-; GCN: ScratchSize: 2064
+; GCN: ScratchSize: 18448{{$}}
 define void @direct_recursion_use_stack(i32 %val) #2 {
   %alloca = alloca [512 x i32], align 4, addrspace(5)
   call void asm sideeffect "; use $0", "v"([512 x i32] addrspace(5)* %alloca) #0
   %cmp = icmp eq i32 %val, 0
   br i1 %cmp, label %ret, label %call
 
 call:
   %val.sub1 = sub i32 %val, 1
   call void @direct_recursion_use_stack(i32 %val.sub1)
   br label %ret
 
 ret:
   ret void
 }
 
 ; GCN-LABEL: {{^}}usage_direct_recursion:
 ; GCN: is_ptr64 = 1
 ; GCN: is_dynamic_callstack = 1
-; GCN: workitem_private_segment_byte_size = 2064
+; GCN: workitem_private_segment_byte_size = 18448{{$}}
 define amdgpu_kernel void @usage_direct_recursion(i32 %n) #0 {
   call void @direct_recursion_use_stack(i32 %n)
   ret void
 }
 
 ; Make sure there's no assert when a sgpr96 is used.
 ; GCN-LABEL: {{^}}count_use_sgpr96_external_call
 ; GCN: ; sgpr96 s[{{[0-9]+}}:{{[0-9]+}}]

llvm/test/CodeGen/AMDGPU/recursion.ll

+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck %s
+
+; CHECK-LABEL: {{^}}recursive:
+; CHECK: ScratchSize: 16
+define void @recursive() {
+  call void @recursive()
+  store volatile i32 0, i32 addrspace(1)* undef
+  ret void
+}
+
+; CHECK-LABEL: {{^}}tail_recursive:
+; CHECK: ScratchSize: 0
+define void @tail_recursive() {
+  tail call void @tail_recursive()
+  ret void
+}
+
+define void @calls_tail_recursive() norecurse {
+  tail call void @tail_recursive()
+  ret void
+}
+
+; CHECK-LABEL: {{^}}tail_recursive_with_stack:
+define void @tail_recursive_with_stack() {
+  %alloca = alloca i32, addrspace(5)
+  store volatile i32 0, i32 addrspace(5)* %alloca
+  tail call void @tail_recursive_with_stack()
+  ret void
+}
+
+; For an arbitrary recursive call, report a large number for unknown stack usage.
+; CHECK-LABEL: {{^}}calls_recursive:
+; CHECK: .amdhsa_private_segment_fixed_size 16400{{$}}
+define amdgpu_kernel void @calls_recursive() {
+  call void @recursive()
+  ret void
+}
+
+; Make sure we do not report a huge stack size for tail recursive
+; functions
+; CHECK-LABEL: {{^}}kernel_indirectly_calls_tail_recursive:
+; CHECK: .amdhsa_private_segment_fixed_size 0{{$}}
+define amdgpu_kernel void @kernel_indirectly_calls_tail_recursive() {
+  call void @calls_tail_recursive()
+  ret void
+}
+
+; TODO: Even though tail_recursive is only called as a tail call, we
+; end up treating it as generally recursive call from the regular call
+; in the kernel.
+
+; CHECK-LABEL: {{^}}kernel_calls_tail_recursive:
+; CHECK: .amdhsa_private_segment_fixed_size 16384{{$}}
+define amdgpu_kernel void @kernel_calls_tail_recursive() {
+  call void @tail_recursive()
+  ret void
+}
+
+; CHECK-LABEL: {{^}}kernel_calls_tail_recursive_with_stack:
+; CHECK: .amdhsa_private_segment_fixed_size 16384{{$}}
+define amdgpu_kernel void @kernel_calls_tail_recursive_with_stack() {
+  call void @tail_recursive_with_stack()
+  ret void
+}