[AMDGPU] SIRemoveShortExecBranches should not remove branches exiting loops

[critson]

Summary

Check that a s_cbranch_execz is not a loop exit before removing it.
As the pass is generating infinite loops.

[cdevadas]

You are right, the s_cbranch_execz shouldn't be removed from this loop exit BB.

But depending on the MachineLoopInfo to fix it is not good in general. There are irreducible control flows and the check may fail.

There is already a special check in function mustRetainExeczBranch, looking for loop code if the opcode is S_CBRANCH_VCCNZ & S_CBRANCH_VCCZ.
We should include S_CBRANCH_SCC0 & S_CBRANCH_SCC1 here. I hope this guarantees to be working for such cases.

If everyone agrees, I can create a patch and fix it.

[critson]

In D72997#1829191, @cdevadas wrote:

There is already a special check in function mustRetainExeczBranch, looking for loop code if the opcode is S_CBRANCH_VCCNZ & S_CBRANCH_VCCZ.
We should include S_CBRANCH_SCC0 & S_CBRANCH_SCC1 here. I hope this guarantees to be working for such cases.

That would work for my test case (essentially a slightly more complex version of lit test).
I am happy with your proposed change, if others agree then I will abandon this one.

[arsenm]

In D72997#1829191, @cdevadas wrote:

You are right, the s_cbranch_execz shouldn't be removed from this loop exit BB.

But depending on the MachineLoopInfo to fix it is not good in general. There are irreducible control flows and the check may fail.

There is already a special check in function mustRetainExeczBranch, looking for loop code if the opcode is S_CBRANCH_VCCNZ & S_CBRANCH_VCCZ.
We should include S_CBRANCH_SCC0 & S_CBRANCH_SCC1 here. I hope this guarantees to be working for such cases.

If everyone agrees, I can create a patch and fix it.

This is probably sufficient

[arsenm]

In D72997#1830227, @arsenm wrote:

In D72997#1829191, @cdevadas wrote:

You are right, the s_cbranch_execz shouldn't be removed from this loop exit BB.

But depending on the MachineLoopInfo to fix it is not good in general. There are irreducible control flows and the check may fail.

There is already a special check in function mustRetainExeczBranch, looking for loop code if the opcode is S_CBRANCH_VCCNZ & S_CBRANCH_VCCZ.
We should include S_CBRANCH_SCC0 & S_CBRANCH_SCC1 here. I hope this guarantees to be working for such cases.

If everyone agrees, I can create a patch and fix it.

This is probably sufficient

Also needs to catch execz/execnz/ Maybe this should juts check isBranch?

[critson]

Update patch to current state of discussion.

[arsenm]

llvm/lib/Target/AMDGPU/SIRemoveShortExecBranches.cpp

110

Special casing execnz just doesn't make sense to me. execz/execnz are morally equivalent

[critson]

Remove C_BRANCH_EXECNZ special case.
This revives a c_branch_execz in one lit test.

Alternatively we special case both exec_nz/exec_z.
However this ceases to feel much better than simply stating all the accepted branch types.

[nhaehnle]

LGTM. You'll have to re-submit D68092 as well then.

