//===- Localizer.cpp ---------------------- Localize some instrs -*- C++ -*-==//
//
// 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 file implements the Localizer class.
//===----------------------------------------------------------------------===//

#include "llvm/CodeGen/GlobalISel/Localizer.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
#include "llvm/CodeGen/GlobalISel/Utils.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
#include "llvm/CodeGen/MachinePassManager.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/TargetLowering.h"
#include "llvm/IR/Analysis.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "localizer"

using namespace llvm;

namespace {

class LocalizerImpl {
  /// MRI contains all the register class/bank information that this
  /// pass uses and updates.
  MachineRegisterInfo *MRI = nullptr;
  /// TTI used for getting remat costs for instructions.
  TargetTransformInfo *TTI = nullptr;

  /// Check if \p MOUse is used in the same basic block as \p Def.
  /// If the use is in the same block, we say it is local.
  /// When the use is not local, \p InsertMBB will contain the basic
  /// block when to insert \p Def to have a local use.
  static bool isLocalUse(MachineOperand &MOUse, const MachineInstr &Def,
                         MachineBasicBlock *&InsertMBB);

  /// Initialize the field members using \p MF.
  void init(MachineFunction &MF, function_ref<TargetTransformInfo *()> GetTTI);

  typedef SmallSetVector<MachineInstr *, 32> LocalizedSetVecT;

  /// If \p Op is a reg operand of a PHI, return the number of total
  /// operands in the PHI that are the same as \p Op, including itself.
  unsigned getNumPhiUses(MachineOperand &Op) const;

  /// Do inter-block localization from the entry block.
  bool localizeInterBlock(MachineFunction &MF,
                          LocalizedSetVecT &LocalizedInstrs);

  /// Do intra-block localization of already localized instructions.
  bool localizeIntraBlock(LocalizedSetVecT &LocalizedInstrs);

public:
  bool runOnMachineFunction(MachineFunction &MF,
                            function_ref<TargetTransformInfo *()> GetTTI);
};

} // namespace

char LocalizerLegacy::ID = 0;
INITIALIZE_PASS_BEGIN(LocalizerLegacy, DEBUG_TYPE,
                      "Move/duplicate certain instructions close to their use",
                      false, false)
INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
INITIALIZE_PASS_END(LocalizerLegacy, DEBUG_TYPE,
                    "Move/duplicate certain instructions close to their use",
                    false, false)

LocalizerLegacy::LocalizerLegacy() : MachineFunctionPass(ID) {}

void LocalizerImpl::init(MachineFunction &MF,
                         function_ref<TargetTransformInfo *()> GetTTI) {
  MRI = &MF.getRegInfo();
  TTI = GetTTI();
}

void LocalizerLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequired<TargetTransformInfoWrapperPass>();
  AU.setPreservesCFG();
  getSelectionDAGFallbackAnalysisUsage(AU);
  MachineFunctionPass::getAnalysisUsage(AU);
}

bool LocalizerImpl::isLocalUse(MachineOperand &MOUse, const MachineInstr &Def,
                               MachineBasicBlock *&InsertMBB) {
  MachineInstr &MIUse = *MOUse.getParent();
  InsertMBB = MIUse.getParent();
  if (MIUse.isPHI())
    InsertMBB = MIUse.getOperand(MOUse.getOperandNo() + 1).getMBB();
  return InsertMBB == Def.getParent();
}

unsigned LocalizerImpl::getNumPhiUses(MachineOperand &Op) const {
  auto *MI = dyn_cast<GPhi>(&*Op.getParent());
  if (!MI)
    return 0;

  Register SrcReg = Op.getReg();
  unsigned NumUses = 0;
  for (unsigned I = 0, NumVals = MI->getNumIncomingValues(); I < NumVals; ++I) {
    if (MI->getIncomingValue(I) == SrcReg)
      ++NumUses;
  }
  return NumUses;
}

