[AMDGPU] Avoid null export insertion when unifying exit blocks

[critson]

Summary

Avoid adding a null export by unifying existing "done" exports
in the unified exit block of a shader with divergent exits.
If a null export is required then place this in a separate block
only visited by divergent exits that require it.

This assumes well-formed IR generated by existing frontends.
Specifically that there is a single "done" export for the shader
which occurs in uniform control flow.

[foad]

Mostly minor comments inline but I'm concerned about the "multiple predecessors with different done exports" case.

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

140

Could be "for (auto &I : reverse(BB)) ..."

143

Why compare with BoolTrue, rather than examine the value? Is this just optimisation to make the test as fast as possible?

154–181

Isn't it possible that different predecessors have different "done" exports? What is this function supposed to do in that case?

158

Don't need the ", 8" (unless you have special knowledge that 8 really is the optimal value).

177

I think return IsCompr ? (...) : (...) would be clearer.

265

Clearer to use .empty() than to use .size() as a boolean, here and 12 lines below.

268

Adding "else Phis[Idx] = Undef" here would simplify the creation of the new export intrinsic below.

313

What does "otherwise" refer back to?

434

Wasn't this supposed to say "exec mask"?

[critson]

In D102830#2770970, @foad wrote:

Mostly minor comments inline but I'm concerned about the "multiple predecessors with different done exports" case.

I have gone into this a bit more in the comment.
However this expects input IR is well-formed w.r.t. exports, as it is produced by existing front-ends.
If we wanted to handle the generic case then I feel like we might as well write a pass that checks all exports and adjusts them so they are well formed and run it after this pass.

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

143

I believe this is standard practice rather than retrieving the constant int values representing true and false repeatedly.

154–181

This does assumes the IR is well-formed to a certain degree.
I have added a comment to try to document this.

Consider the following IR:

if (condition) {
 export done 1
} else {
 export done 2
}
return

This will eventually be compiled to:

set exec mask for if-branch
export done 1
set exec mask for else-branch
export done 2
restore exec mask
return

Multiple export done instructions will be executed which is invalid.
The expected cases for this are:

1. Divergent exits with their own exports, e.g.

if (condition) {
  export done 1
  return
} else {
  export done 2
  return
}

2. Or divergent exits without exports or a uniformly reached export, e.g.

export done
if (condition) {
  return
} else {
  return
}

268

Sure, I had to change the type of Phis and add some casts.

313

I don't know, let's remove it.

[critson]

• Address reviewer comments

[foad]

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

154–181

OK. Thanks for explaining. I find it hard to get my head round the possible cases that we do or don't have to handle. Do you cope with:

if (uniform condition)
  export done 1
else
  export done 2
if (divergent condition)
  return 1
else
  return 2

?

268

Sorry, I didn't realise it would need a cast.

[ruiling]

I don't see why we need to handle such situation that would never happen. Can we simply assert each function has at most one exp_done?

[critson]

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

154–181

Good question, I think the answer is no -- although it is in principle valid IR that we do not generate.
I realize I was thinking primarily of divergent conditions, rather than uniform ones.
Although the comment I added is still accurate, <= 1 reachable export done.

In principle I can rework this to handle the case you raised, but I am not sure it is worth it?

[foad]

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

154–181

I have no idea what cases must be handled or are "worth" handling, so I think I'll leave that part of the review to someone who understands that.

[critson]

Rework based on the assumption there is only one "done" export.
This is true for existing front-ends.

[foad]

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

9–15

This comment seems obsolete now.

31

Same here.

[critson]

• Remove redundant comments

[ruiling]

Hi, I think we can further simplify the logic in our backend regarding export done and implement my idea as D105610, could you take a look if you have concern? the main idea is we just trust frontend will setup exp done correctly and only insert exp-null-done for kill intrinsic which is handled in SILateBranchLowering.cpp.

llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

public:
   static char ID; // Pass identification, replacement for typeid
 
   AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) {
     initializeAMDGPUUnifyDivergentExitNodesPass(*PassRegistry::getPassRegistry());
   }
 
   // We can preserve non-critical-edgeness when we unify function exit nodes
   void getAnalysisUsage(AnalysisUsage &AU) const override;
