//===- LiveRegMatrix.cpp - Track register interference --------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file defines the LiveRegMatrix analysis pass.
//
//===----------------------------------------------------------------------===//

#include "llvm/CodeGen/LiveRegMatrix.h"
#include "RegisterCoalescer.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/CodeGen/LiveInterval.h"
#include "llvm/CodeGen/LiveIntervalUnion.h"
#include "llvm/CodeGen/LiveIntervals.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/CodeGen/TargetSubtargetInfo.h"
#include "llvm/CodeGen/VirtRegMap.h"
#include "llvm/InitializePasses.h"
#include "llvm/MC/LaneBitmask.h"
#include "llvm/MC/MCRegisterInfo.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <cassert>

using namespace llvm;

#define DEBUG_TYPE "regalloc"

STATISTIC(NumAssigned   , "Number of registers assigned");
STATISTIC(NumUnassigned , "Number of registers unassigned");

char LiveRegMatrixWrapperLegacy::ID = 0;
INITIALIZE_PASS_BEGIN(LiveRegMatrixWrapperLegacy, "liveregmatrix",
                      "Live Register Matrix", false, false)
INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
INITIALIZE_PASS_END(LiveRegMatrixWrapperLegacy, "liveregmatrix",
                    "Live Register Matrix", false, true)

void LiveRegMatrixWrapperLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.setPreservesAll();
  AU.addRequiredTransitive<LiveIntervalsWrapperPass>();
  AU.addRequiredTransitive<VirtRegMapWrapperLegacy>();
  MachineFunctionPass::getAnalysisUsage(AU);
}

bool LiveRegMatrixWrapperLegacy::runOnMachineFunction(MachineFunction &MF) {
  auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
  auto &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
  LRM.init(MF, LIS, VRM);
  return false;
}

void LiveRegMatrix::init(MachineFunction &MF, LiveIntervals &pLIS,
                         VirtRegMap &pVRM) {
  TRI = MF.getSubtarget().getRegisterInfo();
  LIS = &pLIS;
  VRM = &pVRM;

  unsigned NumRegUnits = TRI->getNumRegUnits();
  if (NumRegUnits != Matrix.size())
    Queries.reset(new LiveIntervalUnion::Query[NumRegUnits]);
  Matrix.init(*LIUAlloc, NumRegUnits);

  // Make sure no stale queries get reused.
  invalidateVirtRegs();
}

void LiveRegMatrixWrapperLegacy::releaseMemory() { LRM.releaseMemory(); }

void LiveRegMatrix::releaseMemory() {
  for (unsigned i = 0, e = Matrix.size(); i != e; ++i) {
    Matrix[static_cast<MCRegUnit>(i)].clear();
    // No need to clear Queries here, since LiveIntervalUnion::Query doesn't
    // have anything important to clear and LiveRegMatrix's runOnFunction()
    // does a std::unique_ptr::reset anyways.
  }
}

template <typename Callable>
static bool foreachUnit(const TargetRegisterInfo *TRI,
                        const LiveInterval &VRegInterval, MCRegister PhysReg,
                        Callable Func) {
  if (VRegInterval.hasSubRanges()) {
    for (MCRegUnitMaskIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
      MCRegUnit Unit = (*Units).first;
      LaneBitmask Mask = (*Units).second;
      for (const LiveInterval::SubRange &S : VRegInterval.subranges()) {
        if ((S.LaneMask & Mask).any()) {
          if (Func(Unit, S))
            return true;
          break;
        }
      }
    }
  } else {
    for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
      if (Func(Unit, VRegInterval))
        return true;
    }
  }
  return false;
}

void LiveRegMatrix::assign(const LiveInterval &VirtReg, MCRegister PhysReg) {
  LLVM_DEBUG(dbgs() << "assigning " << printReg(VirtReg.reg(), TRI) << " to "
                    << printReg(PhysReg, TRI) << ':');
  assert(!VRM->hasPhys(VirtReg.reg()) && "Duplicate VirtReg assignment");
  VRM->assignVirt2Phys(VirtReg.reg(), PhysReg);

  foreachUnit(
      TRI, VirtReg, PhysReg, [&](MCRegUnit Unit, const LiveRange &Range) {
        LLVM_DEBUG(dbgs() << ' ' << printRegUnit(Unit, TRI) << ' ' << Range);
        Matrix[Unit].unify(VirtReg, Range);
        return false;
      });

  ++NumAssigned;
  LLVM_DEBUG(dbgs() << '\n');
}