bool LocalizerImpl::localizeInterBlock(MachineFunction &MF,
                                       LocalizedSetVecT &LocalizedInstrs) {
  bool Changed = false;
  DenseMap<std::pair<MachineBasicBlock *, Register>, Register> MBBWithLocalDef;

  // Since the IRTranslator only emits constants into the entry block, and the
  // rest of the GISel pipeline generally emits constants close to their users,
  // we only localize instructions in the entry block here. This might change if
  // we start doing CSE across blocks.
  auto &MBB = MF.front();
  auto &TL = *MF.getSubtarget().getTargetLowering();
  for (MachineInstr &MI : llvm::reverse(MBB)) {
    if (!TL.shouldLocalize(MI, TTI))
      continue;
    LLVM_DEBUG(dbgs() << "Should localize: " << MI);
    assert(MI.getDesc().getNumDefs() == 1 &&
           "More than one definition not supported yet");
    Register Reg = MI.getOperand(0).getReg();
    // Check if all the users of MI are local. We are going to invalidate the
    // list of use operands, so we can't use a range iterator. Debug users stay
    // with the original definition.
    for (MachineOperand &MOUse :
         llvm::make_early_inc_range(MRI->use_nodbg_operands(Reg))) {
      // Check if the use is already local.
      MachineBasicBlock *InsertMBB;
      LLVM_DEBUG(MachineInstr &MIUse = *MOUse.getParent();
                 dbgs() << "Checking use: " << MIUse
                        << " #Opd: " << MOUse.getOperandNo() << '\n');
      if (isLocalUse(MOUse, MI, InsertMBB)) {
        // Even if we're in the same block, if the block is very large we could
        // still have many long live ranges. Try to do intra-block localization
        // too.
        LocalizedInstrs.insert(&MI);
        continue;
      }

      // PHIs look like a single user but can use the same register in multiple
      // edges, causing remat into each predecessor. Allow this to a certain
      // extent.
      unsigned NumPhiUses = getNumPhiUses(MOUse);
      const unsigned PhiThreshold = 2; // FIXME: Tune this more.
      if (NumPhiUses > PhiThreshold)
        continue;

      LLVM_DEBUG(dbgs() << "Fixing non-local use\n");
      Changed = true;
      auto MBBAndReg = std::make_pair(InsertMBB, Reg);
      auto NewVRegIt = MBBWithLocalDef.find(MBBAndReg);
      if (NewVRegIt == MBBWithLocalDef.end()) {
        // Create the localized instruction.
        MachineInstr *LocalizedMI = MF.CloneMachineInstr(&MI);
        LocalizedInstrs.insert(LocalizedMI);
        MachineInstr &UseMI = *MOUse.getParent();
        if (MRI->hasOneNonDBGUse(Reg) && !UseMI.isPHI())
          InsertMBB->insert(UseMI, LocalizedMI);
        else
          InsertMBB->insert(InsertMBB->SkipPHIsAndLabels(InsertMBB->begin()),
                            LocalizedMI);

        // Set a new register for the definition.
        Register NewReg = MRI->cloneVirtualRegister(Reg);
        LocalizedMI->getOperand(0).setReg(NewReg);
        NewVRegIt =
            MBBWithLocalDef.try_emplace(MBBAndReg, NewReg).first;
        LLVM_DEBUG(dbgs() << "Inserted: " << *LocalizedMI);
      }
      LLVM_DEBUG(dbgs() << "Update use with: " << printReg(NewVRegIt->second)
                        << '\n');
      // Update the user reg.
      MOUse.setReg(NewVRegIt->second);
    }
  }
  return Changed;
}

