[AMDGPU] Unify early PS termination blocks

[critson]

Summary

Depends on D77544.
Generate a single early exit block out-of-line and branch to this
if all lanes are killed. This avoids branching if lanes are active.

[critson]

This is a stepping stop to adding early termination in more places because it means termination can be added anywhere in a basic block (without splitting basic blocks).

Will need to be revised with D77544 is committed.

[arsenm]

Is this the right place for this? We're trying to remove this pass

[cdevadas]

We are trying to have the kill intrinsic handling early during the instruction selection.
I can try to incorporate it all entirely there including the unify block code (if can't find a better place now to have it).
I will also update D77544 if this patch goes upstream.

[critson]

In D82641#2118340, @cdevadas wrote:

We are trying to have the kill intrinsic handling early during the instruction selection.
I can try to incorporate it all entirely there including the unify block code (if can't find a better place now to have it).
I will also update D77544 if this patch goes upstream.

This pass seems to be the logic place for this right now.
Is there any code or ongoing discussion I can refer to for moving kill handling to instruction selection?
I can see a few problems with doing that, but that might just be because I do not know what is planned.

I am fine with this getting blown away when SIInsertSkips is replaced by a better solution, and contributing again to that solution.

I think this patch can be simplified if D77544 is submitted first and the order of SIInsertSkips and SIPreEmitPeephole are swapped.

[critson]

Rework this on top of D77544.
Make insertion of skips near the beginning of blocks work correctly by splitting blocks.

[nhaehnle]

In D82641#2118340, @cdevadas wrote:

We are trying to have the kill intrinsic handling early during the instruction selection.

Whether this is a good idea depends on what is meant by it :) Kill intrinsics interact with WQM for how pixel shaders are defined, and we really need to move the WQM pass later because it adds instructions that interfere with scheduling in a bad way. Setting up the CFG to be the right shape earlier makes sense to me.

Mostly looks good to me, some minor inline comments, but most importantly: please add a test case in skip-if-dead.ll with a function that triggers the early-exit and returns a value (the only test return right now is @test_kill_control_flow, but it doesn't trigger this early exit).

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

167

Should probably be renamed to generatePsEndPgm.

222–223

Why the update to NextBBI? The old code didn't do that, I believe on purpose because while it doesn't break correctness, there simply shouldn't be a need to inspect the EarlyExitBlock.

[nhaehnle]

In D82641#2119773, @nhaehnle wrote:

Mostly looks good to me, some minor inline comments, but most importantly: please add a test case in skip-if-dead.ll with a function that triggers the early-exit and returns a value (the only test return right now is @test_kill_control_flow, but it doesn't trigger this early exit).

I see this may be covered by llvm/test/CodeGen/AMDGPU/transform-block-with-return-to-epilog.ll?

[critson]

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

222–223

This update is not part of the iteration.
EarlyExitBlock is not inspected because this code is executed after the iteration of blocks.
This update is required so the SplitBB is inserted in the correct place (immediately following MBB). Technically this is only required when NextBBI is MBB->end(), in which case the EarlyExitBlock will become the next block; however, simply resetting the iterator keeps it more generic w.r.t. EarlyExitBlock placement.

[critson]

• Rebase
• Rename generateEndPgm
• Edit comment for clarity
• Add explicit test to skip-if-dead.ll

[nhaehnle]

Thanks, LGTM

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

void GCNPassConfig::addPreEmitPass() {
   //
   // FIXME: This stand-alone pass will emit indiv. S_NOP 0, as needed. It would
   // be better for it to emit S_NOP <N> when possible.
   addPass(&PostRAHazardRecognizerID);
   if (getOptLevel() > CodeGenOpt::None)
     addPass(&SIInsertHardClausesID);
 
   addPass(&SIRemoveShortExecBranchesID);
-  addPass(&SIPreEmitPeepholeID);
   addPass(&SIInsertSkipsPassID);
+  addPass(&SIPreEmitPeepholeID);
   addPass(&BranchRelaxationPassID);
 }
 
 TargetPassConfig *GCNTargetMachine::createPassConfig(PassManagerBase &PM) {
   return new GCNPassConfig(*this, PM);
 }
 
 yaml::MachineFunctionInfo *GCNTargetMachine::createDefaultFuncInfoYAML() const {

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

 
 class SIInsertSkips : public MachineFunctionPass {
 private:
   const SIRegisterInfo *TRI = nullptr;
   const SIInstrInfo *TII = nullptr;
   unsigned SkipThreshold = 0;
   MachineDominatorTree *MDT = nullptr;
 
+  MachineBasicBlock *EarlyExitBlock = nullptr;
+
   bool shouldSkip(const MachineBasicBlock &From,
                   const MachineBasicBlock &To) const;
 
   bool dominatesAllReachable(MachineBasicBlock &MBB);
+  void createEarlyExitBlock(MachineBasicBlock &MBB);
   void skipIfDead(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
                   DebugLoc DL);
 
   bool kill(MachineInstr &MI);
 
   bool skipMaskBranch(MachineInstr &MI, MachineBasicBlock &MBB);
 
 public:

 bool SIInsertSkips::dominatesAllReachable(MachineBasicBlock &MBB) {
   for (MachineBasicBlock *Other : depth_first(&MBB)) {
     if (!MDT->dominates(&MBB, Other))
       return false;
   }
   return true;
 }
 
+static void generateEndPgm(MachineBasicBlock &MBB,
+                           MachineBasicBlock::iterator I, DebugLoc DL,
+                           const SIInstrInfo *TII) {
+  // Generate "null export; s_endpgm".
+  BuildMI(MBB, I, DL, TII->get(AMDGPU::EXP_DONE))
+      .addImm(0x09) // V_008DFC_SQ_EXP_NULL
+      .addReg(AMDGPU::VGPR0, RegState::Undef)
+      .addReg(AMDGPU::VGPR0, RegState::Undef)
+      .addReg(AMDGPU::VGPR0, RegState::Undef)
+      .addReg(AMDGPU::VGPR0, RegState::Undef)
+      .addImm(1)  // vm
+      .addImm(0)  // compr
+      .addImm(0); // en
+  BuildMI(MBB, I, DL, TII->get(AMDGPU::S_ENDPGM)).addImm(0);
+}
+
+void SIInsertSkips::createEarlyExitBlock(MachineBasicBlock &MBB) {
+  MachineFunction *MF = MBB.getParent();
+  DebugLoc DL;
+
+  assert(!EarlyExitBlock);
+  EarlyExitBlock = MF->CreateMachineBasicBlock();
+  MF->insert(MF->end(), EarlyExitBlock);
+
+  generateEndPgm(*EarlyExitBlock, EarlyExitBlock->end(), DL, TII);
+}
+
 /// Insert an "if exec=0 { null export; s_endpgm }" sequence before the given
 /// iterator. Only applies to pixel shaders.
 void SIInsertSkips::skipIfDead(MachineBasicBlock &MBB,
                                MachineBasicBlock::iterator I, DebugLoc DL) {
   MachineFunction *MF = MBB.getParent();
   assert(MF->getFunction().getCallingConv() == CallingConv::AMDGPU_PS);
 
-  // Currently, SI_KILL_*_TERMINATOR is expected to occur only as the last
-  // terminator of a basic block. If this ever changes, we need to optionally
-  // split MBB here.
-  assert(I == MBB.end());
-
   // It is possible for an SI_KILL_*_TERMINATOR to sit at the bottom of a
   // basic block that has no further successors (e.g., there was an
   // `unreachable` there in IR). This can happen with original source of the
   // form:
   //
   //   if (uniform_condition) {
   //     write_to_memory();
   //     discard;
   //   }
   //
   // In this case, we write the "null_export; s_endpgm" skip code in the
   // already-existing basic block.
   auto NextBBI = std::next(MBB.getIterator());
-  bool NoSuccessor = llvm::find(MBB.successors(), &*NextBBI) == MBB.succ_end();
-  MachineBasicBlock *SkipBB;
+  bool NoSuccessor = NextBBI != MF->end() &&
+                     llvm::find(MBB.successors(), &*NextBBI) == MBB.succ_end();
 
   if (NoSuccessor) {
-    SkipBB = &MBB;
+    generateEndPgm(MBB, I, DL, TII);
   } else {
-    // Create a new basic block that will contain the "null export; s_endpgm"
-    // and set up the branching to go around it.
-    SkipBB = MF->CreateMachineBasicBlock();
-    MF->insert(NextBBI, SkipBB);
-
-    BuildMI(&MBB, DL, TII->get(AMDGPU::S_CBRANCH_EXECNZ)).addMBB(&*NextBBI);
-    MBB.addSuccessor(SkipBB);
-
-    MDT->addNewBlock(SkipBB, &MBB);
+    if (!EarlyExitBlock) {
+      createEarlyExitBlock(MBB);
+      // Reset iterator as next block may have changed
+      NextBBI = std::next(MBB.getIterator());
+    }
+
+    auto BranchMI = BuildMI(MBB, I, DL, TII->get(AMDGPU::S_CBRANCH_EXECZ))
+                        .addMBB(EarlyExitBlock);
+
+    // Split the block if the branch will not come at the end.
+    auto Next = std::next(BranchMI->getIterator());
+    if (Next != MBB.end() && !Next->isTerminator()) {
+      MachineBasicBlock *SplitBB =
+          MF->CreateMachineBasicBlock(MBB.getBasicBlock());
+      MF->insert(NextBBI, SplitBB);
+      SplitBB->splice(SplitBB->begin(), &MBB, I, MBB.end());
+      SplitBB->transferSuccessorsAndUpdatePHIs(&MBB);
+      // FIXME: the expectation is that this will be used near the beginning
+      //        of a block so just assume all registers are still live.
+      for (auto LiveIn : MBB.liveins())
+        SplitBB->addLiveIn(LiveIn);
+      MBB.addSuccessor(SplitBB);
+      MDT->addNewBlock(SplitBB, &MBB);
     }
 
-  // Generate "null export; s_endpgm".
-  BuildMI(SkipBB, DL, TII->get(AMDGPU::EXP_DONE))
-      .addImm(0x09) // V_008DFC_SQ_EXP_NULL
-      .addReg(AMDGPU::VGPR0, RegState::Undef)
-      .addReg(AMDGPU::VGPR0, RegState::Undef)
-      .addReg(AMDGPU::VGPR0, RegState::Undef)
-      .addReg(AMDGPU::VGPR0, RegState::Undef)
-      .addImm(1)  // vm
-      .addImm(0)  // compr
-      .addImm(0); // en
-  BuildMI(SkipBB, DL, TII->get(AMDGPU::S_ENDPGM)).addImm(0);
+    MBB.addSuccessor(EarlyExitBlock);
+    MDT->getBase().insertEdge(&MBB, EarlyExitBlock);
+  }
 }
 
 /// Translate a SI_KILL_*_TERMINATOR into exec-manipulating instructions.
 /// Return true unless the terminator is a no-op.
 bool SIInsertSkips::kill(MachineInstr &MI) {
   MachineBasicBlock &MBB = *MI.getParent();
   DebugLoc DL = MI.getDebugLoc();
 
bool SIInsertSkips::runOnMachineFunction(MachineFunction &MF) {
   }
 
   for (MachineInstr *Kill : KillInstrs) {
     skipIfDead(*Kill->getParent(), std::next(Kill->getIterator()),
                Kill->getDebugLoc());
     Kill->eraseFromParent();
   }
   KillInstrs.clear();
+  EarlyExitBlock = nullptr;
 
   return MadeChange;
 }

llvm/test/CodeGen/AMDGPU/insert-skips-kill-uncond.mir

 # RUN: llc -march=amdgcn -mcpu=polaris10 -run-pass si-insert-skips -amdgpu-skip-threshold-legacy=1 %s -o - | FileCheck %s
 # https://bugs.freedesktop.org/show_bug.cgi?id=99019
 --- |
   define amdgpu_ps void @kill_uncond_branch() {
     ret void
   }
 ...
 ---
 
 # CHECK-LABEL: name: kill_uncond_branch
 
 # CHECK: bb.0:
 # CHECK: S_CBRANCH_VCCNZ %bb.1, implicit $vcc
 
 # CHECK: bb.1:
 # CHECK: V_CMPX_LE_F32_e32
-# CHECK-NEXT: S_CBRANCH_EXECNZ %bb.2, implicit $exec
+# CHECK-NEXT: S_CBRANCH_EXECZ %bb.3, implicit $exec
 
-# CHECK: bb.3:
-# CHECK-NEXT: EXP_DONE
+# CHECK: bb.2:
 # CHECK: S_ENDPGM 0
 
-# CHECK: bb.2:
+# CHECK: bb.3:
+# CHECK-NEXT: EXP_DONE
 # CHECK: S_ENDPGM 0
 
 name: kill_uncond_branch
 
 body: |
   bb.0:
     successors: %bb.1
     S_CBRANCH_VCCNZ %bb.1, implicit $vcc

llvm/test/CodeGen/AMDGPU/skip-if-dead.ll

 ; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck %s
 
 ; CHECK-LABEL: {{^}}test_kill_depth_0_imm_pos:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: ; %bb.1:
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_0_imm_pos() #0 {
   call void @llvm.amdgcn.kill(i1 true)
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_kill_depth_0_imm_neg:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: s_mov_b64 exec, 0
-; CHECK-NEXT: s_cbranch_execnz BB1_2
+; CHECK-NEXT: s_cbranch_execz BB1_2
 ; CHECK-NEXT: ; %bb.1:
-; CHECK-NEXT: exp null off, off, off, off done vm
 ; CHECK-NEXT: s_endpgm
 ; CHECK-NEXT: BB1_2:
+; CHECK-NEXT: exp null off, off, off, off done vm
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_0_imm_neg() #0 {
   call void @llvm.amdgcn.kill(i1 false)
   ret void
 }
 
 ; FIXME: Ideally only one would be emitted
 ; CHECK-LABEL: {{^}}test_kill_depth_0_imm_neg_x2:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: s_mov_b64 exec, 0
-; CHECK-NEXT: s_cbranch_execnz BB2_2
-; CHECK:      exp null
-; CHECK-NEXT: s_endpgm
-; CHECK:      BB2_2:
+; CHECK-NEXT: s_cbranch_execz BB2_3
+; CHECK-NEXT: ; %bb.1:
 ; CHECK-NEXT: s_mov_b64 exec, 0
-; CHECK-NEXT: s_cbranch_execnz BB2_4
-; CHECK:      exp null
+; CHECK-NEXT: s_cbranch_execz BB2_3
+; CHECK-NEXT: ; %bb.2:
 ; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB2_4:
+; CHECK-NEXT: BB2_3:
+; CHECK:      exp null
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_0_imm_neg_x2() #0 {
   call void @llvm.amdgcn.kill(i1 false)
   call void @llvm.amdgcn.kill(i1 false)
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_kill_depth_var:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v0
-; CHECK-NEXT: s_cbranch_execnz BB3_2
-; CHECK:      exp null
+; CHECK-NEXT: s_cbranch_execz BB3_2
+; CHECK-NEXT: ; %bb.1:
 ; CHECK-NEXT: s_endpgm
 ; CHECK-NEXT: BB3_2:
+; CHECK:      exp null
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_var(float %x) #0 {
   %cmp = fcmp olt float %x, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp)
   ret void
 }
 
 ; FIXME: Ideally only one would be emitted
 ; CHECK-LABEL: {{^}}test_kill_depth_var_x2_same:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v0
-; CHECK-NEXT: s_cbranch_execnz BB4_2
-; CHECK:      exp null
-; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB4_2:
+; CHECK-NEXT: s_cbranch_execz BB4_3
+; CHECK-NEXT: ; %bb.1:
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v0
-; CHECK-NEXT: s_cbranch_execnz BB4_4
-; CHECK:      exp null
+; CHECK-NEXT: s_cbranch_execz BB4_3
+; CHECK-NEXT: ; %bb.2:
 ; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB4_4:
+; CHECK-NEXT: BB4_3:
+; CHECK:      exp null
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_var_x2_same(float %x) #0 {
   %cmp = fcmp olt float %x, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp)
   call void @llvm.amdgcn.kill(i1 %cmp)
   ret void
 }
 
 ; FIXME: Ideally only one early-exit would be emitted
 ; CHECK-LABEL: {{^}}test_kill_depth_var_x2:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v0
