[llvm] [AMDGPU] Sink async DMA out of s_cbranch_execz then-blocks (PR #196374)

[Vigneshwar Jayakumar via llvm-commits]

================
@@ -654,6 +663,128 @@ void SILowerControlFlow::optimizeEndCf() {
   }
 }
 
+static MachineInstr *findExecRestoreAtBlockStart(MachineBasicBlock &JoinBB,
+                                                 const SIRegisterInfo *TRI) {
+  for (MachineInstr &MI : JoinBB) {
+    if (MI.isMetaInstruction() || MI.isDebugInstr())
+      continue;
+    return MI.modifiesRegister(AMDGPU::EXEC, TRI) ? &MI : nullptr;
+  }
+  return nullptr;
+}
+
+bool SILowerControlFlow::sinkAsyncDMAOutOfExeczBlocks(MachineFunction &MF) {
+  bool Changed = false;
+
+  for (MachineBasicBlock &MBB : MF) {
+    // Header must end with an S_CBRANCH_EXECZ to JoinBB.
+    auto BrIt = llvm::find_if(MBB.terminators(), [](const MachineInstr &MI) {
+      return MI.getOpcode() == AMDGPU::S_CBRANCH_EXECZ;
+    });
+    if (BrIt == MBB.terminators().end())
+      continue;
+    MachineBasicBlock *JoinBB = BrIt->getOperand(0).getMBB();
+
+    if (MBB.succ_size() != 2)
+      continue;
+    MachineBasicBlock *S0 = *MBB.succ_begin();
+    MachineBasicBlock *S1 = *std::next(MBB.succ_begin());
+    MachineBasicBlock *ThenBB = (S0 == JoinBB) ? S1 : S0;
+    if (ThenBB == JoinBB || ThenBB->succ_size() != 1 ||
+        *ThenBB->succ_begin() != JoinBB)
+      continue;
+
+    // Sink before the EXEC restore at the top of JoinBB.
+    MachineInstr *ExecRestore = findExecRestoreAtBlockStart(*JoinBB, TRI);
+    if (!ExecRestore)
+      continue;
+
+    SmallVector<MachineInstr *, 4> ToSink;
+    bool Eligible = true;
+    for (MachineInstr &TMI : *ThenBB) {
+      if (TMI.isMetaInstruction() || TMI.isDebugInstr() || TMI.isTerminator())
+        continue;
+      if (TII->hasUnwantedEffectsWhenEXECEmpty(TMI)) {
+        Eligible = false;
+        break;
+      }
+      if (SIInstrInfo::usesASYNC_CNT(TMI)) {
+        ToSink.push_back(&TMI);
+        continue;
+      }
+      // Bail on stores that could alias the sunk DMAs.
+      if (TMI.mayStore()) {
+        Eligible = false;
+        break;
+      }
+      // Bail on non-invariant loads that could alias the sunk DMAs.
+      if (TMI.mayLoad() && !TMI.isDereferenceableInvariantLoad()) {
----------------
VigneshwarJ wrote:

when EXEC is zero, issuing the sunk DMA has no per-lane memory effect. It does not make reordering safe for waves where EXEC is nonzero.
If a non-invariant load follows an async load-to-LDS, sinking reverses their issue order. If the load may read the LDS destination written by the DMA

https://github.com/llvm/llvm-project/pull/196374