//===-- 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:
///  - 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 {
  MachineFunction &MF;
  const GCNSubtarget &ST;
  const SIInstrInfo &TII;
  const SIRegisterInfo &TRI;
  MachineRegisterInfo &MRI;
  VirtRegMap &VRM;
  LiveRegMatrix ‎
  LiveIntervals &LIS;
  const RegisterClassInfo &RegClassInfo;

  bool attemptReassignmentsToAGPR(SmallSetVector<Register, 4> &InterferingRegs,
                                  MCPhysReg PrefPhysReg) const;

public:
  AMDGPURewriteAGPRCopyMFMAImpl(MachineFunction &MF, VirtRegMap &VRM,
                                LiveRegMatrix &LRM, LiveIntervals &LIS,
                                const RegisterClassInfo &RegClassInfo)
      : MF(MF), ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()),
        TRI(*ST.getRegisterInfo()), MRI(MF.getRegInfo()), VRM(VRM), LRM(LRM),
        LIS(LIS), RegClassInfo(RegClassInfo) {}

  bool isRewriteCandidate(const MachineInstr &MI) const {
    return TII.isMAI(MI) && AMDGPU::getMFMASrcCVDstAGPROp(MI.getOpcode()) != -1;
  }

  /// Compute the register class constraints based on the uses of \p Reg,
  /// excluding MFMA uses from which can be rewritten to change the register
  /// class constraint. This should be nearly identical to
  /// MachineRegisterInfo::recomputeRegClass.

  /// \p RewriteCandidates will collect the set of MFMA instructions that need
  /// to have the opcode mutated to perform the replacement.
  ///
  /// \p RewriteRegs will accumulate the set of register used by those MFMAs
  /// that need to have the register classes adjusted.
  bool recomputeRegClassExceptRewritable(
      Register Reg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
      SmallSetVector<Register, 4> &RewriteRegs) const;

  bool run(MachineFunction &MF) const;
};

bool AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
    Register StartReg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
    SmallSetVector<Register, 4> &RewriteRegs) const {
  SmallVector<Register, 8> Worklist = {StartReg};

  // Recursively visit all transitive MFMA users
  while (!Worklist.empty()) {
    Register Reg = Worklist.pop_back_val();
    const TargetRegisterClass *OldRC = MRI.getRegClass(Reg);

    // Inflate to the equivalent AV_* class.
    const TargetRegisterClass *NewRC = TRI.getLargestLegalSuperClass(OldRC, MF);
    if (OldRC == NewRC)
      return false;

    // 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();

      // We can swap the classes of dst + src2 as a pair to AGPR, so ignore the
      // effects of rewrite candidates. It just so happens that we can use
      // either AGPR or VGPR in src0/src1, so don't bother checking the
      // constraint effects of the individual operands.
      if (isRewriteCandidate(*MI)) {
        const MachineOperand *VDst =
            TII.getNamedOperand(*MI, AMDGPU::OpName::vdst);
        const MachineOperand *Src2 =
            TII.getNamedOperand(*MI, AMDGPU::OpName::src2);
        for (const MachineOperand *Op : {VDst, Src2}) {
          if (!Op->isReg())
            continue;

          Register OtherReg = Op->getReg();
          if (OtherReg.isPhysical())
            return false;

          if (OtherReg != Reg && RewriteRegs.insert(OtherReg))
            Worklist.push_back(OtherReg);
        }

        if (!is_contained(RewriteCandidates, MI)) {
          LLVM_DEBUG({
            Register VDstPhysReg = VRM.getPhys(VDst->getReg());
            dbgs() << "Attempting to replace VGPR MFMA with AGPR version:"
                   << " Dst=[" << printReg(VDst->getReg()) << " => "
                   << printReg(VDstPhysReg, &TRI);

            if (Src2->isReg()) {
              Register Src2PhysReg = VRM.getPhys(Src2->getReg());
              dbgs() << "], Src2=[" << printReg(Src2->getReg(), &TRI) << " => "
                     << printReg(Src2PhysReg, &TRI);
            }

            dbgs() << "]: " << MI;
          });

          RewriteCandidates.push_back(MI);
        }

        continue;
      }

      unsigned OpNo = &MO - &MI->getOperand(0);
      NewRC = MI->getRegClassConstraintEffect(OpNo, NewRC, &TII, &TRI);
      if (!NewRC || NewRC == OldRC) {
        LLVM_DEBUG(dbgs() << "User of " << printReg(Reg, &TRI)
                          << " cannot be reassigned to "
                          << TRI.getRegClassName(NewRC) << ": " << *MI);
        return false;
      }
    }
  }

  return true;
}