-; CHECK-NEXT: s_cbranch_execnz BB5_2
-; CHECK:      exp null
-; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB5_2:
+; CHECK-NEXT: s_cbranch_execz BB5_3
+; CHECK-NEXT: ; %bb.1
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v1
-; CHECK-NEXT: s_cbranch_execnz BB5_4
-; CHECK:      exp null
+; CHECK-NEXT: s_cbranch_execz BB5_3
+; CHECK-NEXT: ; %bb.2
 ; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB5_4:
+; CHECK-NEXT: BB5_3:
+; CHECK:      exp null
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_var_x2(float %x, float %y) #0 {
   %cmp.x = fcmp olt float %x, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp.x)
   %cmp.y = fcmp olt float %y, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp.y)
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_kill_depth_var_x2_instructions:
 ; CHECK-NEXT: ; %bb.0:
 ; CHECK-NEXT: v_cmpx_gt_f32_e32 vcc, 0, v0
-; CHECK-NEXT: s_cbranch_execnz BB6_2
+; CHECK-NEXT: s_cbranch_execz BB6_3
 ; CHECK-NEXT: ; %bb.1:
-; CHECK-NEXT: exp
-; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB6_2:
 ; CHECK: v_mov_b32_e64 v7, -1
 ; CHECK: v_cmpx_gt_f32_e32 vcc, 0, v7