llvm/lib/Target/AMDGPU/SIRemoveShortExecBranches.cpp

 ///
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIInstrInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/InitializePasses.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "si-remove-short-exec-branches"
 
 static unsigned SkipThreshold;
 
 static cl::opt<unsigned, true> SkipThresholdFlag(
     "amdgpu-skip-threshold", cl::Hidden,
     cl::desc(
         "Number of instructions before jumping over divergent control flow"),
     cl::location(SkipThreshold), cl::init(12));
 
 namespace {
 
 class SIRemoveShortExecBranches : public MachineFunctionPass {
 private:
   const SIInstrInfo *TII = nullptr;
+  MachineLoopInfo *MLI = nullptr;
+
   bool getBlockDestinations(MachineBasicBlock &SrcMBB,
                             MachineBasicBlock *&TrueMBB,
                             MachineBasicBlock *&FalseMBB,
                             SmallVectorImpl<MachineOperand> &Cond);
   bool mustRetainExeczBranch(const MachineBasicBlock &From,
                              const MachineBasicBlock &To) const;
   bool removeExeczBranch(MachineInstr &MI, MachineBasicBlock &SrcMBB);
 
 public:
   static char ID;
 
   SIRemoveShortExecBranches() : MachineFunctionPass(ID) {
     initializeSIRemoveShortExecBranchesPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineLoopInfo>();
+    AU.addPreserved<MachineLoopInfo>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
 };
 
 } // End anonymous namespace.
 
-INITIALIZE_PASS(SIRemoveShortExecBranches, DEBUG_TYPE,
+INITIALIZE_PASS_BEGIN(SIRemoveShortExecBranches, DEBUG_TYPE,
+                "SI remove short exec branches", false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
+INITIALIZE_PASS_END(SIRemoveShortExecBranches, DEBUG_TYPE,
                 "SI remove short exec branches", false, false)
 
 char SIRemoveShortExecBranches::ID = 0;
 
 char &llvm::SIRemoveShortExecBranchesID = SIRemoveShortExecBranches::ID;
 
 bool SIRemoveShortExecBranches::getBlockDestinations(
     MachineBasicBlock &SrcMBB, MachineBasicBlock *&TrueMBB,
for (MachineBasicBlock::const_iterator I = MBB.begin(), E = MBB.end();
          I != E; ++I) {
       // When a uniform loop is inside non-uniform control flow, the branch
       // leaving the loop might be an S_CBRANCH_VCCNZ, which is never taken
       // when EXEC = 0. We should skip the loop lest it becomes infinite.
       if (I->getOpcode() == AMDGPU::S_CBRANCH_VCCNZ ||
           I->getOpcode() == AMDGPU::S_CBRANCH_VCCZ)
         return true;
 
       if (TII->hasUnwantedEffectsWhenEXECEmpty(*I))
         return true;
 
       // These instructions are potentially expensive even if EXEC = 0.
       if (TII->isSMRD(*I) || TII->isVMEM(*I) || TII->isFLAT(*I) ||
           I->getOpcode() == AMDGPU::S_WAITCNT)
         return true;
 
       ++NumInstr;
bool SIRemoveShortExecBranches::removeExeczBranch(MachineInstr &MI,
   MachineBasicBlock *FalseMBB = nullptr;
   SmallVector<MachineOperand, 1> Cond;
 
   if (!getBlockDestinations(SrcMBB, TrueMBB, FalseMBB, Cond))
     return false;
 
   // Consider only the forward branches.
   if ((SrcMBB.getNumber() >= TrueMBB->getNumber()) ||
+      (MLI->getLoopFor(&SrcMBB) != MLI->getLoopFor(TrueMBB)) ||
       mustRetainExeczBranch(*FalseMBB, *TrueMBB))
     return false;
 
   LLVM_DEBUG(dbgs() << "Removing the execz branch: " << MI);
   MI.eraseFromParent();
   SrcMBB.removeSuccessor(TrueMBB);
 
   return true;
 }
 
 bool SIRemoveShortExecBranches::runOnMachineFunction(MachineFunction &MF) {
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   TII = ST.getInstrInfo();
+  MLI = &getAnalysis<MachineLoopInfo>();
   MF.RenumberBlocks();
   bool Changed = false;
 
   for (MachineBasicBlock &MBB : MF) {
     MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
     if (MBBI == MBB.end())
       continue;
 

llvm/test/CodeGen/AMDGPU/divergent-branch-uniform-condition.ll

 ; CHECK-NEXT:  BB0_2: ; %Flow
 ; CHECK-NEXT:    ; in Loop: Header=BB0_3 Depth=1
 ; CHECK-NEXT:    s_and_b64 s[10:11], exec, s[6:7]
 ; CHECK-NEXT:    s_or_b64 s[2:3], s[10:11], s[2:3]
 ; CHECK-NEXT:    s_andn2_b64 s[4:5], s[4:5], exec
 ; CHECK-NEXT:    s_and_b64 s[8:9], s[8:9], exec
 ; CHECK-NEXT:    s_or_b64 s[4:5], s[4:5], s[8:9]
 ; CHECK-NEXT:    s_andn2_b64 exec, exec, s[2:3]
+; CHECK-NEXT:    s_cbranch_execz BB0_6
 ; CHECK-NEXT:  BB0_3: ; %loop
 ; CHECK-NEXT:    ; =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    s_or_b64 s[6:7], s[6:7], exec
 ; CHECK-NEXT:    s_cmp_lt_u32 s0, 32
 ; CHECK-NEXT:    s_mov_b64 s[8:9], -1
 ; CHECK-NEXT:    s_cbranch_scc0 BB0_2
 ; CHECK-NEXT:  ; %bb.4: ; %endif1
 ; CHECK-NEXT:    ; in Loop: Header=BB0_3 Depth=1
 ; CHECK-NEXT:    s_mov_b64 s[6:7], -1
 ; CHECK-NEXT:    s_and_saveexec_b64 s[8:9], vcc
 ; CHECK-NEXT:    s_xor_b64 s[8:9], exec, s[8:9]
 ; CHECK-NEXT:    s_cbranch_execz BB0_1
 ; CHECK-NEXT:  ; %bb.5: ; %endif2
 ; CHECK-NEXT:    ; in Loop: Header=BB0_3 Depth=1
 ; CHECK-NEXT:    s_add_i32 s0, s0, 1
 ; CHECK-NEXT:    s_xor_b64 s[6:7], exec, -1
 ; CHECK-NEXT:    s_branch BB0_1
-; CHECK-NEXT:  ; %bb.6: ; %Flow2
+; CHECK-NEXT:  BB0_6: ; %Flow2
 ; CHECK-NEXT:    s_or_b64 exec, exec, s[2:3]
 ; CHECK-NEXT:    v_mov_b32_e32 v1, 0
 ; CHECK-NEXT:    s_and_saveexec_b64 s[0:1], s[4:5]
 ; CHECK-NEXT:  ; %bb.7: ; %if1
 ; CHECK-NEXT:    v_sqrt_f32_e32 v1, v0
 ; CHECK-NEXT:  ; %bb.8: ; %endloop
 ; CHECK-NEXT:    s_or_b64 exec, exec, s[0:1]
 ; CHECK-NEXT:    exp mrt0 v1, v1, v1, v1 done vm