void LiveRegMatrix::unassign(const LiveInterval &VirtReg,
                             bool ClearAllReferencingSegments) {
  Register PhysReg = VRM->getPhys(VirtReg.reg());
  LLVM_DEBUG(dbgs() << "unassigning " << printReg(VirtReg.reg(), TRI)
                    << " from " << printReg(PhysReg, TRI) << ':');
  VRM->clearVirt(VirtReg.reg());

  if (!ClearAllReferencingSegments) {
    foreachUnit(TRI, VirtReg, PhysReg,
                [&](MCRegUnit Unit, const LiveRange &Range) {
                  LLVM_DEBUG(dbgs() << ' ' << printRegUnit(Unit, TRI));
                  Matrix[Unit].extract(VirtReg, Range);
                  return false;
                });
  } else {
    for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
      Matrix[Unit].clearAllSegmentsReferencing(VirtReg);
    }
  }

  ++NumUnassigned;
  LLVM_DEBUG(dbgs() << '\n');
}

bool LiveRegMatrix::isPhysRegUsed(MCRegister PhysReg) const {
  for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
    if (!Matrix[Unit].empty())
      return true;
  }
  return false;
}

bool LiveRegMatrix::checkRegMaskInterference(const LiveInterval &VirtReg,
                                             MCRegister PhysReg) {
  // Check if the cached information is valid.
  // The same BitVector can be reused for all PhysRegs.
  // We could cache multiple VirtRegs if it becomes necessary.
  if (RegMaskVirtReg != VirtReg.reg() || RegMaskTag != UserTag) {
    RegMaskVirtReg = VirtReg.reg();
    RegMaskTag = UserTag;
    RegMaskUsable.clear();
    LIS->checkRegMaskInterference(VirtReg, RegMaskUsable);
  }

  // The BitVector is indexed by PhysReg, not register unit.
  // Regmask interference is more fine grained than regunits.
  // For example, a Win64 call can clobber %ymm8 yet preserve %xmm8.
  return !RegMaskUsable.empty() &&
         (!PhysReg || !RegMaskUsable.test(PhysReg.id()));
}

bool LiveRegMatrix::checkRegUnitInterference(const LiveInterval &VirtReg,
                                             MCRegister PhysReg) {
  if (VirtReg.empty())
    return false;
  CoalescerPair CP(VirtReg.reg(), PhysReg, *TRI);

  bool Result = foreachUnit(
      TRI, VirtReg, PhysReg, [&](MCRegUnit Unit, const LiveRange &Range) {
        const LiveRange &UnitRange = LIS->getRegUnit(Unit);
        return Range.overlaps(UnitRange, CP, *LIS->getSlotIndexes());
      });
  return Result;
}

bool LiveRegMatrix::checkRegMaskInterference(SlotIndex Start, SlotIndex End,
                                             MCRegister PhysReg) {
  ArrayRef<SlotIndex> Slots = LIS->getRegMaskSlots();
  ArrayRef<const uint32_t *> Bits = LIS->getRegMaskBits();

  // Find the first regmask slot that is not before Start.
  auto SlotI = llvm::lower_bound(Slots, Start);
  for (; SlotI != Slots.end() && *SlotI < End; ++SlotI) {
    if (MachineOperand::clobbersPhysReg(Bits[SlotI - Slots.begin()], PhysReg))
      return true;
  }
  return false;
}

bool LiveRegMatrix::checkRegUnitInterference(SlotIndex Start, SlotIndex End,
                                             MCRegister PhysReg) {
  for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
    if (LIS->getRegUnit(Unit).overlaps(Start, End))
      return true;
  }
  return false;
}

