[AMDGPU] Add SI_EARLY_TERMINATE_SCC0 for early terminating shader

[critson]

Summary

Add pseudo instruction to allow early termination of pixel shader
anywhere based on the value of SCC. The intention is to use this
when a mask of live lanes is updated, e.g. live lanes in WQM pass.
This facilitates early termination of shaders even when EXEC is
incomplete, e.g. in non-uniform control flow.

[critson]

• Rebase

[critson]

• Remove restrictions on types of shader where early terminate can occur.

[arsenm]

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

394

const reference

[nhaehnle]

LGTM modulo the inline comment.

[critson]

• Address reviewer comments.
• Rebase.

[critson]

Rebase

[foad]

Patch looks fine to me. Why is it useful to terminate on scc0 instead of scc1, or is it an arbitrary choice? Could you give a slightly more realistic example of how it would be used? Your tests all have:

S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec

which will terminate early if sgpr0 is not 0. Is the idea that this would be used when we're just about to AND sgpr0 into exec?

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

511

Could erase here unconditionally, to avoid having an erase in earlyTerm as well.

llvm/lib/Target/AMDGPU/SIInstructions.td

364

Am I right in thinking that this is not a terminator, so it can appear in the middle of an MBB? Is that normal for this kind of pseudo?

[arsenm]

llvm/lib/Target/AMDGPU/SIInstructions.td

364

If this does any branching to any point in the program, this should be a terminator. If it behaves more like a trap/abort, it's OK if it's not a terminator

[critson]

In D88777#2486221, @foad wrote:

Patch looks fine to me. Why is it useful to terminate on scc0 instead of scc1, or is it an arbitrary choice? Could you give a slightly more realistic example of how it would be used? Your tests all have:

S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec

which will terminate early if sgpr0 is not 0. Is the idea that this would be used when we're just about to AND sgpr0 into exec?

The expectation is that this is used after updating a stored EXEC mask (not EXEC itself).
So SCC will be set from S_AND or S_ANDN2 in most cases.
SCC0 means mask is empty and program can be terminated.

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

511

Yes.

llvm/lib/Target/AMDGPU/SIInstructions.td

364

I think it is reasonable to consider this an abort. It always branches to terminate the program.

[critson]

• Address reviewer comments
• Change tests to use more indicative S_AND_B32

[foad]

LGTM.

llvm/lib/Target/AMDGPU/SIInsertSkips.cpp

bool shouldSkip(const MachineBasicBlock &From,
                   const MachineBasicBlock &To) const;
 
   bool dominatesAllReachable(MachineBasicBlock &MBB);
   void ensureEarlyExitBlock(MachineBasicBlock &MBB, bool ClearExec);
   void skipIfDead(MachineBasicBlock &MBB, MachineBasicBlock::iterator I,
                   DebugLoc DL);
 
   bool kill(MachineInstr &MI);
+  void earlyTerm(MachineInstr &MI);
 
   bool skipMaskBranch(MachineInstr &MI, MachineBasicBlock &MBB);
 
 public:
   static char ID;
 
   SIInsertSkips() : MachineFunctionPass(ID) {}
 
BuildMI(MBB, &MI, DL, TII->get(Opcode), Exec)
         .add(Op);
     return true;
   }
   default:
     llvm_unreachable("invalid opcode, expected SI_KILL_*_TERMINATOR");
   }
 }
 