-; CHECK-NEXT: s_cbranch_execnz BB6_4
-; CHECK-NEXT: ; %bb.3:
-; CHECK-NEXT: exp
+; CHECK-NEXT: s_cbranch_execz BB6_3
+; CHECK-NEXT: ; %bb.2:
 ; CHECK-NEXT: s_endpgm
-; CHECK-NEXT: BB6_4:
+; CHECK-NEXT: BB6_3:
+; CHECK-NEXT: exp null
 ; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @test_kill_depth_var_x2_instructions(float %x) #0 {
   %cmp.x = fcmp olt float %x, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp.x)
   %y = call float asm sideeffect "v_mov_b32_e64 v7, -1", "={v7}"()
   %cmp.y = fcmp olt float %y, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp.y)
   ret void
exit:
   store volatile i32 8, i32 addrspace(1)* undef
   ret void
 }
 
 ; bug 28550
 ; CHECK-LABEL: {{^}}phi_use_def_before_kill:
 ; CHECK: v_cndmask_b32_e64 [[PHIREG:v[0-9]+]], 0, -1.0,
 ; CHECK: v_cmpx_lt_f32_e32 vcc, 0,
-; CHECK-NEXT: s_cbranch_execnz [[BB4:BB[0-9]+_[0-9]+]]
-
-; CHECK: exp
-; CHECK-NEXT: s_endpgm
+; CHECK-NEXT: s_cbranch_execz [[EXITBB:BB[0-9]+_[0-9]+]]
 
