//===-- X86PreTileConfig.cpp - Tile Register Pre-configure-----------------===//
//
// 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 Pass to pre-config the shapes of AMX registers
/// AMX register needs to be configured before use. The shapes of AMX register
/// are encoded in the 1st and 2nd machine operand of AMX pseudo instructions.
///
/// The instruction ldtilecfg is used to config the shapes. It must be reachable
/// for all variable shapes. ldtilecfg will be inserted more than once if we
/// cannot find a dominating point for all AMX instructions.
///
/// The configure register is caller saved according to ABI. We need to insert
/// ldtilecfg again after the call instruction if callee clobbers any AMX
/// registers.
///
/// This pass calculates all points that ldtilecfg need to be inserted to and
/// insert them. It reports error if the reachability conditions aren't met.
//
//===----------------------------------------------------------------------===//

#include "X86.h"
#include "X86InstrBuilder.h"
#include "X86MachineFunctionInfo.h"
#include "X86RegisterInfo.h"
#include "X86Subtarget.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/IR/Module.h"
#include "llvm/InitializePasses.h"

using namespace llvm;

#define DEBUG_TYPE "x86-pre-tile-config"

static void emitErrorMsg(MachineFunction &MF) {
  LLVMContext &Context = MF.getFunction().getContext();
  Context.emitError(
      MF.getName() +
      ": Failed to config tile register, please define the shape earlier");
}

namespace {

struct MIRef {
  MachineInstr *MI = nullptr;
  MachineBasicBlock *MBB = nullptr;
  // A virtual position for instruction that will be inserted after MI.
  size_t Pos = 0;
  MIRef() = default;
  MIRef(MachineBasicBlock *MBB) : MBB(MBB) {
    for (auto I = MBB->begin(), E = MBB->end(); I != E && I->isPHI();
         ++I, ++Pos)
      MI = &*I;
  }
  MIRef(MachineInstr *MI)
      : MI(MI), MBB(MI->getParent()),
        Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {}
  MIRef(MachineInstr *MI, MachineBasicBlock *MBB)
      : MI(MI), MBB(MBB),
        Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {}
  MIRef(MachineInstr *MI, MachineBasicBlock *MBB, size_t Pos)
      : MI(MI), MBB(MBB), Pos(Pos) {}
  operator bool() const { return MBB != nullptr; }
  bool operator==(const MIRef &RHS) const {
    return MI == RHS.MI && MBB == RHS.MBB;
  }
  bool operator!=(const MIRef &RHS) const { return !(*this == RHS); }
  bool operator<(const MIRef &RHS) const {
    // Comparison between different BBs happens when inserting a MIRef into set.
    // So we compare MBB first to make the insertion happy.
    return std::tie(MBB, Pos) < std::tie(RHS.MBB, RHS.Pos);
  }
  bool operator>(const MIRef &RHS) const {
    // Comparison between different BBs happens when inserting a MIRef into set.
    // So we compare MBB first to make the insertion happy.
    return std::tie(MBB, Pos) > std::tie(RHS.MBB, RHS.Pos);
  }
};

struct BBInfo {
  MIRef FirstAMX;
  MIRef LastCall;
  bool HasAMXRegLiveIn = false;
  bool TileCfgForbidden = false;
  bool NeedTileCfgLiveIn = false;
};

class X86PreTileConfigImpl {
  std::function<MachineLoopInfo *()> GetMLI;
  MachineRegisterInfo *MRI = nullptr;
  const MachineLoopInfo *MLI = nullptr;
  SmallPtrSet<MachineInstr *, 8> DefVisited;
  DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo;
  DenseMap<MachineBasicBlock *, SmallVector<MIRef, 8>> ShapeBBs;

  /// Check if the callee will clobber AMX registers.
  bool isDestructiveCall(MachineInstr &MI, BitVector UsableRegs) {
    auto Iter = llvm::find_if(
        MI.operands(), [](MachineOperand &MO) { return MO.isRegMask(); });
    if (Iter == MI.operands_end())
      return false;
    UsableRegs.clearBitsInMask(Iter->getRegMask());
    return !UsableRegs.none();
  }

