//===-- SIPostRABundler.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
/// This pass creates bundles of memory instructions to protect adjacent loads
/// and stores from being rescheduled apart from each other post-RA.
///
//===----------------------------------------------------------------------===//

#include "SIPostRABundler.h"
#include "AMDGPU.h"
#include "GCNSubtarget.h"
#include "llvm/CodeGen/MachineFunctionPass.h"

using namespace llvm;

#define DEBUG_TYPE "si-post-ra-bundler"

namespace {

class SIPostRABundlerLegacy : public MachineFunctionPass {
public:
  static char ID;

public:
  SIPostRABundlerLegacy() : MachineFunctionPass(ID) {}

  bool runOnMachineFunction(MachineFunction &MF) override;

  StringRef getPassName() const override {
    return "SI post-RA bundler";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesAll();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
};

class SIPostRABundler {
public:
  bool run(MachineFunction &MF);

private:
  const SIRegisterInfo *TRI;

  SmallSet<Register, 16> Defs;

  void collectUsedRegUnits(const MachineInstr &MI,
                           BitVector &UsedRegUnits) const;

  bool isBundleCandidate(const MachineInstr &MI) const;
  bool isDependentLoad(const MachineInstr &MI) const;
  bool canBundle(const MachineInstr &MI, const MachineInstr &NextMI) const;
};

} // End anonymous namespace.

INITIALIZE_PASS(SIPostRABundlerLegacy, DEBUG_TYPE, "SI post-RA bundler", false,
                false)

char SIPostRABundlerLegacy::ID = 0;

char &llvm::SIPostRABundlerLegacyID = SIPostRABundlerLegacy::ID;

FunctionPass *llvm::createSIPostRABundlerPass() {
  return new SIPostRABundlerLegacy();
}

bool SIPostRABundler::isDependentLoad(const MachineInstr &MI) const {
  if (!MI.mayLoad())
    return false;

  for (const MachineOperand &Op : MI.explicit_operands()) {
    if (!Op.isReg())
      continue;
    Register Reg = Op.getReg();
    for (Register Def : Defs)
      if (TRI->regsOverlap(Reg, Def))
        return true;
  }

  return false;
}

void SIPostRABundler::collectUsedRegUnits(const MachineInstr &MI,
                                          BitVector &UsedRegUnits) const {
  if (MI.isDebugInstr())
    return;

  for (const MachineOperand &Op : MI.operands()) {
    if (!Op.isReg() || !Op.readsReg())
      continue;

    Register Reg = Op.getReg();
    assert(!Op.getSubReg() &&
           "subregister indexes should not be present after RA");

    for (MCRegUnit Unit : TRI->regunits(Reg))
      UsedRegUnits.set(static_cast<unsigned>(Unit));
  }
}

static bool isMemoryInst(const MachineInstr &MI) {
  return SIInstrFlags::isMUBUF(MI) || SIInstrFlags::isMTBUF(MI) ||
         SIInstrFlags::isSMRD(MI) || SIInstrFlags::isDS(MI) ||
         SIInstrFlags::isFLAT(MI) || SIInstrFlags::isMIMG(MI) ||
         SIInstrFlags::isVIMAGE(MI) || SIInstrFlags::isVSAMPLE(MI);
}

static bool hasSameMemFormat(const MachineInstr &A, const MachineInstr &B) {
  return SIInstrFlags::isMUBUF(A) == SIInstrFlags::isMUBUF(B) &&
         SIInstrFlags::isMTBUF(A) == SIInstrFlags::isMTBUF(B) &&
         SIInstrFlags::isSMRD(A) == SIInstrFlags::isSMRD(B) &&
         SIInstrFlags::isDS(A) == SIInstrFlags::isDS(B) &&
         SIInstrFlags::isFLAT(A) == SIInstrFlags::isFLAT(B) &&
         SIInstrFlags::isMIMG(A) == SIInstrFlags::isMIMG(B) &&
         SIInstrFlags::isVIMAGE(A) == SIInstrFlags::isVIMAGE(B) &&
         SIInstrFlags::isVSAMPLE(A) == SIInstrFlags::isVSAMPLE(B);
}

bool SIPostRABundler::isBundleCandidate(const MachineInstr &MI) const {
  return isMemoryInst(MI) && MI.mayLoadOrStore() && !MI.isBundled();
}

bool SIPostRABundler::canBundle(const MachineInstr &MI,
                                const MachineInstr &NextMI) const {
  return isMemoryInst(MI) && MI.mayLoadOrStore() && !NextMI.isBundled() &&
         NextMI.mayLoad() == MI.mayLoad() &&
         NextMI.mayStore() == MI.mayStore() && hasSameMemFormat(MI, NextMI) &&
         !isDependentLoad(NextMI);
}

bool SIPostRABundlerLegacy::runOnMachineFunction(MachineFunction &MF) {
  if (skipFunction(MF.getFunction()))
    return false;
  return SIPostRABundler().run(MF);
}

PreservedAnalyses SIPostRABundlerPass::run(MachineFunction &MF,
                                           MachineFunctionAnalysisManager &) {
  SIPostRABundler().run(MF);
  return PreservedAnalyses::all();
}

bool SIPostRABundler::run(MachineFunction &MF) {

  TRI = MF.getSubtarget<GCNSubtarget>().getRegisterInfo();
  BitVector BundleUsedRegUnits(TRI->getNumRegUnits());
  BitVector KillUsedRegUnits(TRI->getNumRegUnits());

  bool Changed = false;
  for (MachineBasicBlock &MBB : MF) {
    bool HasIGLPInstrs = llvm::any_of(MBB.instrs(), [](MachineInstr &MI) {
      unsigned Opc = MI.getOpcode();
      return Opc == AMDGPU::SCHED_GROUP_BARRIER || Opc == AMDGPU::IGLP_OPT;
    });

    // Don't cluster with IGLP instructions.
    if (HasIGLPInstrs)
      continue;

    MachineBasicBlock::instr_iterator Next;
    MachineBasicBlock::instr_iterator B = MBB.instr_begin();
    MachineBasicBlock::instr_iterator E = MBB.instr_end();

    for (auto I = B; I != E; I = Next) {
      Next = std::next(I);
      if (!isBundleCandidate(*I))
        continue;

      assert(Defs.empty());

      if (I->getNumExplicitDefs() != 0)
        Defs.insert(I->defs().begin()->getReg());

      MachineBasicBlock::instr_iterator BundleStart = I;
      MachineBasicBlock::instr_iterator BundleEnd = I;
      unsigned ClauseLength = 1;
      for (I = Next; I != E; I = Next) {
        Next = std::next(I);

        assert(BundleEnd != I);
        if (canBundle(*BundleEnd, *I)) {
          BundleEnd = I;
          if (I->getNumExplicitDefs() != 0)
            Defs.insert(I->defs().begin()->getReg());
          ++ClauseLength;
        } else if (!I->isMetaInstruction() ||
                   I->getOpcode() == AMDGPU::SCHED_BARRIER) {
          // SCHED_BARRIER is not bundled to be honored by scheduler later.
          // Allow other meta instructions in between bundle candidates, but do
          // not start or end a bundle on one.
          //
          // TODO: It may be better to move meta instructions like dbg_value
          // after the bundle. We're relying on the memory legalizer to unbundle
          // these.
          break;
        }
      }

      Next = std::next(BundleEnd);
      if (ClauseLength > 1) {
        Changed = true;

        // Before register allocation, kills are inserted after potential soft
        // clauses to hint register allocation. Look for kills that look like
        // this, and erase them.
        if (Next != E && Next->isKill()) {

          // TODO: Should maybe back-propagate kill flags to the bundle.
          for (const MachineInstr &BundleMI : make_range(BundleStart, Next))
            collectUsedRegUnits(BundleMI, BundleUsedRegUnits);

          BundleUsedRegUnits.flip();

          while (Next != E && Next->isKill()) {
            MachineInstr &Kill = *Next;
            collectUsedRegUnits(Kill, KillUsedRegUnits);

            KillUsedRegUnits &= BundleUsedRegUnits;

            // Erase the kill if it's a subset of the used registers.
            //
            // TODO: Should we just remove all kills? Is there any real reason to
            // keep them after RA?
            if (KillUsedRegUnits.none()) {
              ++Next;
              Kill.eraseFromParent();
            } else
              break;

            KillUsedRegUnits.reset();
          }

          BundleUsedRegUnits.reset();
        }

        finalizeBundle(MBB, BundleStart, Next);
      }

      Defs.clear();
    }
  }

  return Changed;
}