/// Attempt to reassign the registers in \p InterferingRegs to be AGPRs, with a
/// preference to use \p PhysReg first. Returns false if the reassignments
/// cannot be trivially performed.
bool AMDGPURewriteAGPRCopyMFMAImpl::attemptReassignmentsToAGPR(
    SmallSetVector<Register, 4> &InterferingRegs, MCPhysReg PrefPhysReg) const {
  // FIXME: The ordering may matter here, but we're just taking uselistorder
  // with the special case of ensuring to process the starting instruction
  // first. We probably should extract the priority advisor out of greedy and
  // use that ordering.
  for (Register InterferingReg : InterferingRegs) {
    LiveInterval &ReassignLI = LIS.getInterval(InterferingReg);
    const TargetRegisterClass *EquivalentAGPRRegClass =
        TRI.getEquivalentAGPRClass(MRI.getRegClass(InterferingReg));

    MCPhysReg Assignable = AMDGPU::NoRegister;
    if (EquivalentAGPRRegClass->contains(PrefPhysReg) &&
        LRM.checkInterference(ReassignLI, PrefPhysReg) ==
            LiveRegMatrix::IK_Free) {
      // First try to assign to the AGPR we were already copying to. This
      // should be the first assignment we attempt. We have to guard
      // against the use being a subregister (which doesn't have an exact
      // class match).

      // TODO: If this does happen to be a subregister use, we should
      // still try to assign to a subregister of the original copy result.
      Assignable = PrefPhysReg;
    } else {
      ArrayRef<MCPhysReg> AllocOrder =
          RegClassInfo.getOrder(EquivalentAGPRRegClass);
      for (MCPhysReg Reg : AllocOrder) {
        if (LRM.checkInterference(ReassignLI, Reg) == LiveRegMatrix::IK_Free) {
          Assignable = Reg;
          break;
        }
      }
    }

    if (!Assignable) {
      LLVM_DEBUG(dbgs() << "Unable to reassign VGPR "
                        << printReg(InterferingReg, &TRI)
                        << " to a free AGPR\n");
      return false;
    }

    LLVM_DEBUG(dbgs() << "Reassigning VGPR " << printReg(InterferingReg, &TRI)
                      << " to " << printReg(Assignable, &TRI) << '\n');
    LRM.assign(ReassignLI, Assignable);
  }

  return true;
}

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)) {
    LLVM_DEBUG(dbgs() << "skipping function that did not allocate AGPRs\n");
    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.hasAGPRs(VirtRegRC))
      continue;

    const TargetRegisterClass *AssignedRC = VirtRegRC;
    if (TRI.hasVGPRs(VirtRegRC)) {
      // If this is an AV register, we have to check if the actual assignment is
      // to an AGPR
      AssignedRC = TRI.getPhysRegBaseClass(PhysReg);
      if (!TRI.isAGPRClass(AssignedRC))
        continue;
    }

    LiveInterval &LI = LIS.getInterval(VReg);

    for (VNInfo *VNI : LI.vnis()) {
      if (VNI->isPHIDef() || VNI->isUnused())
        continue;

      MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
      if (!DefMI || !DefMI->isCopy())
        continue;

      Register MFMADstReg = DefMI->getOperand(1).getReg();
      if (!MFMADstReg.isVirtual())
        continue;

      LiveInterval &CopySrcLI = LIS.getInterval(MFMADstReg);
      LiveQueryResult LRQ = CopySrcLI.Query(VNI->def.getRegSlot());
      MachineInstr *MFMA = LIS.getInstructionFromIndex(LRQ.valueIn()->def);
      if (!MFMA || !isRewriteCandidate(*MFMA))
        continue;

      // src2 and dst have the same physical class constraint; try to preserve
      // the original src2 subclass if one were to exist.
      SmallVector<MachineInstr *, 4> RewriteCandidates = {MFMA};
      SmallSetVector<Register, 4> RewriteRegs;

      // Make sure we reassign the MFMA we found the copy from first. We want
      // to ensure dst ends up in the physreg we were originally copying to.
      RewriteRegs.insert(MFMADstReg);

      // 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.
      //
      // Note recomputeRegClassExceptRewritable will consider the constraints of
      // this MFMA's src2 as well as the src2/dst of any transitive MFMA users.
      if (!recomputeRegClassExceptRewritable(MFMADstReg, RewriteCandidates,
                                             RewriteRegs)) {
        LLVM_DEBUG(dbgs() << "Could not recompute the regclass of dst reg "
                          << printReg(MFMADstReg, &TRI) << '\n');
        continue;
      }

      // If src2 and dst are different registers, we need to also reassign the
      // input to an available AGPR if it 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.
      //
      // It's likely that the MFMA is used in sequence with other MFMAs; if we
      // cannot migrate the full use/def chain of MFMAs, we would need to
      // introduce intermediate copies somewhere. So we only make the
      // transform if all the interfering MFMAs can also be migrated. Collect
      // the set of rewritable MFMAs and check if we can assign an AGPR at
      // that point.
      //
      // If any of the MFMAs aren't reassignable, we give up and rollback to
      // the original register assignments.

      using RecoloringStack =
          SmallVector<std::pair<const LiveInterval *, MCRegister>, 8>;
      RecoloringStack TentativeReassignments;

      for (Register RewriteReg : RewriteRegs) {
        LiveInterval &LI = LIS.getInterval(RewriteReg);
        TentativeReassignments.push_back({&LI, VRM.getPhys(RewriteReg)});
        LRM.unassign(LI);
      }

      if (!attemptReassignmentsToAGPR(RewriteRegs, PhysReg)) {
        // Roll back the register assignments to the original state.
        for (auto [LI, OldAssign] : TentativeReassignments) {
          if (VRM.hasPhys(LI->reg()))
            LRM.unassign(*LI);
          LRM.assign(*LI, OldAssign);
        }

        continue;
      }

      // Fixup the register classes of the virtual registers now that we've
      // committed to the reassignments.
      for (Register InterferingReg : RewriteRegs) {
        const TargetRegisterClass *EquivalentAGPRRegClass =
            TRI.getEquivalentAGPRClass(MRI.getRegClass(InterferingReg));
        MRI.setRegClass(InterferingReg, EquivalentAGPRRegClass);
      }

      for (MachineInstr *RewriteCandidate : RewriteCandidates) {
        int NewMFMAOp =
            AMDGPU::getMFMASrcCVDstAGPROp(RewriteCandidate->getOpcode());
        RewriteCandidate->setDesc(TII.get(NewMFMAOp));
      }

      // We likely left an identity copy behind after assignment; let
      // VirtRegRewriter deal with it later.
      MadeChange = true;
    }
  }

  return MadeChange;
}

class AMDGPURewriteAGPRCopyMFMALegacy : public MachineFunctionPass {
public:
  static char ID;
  RegisterClassInfo RegClassInfo;

  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;

  RegClassInfo.runOnMachineFunction(MF);

  auto &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
  auto &LRM = getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
  auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();

  AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS, RegClassInfo);
  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);
  RegisterClassInfo RegClassInfo;
  RegClassInfo.runOnMachineFunction(MF);

  AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS, RegClassInfo);
  if (!Impl.run(MF))
    return PreservedAnalyses::all();
  auto PA = getMachineFunctionPassPreservedAnalyses();
  PA.preserveSet<CFGAnalyses>();
  return PA;
}