+void SIInsertSkips::earlyTerm(MachineInstr &MI) {
+  MachineBasicBlock &MBB = *MI.getParent();
+  DebugLoc DL = MI.getDebugLoc();
+
+  ensureEarlyExitBlock(MBB, true);
+
+  auto BranchMI = BuildMI(MBB, MI, DL, TII->get(AMDGPU::S_CBRANCH_SCC0))
+                      .addMBB(EarlyExitBlock);
+  auto Next = std::next(MI.getIterator());
+
+  if (Next != MBB.end() && !Next->isTerminator())
+    splitBlock(MBB, *BranchMI, MDT);
+
+  MBB.addSuccessor(EarlyExitBlock);
+  MDT->getBase().insertEdge(&MBB, EarlyExitBlock);
+
+  MI.eraseFromParent();
+}
+
 // Returns true if a branch over the block was inserted.
 bool SIInsertSkips::skipMaskBranch(MachineInstr &MI,
                                    MachineBasicBlock &SrcMBB) {
   MachineBasicBlock *DestBB = MI.getOperand(0).getMBB();
 
   if (!shouldSkip(**SrcMBB.succ_begin(), *DestBB))
     return false;
 

 bool SIInsertSkips::runOnMachineFunction(MachineFunction &MF) {
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   TII = ST.getInstrInfo();
   TRI = &TII->getRegisterInfo();
   MDT = &getAnalysis<MachineDominatorTree>();
   SkipThreshold = SkipThresholdFlag;
 
   SmallVector<MachineInstr *, 4> KillInstrs;
+  SmallVector<MachineInstr *, 4> EarlyTermInstrs;
   bool MadeChange = false;
 
   for (MachineBasicBlock &MBB : MF) {
     MachineBasicBlock::iterator I, Next;
     for (I = MBB.begin(); I != MBB.end(); I = Next) {
       Next = std::next(I);
       MachineInstr &MI = *I;
 
for (I = MBB.begin(); I != MBB.end(); I = Next) {
         if (MF.getFunction().getCallingConv() == CallingConv::AMDGPU_PS &&
             dominatesAllReachable(MBB)) {
           KillInstrs.push_back(&MI);
         } else {
           MI.eraseFromParent();
         }
         break;
 
+      case AMDGPU::SI_EARLY_TERMINATE_SCC0:
+        assert(MF.getFunction().getCallingConv() == CallingConv::AMDGPU_PS);
+        EarlyTermInstrs.push_back(&MI);
+        break;
+
       default:
         break;
       }
     }
   }
 
+  for (MachineInstr *Instr : EarlyTermInstrs)
+    earlyTerm(*Instr);
   for (MachineInstr *Kill : KillInstrs) {
     skipIfDead(*Kill->getParent(), std::next(Kill->getIterator()),
                Kill->getDebugLoc());
     Kill->eraseFromParent();
   }
   KillInstrs.clear();
+  EarlyTermInstrs.clear();
   EarlyExitBlock = nullptr;
 
   return MadeChange;
 }

llvm/lib/Target/AMDGPU/SIInstructions.td

 
 def SI_IF_BREAK : CFPseudoInstSI <
   (outs SReg_1:$dst), (ins SReg_1:$vcc, SReg_1:$src), []> {
   let Size = 4;
   let isAsCheapAsAMove = 1;
   let isReMaterializable = 1;
 }
 
+// Branch to the early termination block of the shader if SCC is 0.
+// This uses SCC from a previous SALU operation, i.e. the update of
+// a mask of live lanes after a kill/demote operation.
+// Only valid in pixel shaders.
+def SI_EARLY_TERMINATE_SCC0 : SPseudoInstSI <(outs), (ins)> {
+  let Uses = [EXEC,SCC];
+}
+
 let Uses = [EXEC] in {
 
 multiclass PseudoInstKill <dag ins> {
   // Even though this pseudo can usually be expanded without an SCC def, we
   // conservatively assume that it has an SCC def, both because it is sometimes
   // required in degenerate cases (when V_CMPX cannot be used due to constant
   // bus limitations) and because it allows us to avoid having to track SCC
   // liveness across basic blocks.

llvm/test/CodeGen/AMDGPU/early-term.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx1010 -run-pass=si-insert-skips -verify-machineinstrs  %s -o - | FileCheck %s
+
+--- |
+  define amdgpu_ps void @early_term_scc0_end_block() {
+    ret void
+  }
+
+  define amdgpu_ps void @early_term_scc0_next_terminator() {
+    ret void
+  }
+
+  define amdgpu_ps void @early_term_scc0_in_block() {
+    ret void
+  }
+
+  define amdgpu_ps void @early_term_scc0_with_kill() {
+    ret void
+  }
+...
+---
+
+name: early_term_scc0_end_block
+tracksRegLiveness: true
+liveins:
+  - { reg: '$sgpr0' }
+body: |
+  ; CHECK-LABEL: name: early_term_scc0_end_block
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000), %bb.2(0x00000000)
+  ; CHECK:   liveins: $sgpr0
+  ; CHECK:   $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+  ; CHECK:   S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+  ; CHECK:   S_CBRANCH_SCC0 %bb.2, implicit $scc
+  ; CHECK: bb.1:
+  ; CHECK:   liveins: $vgpr0
+  ; CHECK:   EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  ; CHECK: bb.2:
+  ; CHECK:   $exec_lo = S_MOV_B32 0
+  ; CHECK:   EXP_DONE 9, undef $vgpr0, undef $vgpr0, undef $vgpr0, undef $vgpr0, 1, 0, 0, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  bb.0:
+    liveins: $sgpr0
+    successors: %bb.1
+
+    $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+    S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+    SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec
+
+  bb.1:
+    liveins: $vgpr0
+    EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+    S_ENDPGM 0
+...
+
+---
+
+
+name: early_term_scc0_next_terminator
+tracksRegLiveness: true
+liveins:
+  - { reg: '$sgpr0' }
+body: |
+  ; CHECK-LABEL: name: early_term_scc0_next_terminator
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.2(0x80000000), %bb.3(0x00000000)
+  ; CHECK:   liveins: $sgpr0
+  ; CHECK:   $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+  ; CHECK:   S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+  ; CHECK:   S_CBRANCH_SCC0 %bb.3, implicit $scc
+  ; CHECK:   S_BRANCH %bb.2
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.2(0x80000000)
+  ; CHECK:   $vgpr0 = V_MOV_B32_e32 1, implicit $exec
+  ; CHECK: bb.2:
+  ; CHECK:   liveins: $vgpr0
+  ; CHECK:   EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  ; CHECK: bb.3:
+  ; CHECK:   $exec_lo = S_MOV_B32 0
+  ; CHECK:   EXP_DONE 9, undef $vgpr0, undef $vgpr0, undef $vgpr0, undef $vgpr0, 1, 0, 0, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  bb.0:
+    liveins: $sgpr0
+    successors: %bb.2
+
+    $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+    S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+    SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec
+    S_BRANCH %bb.2
+
+  bb.1:
+    successors: %bb.2
+    $vgpr0 = V_MOV_B32_e32 1, implicit $exec
+    S_BRANCH %bb.2
+
+  bb.2:
+    liveins: $vgpr0
+    EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+    S_ENDPGM 0
+...
+
+---
+
+name: early_term_scc0_in_block
+tracksRegLiveness: true
+liveins:
+  - { reg: '$sgpr0' }
+body: |
+  ; CHECK-LABEL: name: early_term_scc0_in_block
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.3(0x40000000), %bb.2(0x40000000)
+  ; CHECK:   liveins: $sgpr0
+  ; CHECK:   $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+  ; CHECK:   S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+  ; CHECK:   S_CBRANCH_SCC0 %bb.2, implicit $scc
+  ; CHECK: bb.3:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $vgpr0, $scc
+  ; CHECK:   $vgpr1 = V_MOV_B32_e32 1, implicit $exec
+  ; CHECK: bb.1:
+  ; CHECK:   liveins: $vgpr0, $vgpr1
+  ; CHECK:   EXP 1, $vgpr1, $vgpr1, $vgpr1, $vgpr1, -1, -1, 15, implicit $exec
+  ; CHECK:   EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  ; CHECK: bb.2:
+  ; CHECK:   $exec_lo = S_MOV_B32 0
+  ; CHECK:   EXP_DONE 9, undef $vgpr0, undef $vgpr0, undef $vgpr0, undef $vgpr0, 1, 0, 0, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  bb.0:
+    liveins: $sgpr0
+    successors: %bb.1
+
+    $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+    S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+    SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec
+    $vgpr1 = V_MOV_B32_e32 1, implicit $exec
+
+  bb.1:
+    liveins: $vgpr0, $vgpr1
+    EXP 1, $vgpr1, $vgpr1, $vgpr1, $vgpr1, -1, -1, 15, implicit $exec
+    EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+    S_ENDPGM 0
+...
+
+---
+
+name: early_term_scc0_with_kill
+tracksRegLiveness: true
+liveins:
+  - { reg: '$sgpr0' }
+  - { reg: '$vgpr2' }
+body: |
+  ; CHECK-LABEL: name: early_term_scc0_with_kill
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000), %bb.3(0x00000000)
+  ; CHECK:   liveins: $sgpr0, $vgpr2
+  ; CHECK:   $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+  ; CHECK:   V_CMPX_LE_F32_nosdst_e32 0, killed $vgpr2, implicit-def $exec, implicit $mode, implicit $exec
+  ; CHECK:   S_CBRANCH_EXECZ %bb.3, implicit $exec
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.4(0x40000000), %bb.3(0x40000000)
+  ; CHECK:   liveins: $sgpr0, $vgpr0
+  ; CHECK:   S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+  ; CHECK:   S_CBRANCH_SCC0 %bb.3, implicit $scc
+  ; CHECK: bb.4:
+  ; CHECK:   successors: %bb.2(0x80000000)
+  ; CHECK:   liveins: $vgpr0, $scc
+  ; CHECK:   $vgpr1 = V_MOV_B32_e32 1, implicit $exec
+  ; CHECK: bb.2:
+  ; CHECK:   liveins: $vgpr0, $vgpr1
+  ; CHECK:   EXP 1, $vgpr1, $vgpr1, $vgpr1, $vgpr1, -1, -1, 15, implicit $exec
+  ; CHECK:   EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  ; CHECK: bb.3:
+  ; CHECK:   $exec_lo = S_MOV_B32 0
+  ; CHECK:   EXP_DONE 9, undef $vgpr0, undef $vgpr0, undef $vgpr0, undef $vgpr0, 1, 0, 0, implicit $exec
+  ; CHECK:   S_ENDPGM 0
+  bb.0:
+    liveins: $sgpr0, $vgpr2
+    successors: %bb.1
+    $vgpr0 = V_MOV_B32_e32 0, implicit $exec
+    SI_KILL_F32_COND_IMM_TERMINATOR killed $vgpr2, 0, 3, implicit-def $exec, implicit-def $vcc, implicit-def $scc, implicit $exec
+
+  bb.1:
+    liveins: $sgpr0, $vgpr0
+    successors: %bb.2
+    S_CMP_EQ_U32 killed $sgpr0, 0, implicit-def $scc
+    SI_EARLY_TERMINATE_SCC0 implicit $scc, implicit $exec
+    $vgpr1 = V_MOV_B32_e32 1, implicit $exec
+
+  bb.2:
+    liveins: $vgpr0, $vgpr1
+    EXP 1, $vgpr1, $vgpr1, $vgpr1, $vgpr1, -1, -1, 15, implicit $exec
+    EXP_DONE 0, $vgpr0, $vgpr0, $vgpr0, $vgpr0, -1, -1, 15, implicit $exec
+    S_ENDPGM 0
+...