LiveIntervalUnion::Query &LiveRegMatrix::query(const LiveRange &LR,
                                               MCRegUnit RegUnit) {
  LiveIntervalUnion::Query &Q = Queries[static_cast<unsigned>(RegUnit)];
  Q.init(UserTag, LR, Matrix[RegUnit]);
  return Q;
}

LiveRegMatrix::InterferenceKind
LiveRegMatrix::checkInterference(const LiveInterval &VirtReg,
                                 MCRegister PhysReg) {
  if (VirtReg.empty())
    return IK_Free;

  // Regmask interference is the fastest check.
  if (checkRegMaskInterference(VirtReg, PhysReg))
    return IK_RegMask;

  // Check for fixed interference.
  if (checkRegUnitInterference(VirtReg, PhysReg))
    return IK_RegUnit;

  // Check the matrix for virtual register interference.
  bool Interference = foreachUnit(TRI, VirtReg, PhysReg,
                                  [&](MCRegUnit Unit, const LiveRange &LR) {
                                    return query(LR, Unit).checkInterference();
                                  });
  if (Interference)
    return IK_VirtReg;

  return IK_Free;
}

bool LiveRegMatrix::checkInterference(SlotIndex Start, SlotIndex End,
                                      MCRegister PhysReg) {
  // Regmask interference is the fastest check.
  if (checkRegMaskInterference(Start, End, PhysReg))
    return true;

  // Check for fixed interference.
  if (checkRegUnitInterference(Start, End, PhysReg))
    return true;

  // Construct artificial live range containing only one segment [Start, End).
  VNInfo valno(0, Start);
  LiveRange::Segment Seg(Start, End, &valno);
  LiveRange LR;
  LR.addSegment(Seg);

  // Check the matrix for virtual register interference with that segment.
  for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
    // LR is stack-allocated. LiveRegMatrix caches queries by a key that
    // includes the address of the live range. If (for the same reg unit) this
    // checkInterference overload is called twice, without any other query()
    // calls in between (on heap-allocated LiveRanges)  - which would invalidate
    // the cached query - the LR address seen the second time may well be the
    // same as that seen the first time, while the Start/End/valno may not - yet
    // the same cached result would be fetched. To avoid that, we don't cache
    // this query.
    //
    // FIXME: the usability of the Query API needs to be improved to avoid
    // subtle bugs due to query identity. Avoiding caching, for example, would
    // greatly simplify things.
    LiveIntervalUnion::Query Q;
    Q.reset(UserTag, LR, Matrix[Unit]);
    if (Q.checkInterference())
      return true;
  }
  return false;
}

LaneBitmask LiveRegMatrix::checkInterferenceLanes(SlotIndex Start,
                                                  SlotIndex End,
                                                  MCRegister PhysReg) {
  // Construct artificial live range containing only one segment [Start, End).
  VNInfo valno(0, Start);
  LiveRange::Segment Seg(Start, End, &valno);
  LiveRange LR;
  LR.addSegment(Seg);

  LaneBitmask InterferingLanes;

  // Check for interference with that segment
  for (MCRegUnitMaskIterator MCRU(PhysReg, TRI); MCRU.isValid(); ++MCRU) {
    auto [Unit, Lanes] = *MCRU;
    // LR is stack-allocated. LiveRegMatrix caches queries by a key that
    // includes the address of the live range. If (for the same reg unit) this
    // checkInterference overload is called twice, without any other query()
    // calls in between (on heap-allocated LiveRanges)  - which would invalidate
    // the cached query - the LR address seen the second time may well be the
    // same as that seen the first time, while the Start/End/valno may not - yet
    // the same cached result would be fetched. To avoid that, we don't cache
    // this query.
    //
    // FIXME: the usability of the Query API needs to be improved to avoid
    // subtle bugs due to query identity. Avoiding caching, for example, would
    // greatly simplify things.
    LiveIntervalUnion::Query Q;
    Q.reset(UserTag, LR, Matrix[Unit]);
    if (Q.checkInterference())
      InterferingLanes |= Lanes;
  }

  return InterferingLanes;
}

