|
| 1 | +//===-- AMDGPURewriteAGPRCopyMFMA.cpp -------------------------------------===// |
| 2 | +// |
| 3 | +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| 4 | +// See https://llvm.org/LICENSE.txt for license information. |
| 5 | +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | +// |
| 7 | +//===----------------------------------------------------------------------===// |
| 8 | +// |
| 9 | +/// \file \brief Try to replace MFMA instructions using VGPRs with MFMA |
| 10 | +/// instructions using AGPRs. We expect MFMAs to be selected using VGPRs, and |
| 11 | +/// only use AGPRs if it helps avoid spilling. In this case, the MFMA will have |
| 12 | +/// copies between AGPRs and VGPRs and the AGPR variant of an MFMA pseudo. This |
| 13 | +/// pass will attempt to delete the cross register bank copy and replace the |
| 14 | +/// MFMA opcode. |
| 15 | +/// |
| 16 | +/// TODO: |
| 17 | +/// - Handle non-tied dst+src2 cases. We need to try to find a copy from an |
| 18 | +/// AGPR from src2, or reassign src2 to an available AGPR (which should work |
| 19 | +/// in the common case of a load). |
| 20 | +/// |
| 21 | +/// - Handle multiple MFMA uses of the same register. e.g. chained MFMAs that |
| 22 | +/// can be rewritten as a set |
| 23 | +/// |
| 24 | +/// - Update LiveIntervals incrementally instead of recomputing from scratch |
| 25 | +/// |
| 26 | +//===----------------------------------------------------------------------===// |
| 27 | + |
| 28 | +#include "AMDGPU.h" |
| 29 | +#include "GCNSubtarget.h" |
| 30 | +#include "SIMachineFunctionInfo.h" |
| 31 | +#include "SIRegisterInfo.h" |
| 32 | +#include "llvm/CodeGen/LiveIntervals.h" |
| 33 | +#include "llvm/CodeGen/LiveRegMatrix.h" |
| 34 | +#include "llvm/CodeGen/MachineFunctionPass.h" |
| 35 | +#include "llvm/CodeGen/VirtRegMap.h" |
| 36 | +#include "llvm/InitializePasses.h" |
| 37 | + |
| 38 | +using namespace llvm; |
| 39 | + |
| 40 | +#define DEBUG_TYPE "amdgpu-rewrite-agpr-copy-mfma" |
| 41 | + |
| 42 | +namespace { |
| 43 | + |
| 44 | +class AMDGPURewriteAGPRCopyMFMAImpl { |
| 45 | + const GCNSubtarget &ST; |
| 46 | + const SIInstrInfo &TII; |
| 47 | + const SIRegisterInfo &TRI; |
| 48 | + MachineRegisterInfo &MRI; |
| 49 | + VirtRegMap &VRM; |
| 50 | + LiveRegMatrix ‎ |
| 51 | + LiveIntervals &LIS; |
| 52 | + |
| 53 | +public: |
| 54 | + AMDGPURewriteAGPRCopyMFMAImpl(MachineFunction &MF, VirtRegMap &VRM, |
| 55 | + LiveRegMatrix &LRM, LiveIntervals &LIS) |
| 56 | + : ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()), |
| 57 | + TRI(*ST.getRegisterInfo()), MRI(MF.getRegInfo()), VRM(VRM), LRM(LRM), |
| 58 | + LIS(LIS) {} |
| 59 | + |
| 60 | + /// Compute the register class constraints based on the uses of \p Reg, |
| 61 | + /// excluding uses from \p ExceptMI. This should be nearly identical to |
| 62 | + /// MachineRegisterInfo::recomputeRegClass. |
| 63 | + const TargetRegisterClass * |
| 64 | + recomputeRegClassExcept(Register Reg, const TargetRegisterClass *OldRC, |
| 65 | + const TargetRegisterClass *NewRC, |
| 66 | + const MachineInstr *ExceptMI) const; |
| 67 | + |
| 68 | + bool run(MachineFunction &MF) const; |
| 69 | +}; |
| 70 | + |
| 71 | +const TargetRegisterClass * |
| 72 | +AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExcept( |
| 73 | + Register Reg, const TargetRegisterClass *OldRC, |
| 74 | + const TargetRegisterClass *NewRC, const MachineInstr *ExceptMI) const { |
| 75 | + |
| 76 | + // Accumulate constraints from all uses. |
| 77 | + for (MachineOperand &MO : MRI.reg_nodbg_operands(Reg)) { |
| 78 | + // Apply the effect of the given operand to NewRC. |
| 79 | + MachineInstr *MI = MO.getParent(); |
| 80 | + if (MI == ExceptMI) |
| 81 | + continue; |
| 82 | + |
| 83 | + unsigned OpNo = &MO - &MI->getOperand(0); |
| 84 | + NewRC = MI->getRegClassConstraintEffect(OpNo, NewRC, &TII, &TRI); |
| 85 | + if (!NewRC || NewRC == OldRC) |
| 86 | + return nullptr; |
| 87 | + } |
| 88 | + |
| 89 | + return NewRC; |
| 90 | +} |
| 91 | + |
| 92 | +bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const { |
| 93 | + // This only applies on subtargets that have a configurable AGPR vs. VGPR |
| 94 | + // allocation. |
| 95 | + if (!ST.hasGFX90AInsts()) |
| 96 | + return false; |
| 97 | + |
| 98 | + // Early exit if no AGPRs were assigned. |
| 99 | + if (!LRM.isPhysRegUsed(AMDGPU::AGPR0)) |
| 100 | + return false; |
| 101 | + |
| 102 | + bool MadeChange = false; |
| 103 | + |
| 104 | + for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) { |
| 105 | + Register VReg = Register::index2VirtReg(I); |
| 106 | + Register PhysReg = VRM.getPhys(VReg); |
| 107 | + if (!PhysReg) |
| 108 | + continue; |
| 109 | + |
| 110 | + // Find AV_* registers assigned to AGPRs. |
| 111 | + const TargetRegisterClass *VirtRegRC = MRI.getRegClass(VReg); |
| 112 | + if (!TRI.isVectorSuperClass(VirtRegRC)) |
| 113 | + continue; |
| 114 | + |
| 115 | + const TargetRegisterClass *AssignedRC = TRI.getPhysRegBaseClass(PhysReg); |
| 116 | + if (!TRI.isAGPRClass(AssignedRC)) |
| 117 | + continue; |
| 118 | + |
| 119 | + LiveInterval &LI = LIS.getInterval(VReg); |
| 120 | + |
| 121 | + // TODO: Test multiple uses |
| 122 | + for (VNInfo *VNI : LI.vnis()) { |
| 123 | + MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def); |
| 124 | + |
| 125 | + // TODO: Handle SplitKit produced copy bundles for partially defined |
| 126 | + // registers. |
| 127 | + if (!DefMI || !DefMI->isFullCopy()) |
| 128 | + continue; |
| 129 | + |
| 130 | + Register CopySrcReg = DefMI->getOperand(1).getReg(); |
| 131 | + LiveInterval &CopySrcLI = LIS.getInterval(CopySrcReg); |
| 132 | + |
| 133 | + LiveQueryResult LRQ = CopySrcLI.Query(VNI->def.getRegSlot()); |
| 134 | + |
| 135 | + MachineInstr *CopySrcMI = LIS.getInstructionFromIndex(LRQ.valueIn()->def); |
| 136 | + if (!CopySrcMI) |
| 137 | + continue; |
| 138 | + |
| 139 | + int AGPROp = AMDGPU::getMFMASrcCVDstAGPROp(CopySrcMI->getOpcode()); |
| 140 | + if (AGPROp == -1) |
| 141 | + continue; |
| 142 | + |
| 143 | + MachineOperand *Src2 = |
| 144 | + TII.getNamedOperand(*CopySrcMI, AMDGPU::OpName::src2); |
| 145 | + |
| 146 | + // FIXME: getMinimalPhysRegClass returns a nonsense AV_* subclass instead |
| 147 | + // of an AGPR or VGPR subclass, so we can't simply use the result on the |
| 148 | + // assignment. |
| 149 | + |
| 150 | + LLVM_DEBUG({ |
| 151 | + Register Src2PhysReg = VRM.getPhys(Src2->getReg()); |
| 152 | + dbgs() << "Attempting to replace VGPR MFMA with AGPR version:" |
| 153 | + << " Dst=[" << printReg(VReg) << " => " |
| 154 | + << printReg(PhysReg, &TRI) << "], Src2=[" |
| 155 | + << printReg(Src2->getReg(), &TRI) << " => " |
| 156 | + << printReg(Src2PhysReg, &TRI) << "]: " << *CopySrcMI; |
| 157 | + }); |
| 158 | + |
| 159 | + // If the inputs are tied and the same register, we can shortcut and |
| 160 | + // directly replace the register. |
| 161 | + if (Src2->getReg() != CopySrcReg) { |
| 162 | + LLVM_DEBUG( |
| 163 | + dbgs() |
| 164 | + << "Replacing untied VGPR MFMAs with AGPR form not yet handled\n"); |
| 165 | + // TODO: Only handles the tied case for now. If the input operand is a |
| 166 | + // different register, we need to also reassign it (either by looking |
| 167 | + // for a compatible copy-from-AGPR, or by seeing if an available AGPR is |
| 168 | + // compatible with all other uses. |
| 169 | + |
| 170 | + // If we can't reassign it, we'd need to introduce a different copy |
| 171 | + // which is likely worse than the copy we'd be saving. |
| 172 | + continue; |
| 173 | + } |
| 174 | + |
| 175 | + const TargetRegisterClass *Src2VirtRegRC = |
| 176 | + MRI.getRegClass(Src2->getReg()); |
| 177 | + |
| 178 | + // We've found av = COPY (MFMA), and need to verify that we can trivially |
| 179 | + // rewrite src2 to use the new AGPR. If we can't trivially replace it, |
| 180 | + // we're going to induce as many copies as we would have emitted in the |
| 181 | + // first place, as well as need to assign another register, and need to |
| 182 | + // figure out where to put them. The live range splitting is smarter than |
| 183 | + // anything we're doing here, so trust it did something reasonable. |
| 184 | + const TargetRegisterClass *Src2ExceptRC = recomputeRegClassExcept( |
| 185 | + Src2->getReg(), Src2VirtRegRC, VirtRegRC, CopySrcMI); |
| 186 | + if (!Src2ExceptRC) |
| 187 | + continue; |
| 188 | + |
| 189 | + const TargetRegisterClass *NewSrc2ConstraintRC = |
| 190 | + TII.getRegClass(TII.get(AGPROp), Src2->getOperandNo(), &TRI, MF); |
| 191 | + |
| 192 | + // Try to constrain src2 to the replacement instruction candidate's |
| 193 | + // register class. |
| 194 | + const TargetRegisterClass *NewSrc2RC = |
| 195 | + TRI.getCommonSubClass(Src2ExceptRC, NewSrc2ConstraintRC); |
| 196 | + if (!NewSrc2RC) { |
| 197 | + // TODO: This is ignoring ther rewritable uses. e.g. a rewritable MFMA |
| 198 | + // using a rewritable MFMA can be rewritten as a pair. |
| 199 | + LLVM_DEBUG(dbgs() << "Other uses of " << printReg(Src2->getReg(), &TRI) |
| 200 | + << " are incompatible with replacement class\n"); |
| 201 | + continue; |
| 202 | + } |
| 203 | + |
| 204 | + MRI.setRegClass(VReg, AssignedRC); |
| 205 | + MRI.setRegClass(Src2->getReg(), NewSrc2RC); |
| 206 | + |
| 207 | + CopySrcMI->setDesc(TII.get(AGPROp)); |
| 208 | + |
| 209 | + // TODO: Is replacing too aggressive, fixup these instructions only? |
| 210 | + MRI.replaceRegWith(CopySrcReg, VReg); |
| 211 | + |
| 212 | + LLVM_DEBUG(dbgs() << "Replaced VGPR MFMA with AGPR: " << *CopySrcMI); |
| 213 | + |
| 214 | + // We left behind an identity copy, so delete it. |
| 215 | + LIS.RemoveMachineInstrFromMaps(*DefMI); |
| 216 | + DefMI->eraseFromParent(); |
| 217 | + |
| 218 | + LRM.unassign(CopySrcLI); |
| 219 | + |
| 220 | + // We don't need the liveness information anymore, so don't bother |
| 221 | + // updating the intervals. Just delete the stale information. |
| 222 | + // TODO: Is it worth preserving these? |
| 223 | + LIS.removeInterval(CopySrcReg); |
| 224 | + LIS.removeInterval(VReg); |
| 225 | + LIS.createAndComputeVirtRegInterval(VReg); |
| 226 | + |
| 227 | + MadeChange = true; |
| 228 | + } |
| 229 | + } |
| 230 | + |
| 231 | + return MadeChange; |
| 232 | +} |
| 233 | + |
| 234 | +class AMDGPURewriteAGPRCopyMFMALegacy : public MachineFunctionPass { |
| 235 | +public: |
| 236 | + static char ID; |
| 237 | + |
| 238 | + AMDGPURewriteAGPRCopyMFMALegacy() : MachineFunctionPass(ID) { |
| 239 | + initializeAMDGPURewriteAGPRCopyMFMALegacyPass( |
| 240 | + *PassRegistry::getPassRegistry()); |
| 241 | + } |
| 242 | + |
| 243 | + bool runOnMachineFunction(MachineFunction &MF) override; |
| 244 | + |
| 245 | + StringRef getPassName() const override { |
| 246 | + return "AMDGPU Rewrite AGPR-Copy-MFMA"; |
| 247 | + } |
| 248 | + |
| 249 | + void getAnalysisUsage(AnalysisUsage &AU) const override { |
| 250 | + AU.addRequired<LiveIntervalsWrapperPass>(); |
| 251 | + AU.addRequired<VirtRegMapWrapperLegacy>(); |
| 252 | + AU.addRequired<LiveRegMatrixWrapperLegacy>(); |
| 253 | + |
| 254 | + AU.addPreserved<LiveIntervalsWrapperPass>(); |
| 255 | + AU.addPreserved<VirtRegMapWrapperLegacy>(); |
| 256 | + AU.addPreserved<LiveRegMatrixWrapperLegacy>(); |
| 257 | + AU.setPreservesAll(); |
| 258 | + MachineFunctionPass::getAnalysisUsage(AU); |
| 259 | + } |
| 260 | +}; |
| 261 | + |
| 262 | +} // End anonymous namespace. |
| 263 | + |
| 264 | +INITIALIZE_PASS_BEGIN(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE, |
| 265 | + "AMDGPU Rewrite AGPR-Copy-MFMA", false, false) |
| 266 | +INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass) |
| 267 | +INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy) |
| 268 | +INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy) |
| 269 | +INITIALIZE_PASS_END(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE, |
| 270 | + "AMDGPU Rewrite AGPR-Copy-MFMA", false, false) |
| 271 | + |
| 272 | +char AMDGPURewriteAGPRCopyMFMALegacy::ID = 0; |
| 273 | + |
| 274 | +char &llvm::AMDGPURewriteAGPRCopyMFMALegacyID = |
| 275 | + AMDGPURewriteAGPRCopyMFMALegacy::ID; |
| 276 | + |
| 277 | +bool AMDGPURewriteAGPRCopyMFMALegacy::runOnMachineFunction( |
| 278 | + MachineFunction &MF) { |
| 279 | + if (skipFunction(MF.getFunction())) |
| 280 | + return false; |
| 281 | + |
| 282 | + auto &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM(); |
| 283 | + auto &LRM = getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM(); |
| 284 | + auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS(); |
| 285 | + |
| 286 | + AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS); |
| 287 | + return Impl.run(MF); |
| 288 | +} |
| 289 | + |
| 290 | +PreservedAnalyses |
| 291 | +AMDGPURewriteAGPRCopyMFMAPass::run(MachineFunction &MF, |
| 292 | + MachineFunctionAnalysisManager &MFAM) { |
| 293 | + VirtRegMap &VRM = MFAM.getResult<VirtRegMapAnalysis>(MF); |
| 294 | + LiveRegMatrix &LRM = MFAM.getResult<LiveRegMatrixAnalysis>(MF); |
| 295 | + LiveIntervals &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF); |
| 296 | + |
| 297 | + AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS); |
| 298 | + if (!Impl.run(MF)) |
| 299 | + return PreservedAnalyses::all(); |
| 300 | + auto PA = getMachineFunctionPassPreservedAnalyses(); |
| 301 | + PA.preserveSet<CFGAnalyses>(); |
| 302 | + return PA; |
| 303 | +} |
0 commit comments