  /// Check if MI is AMX pseudo instruction.
  bool isAMXInstruction(MachineInstr &MI) {
    if (MI.isPHI() || MI.isDebugInstr() || MI.getNumOperands() < 3)
      return false;
    switch (MI.getOpcode()) {
    case X86::PTILESTOREDV:
    case X86::PTCVTROWD2PSrteV:
    case X86::PTCVTROWD2PSrtiV:
    case X86::PTCVTROWPS2BF16HrteV:
    case X86::PTCVTROWPS2BF16HrtiV:
    case X86::PTCVTROWPS2BF16LrteV:
    case X86::PTCVTROWPS2BF16LrtiV:
    case X86::PTCVTROWPS2PHHrteV:
    case X86::PTCVTROWPS2PHHrtiV:
    case X86::PTCVTROWPS2PHLrteV:
    case X86::PTCVTROWPS2PHLrtiV:
    case X86::PTILEMOVROWrteV:
    case X86::PTILEMOVROWrtiV:
      return true;
    }

    // We can simply check if it is AMX instruction by its def.
    // But we should exclude old API which uses physical registers.
    MachineOperand &MO = MI.getOperand(0);
    if (!MO.isReg() || !MO.getReg().isVirtual())
      return false;

    if (MRI->getRegClass(MO.getReg())->getID() != X86::TILERegClassID)
      return false;

    collectShapeInfo(MI);
    return true;
  }

  /// Check if it is an edge from loop bottom to loop head.
  bool isLoopBackEdge(MachineBasicBlock *Header, MachineBasicBlock *Bottom) {
    if (!MLI->isLoopHeader(Header))
      return false;
    auto *ML = MLI->getLoopFor(Header);
    if (ML->contains(Bottom) && ML->isLoopLatch(Bottom))
      return true;

    return false;
  }

  /// Collect the shape def information for later use.
  void collectShapeInfo(MachineInstr &MI);

  /// Try to hoist shapes definded below AMX instructions.
  bool hoistShapesInBB(MachineBasicBlock *MBB, SmallVectorImpl<MIRef> &Shapes) {
    MIRef &FirstAMX = BBVisitedInfo[MBB].FirstAMX;
    auto FirstShapeBelowAMX = llvm::lower_bound(Shapes, FirstAMX);
    auto InsertPoint = FirstAMX.MI->getIterator();
    for (auto I = FirstShapeBelowAMX, E = Shapes.end(); I != E; ++I) {
      // Do not hoist instructions that access memory.
      if (I->MI->mayLoadOrStore())
        return false;
      for (auto &MO : I->MI->operands()) {
        if (MO.isDef())
          continue;
        // Do not hoist instructions if the sources' def under AMX instruction.
        // TODO: We can handle isMoveImmediate MI here.
        if (MO.isReg() && MIRef(MRI->getVRegDef(MO.getReg())) > FirstAMX)
          return false;
        // TODO: Maybe need more checks here.
      }
      MBB->insert(InsertPoint, I->MI->removeFromParent());
    }
    // We only need to mark the last shape in the BB now.
    Shapes.clear();
    Shapes.push_back(MIRef(&*--InsertPoint, MBB));
    return true;
  }

  /// Clear MF related structures.
  void releaseMemory() {
    ShapeBBs.clear();
    DefVisited.clear();
    BBVisitedInfo.clear();
  }

public:
  X86PreTileConfigImpl(std::function<MachineLoopInfo *()> GetMLI)
      : GetMLI(GetMLI) {}
  bool runOnMachineFunction(MachineFunction &MF);
};

class X86PreTileConfigLegacy : public MachineFunctionPass {
public:
  X86PreTileConfigLegacy() : MachineFunctionPass(ID) {}

  /// Return the pass name.
  StringRef getPassName() const override {
    return "Tile Register Pre-configure";
  }

  /// X86PreTileConfig analysis usage.
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesAll();
    AU.addRequired<MachineLoopInfoWrapperPass>();
    MachineFunctionPass::getAnalysisUsage(AU);
  }

  /// Perform ldtilecfg instructions inserting.
  bool runOnMachineFunction(MachineFunction &MF) override;

  static char ID;
};

} // end anonymous namespace

char X86PreTileConfigLegacy::ID = 0;

INITIALIZE_PASS_BEGIN(X86PreTileConfigLegacy, "tilepreconfig",
                      "Tile Register Pre-configure", false, false)
INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
INITIALIZE_PASS_END(X86PreTileConfigLegacy, "tilepreconfig",
                    "Tile Register Pre-configure", false, false)