-; CHECK: [[KILLBB:BB[0-9]+_[0-9]+]]:
+; CHECK: ; %[[KILLBB:bb.[0-9]+]]:
 ; CHECK-NEXT: s_cbranch_scc0 [[PHIBB:BB[0-9]+_[0-9]+]]
 
 ; CHECK: [[PHIBB]]:
 ; CHECK: v_cmp_eq_f32_e32 vcc, 0, [[PHIREG]]
 ; CHECK: s_cbranch_vccz [[ENDBB:BB[0-9]+_[0-9]+]]
 
 ; CHECK: ; %bb10
 ; CHECK: v_mov_b32_e32 v{{[0-9]+}}, 9
 ; CHECK: buffer_store_dword
 
 ; CHECK: [[ENDBB]]:
 ; CHECK-NEXT: s_endpgm
+
+; CHECK: [[EXITBB]]:
+; CHECK: exp null
+; CHECK-NEXT: s_endpgm
 define amdgpu_ps void @phi_use_def_before_kill(float inreg %x) #0 {
 bb:
   %tmp = fadd float %x, 1.000000e+00
   %tmp1 = fcmp olt float 0.000000e+00, %tmp
   %tmp2 = select i1 %tmp1, float -1.000000e+00, float 0.000000e+00
   %cmp.tmp2 = fcmp olt float %tmp2, 0.0
   call void @llvm.amdgcn.kill(i1 %cmp.tmp2)
   br i1 undef, label %phibb, label %bb8