AMDGPU: Introduce a pass to replace VGPR MFMAs with AGPR

In gfx90a-gfx950, it's possible to emit MFMAs which use AGPRs or VGPRs
for vdst and src2. We do not want to do use the AGPR form, unless
required by register pressure as it requires cross bank register
copies from most other instructions. Currently we select the AGPR
or VGPR version depending on a crude heuristic for whether it's possible
AGPRs will be required. We really need the register allocation to
be complete to make a good decision, which is what this pass is for.

This adds the pass, but does not yet remove the selection patterns
for AGPRs. This is a WIP, and NFC-ish. It should be a no-op on any
currently selected code. It also does not yet trigger on the real
examples of interest, which require handling batches of MFMAs at
once.

Author: arsenm

llvm/lib/Target/AMDGPU/AMDGPU.h

@@ -541,6 +541,17 @@ extern char &GCNRewritePartialRegUsesID;
 void initializeAMDGPUWaitSGPRHazardsLegacyPass(PassRegistry &);
 extern char &AMDGPUWaitSGPRHazardsLegacyID;
 
+class AMDGPURewriteAGPRCopyMFMAPass
+    : public PassInfoMixin<AMDGPURewriteAGPRCopyMFMAPass> {
+public:
+  AMDGPURewriteAGPRCopyMFMAPass() = default;
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+void initializeAMDGPURewriteAGPRCopyMFMALegacyPass(PassRegistry &);
+extern char &AMDGPURewriteAGPRCopyMFMALegacyID;
+
 namespace AMDGPU {
 enum TargetIndex {
   TI_CONSTDATA_START,

llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def

@@ -102,6 +102,7 @@ MACHINE_FUNCTION_PASS("amdgpu-isel", AMDGPUISelDAGToDAGPass(*this))
 MACHINE_FUNCTION_PASS("amdgpu-mark-last-scratch-load", AMDGPUMarkLastScratchLoadPass())
 MACHINE_FUNCTION_PASS("amdgpu-pre-ra-long-branch-reg", GCNPreRALongBranchRegPass())
 MACHINE_FUNCTION_PASS("amdgpu-reserve-wwm-regs", AMDGPUReserveWWMRegsPass())
+MACHINE_FUNCTION_PASS("amdgpu-rewrite-agpr-copy-mfma", AMDGPURewriteAGPRCopyMFMAPass())
 MACHINE_FUNCTION_PASS("amdgpu-rewrite-partial-reg-uses", GCNRewritePartialRegUsesPass())
 MACHINE_FUNCTION_PASS("amdgpu-set-wave-priority", AMDGPUSetWavePriorityPass())
 MACHINE_FUNCTION_PASS("amdgpu-pre-ra-optimizations", GCNPreRAOptimizationsPass())

llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp

@@ -0,0 +1,300 @@
+//===-- AMDGPURewriteAGPRCopyMFMA.cpp -------------------------------------===//
+//
+// 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 \brief Try to replace MFMA instructions using VGPRs with MFMA
+/// instructions using AGPRs. We expect MFMAs to be selected using VGPRs, and
+/// only use AGPRs if it helps avoid spilling. In this case, the MFMA will have
+/// copies between AGPRs and VGPRs and the AGPR variant of an MFMA pseudo. This
+/// pass will attempt to delete the cross register bank copy and replace the
+/// MFMA opcode.
+///
+/// TODO:
+///  - Handle non-tied dst+src2 cases. We need to try to find a copy from an
+///    AGPR from src2, or reassign src2 to an available AGPR (which should work
+///    in the common case of a load).
+///
+///  - Handle multiple MFMA uses of the same register. e.g. chained MFMAs that
+///    can be rewritten as a set
+///
+///  - Update LiveIntervals incrementally instead of recomputing from scratch
+///
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "SIMachineFunctionInfo.h"
+#include "SIRegisterInfo.h"
+#include "llvm/CodeGen/LiveIntervals.h"
+#include "llvm/CodeGen/LiveRegMatrix.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/VirtRegMap.h"
+#include "llvm/InitializePasses.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-rewrite-agpr-copy-mfma"
+
+namespace {
+
+class AMDGPURewriteAGPRCopyMFMAImpl {
+  const GCNSubtarget &ST;
+  const SIInstrInfo &TII;
+  const SIRegisterInfo &TRI;
+  MachineRegisterInfo &MRI;
+  VirtRegMap &VRM;
+  LiveRegMatrix ‎
+  LiveIntervals &LIS;
+
+public:
+  AMDGPURewriteAGPRCopyMFMAImpl(MachineFunction &MF, VirtRegMap &VRM,
+                                LiveRegMatrix &LRM, LiveIntervals &LIS)
+      : ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()),
+        TRI(*ST.getRegisterInfo()), MRI(MF.getRegInfo()), VRM(VRM), LRM(LRM),
+        LIS(LIS) {}
+
+  /// Compute the register class constraints based on the uses of \p Reg,
+  /// excluding uses from \p ExceptMI. This should be nearly identical to
+  /// MachineRegisterInfo::recomputeRegClass.
+  const TargetRegisterClass *
+  recomputeRegClassExcept(Register Reg, const TargetRegisterClass *OldRC,
+                          const TargetRegisterClass *NewRC,
+                          const MachineInstr *ExceptMI) const;
+
+  bool run(MachineFunction &MF) const;
+};
+
+const TargetRegisterClass *
+AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExcept(
+    Register Reg, const TargetRegisterClass *OldRC,
+    const TargetRegisterClass *NewRC, const MachineInstr *ExceptMI) const {
+
+  // Accumulate constraints from all uses.
+  for (MachineOperand &MO : MRI.reg_nodbg_operands(Reg)) {
+    // Apply the effect of the given operand to NewRC.
+    MachineInstr *MI = MO.getParent();
+    if (MI == ExceptMI)
+      continue;
+
+    unsigned OpNo = &MO - &MI->getOperand(0);
+    NewRC = MI->getRegClassConstraintEffect(OpNo, NewRC, &TII, &TRI);
+    if (!NewRC || NewRC == OldRC)
+      return nullptr;
+  }
+
+  return NewRC;
+}
+
+bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
+  // This only applies on subtargets that have a configurable AGPR vs. VGPR
+  // allocation.
+  if (!ST.hasGFX90AInsts())
+    return false;
+
+  // Early exit if no AGPRs were assigned.
+  if (!LRM.isPhysRegUsed(AMDGPU::AGPR0))
+    return false;
+
+  bool MadeChange = false;
+
+  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
+    Register VReg = Register::index2VirtReg(I);
+    Register PhysReg = VRM.getPhys(VReg);
+    if (!PhysReg)
+      continue;
+
+    // Find AV_* registers assigned to AGPRs.
+    const TargetRegisterClass *VirtRegRC = MRI.getRegClass(VReg);
+    if (!TRI.isVectorSuperClass(VirtRegRC))
+      continue;
+
+    const TargetRegisterClass *AssignedRC = TRI.getPhysRegBaseClass(PhysReg);
+    if (!TRI.isAGPRClass(AssignedRC))
+      continue;
+
+    LiveInterval &LI = LIS.getInterval(VReg);
+
+    // TODO: Test multiple uses
+    for (VNInfo *VNI : LI.vnis()) {
+      MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
+
+      // TODO: Handle SplitKit produced copy bundles for partially defined
+      // registers.
+      if (!DefMI || !DefMI->isFullCopy())
+        continue;
+
+      Register CopySrcReg = DefMI->getOperand(1).getReg();
+      LiveInterval &CopySrcLI = LIS.getInterval(CopySrcReg);
+
+      LiveQueryResult LRQ = CopySrcLI.Query(VNI->def.getRegSlot());
+
+      MachineInstr *CopySrcMI = LIS.getInstructionFromIndex(LRQ.valueIn()->def);
+      int AGPROp = AMDGPU::getMFMASrcCVDstAGPROp(CopySrcMI->getOpcode());
+      if (AGPROp == -1)
+        continue;
+
+      MachineOperand *Src2 =
+          TII.getNamedOperand(*CopySrcMI, AMDGPU::OpName::src2);
+
+      // FIXME: getMinimalPhysRegClass returns a nonsense AV_* subclass instead
+      // of an AGPR or VGPR subclass, so we can't simply use the result on the
+      // assignment.
+
+      LLVM_DEBUG({
+        Register Src2PhysReg = VRM.getPhys(Src2->getReg());
+        dbgs() << "Attempting to replace VGPR MFMA with AGPR version:"
+               << " Dst=[" << printReg(VReg) << " => "
+               << printReg(PhysReg, &TRI) << "], Src2=["
+               << printReg(Src2->getReg(), &TRI) << " => "
+               << printReg(Src2PhysReg, &TRI) << "]: " << *CopySrcMI;
+      });
+
+      // If the inputs are tied and the same register, we can shortcut and
+      // directly replace the register.
+      if (Src2->getReg() != CopySrcReg) {
+        LLVM_DEBUG(
+            dbgs()
+            << "Replacing untied VGPR MFMAs with AGPR form not yet handled\n");
+        // TODO: Only handles the tied case for now. If the input operand is a
+        // different register, we need to also reassign it (either by looking
+        // for a compatible copy-from-AGPR, or by seeing if an available AGPR is
+        // compatible with all other uses.
+
+        // If we can't reassign it, we'd need to introduce a different copy
+        // which is likely worse than the copy we'd be saving.
+        continue;
+      }
+
+      const TargetRegisterClass *Src2VirtRegRC =
+          MRI.getRegClass(Src2->getReg());
+
+      // We've found av = COPY (MFMA), and need to verify that we can trivially
+      // rewrite src2 to use the new AGPR. If we can't trivially replace it,
+      // we're going to induce as many copies as we would have emitted in the
+      // first place, as well as need to assign another register, and need to
+      // figure out where to put them. The live range splitting is smarter than
+      // anything we're doing here, so trust it did something reasonable.
+      const TargetRegisterClass *Src2ExceptRC = recomputeRegClassExcept(
+          Src2->getReg(), Src2VirtRegRC, VirtRegRC, CopySrcMI);
+      if (!Src2ExceptRC)
+        continue;
+
+      const TargetRegisterClass *NewSrc2ConstraintRC =
+          TII.getRegClass(TII.get(AGPROp), Src2->getOperandNo(), &TRI, MF);
+
+      // Try to constrain src2 to the replacement instruction candidate's
+      // register class.
+      const TargetRegisterClass *NewSrc2RC =
+          TRI.getCommonSubClass(Src2ExceptRC, NewSrc2ConstraintRC);
+      if (!NewSrc2RC) {
+        // TODO: This is ignoring ther rewritable uses. e.g. a rewritable MFMA
+        // using a rewritable MFMA can be rewritten as a pair.
+        LLVM_DEBUG(dbgs() << "Other uses of " << printReg(Src2->getReg(), &TRI)
+                          << " are incompatible with replacement class\n");
+        continue;
+      }
+
+      MRI.setRegClass(VReg, AssignedRC);
+      MRI.setRegClass(Src2->getReg(), NewSrc2RC);
+
+      CopySrcMI->setDesc(TII.get(AGPROp));
+
+      // TODO: Is replacing too aggressive, fixup these instructions only?
+      MRI.replaceRegWith(CopySrcReg, VReg);
+
+      LLVM_DEBUG(dbgs() << "Replaced VGPR MFMA with AGPR: " << *CopySrcMI);
+
+      // We left behind an identity copy, so delete it.
+      LIS.RemoveMachineInstrFromMaps(*DefMI);
+      DefMI->eraseFromParent();
+
+      LRM.unassign(CopySrcLI);
+
+      // We don't need the liveness information anymore, so don't bother
+      // updating the intervals. Just delete the stale information.
+      // TODO: Is it worth preserving these?
+      LIS.removeInterval(CopySrcReg);
+      LIS.removeInterval(VReg);
+      LIS.createAndComputeVirtRegInterval(VReg);
+
+      MadeChange = true;
+    }
+  }
+
+  return MadeChange;
+}
+
+class AMDGPURewriteAGPRCopyMFMALegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  AMDGPURewriteAGPRCopyMFMALegacy() : MachineFunctionPass(ID) {
+    initializeAMDGPURewriteAGPRCopyMFMALegacyPass(
+        *PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override {
+    return "AMDGPU Rewrite AGPR-Copy-MFMA";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<LiveIntervalsWrapperPass>();
+    AU.addRequired<VirtRegMapWrapperLegacy>();
+    AU.addRequired<LiveRegMatrixWrapperLegacy>();
+
+    AU.addPreserved<LiveIntervalsWrapperPass>();
+    AU.addPreserved<VirtRegMapWrapperLegacy>();
+    AU.addPreserved<LiveRegMatrixWrapperLegacy>();
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+} // End anonymous namespace.
+
+INITIALIZE_PASS_BEGIN(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE,
+                      "AMDGPU Rewrite AGPR-Copy-MFMA", false, false)
+INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
+INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
+INITIALIZE_PASS_END(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE,
+                    "AMDGPU Rewrite AGPR-Copy-MFMA", false, false)
+
+char AMDGPURewriteAGPRCopyMFMALegacy::ID = 0;
+
+char &llvm::AMDGPURewriteAGPRCopyMFMALegacyID =
+    AMDGPURewriteAGPRCopyMFMALegacy::ID;
+
+bool AMDGPURewriteAGPRCopyMFMALegacy::runOnMachineFunction(
+    MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  auto &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
+  auto &LRM = getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
+  auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
+
+  AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS);
+  return Impl.run(MF);
+}
+
+PreservedAnalyses
+AMDGPURewriteAGPRCopyMFMAPass::run(MachineFunction &MF,
+                                   MachineFunctionAnalysisManager &MFAM) {
+  VirtRegMap &VRM = MFAM.getResult<VirtRegMapAnalysis>(MF);
+  LiveRegMatrix &LRM = MFAM.getResult<LiveRegMatrixAnalysis>(MF);
+  LiveIntervals &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
+
+  AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS);
+  if (!Impl.run(MF))
+    return PreservedAnalyses::all();
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -535,6 +535,7 @@ extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeAMDGPULowerModuleLDSLegacyPass(*PR);
   initializeAMDGPULowerBufferFatPointersPass(*PR);
   initializeAMDGPUReserveWWMRegsLegacyPass(*PR);
+  initializeAMDGPURewriteAGPRCopyMFMALegacyPass(*PR);
   initializeAMDGPURewriteOutArgumentsPass(*PR);
   initializeAMDGPURewriteUndefForPHILegacyPass(*PR);
   initializeSIAnnotateControlFlowLegacyPass(*PR);
@@ -1578,6 +1579,8 @@ void GCNPassConfig::addOptimizedRegAlloc() {
 bool GCNPassConfig::addPreRewrite() {
   if (EnableRegReassign)
     addPass(&GCNNSAReassignID);
+
+  addPass(&AMDGPURewriteAGPRCopyMFMALegacyID);
   return true;
 }
 

llvm/lib/Target/AMDGPU/CMakeLists.txt

@@ -102,6 +102,7 @@ add_llvm_target(AMDGPUCodeGen
   AMDGPURemoveIncompatibleFunctions.cpp
   AMDGPUReserveWWMRegs.cpp
   AMDGPUResourceUsageAnalysis.cpp
+  AMDGPURewriteAGPRCopyMFMA.cpp
   AMDGPURewriteOutArguments.cpp
   AMDGPURewriteUndefForPHI.cpp
   AMDGPUSelectionDAGInfo.cpp

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

@@ -237,6 +237,10 @@ class SIRegisterInfo final : public AMDGPUGenRegisterInfo {
     return isSGPRClass(getPhysRegBaseClass(Reg));
   }
 
+  bool isVGPRPhysReg(Register Reg) const {
+    return isVGPRClass(getPhysRegBaseClass(Reg));
+  }
+
   /// \returns true if this class contains only VGPR registers
   static bool isVGPRClass(const TargetRegisterClass *RC) {
     return hasVGPRs(RC) && !hasAGPRs(RC) && !hasSGPRs(RC);