AMDGPU: Fix AMDGPUUnifyDivergentExitNodes with no normal returns

[cwabbott]

Summary

The code was assuming in a few places that if there was only one exit
from the function that it was a normal return, which is invalid. It
could be an infinite loop, in which case we still need to insert the
usual fake edge so that the null export happens. This fixes shaders that
end with an infinite loop that discards.

[nhaehnle]

LGTM

[cwabbott]

It seems I forgot to test this against the entire testsuite, which made a bunch
of buildbots unhappy. This version fixes the issues:

• Skip the entire thing for non-pixel shaders that only end with an infinite

loop. This was causing slightly different assembly for four other tests, due to
the insertion of the fake branch, and it generates extra work for normal
shaders that's pointless.

• Update the expected output in update-phi.ll to account for the null export.

[nhaehnle]

LGTM. Did you manage to sort out your commit access issues?

[cwabbott]

In D71192#1848832, @nhaehnle wrote:

LGTM. Did you manage to sort out your commit access issues?

Yes, I just got added yesterday.

[piotr]

Hi @cwabbott,
This commit causes a GPU hang on amdvlk in one test, due to the missing "done" bit on the normal export. I think in the attached case the patch incorrectly classifies the export as being in an infinite loop.

before_AMDGPUUnifyDivergentExitNodes.ll2 KBDownload

[cwabbott]

In D71192#1886101, @piotr wrote:

Hi @cwabbott,
This commit causes a GPU hang on amdvlk in one test, due to the missing "done" bit on the normal export. I think in the attached case the patch incorrectly classifies the export as being in an infinite loop.

before_AMDGPUUnifyDivergentExitNodes.ll2 KBDownload

Hi @piotr,

I think what's happening is that the "normal" return, that already has an export before it, is uniformly reached (i.e. only reached via uniform branches). Normally this means that we don't have to do anything with it, so the pass ignores it, but then we remove the "done" bit even though the final dummy export, which is supposed to replace it, isn't reached when returning normally. There are two ways I can see to solve it:

• When InsertExport is true, unify *all* the return blocks, even the uniformly-reached ones. For this shader, we would still remove the done bit on the original export, but we'd replace the return with a branch to the UnifiedReturnBlock with the dummy export.
• Do the above, but when we can prove that all the existing exports are post-dominated by the normal, uniformly-reached returns, we can avoid clearing their done bits or unifying them.

The second option would prevent the regression in code quality, but it would be more complicated. We've deliberately made this as simple as possible already, because it only kicks in when there are infinite loops (that might have kills in them) and this seems pretty uncommon in real-world shaders. Is this coming from a test or an actual game?

[piotr]

The reported example comes from one of graphicsfuzz CTS tests, so I think it is fine to treat it as an unrealistic corner case where performance is not critical.

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

for (BasicBlock *BB : ReturningBlocks) {
     simplifyCFG(BB, TTI, {2});
   }
 
   return NewRetBlock;
 }
 
 bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
   auto &PDT = getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
-  if (PDT.getRoots().size() <= 1)
+
+  // If there's only one exit, we don't need to do anything, unless this is a
+  // pixel shader and that exit is an infinite loop, since we still have to
+  // insert an export in that case.
+  if (PDT.getRoots().size() <= 1 &&
+      F.getCallingConv() != CallingConv::AMDGPU_PS)
     return false;
 
   LegacyDivergenceAnalysis &DA = getAnalysis<LegacyDivergenceAnalysis>();
 
   // Loop over all of the blocks in a function, tracking all of the blocks that
   // return.
   SmallVector<BasicBlock *, 4> ReturningBlocks;
   SmallVector<BasicBlock *, 4> UnreachableBlocks;
if (!ReturningBlocks.empty()) {
       ReturningBlocks.push_back(UnreachableBlock);
     }
   }
 
   // Now handle return blocks.
   if (ReturningBlocks.empty())
     return false; // No blocks return
 