+  void unifyExports(Function &F, ArrayRef<BasicBlock *> ReturningBlocks,
+                    IRBuilder<> &B);
   BasicBlock *unifyReturnBlockSet(Function &F, DomTreeUpdater &DTU,
                                   ArrayRef<BasicBlock *> ReturningBlocks,
-                                  bool InsertExport, StringRef Name);
+                                  bool UpdateExports, StringRef Name);
   bool runOnFunction(Function &F) override;
 };
 
 } // end anonymous namespace
 
 char AMDGPUUnifyDivergentExitNodes::ID = 0;
 
 char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID;
for (BasicBlock *Pred : predecessors(Top)) {
       if (Visited.insert(Pred).second)
         Stack.push_back(Pred);
     }
   }
 
   return true;
 }
 
-static void removeDoneExport(Function &F) {
-  ConstantInt *BoolFalse = ConstantInt::getFalse(F.getContext());
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      if (IntrinsicInst *Intrin = llvm::dyn_cast<IntrinsicInst>(&I)) {
-        if (Intrin->getIntrinsicID() == Intrinsic::amdgcn_exp) {
-          Intrin->setArgOperand(6, BoolFalse); // done
-        } else if (Intrin->getIntrinsicID() == Intrinsic::amdgcn_exp_compr) {
-          Intrin->setArgOperand(4, BoolFalse); // done
-        }
-      }
-    }
-  }
+// Reverse search block to find a "done" export
+static IntrinsicInst *findExportDone(BasicBlock *BB, ConstantInt *BoolTrue) {
+  for (auto &I : reverse(*BB)) {
+        if (Intrin->getArgOperand(6) == BoolTrue)
+          return Intrin;
+        if (Intrin->getArgOperand(4) == BoolTrue)
+          return Intrin;
+  return nullptr;
+
+// Search a block and its predecessors to find a "done" export.
+// Note: this assumes the IR is already well-formed and has
+// at most one reachable "done" export for a given return block.
+static IntrinsicInst *findExportDone(BasicBlock *BB,
+                                     DenseSet<const BasicBlock *> &Visited,
+                                     ConstantInt *BoolTrue) {
+  SmallVector<BasicBlock *> Stack;
+
+  Stack.push_back(BB);
+  do {
+    BasicBlock *Top = Stack.pop_back_val();
+    IntrinsicInst *Export = findExportDone(Top, BoolTrue);
+    if (Export)
+      return Export;
+
+    for (BasicBlock *Pred : predecessors(Top)) {
+      if (Visited.insert(Pred).second)
+        Stack.push_back(Pred);
     }
+  } while (!Stack.empty());
 
-BasicBlock *AMDGPUUnifyDivergentExitNodes::unifyReturnBlockSet(
-    Function &F, DomTreeUpdater &DTU, ArrayRef<BasicBlock *> ReturningBlocks,
-    bool InsertExport, StringRef Name) {
-  // Otherwise, we need to insert a new basic block into the function, add a PHI
-  // nodes (if the function returns values), and convert all of the return
-  // instructions into unconditional branches.
-  BasicBlock *NewRetBlock = BasicBlock::Create(F.getContext(), Name, &F);
-  IRBuilder<> B(NewRetBlock);
+  return nullptr;
+}
 
-  if (InsertExport) {
-    // Ensure that there's only one "done" export in the shader by removing the
-    // "done" bit set on the original final export. More than one "done" export
-    // can lead to undefined behavior.
-    removeDoneExport(F);
+static bool isChannelEnabled(unsigned Channel, uint64_t EnBits, bool IsCompr) {
+  return IsCompr ? (EnBits & (1 << Channel))
+                 : (EnBits & (0x3 << (2 * Channel)));
+}
+
+static StringRef unifiedChannelName(unsigned Channel) {
+  switch (Channel) {
+  case 0:
+    return "UnifiedExportVSRC0";
+  case 1:
+    return "UnifiedExportVSRC1";
+  case 2:
+    return "UnifiedExportVSRC2";
+  case 3:
+    return "UnifiedExportVSRC3";
+  default:
+    return "";
+  }
+}
 
+void AMDGPUUnifyDivergentExitNodes::unifyExports(
+    Function &F, ArrayRef<BasicBlock *> ReturningBlocks, IRBuilder<> &B) {
+  // Ensure that there's only one "done" export in the shader by unifying
+  // all "done" exports into a single instruction in the unified return block.
+  // More than one "done" export can lead to undefined behavior.
+
+  // Find exports to unify
+  ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
+  DenseMap<BasicBlock *, IntrinsicInst *> Exports;
+  DenseSet<const BasicBlock *> Visited;
+  Type *ExportType = nullptr;
+  uint64_t EnBits = 0;
+  uint64_t Target = 0;
+  bool IsCompr = false;
+  bool MatchingExports = true;
+  for (BasicBlock *BB : ReturningBlocks) {
+    IntrinsicInst *Intrin = findExportDone(BB, Visited, BoolTrue);
+    if (!Intrin)
+      continue;
+
+    // Check all exports we are unifying are consistent
+    uint64_t TargetVal =
+        cast<ConstantInt>(Intrin->getArgOperand(0))->getZExtValue();
+    uint64_t EnBitsVal =
+        cast<ConstantInt>(Intrin->getArgOperand(1))->getZExtValue();
+    Type *TypeVal = Intrin->getFunctionType()->getParamType(2);
+    bool ComprVal = Intrin->getIntrinsicID() == Intrinsic::amdgcn_exp_compr;
+    if (!Exports.empty()) {
+      MatchingExports = MatchingExports && (Target == TargetVal) &&
+                        (EnBits == EnBitsVal) && (ExportType == TypeVal) &&
+                        (IsCompr == ComprVal);
+    } else {
+      Target = TargetVal;
+      EnBits = EnBitsVal;
+      ExportType = TypeVal;
+      IsCompr = ComprVal;
+    }
+    Exports[BB] = Intrin;
+  }
+
+  if (Exports.empty() || !MatchingExports) {
+    LLVM_DEBUG(dbgs() << "Unifying exports using null export\n");
+
+    // Clear "done" bit from exports
+    ConstantInt *BoolFalse = ConstantInt::getFalse(F.getContext());
+    for (auto Export : Exports) {
+      IntrinsicInst *Intrin = Export.second;
+      if (Intrin->getIntrinsicID() == Intrinsic::amdgcn_exp) {
+        Intrin->setArgOperand(6, BoolFalse); // done
+      } else if (Intrin->getIntrinsicID() == Intrinsic::amdgcn_exp_compr) {
+        Intrin->setArgOperand(4, BoolFalse); // done
+      }
+    }
+    // Insert null export for unified return block
     Value *Undef = UndefValue::get(B.getFloatTy());
     B.CreateIntrinsic(Intrinsic::amdgcn_exp, { B.getFloatTy() },
                       {
                         B.getInt32(AMDGPU::Exp::ET_NULL),
                         B.getInt32(0), // enabled channels
                         Undef, Undef, Undef, Undef, // values
                         B.getTrue(), // done
                         B.getTrue(), // valid mask
                       });
-  }
+  } else {
+    LLVM_DEBUG(dbgs() << "Unifying exports by merging\n");
+
+    // Create one phi node per enabled export channel
+    const unsigned Channels = IsCompr ? 2 : 4;
+    Value *Phis[4] = {nullptr};
+    for (unsigned Idx = 0; Idx < Channels; ++Idx) {
+      if (isChannelEnabled(Idx, EnBits, IsCompr))
+        Phis[Idx] = B.CreatePHI(ExportType, ReturningBlocks.size(),
+                                unifiedChannelName(Idx));
+    }
+    // Add export values to phi nodes and remove old export
+    Value *Undef = UndefValue::get(ExportType);
+    for (BasicBlock *BB : ReturningBlocks) {
+      IntrinsicInst *Intrin = nullptr;
+      if (Exports.count(BB))
+        Intrin = Exports[BB];
+      for (unsigned Idx = 0; Idx < Channels; ++Idx) {
+        if (isChannelEnabled(Idx, EnBits, IsCompr)) {
+          cast<PHINode>(Phis[Idx])->addIncoming(
+              Intrin ? Intrin->getOperand(2 + Idx) : Undef, BB);
+        } else {
+          Phis[Idx] = Undef;
+        }
+      }
+      if (Intrin)
+        Intrin->eraseFromParent();
+    }
+    // Insert new unified export based on phi nodes
+    if (IsCompr) {
+      B.CreateIntrinsic(Intrinsic::amdgcn_exp_compr, {ExportType},
+                        {
+                            B.getInt32(Target), B.getInt32(EnBits), Phis[0],
+                            Phis[1],
+                            B.getTrue(), // done
+                            B.getTrue(), // valid mask
+                        });
+    } else {
+      B.CreateIntrinsic(Intrinsic::amdgcn_exp, {ExportType},
+                        {
+                            B.getInt32(Target), B.getInt32(EnBits), Phis[0],
+                            Phis[1], Phis[2], Phis[3],
+                            B.getTrue(), // done
+                            B.getTrue(), // valid mask
+                        });
+    }
+  }
 
+BasicBlock *AMDGPUUnifyDivergentExitNodes::unifyReturnBlockSet(
+    Function &F, DomTreeUpdater &DTU, ArrayRef<BasicBlock *> ReturningBlocks,
+    bool UpdateExports, StringRef Name) {
+  // We need to insert a new basic block into the function,
+  // add a PHI nodes (if the function returns values),
+  // and convert all of the return instructions into unconditional branches.
+  BasicBlock *NewRetBlock = BasicBlock::Create(F.getContext(), Name, &F);
+  IRBuilder<> B(NewRetBlock);
+
+  if (UpdateExports)
+    unifyExports(F, ReturningBlocks, B);
+
   PHINode *PN = nullptr;
   if (F.getReturnType()->isVoidTy()) {
     B.CreateRetVoid();
   } else {
     // If the function doesn't return void... add a PHI node to the block...
     PN = B.CreatePHI(F.getReturnType(), ReturningBlocks.size(),
                      "UnifiedRetVal");
-    assert(!InsertExport);
+    assert(!UpdateExports);
     B.CreateRet(PN);
   }
 
   // Loop over all of the blocks, replacing the return instruction with an
   // unconditional branch.
   std::vector<DominatorTree::UpdateType> Updates;
   Updates.reserve(ReturningBlocks.size());
   for (BasicBlock *BB : ReturningBlocks) {
bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
   // 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;
+  bool UpdateExports = false;
 
   bool Changed = false;
   std::vector<DominatorTree::UpdateType> Updates;
 
   for (BasicBlock *BB : PDT.roots()) {
     if (isa<ReturnInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         ReturningBlocks.push_back(BB);
if (isa<ReturnInst>(BB->getTerminator())) {
         DummyReturnBB = BasicBlock::Create(F.getContext(),
                                            "DummyReturnBlock", &F);
         Type *RetTy = F.getReturnType();
         Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy);
 
         // For pixel shaders, the producer guarantees that an export is
         // executed before each return instruction. However, if there is an
         // infinite loop and we insert a return ourselves, we need to uphold
-        // that guarantee by inserting a null export. This can happen e.g. in
-        // an infinite loop with kill instructions, which is supposed to
-        // terminate. However, we don't need to do this if there is a non-void
-        // return value, since then there is an epilog afterwards which will
-        // still export.
+        // that guarantee by inserting a null export, or moving an existing
+        // export. This can happen e.g. in an infinite loop with kill
+        // instructions, which is supposed to terminate. However, we don't need
+        // to do this if there is a non-void return value, since then there is
+        // an epilog afterwards which will still export.
+        //
+        // We try not to insert new exports by moving an existing export to
+        // the unified return block.  When this is not possible, due to
+        // multiple incompatible final exports, then the "done" flag on
+        // all other exports is clear and a final null export is added.
         //
         // Note: In the case where only some threads enter the infinite loop,
         // this can result in the null export happening redundantly after the
         // original exports. However, The last "real" export happens after all
         // the threads that didn't enter an infinite loop converged, which
         // means that the only extra threads to execute the null export are
         // threads that entered the infinite loop, and they only could've
         // exited through being killed which sets their exec bit to 0.
         // Therefore, unless there's an actual infinite loop, which can have
         // invalid results, or there's a kill after the last export, which we
         // assume the frontend won't do, this export will have the same exec
         // mask as the last "real" export, and therefore the valid mask will be
         // overwritten with the same value and will still be correct. Also,
         // even though this forces an extra unnecessary export wait, we assume
         // that this happens rare enough in practice to that we don't have to
         // worry about performance.
         if (F.getCallingConv() == CallingConv::AMDGPU_PS &&
             RetTy->isVoidTy()) {
-          InsertExport = true;
+          UpdateExports = true;
         }
 
         ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
         ReturningBlocks.push_back(DummyReturnBB);
       }
 
       if (BI->isUnconditional()) {
         BasicBlock *LoopHeaderBB = BI->getSuccessor(0);
bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
   if (RequireAndPreserveDomTree)
     DTU.applyUpdates(Updates);
   Updates.clear();
 
   // Now handle return blocks.
   if (ReturningBlocks.empty())
     return Changed; // No blocks return
 
-  if (ReturningBlocks.size() == 1 && !InsertExport)
+  if (ReturningBlocks.size() == 1 && !UpdateExports)
     return Changed; // Already has a single return block
 
   // 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) {
+  if (UpdateExports) {
     llvm::append_range(BlocksToUnify, UniformlyReachedRetBlocks);
   }
 
-  unifyReturnBlockSet(F, DTU, BlocksToUnify, InsertExport,
+  unifyReturnBlockSet(F, DTU, BlocksToUnify, UpdateExports,
                       "UnifiedReturnBlock");
   return true;
 }

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

 ; RUN: llc -march=amdgcn -verify-machineinstrs -simplifycfg-require-and-preserve-domtree=1 < %s | FileCheck -enable-var-scope %s
 
 ; Although it's modeled without any control flow in order to get better code
 ; out of the structurizer, @llvm.amdgcn.kill actually ends the thread that calls
 ; it with "true". In case it's called in a provably infinite loop, we still
 ; need to successfully exit and export something, even if we can't know where
-; to jump to in the LLVM IR. Therefore we insert a null export ourselves in
-; this case right before the s_endpgm to avoid GPU hangs, which is what this
-; tests.
+; to jump to in the LLVM IR. Therefore we move the export to the shader end.
 
 ; CHECK-LABEL: return_void
-; Make sure that we remove the done bit from the original export
-; CHECK: exp mrt0 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} vm
+; Make sure there is a single export
+; CHECK-NOT: s_endpgm
+; CHECK: exp mrt0 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} done vm
+; CHECK-NEXT: s_endpgm
+; Early termination block follows
 ; CHECK: exp null off, off, off, off done vm
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @return_void(float %0) #0 {
 main_body:
   %cmp = fcmp olt float %0, 1.000000e+01
   br i1 %cmp, label %end, label %loop
 
 loop:
   call void @llvm.amdgcn.kill(i1 false) #3
   br label %loop
 
 end:
   call void @llvm.amdgcn.exp.f32(i32 0, i32 15, float 0., float 0., float 0., float 1., i1 true, i1 true) #3
   ret void
 }
 
 ; Check that we also remove the done bit from compressed exports correctly.
 ; CHECK-LABEL: return_void_compr
-; CHECK: exp mrt0 v{{[0-9]+}}, off, v{{[0-9]+}}, off compr vm
+; CHECK-NOT: s_endpgm
+; CHECK: exp mrt0 v{{[0-9]+}}, off, v{{[0-9]+}}, off done compr vm
+; CHECK-NEXT: s_endpgm
+; Early termination block follows
 ; CHECK: exp null off, off, off, off done vm
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @return_void_compr(float %0) #0 {
 main_body:
   %cmp = fcmp olt float %0, 1.000000e+01
   br i1 %cmp, label %end, label %loop
 
 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-NOT: s_endpgm
 ; CHECK: exp null off, off, off, off done vm
 ; CHECK-NEXT: s_endpgm
-; SILateBranchLowering inserts an extra null export here, but it should be harmless.
+; Early termination block follows
 ; 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

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.
+; Test that divergent exit is unified with uniform export exit.
 
 ; IR-LABEL: @uniformly_reached_export
 ; IR-NEXT: .entry:
-; IR: br i1 [[CND:%.*]], label %[[EXP:.*]], label %[[FLOW:.*]]
+; IR: br i1 [[CND:%.*]], label %[[LOOP:.*]], label %UnifiedReturnBlock
+
+; IR: [[LOOP]]:
+; IR-NEXT: br i1 false, label %[[FLOW:.*]], label %[[LOOP]]
 
 ; IR: [[FLOW]]:
-; IR-NEXT: phi
-; IR-NEXT: br i1 [[CND2:%.*]], label %[[LOOP:.*]], label %UnifiedReturnBlock
+; IR-NEXT: br label %UnifiedReturnBlock
 
-; IR: [[LOOP]]:
-; IR-NEXT: br i1 false, label %[[FLOW1:.*]], label %[[LOOP]]
+; IR: UnifiedReturnBlock:
+; IR-NEXT: call void @llvm.amdgcn.exp.compr.v2f16(i32 0, i32 15, <2 x half> <half 0xH3C00, half 0xH0000>, <2 x half> <half 0xH0000, half 0xH3C00>, i1 true, i1 true)
+; IR-NEXT: ret void
 
-; 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]]
+define amdgpu_ps void @uniformly_reached_export(float inreg %tmp25) {
+.entry:
+  %tmp26 = fcmp olt float %tmp25, 0.000000e+00
+  br i1 %tmp26, label %loop, label %bb27
+
+loop:                                             ; preds = %loop, %.entry
+  br label %loop
+
+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
+}
+
+; Test that early uniform export is moved to unified exit block.
+
+; IR-LABEL: @uniformly_reached_early_export
+; IR-NEXT: .entry:
+; IR: br i1 [[CND:%.*]], label %[[LOOP:.*]], label %UnifiedReturnBlock
+
+; IR: [[LOOP]]:
+; IR-NEXT: br i1 false, label %[[FLOW:.*]], label %[[LOOP]]
 
-; IR: [[FLOW1]]:
+; IR: [[FLOW]]:
 ; IR-NEXT: br label %UnifiedReturnBlock
 
 ; 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: call void @llvm.amdgcn.exp.compr.v2f16(i32 0, i32 15, <2 x half> <half 0xH3C00, half 0xH0000>, <2 x half> <half 0xH0000, half 0xH3C00>, i1 true, i1 true)
 ; IR-NEXT: ret void
 
-define amdgpu_ps void @uniformly_reached_export(float inreg %tmp25) {
+define amdgpu_ps void @uniformly_reached_early_export(float inreg %tmp25) {
 .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)
   %tmp26 = fcmp olt float %tmp25, 0.000000e+00
   br i1 %tmp26, label %loop, label %bb27
 
-loop:                                               ; preds = %loop, %.entry
+loop:                                             ; preds = %loop, %.entry
   br label %loop
 
 bb27:                                             ; preds = %.entry
+  ret void
+}
+
+; Test that incompatible divergent exports are not unified.
+
+; IR-LABEL: @incompatible_divergent_exports
+
+; IR: bb26:
+; IR-NEXT: call void @llvm.amdgcn.exp.compr.v2f16(i32 immarg 1, 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
+
+; IR: bb27:
+; 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
+
+; 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 @incompatible_divergent_exports(float inreg %tmp25) {
+.entry:
+  %tmp26 = fcmp olt float %tmp25, 1.000000e+00
+  br i1 %tmp26, label %loop, label %bb25
+
+loop:                                             ; preds = %loop, %.entry
+  br label %loop
+
+bb25:                                             ; preds = %.entry
+  %tmp27 = fcmp olt float %tmp25, 0.000000e+00
+  br i1 %tmp27, label %bb26, label %bb27
+
+bb26:                                             ; preds = %.bb25
+  call void @llvm.amdgcn.exp.compr.v2f16(i32 immarg 1, 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
+
+bb27:                                             ; preds = %.bb25
   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
 }
 
 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 }