[AMDGPU] Skip non-first termintors when forcing emit zero flag

Author: shiltian

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

@@ -1824,7 +1824,14 @@ bool SIInsertWaitcnts::generateWaitcntInstBefore(MachineInstr &MI,
   // Verify that the wait is actually needed.
   ScoreBrackets.simplifyWaitcnt(Wait);
 
-  if (ForceEmitZeroFlag)
+  // When forcing emit, we need to skip non-first terminators of a MBB because
+  // that would break the terminators of the MBB.
+  auto CheckIfMBBNonFirstTerminators = [](MachineInstr &MI) {
+    if (!MI.isTerminator())
+      return false;
+    return MI.getParent()->getFirstTerminator() != &MI;
+  };
+  if (ForceEmitZeroFlag && !CheckIfMBBNonFirstTerminators(MI))
     Wait = WCG->getAllZeroWaitcnt(/*IncludeVSCnt=*/false);
 
   if (ForceEmitWaitcnt[LOAD_CNT])

llvm/test/CodeGen/AMDGPU/waitcnt-debug-non-first-terminators.mir

@@ -0,0 +1,22 @@
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass si-insert-waitcnts -amdgpu-waitcnt-forcezero=1 %s -o - | FileCheck %s
+
+...
+
+# CHECK-LABEL: waitcnt-debug-non-first-terminators
+# CHECK: S_WAITCNT 0
+# CHECK-NEXT: S_CBRANCH_SCC1 %bb.1, implicit $scc
+# CHECK-NEXT: S_BRANCH %bb.2, implicit $scc
+
+name: waitcnt-debug-non-first-terminators
+liveins:
+machineFunctionInfo:
+  isEntryFunction: true
+body:             |
+  bb.0:
+    S_CBRANCH_SCC1 %bb.1, implicit $scc
+    S_BRANCH %bb.2, implicit $scc
+  bb.1:
+    S_NOP 0
+  bb.2:
+    S_NOP 0
+...