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

[Matt Arsenault via llvm-commits]

================
@@ -0,0 +1,229 @@
+//===-- SISinkAsyncDMA.cpp - Sink async DMA out of execz then-blocks ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// LLVM lowers a divergent branch around global_load_async_to_lds /
+/// global_store_async_from_lds with an S_CBRANCH_EXECZ, so fully-masked waves
+/// skip the DMA entirely and the ASYNCcnt observed at the join depends on
+/// whether the wave took the branch. Software-pipelined kernels then have to
+/// use a conservative async waitcnt.
+///
+/// This pass sinks each such DMA into the join, immediately before SI_ELSE or
+/// SI_END_CF:
+///
+///        MBB                     MBB        SI_IF sets EXEC to the then-block
+///       /   \                   /   \       mask before the branch, so both
+///   ThenBB   |               ThenBB  |      edges carry it and per-lane
+///    [DMA]   |      ==>          \  /       behavior is unchanged. But every
+///       \   /                   JoinBB      wave now issues the DMA, so
+///      JoinBB                    [DMA]      ASYNCcnt at the join no longer
+///    [SI_END_CF]                   |        depends on the branch.
+///                             [SI_END_CF]
+///
+/// The join is split so that SI_END_CF starts a block of its own, because
+/// SILowerControlFlow emits the EXEC restore at the top of the block holding
+/// it, which would otherwise place it above the sunk DMAs.
+
+//
+//===----------------------------------------------------------------------===//
+
+#include "SISinkAsyncDMA.h"
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIInstrInfo.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineSSAUpdater.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "si-sink-async-dma"
+
+namespace {
+
+class SISinkAsyncDMA {
+  const SIInstrInfo *TII = nullptr;
+  const SIRegisterInfo *TRI = nullptr;
+  MachineRegisterInfo *MRI = nullptr;
+
+  bool sinkFromBlock(MachineBasicBlock &MBB);
+
+public:
+  bool run(MachineFunction &MF);
+};
+
+class SISinkAsyncDMALegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  SISinkAsyncDMALegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override {
+    return "SI sink async DMA out of execz then-blocks";
+  }
+
+  MachineFunctionProperties getRequiredProperties() const override {
+    return MachineFunctionProperties().setIsSSA();
+  }
+
+  MachineFunctionProperties getClearedProperties() const override {
+    return MachineFunctionProperties().setNoPHIs();
+  }
+};
+
+} // namespace
+
+char SISinkAsyncDMALegacy::ID = 0;
+
+INITIALIZE_PASS(SISinkAsyncDMALegacy, DEBUG_TYPE,
+                "SI sink async DMA out of execz then-blocks", false, false)
+
+char &llvm::SISinkAsyncDMALegacyID = SISinkAsyncDMALegacy::ID;
+
+static bool isAsyncDMA(const MachineInstr &MI) {
+  return SIInstrInfo::isLDSDMA(MI) && SIInstrInfo::usesASYNC_CNT(MI);
+}
+
+static bool isAsyncMarker(const MachineInstr &MI) {
+  return MI.getOpcode() == AMDGPU::ASYNCMARK ||
+         MI.getOpcode() == AMDGPU::WAIT_ASYNCMARK;
+}
+
+bool SISinkAsyncDMA::sinkFromBlock(MachineBasicBlock &MBB) {
+  if (MBB.succ_size() != 2)
+    return false;
+
+  // A region head ends in SI_IF or SI_ELSE ($dst, $cond, $target), which define
+  // the mask the region restores in $dst and the join block in $target.
+  auto ControlMI = MBB.getFirstTerminator();
+  if (ControlMI == MBB.end() || (ControlMI->getOpcode() != AMDGPU::SI_IF &&
+                                 ControlMI->getOpcode() != AMDGPU::SI_ELSE))
+    return false;
+
+  Register SavedExec = ControlMI->getOperand(0).getReg();
+  MachineBasicBlock *JoinBB = ControlMI->getOperand(2).getMBB();
+
+  if (!MBB.isSuccessor(JoinBB) || JoinBB->pred_size() != 2)
+    return false;
+
+  auto ThenIt = find_if(MBB.successors(),
+                        [JoinBB](MachineBasicBlock *S) { return S != JoinBB; });
+  if (ThenIt == MBB.succ_end() || (*ThenIt)->getSingleSuccessor() != JoinBB)
+    return false;
+  MachineBasicBlock *ThenBB = *ThenIt;
+
+  auto Boundary = JoinBB->getFirstNonPHI();
+  while (Boundary != JoinBB->end() && Boundary->isMetaInstruction())
+    ++Boundary;
+  if (Boundary == JoinBB->end())
+    return false;
----------------
arsenm wrote:

Manual control flow parsing is annoying, can this use AnalyzeBranch?

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