void X86PreTileConfigImpl::collectShapeInfo(MachineInstr &MI) {
  auto RecordShape = [&](MachineInstr *MI, MachineBasicBlock *MBB) {
    MIRef MIR(MI, MBB);
    auto &Refs = ShapeBBs[MBB];
    auto I = llvm::lower_bound(Refs, MIR);
    if (I == Refs.end() || *I != MIR)
      Refs.insert(I, MIR);
  };

  SmallVector<Register, 8> WorkList(
      {MI.getOperand(1).getReg(), MI.getOperand(2).getReg()});
  while (!WorkList.empty()) {
    Register R = WorkList.pop_back_val();
    MachineInstr *DefMI = MRI->getVRegDef(R);
    assert(DefMI && "R must has one define instruction");
    MachineBasicBlock *DefMBB = DefMI->getParent();
    if (DefMI->isMoveImmediate() || !DefVisited.insert(DefMI).second)
      continue;

    if (DefMI->isPHI()) {
      for (unsigned I = 1; I < DefMI->getNumOperands(); I += 2)
        if (isLoopBackEdge(DefMBB, DefMI->getOperand(I + 1).getMBB()))
          RecordShape(DefMI, DefMBB); // In this case, PHI is also a shape def.
        else
          WorkList.push_back(DefMI->getOperand(I).getReg());
    } else {
      RecordShape(DefMI, DefMBB);
    }
  }
}