bool LocalizerImpl::localizeIntraBlock(LocalizedSetVecT &LocalizedInstrs) {
  bool Changed = false;

  // For each already-localized instruction which has multiple users, then we
  // scan the block top down from the current position until we hit one of them.

  // FIXME: Consider doing inst duplication if live ranges are very long due to
  // many users, but this case may be better served by regalloc improvements.

  for (MachineInstr *MI : LocalizedInstrs) {
    Register Reg = MI->getOperand(0).getReg();
    MachineBasicBlock &MBB = *MI->getParent();
    // All of the user MIs of this reg.
    SmallPtrSet<MachineInstr *, 32> Users;
    for (MachineInstr &UseMI : MRI->use_nodbg_instructions(Reg)) {
      if (!UseMI.isPHI())
        Users.insert(&UseMI);
    }
    MachineBasicBlock::iterator II(MI);
    // If all the users were PHIs then they're not going to be in our block, we
    // may still benefit from sinking, especially since the value might be live
    // across a call.
    if (Users.empty()) {
      // Make sure we don't sink in between two terminator sequences by scanning
      // forward, not backward.
      II = MBB.getFirstTerminatorForward();
      LLVM_DEBUG(dbgs() << "Only phi users: moving inst to end: " << *MI);
    } else {
      ++II;
      while (II != MBB.end() && !Users.count(&*II))
        ++II;
      assert(II != MBB.end() && "Didn't find the user in the MBB");
      LLVM_DEBUG(dbgs() << "Intra-block: moving " << *MI << " before " << *II);
    }

    // Leave each DBG_VALUE above MI that uses one of its results behind as
    // undef, so that it still ends any earlier location of the variable, and
    // put a copy right after MI, like MachineSink does. No copy is needed if a
    // later DBG_VALUE of the same variable replaces it before that.
    SmallVector<MachineInstr *, 4> DbgUsers;
    SmallVector<MachineInstr *, 4> DbgUsersToCopy;
    for (MachineInstr &DbgMI :
         make_range(std::next(MachineBasicBlock::iterator(MI)), II)) {
      if (!DbgMI.isDebugValue())
        continue;
      erase_if(DbgUsersToCopy, [&](const MachineInstr *Earlier) {
        return Earlier->getDebugVariable() == DbgMI.getDebugVariable() &&
               Earlier->getDebugLoc().getInlinedAt() ==
                   DbgMI.getDebugLoc().getInlinedAt() &&
               Earlier->getDebugExpression()->fragmentsOverlap(
                   DbgMI.getDebugExpression());
      });
      if (DbgMI.hasDebugOperandForReg(Reg)) {
        DbgUsers.push_back(&DbgMI);
        DbgUsersToCopy.push_back(&DbgMI);
      }
    }

    MI->removeFromParent();
    MBB.insert(II, MI);
    Changed = true;

    for (MachineInstr *DbgMI : DbgUsersToCopy)
      MBB.insert(II, MBB.getParent()->CloneMachineInstr(DbgMI));
    for (MachineInstr *DbgMI : DbgUsers)
      DbgMI->setDebugValueUndef();

    // If the instruction (constant) being localized has single user, we can
    // propagate debug location from user.
    if (Users.size() == 1) {
      const auto &DefDL = MI->getDebugLoc();
      const auto &UserDL = (*Users.begin())->getDebugLoc();

      if ((!DefDL || DefDL.getLine() == 0) && UserDL && UserDL.getLine() != 0) {
        MI->setDebugLoc(UserDL);
      }
    }
  }
  return Changed;
}

bool LocalizerImpl::runOnMachineFunction(
    MachineFunction &MF, function_ref<TargetTransformInfo *()> GetTTI) {
  // If the ISel pipeline failed, do not bother running that pass.
  if (MF.getProperties().hasFailedISel())
    return false;

  LLVM_DEBUG(dbgs() << "Localize instructions for: " << MF.getName() << '\n');

  init(MF, GetTTI);

  // Keep track of the instructions we localized. We'll do a second pass of
  // intra-block localization to further reduce live ranges.
  LocalizedSetVecT LocalizedInstrs;

  bool Changed = localizeInterBlock(MF, LocalizedInstrs);
  Changed |= localizeIntraBlock(LocalizedInstrs);
  return Changed;
}

bool LocalizerLegacy::runOnMachineFunction(MachineFunction &MF) {
  LocalizerImpl Impl;
  return Impl.runOnMachineFunction(MF, [&]() {
    return &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(
        MF.getFunction());
  });
}

PreservedAnalyses LocalizerPass::run(MachineFunction &MF,
                                     MachineFunctionAnalysisManager &MFAM) {
  MFPropsModifier<LocalizerPass> _(*this, MF);
  LocalizerImpl Impl;
  bool Changed = Impl.runOnMachineFunction(MF, [&]() {
    Function &F = MF.getFunction();
    FunctionAnalysisManager &FAM =
        MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
            .getManager();
    return &FAM.getResult<TargetIRAnalysis>(F);
  });
  return Changed ? getMachineFunctionPassPreservedAnalyses()
                       .preserveSet<CFGAnalyses>()
                 : PreservedAnalyses::all();
}