[AMDGPU] Fix AMDGPUUnifyDivergentExitNodes

[piotr]

Summary

For the case where "done" bits on existing exports are removed
by unifyReturnBlockSet(), unify all return blocks - even the
uniformly reached ones. We do not want to end up with a non-unified,
uniformly reached block containing a normal export with the "done"
bit cleared.

That case is believed to be rare - possible with infinite loops
in pixel shaders.

This is a fix for D71192.

[piotr]

See D71192#1895125 for more context.

[piotr]

Fix formatting.

[arsenm]

llvm/test/CodeGen/AMDGPU/multi-divergent-exit-region.ll

747

Extra newlines

763

Can drop attribute comment

[arsenm]

I think this makes sense

[cwabbott]

llvm/test/CodeGen/AMDGPU/multi-divergent-exit-region.ll

752

You don't need .preheader.1 anymore and can delete it, right?

[piotr]

llvm/test/CodeGen/AMDGPU/multi-divergent-exit-region.ll

752

Yes, the test could be simplified i guess. I just took the version I had generated. (I just pressed "submit", so can do it in another commit).

[piotr]

llvm/test/CodeGen/AMDGPU/multi-divergent-exit-region.ll

752

Test updated in 4d8a72027706883a14f8ea0b9e0986c6e44da496.

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

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> UniformlyReachedRetBlocks;
   SmallVector<BasicBlock *, 4> UnreachableBlocks;
 
   // Dummy return block for infinite loop.
   BasicBlock *DummyReturnBB = nullptr;
 
   bool InsertExport = false;
 
   for (BasicBlock *BB : PDT.getRoots()) {
     if (isa<ReturnInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         ReturningBlocks.push_back(BB);
+      else
+        UniformlyReachedRetBlocks.push_back(BB);
     } else if (isa<UnreachableInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         UnreachableBlocks.push_back(BB);
     } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {
 
       ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
       if (DummyReturnBB == nullptr) {
         DummyReturnBB = BasicBlock::Create(F.getContext(),
if (ReturningBlocks.empty())
     return false; // No blocks return
 
   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");
+  // Unify returning blocks. If we are going to insert the export it is also
+  // necessary to include blocks that are uniformly reached, because in addition
+  // to inserting the export the "done" bits on existing exports will be cleared and
+  // we do not want to end up with the normal export in a non-unified, uniformly
+  // reached block with the "done" bit cleared.
+  auto BlocksToUnify = std::move(ReturningBlocks);
+  if (InsertExport) {
+    BlocksToUnify.insert(BlocksToUnify.end(), UniformlyReachedRetBlocks.begin(),
+                         UniformlyReachedRetBlocks.end());
+  }
+
+  unifyReturnBlockSet(F, BlocksToUnify, InsertExport, TTI,
+                      "UnifiedReturnBlock");
   return true;
 }

llvm/test/CodeGen/AMDGPU/multi-divergent-exit-region.ll

 
 bb4:                                              ; preds = %bb3
   ret void
 
 bb5:                                              ; preds = %bb3
   unreachable
 }
 
+; Test that there is an extra export inserted after the normal export,
+; if the normal export is inside a uniformly reached block and there is
+; an infinite loop in the pixel shader.
+
+; IR-LABEL: @uniformly_reached_export
+; IR-NEXT: .entry:
+; IR: br i1 [[CND:%.*]], label %[[EXP:.*]], label %[[FLOW:.*]]
+
+; IR: [[FLOW]]:
+; IR-NEXT: phi
+; IR-NEXT: br i1 [[CND2:%.*]], label %[[PREHEADER:.*]], label %[[FLOW2:.*]]
+
+; IR: [[FLOW2]]:
+; IR-NEXT: br label %UnifiedReturnBlock
+
+; IR: [[EXP]]:
+; IR-NEXT: call void @llvm.amdgcn.exp.compr.v2f16(i32 immarg 0, i32 immarg 15, <2 x half> <half 0xH3C00, half 0xH0000>, <2 x half> <half 0xH0000, half 0xH3C00>, i1 immarg false, i1 immarg true)
+; IR-NEXT: br label %[[FLOW]]
+
+; 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
+
+define amdgpu_ps void @uniformly_reached_export(float inreg %tmp25) {
+.entry:
+
+  %tmp26 = fcmp olt float %tmp25, 0.000000e+00
+
+  br i1 %tmp26, label %.preheader.1, label %bb27
+
+.preheader.1:                                     ; preds = %.entry
+  br label %bb
+
+bb:                                               ; preds = %bb, %.preheader.1
+  br label %bb
+
+bb27:                                             ; preds = %.entry
+  call void @llvm.amdgcn.exp.compr.v2f16(i32 immarg 0, i32 immarg 15, <2 x half> <half 0xH3C00, half 0xH0000>, <2 x half> <half 0xH0000, half 0xH3C00>, i1 immarg true, i1 immarg true)
+  ret void
+}
+
+; Function Attrs: inaccessiblememonly nounwind writeonly
+declare void @llvm.amdgcn.exp.compr.v2f16(i32 immarg, i32 immarg, <2 x half>, <2 x half>, i1 immarg, i1 immarg) #0
 declare i32 @llvm.amdgcn.workitem.id.x() #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }

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

 ; 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 [[UNIFIEDRETURNBLOCK:%.*]]
 ; IR:       TransitionBlock:
 ; IR-NEXT:    br i1 [[N30]], label [[DOTLOOPEXIT]], label [[N28]]
-; IR:       n31:
-; IR-NEXT:    ret void
 ; 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