Register LiveRegMatrix::getOneVReg(unsigned PhysReg) const {
  const LiveInterval *VRegInterval = nullptr;
  for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
    if ((VRegInterval = Matrix[Unit].getOneVReg()))
      return VRegInterval->reg();
  }

  return MCRegister::NoRegister;
}

#ifndef NDEBUG
bool LiveRegMatrix::isValid() const {
  // Build set of all valid LiveInterval pointers from LiveIntervals.
  DenseSet<const LiveInterval *> ValidIntervals;
  for (unsigned RegIdx = 0, NumRegs = VRM->getRegInfo().getNumVirtRegs();
       RegIdx < NumRegs; ++RegIdx) {
    Register VReg = Register::index2VirtReg(RegIdx);
    // Only track assigned registers since unassigned ones won't be in Matrix
    if (VRM->hasPhys(VReg) && LIS->hasInterval(VReg))
      ValidIntervals.insert(&LIS->getInterval(VReg));
  }

  // Now scan all LiveIntervalUnions in the matrix and verify each pointer
  unsigned NumDanglingPointers = 0;
  for (unsigned I = 0, Size = Matrix.size(); I < Size; ++I) {
    MCRegUnit Unit = static_cast<MCRegUnit>(I);
    for (const LiveInterval *LI : Matrix[Unit]) {
      if (!ValidIntervals.contains(LI)) {
        ++NumDanglingPointers;
        dbgs() << "ERROR: LiveInterval pointer is not found in LiveIntervals:\n"
               << "  Register Unit: " << printRegUnit(Unit, TRI) << '\n'
               << "  LiveInterval pointer: " << LI << '\n';
      }
    }
  }

  // Reverse check: every VRM-assigned vreg with a non-empty live interval
  // must have its segments present in the Matrix for its assigned phys reg.
  unsigned NumMissing = 0;
  for (unsigned RegIdx = 0, NumRegs = VRM->getRegInfo().getNumVirtRegs();
       RegIdx < NumRegs; ++RegIdx) {
    Register VReg = Register::index2VirtReg(RegIdx);
    if (!VRM->hasPhys(VReg) || !LIS->hasInterval(VReg))
      continue;
    const LiveInterval &LI = LIS->getInterval(VReg);
    if (LI.empty())
      continue;
    MCRegister PhysReg = VRM->getPhys(VReg);
    // Check that the first segment of LI is present in the LiveUnion for
    // at least one reg unit of PhysReg.
    SlotIndex FirstStart = LI.beginIndex();
    bool Found = false;
    for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
      auto It = Matrix[Unit].find(FirstStart);
      if (It.valid() && It.start() == FirstStart && It.value() == &LI) {
        Found = true;
        break;
      }
    }
    if (!Found) {
      ++NumMissing;
      dbgs() << "ERROR: VirtReg " << printReg(VReg, TRI) << " assigned to "
             << printReg(PhysReg, TRI)
             << " in VirtRegMap but not found in LiveRegMatrix\n";
      dbgs() << "  LiveInterval: " << LI << "\n";
      dbgs() << "  FirstStart: " << FirstStart << "\n";
      for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
        dbgs() << "  RegUnit " << printRegUnit(Unit, TRI) << " segments: ";
        auto It = Matrix[Unit].find(FirstStart);
        if (It.valid())
          dbgs() << "[" << It.start() << "," << It.stop() << ") -> "
                 << printReg(It.value()->reg(), TRI);
        else
          dbgs() << "(none found)";
        dbgs() << "\n";
      }
    }
  }

  return NumDanglingPointers == 0 && NumMissing == 0;
}
#endif

AnalysisKey LiveRegMatrixAnalysis::Key;

LiveRegMatrix LiveRegMatrixAnalysis::run(MachineFunction &MF,
                                         MachineFunctionAnalysisManager &MFAM) {
  auto &LIS = MFAM.getResult<LiveIntervalsAnalysis>(MF);
  auto &VRM = MFAM.getResult<VirtRegMapAnalysis>(MF);
  LiveRegMatrix LRM;
  LRM.init(MF, LIS, VRM);
  return LRM;
}