bool X86PreTileConfigImpl::runOnMachineFunction(MachineFunction &MF) {
  scope_exit ClearStateOnExit([this] { releaseMemory(); });

  X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
  // Early exit in the common case of non-AMX code.
  if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)
    return false;

  const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();
  const TargetInstrInfo *TII = ST.getInstrInfo();
  const TargetRegisterInfo *TRI = ST.getRegisterInfo();
  const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID);

  BitVector AMXRegs(TRI->getNumRegs());
  for (unsigned I = 0; I < RC->getNumRegs(); I++)
    AMXRegs.set(X86::TMM0 + I);

  // Iterate MF to collect information.
  MRI = &MF.getRegInfo();
  MLI = GetMLI();
  SmallSet<MIRef, 8> CfgNeedInsert;
  SmallVector<MachineBasicBlock *, 8> CfgLiveInBBs;
  for (auto &MBB : MF) {
    size_t Pos = 0;
    auto &Info = BBVisitedInfo[&MBB];
    for (auto &MI : MBB) {
      ++Pos;
      if (isAMXInstruction(MI)) {
        // If there's call before the AMX, we need to reload tile config.
        if (Info.LastCall)
          CfgNeedInsert.insert(Info.LastCall);
        else // Otherwise, we need tile config to live in this BB.
          Info.NeedTileCfgLiveIn = true;
        // Always record the first AMX in case there's shape def after it.
        if (!Info.FirstAMX)
          Info.FirstAMX = MIRef(&MI, &MBB, Pos);
      } else if (MI.isCall() && isDestructiveCall(MI, AMXRegs)) {
        // Record the call only if the callee clobbers all AMX registers.
        Info.LastCall = MIRef(&MI, &MBB, Pos);
      }
    }
    if (Info.NeedTileCfgLiveIn) {
      if (&MBB == &MF.front())
        CfgNeedInsert.insert(MIRef(&MBB));
      else
        CfgLiveInBBs.push_back(&MBB);
    }
    if (Info.FirstAMX || Info.HasAMXRegLiveIn)
      for (auto *Succ : MBB.successors())
        if (!isLoopBackEdge(Succ, &MBB))
          BBVisitedInfo[Succ].HasAMXRegLiveIn = true;
  }

  // Update NeedTileCfgLiveIn for predecessors.
  while (!CfgLiveInBBs.empty()) {
    MachineBasicBlock *MBB = CfgLiveInBBs.pop_back_val();
    for (auto *Pred : MBB->predecessors()) {
      auto &Info = BBVisitedInfo[Pred];
      if (Info.LastCall) {
        CfgNeedInsert.insert(Info.LastCall);
      } else if (!Info.NeedTileCfgLiveIn) {
        Info.NeedTileCfgLiveIn = true;
        if (Pred == &MF.front())
          CfgNeedInsert.insert(MIRef(Pred));
        else
          CfgLiveInBBs.push_back(Pred);
      }
    }
  }

  // There's no AMX instruction if we didn't find a tile config live in point.
  if (CfgNeedInsert.empty())
    return false;

  // Avoid to insert ldtilecfg before any shape defs.
  SmallVector<MachineBasicBlock *, 8> WorkList;
  for (auto &I : ShapeBBs) {
    auto &Info = BBVisitedInfo[I.first];
    // TODO: We can hoist shapes across BBs here.
    if (Info.HasAMXRegLiveIn) {
      // We are not able to config tile registers since the shape to config
      // is not defined yet. Emit error message and continue. The function
      // would not config tile registers.
      emitErrorMsg(MF);
      return false;
    }
    if (Info.FirstAMX && Info.FirstAMX < I.second.back() &&
        !hoistShapesInBB(I.first, I.second)) {
      emitErrorMsg(MF);
      return false;
    }
    WorkList.push_back(I.first);
  }
  while (!WorkList.empty()) {
    MachineBasicBlock *MBB = WorkList.pop_back_val();
    for (auto *Pred : MBB->predecessors()) {
      auto &Info = BBVisitedInfo[Pred];
      if (!Info.TileCfgForbidden && !isLoopBackEdge(MBB, Pred)) {
        Info.TileCfgForbidden = true;
        WorkList.push_back(Pred);
      }
    }
  }

  DebugLoc DL;
  SmallSet<MIRef, 8> VisitedOrInserted;
  int SS = MF.getFrameInfo().CreateStackObject(
      ST.getTileConfigSize(), ST.getTileConfigAlignment(), false);

  // Try to insert for the tile config live in points.
  for (const auto &I : CfgNeedInsert) {
    SmallSet<MIRef, 8> InsertPoints;
    SmallVector<MIRef, 8> WorkList({I});
    while (!WorkList.empty()) {
      MIRef I = WorkList.pop_back_val();
      if (!VisitedOrInserted.count(I)) {
        if (!BBVisitedInfo[I.MBB].TileCfgForbidden) {
          // If the BB is all shapes reachable, stop sink and try to insert.
          InsertPoints.insert(I);
        } else {
          // Avoid the BB to be multi visited.
          VisitedOrInserted.insert(I);
          // Sink the inserting point along the chain with NeedTileCfgLiveIn =
          // true when MBB isn't all shapes reachable.
          for (auto *Succ : I.MBB->successors())
            if (BBVisitedInfo[Succ].NeedTileCfgLiveIn)
              WorkList.push_back(MIRef(Succ));
        }
      }
    }

    // A given point might be forked due to shape conditions are not met.
    for (MIRef I : InsertPoints) {
      // Make sure we insert ldtilecfg after the last shape def in MBB.
      auto It = ShapeBBs.find(I.MBB);
      if (It != ShapeBBs.end() && I < It->second.back())
        I = It->second.back();
      // There're chances the MBB is sunk more than once. Record it to avoid
      // multi insert.
      if (VisitedOrInserted.insert(I).second) {
        auto II = I.MI ? I.MI->getIterator() : I.MBB->instr_begin();
        addFrameReference(BuildMI(*I.MBB, ++II, DL, TII->get(X86::PLDTILECFGV)),
                          SS);
      }
    }
  }

  // Zero stack slot.
  MachineBasicBlock &MBB = MF.front();
  MachineInstr *MI = &*MBB.begin();
  if (ST.hasAVX512()) {
    Register Zmm = MRI->createVirtualRegister(&X86::VR512RegClass);
    BuildMI(MBB, MI, DL, TII->get(X86::AVX512_512_SET0), Zmm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSZmr)), SS)
        .addReg(Zmm);
  } else if (ST.hasAVX2()) {
    Register Ymm = MRI->createVirtualRegister(&X86::VR256RegClass);
    BuildMI(MBB, MI, DL, TII->get(X86::AVX_SET0), Ymm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS)
        .addReg(Ymm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS, 32)
        .addReg(Ymm);
  } else {
    assert(ST.hasSSE2() && "AMX should assume SSE2 enabled");
    unsigned StoreOpc = ST.hasAVX() ? X86::VMOVUPSmr : X86::MOVUPSmr;
    Register Xmm = MRI->createVirtualRegister(&X86::VR128RegClass);
    BuildMI(MBB, MI, DL, TII->get(X86::V_SET0), Xmm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(StoreOpc)), SS).addReg(Xmm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(StoreOpc)), SS, 16)
        .addReg(Xmm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(StoreOpc)), SS, 32)
        .addReg(Xmm);
    addFrameReference(BuildMI(MBB, MI, DL, TII->get(StoreOpc)), SS, 48)
        .addReg(Xmm);
  }
  // Fill in the palette first.
  addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOV8mi)), SS).addImm(1);

  return true;
}

FunctionPass *llvm::createX86PreTileConfigLegacyPass() {
  return new X86PreTileConfigLegacy();
}

bool X86PreTileConfigLegacy::runOnMachineFunction(MachineFunction &MF) {
  X86PreTileConfigImpl Impl(
      [this]() { return &getAnalysis<MachineLoopInfoWrapperPass>().getLI(); });
  return Impl.runOnMachineFunction(MF);
}

PreservedAnalyses
X86PreTileConfigPass::run(MachineFunction &MF,
                          MachineFunctionAnalysisManager &MFAM) {
  X86PreTileConfigImpl Impl(
      [&MFAM, &MF]() { return &MFAM.getResult<MachineLoopAnalysis>(MF); });
  return Impl.runOnMachineFunction(MF)
             ? getMachineFunctionPassPreservedAnalyses()
                   .preserveSet<CFGAnalyses>()
             : PreservedAnalyses::all();
}