-  if (ReturningBlocks.size() == 1)
+  if (ReturningBlocks.size() == 1 && !InsertExport)
     return false; // Already has a single return block
 
   const TargetTransformInfo &TTI
     = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
   unifyReturnBlockSet(F, ReturningBlocks, InsertExport, TTI, "UnifiedReturnBlock");
   return true;
 }

llvm/test/CodeGen/AMDGPU/kill-infinite-loop.ll

loop:
   call void @llvm.amdgcn.kill(i1 false) #3
   br label %loop
 
 end:
   call void @llvm.amdgcn.exp.compr.v2i16(i32 0, i32 5, <2 x i16> < i16 0, i16 0 >, <2 x i16> < i16 0, i16 0 >, i1 true, i1 true) #3
   ret void
 }
 
+; test the case where there's only a kill in an infinite loop
+; CHECK-LABEL: only_kill
+; CHECK: exp null off, off, off, off done vm
+; CHECK-NEXT: s_endpgm
+; SIInsertSkips inserts an extra null export here, but it should be harmless.
+; CHECK: exp null off, off, off, off done vm
+; CHECK-NEXT: s_endpgm
+define amdgpu_ps void @only_kill() #0 {
+main_body:
+  br label %loop
+
+loop:
+  call void @llvm.amdgcn.kill(i1 false) #3
+  br label %loop
+}
+
 ; In case there's an epilog, we shouldn't have to do this.
 ; CHECK-LABEL: return_nonvoid
 ; CHECK-NOT: exp null off, off, off, off done vm
 define amdgpu_ps float @return_nonvoid(float %0) #0 {
 main_body:
   %cmp = fcmp olt float %0, 1.000000e+01
   br i1 %cmp, label %end, label %loop
 

llvm/test/CodeGen/AMDGPU/update-phi.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -mtriple=amdgcn-- -S -amdgpu-unify-divergent-exit-nodes -verify %s | FileCheck -check-prefix=IR %s
 
 ; Make sure that the phi in n28 is updated when the block is split by unify
 ; divergent exit nodes.
 
 define amdgpu_ps void @_amdgpu_ps_main() local_unnamed_addr #3 {
 ; IR-LABEL: @_amdgpu_ps_main(
 ; IR-NEXT:  .entry:
 ; IR-NEXT:    br label [[DOTLOOPEXIT:%.*]]
 ; IR:       .loopexit:
 ; IR-NEXT:    br label [[N28:%.*]]
 ; IR:       n28:
 ; IR-NEXT:    [[DOT01:%.*]] = phi float [ 0.000000e+00, [[DOTLOOPEXIT]] ], [ [[N29:%.*]], [[TRANSITIONBLOCK:%.*]] ]
 ; IR-NEXT:    [[N29]] = fadd float [[DOT01]], 1.000000e+00
 ; IR-NEXT:    [[N30:%.*]] = fcmp ogt float [[N29]], 4.000000e+00
-; IR-NEXT:    br i1 true, label [[TRANSITIONBLOCK]], label [[DUMMYRETURNBLOCK:%.*]]
+; IR-NEXT:    br i1 true, label [[TRANSITIONBLOCK]], label [[UNIFIEDRETURNBLOCK:%.*]]
 ; IR:       TransitionBlock:
 ; IR-NEXT:    br i1 [[N30]], label [[DOTLOOPEXIT]], label [[N28]]
 ; IR:       n31:
 ; IR-NEXT:    ret void
-; IR:       DummyReturnBlock:
+; IR:       UnifiedReturnBlock:
+; IR-NEXT:    call void @llvm.amdgcn.exp.f32(i32 9, i32 0, float undef, float undef, float undef, float undef, i1 true, i1 true)
 ; IR-NEXT:    ret void
 ;
 .entry:
   br label %.loopexit
 
 .loopexit:                                        ; preds = %n28, %.entry
   br label %n28