/*========================== begin_copyright_notice ============================

Copyright (C) 2023 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/
#include "IGCLivenessAnalysis.h"
#include "AdaptorCommon/ImplicitArgs.hpp"
#include "Compiler/CISACodeGen/helper.h"
#include "Compiler/CodeGenContextWrapper.hpp"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "common/debug/Debug.hpp"

#include "Compiler/IGCPassSupport.h"
#include "common/debug/Debug.hpp"
#include "common/igc_regkeys.hpp"
#include "llvmWrapper/IR/Function.h"

#include <fstream>
#include <queue>
#include <regex>

char IGCLivenessAnalysis::ID = 0;
// Register pass to igc-opt
#define PASS_FLAG2 "igc-df-liveness"
#define PASS_DESCRIPTION2                                                                                              \
  "computes full dataflow liveness analysis & and register pressure "                                                  \
  "estimator"
#define PASS_CFG_ONLY2 false
#define PASS_ANALYSIS2 true
IGC_INITIALIZE_PASS_BEGIN(IGCLivenessAnalysis, PASS_FLAG2, PASS_DESCRIPTION2, PASS_CFG_ONLY2, PASS_ANALYSIS2)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_END(IGCLivenessAnalysis, PASS_FLAG2, PASS_DESCRIPTION2, PASS_CFG_ONLY2, PASS_ANALYSIS2)
#define PRINT(str) llvm::errs() << str

IGCLivenessAnalysis::IGCLivenessAnalysis() : FunctionPass(ID) {
  initializeIGCLivenessAnalysisPass(*PassRegistry::getPassRegistry());
};

unsigned int IGCLivenessAnalysisBase::registerSizeInBytes() {
  if (CGCtx->platform.isProductChildOf(IGFX_PVC))
    return 64;
  return 32;
}

SIMDMode IGCLivenessAnalysisBase::bestGuessSIMDSize(Function *F, GenXFunctionGroupAnalysis *FGA) {
  switch (IGC_GET_FLAG_VALUE(ForceOCLSIMDWidth)) {
  case 0:
    break;
  case 8:
    return SIMDMode::SIMD8;
  case 16:
    return SIMDMode::SIMD16;
  case 32:
    return SIMDMode::SIMD32;
  }

  // if we can find metadata with stipulations how we should compile we use it
  if (F && MDUtils->findFunctionsInfoItem(F) != MDUtils->end_FunctionsInfo()) {
    IGC::IGCMD::FunctionInfoMetaDataHandle funcInfoMD = MDUtils->getFunctionsInfoItem(F);
    unsigned SimdSize = funcInfoMD->getSubGroupSize()->getSIMDSize();
    if (FGA) {
      llvm::Function *Kernel = F;
      auto *FG = FGA->getGroup(F);
      Kernel = FG ? FG->getHead() : nullptr;
      if (Kernel && MDUtils->findFunctionsInfoItem(F) != MDUtils->end_FunctionsInfo()) {
        funcInfoMD = MDUtils->getFunctionsInfoItem(Kernel);
        SimdSize = funcInfoMD->getSubGroupSize()->getSIMDSize();
      }
    }
    if (SimdSize) {
      return lanesToSIMDMode(SimdSize);
    }
  }

  // rule for pvc
  if (CGCtx->platform.isProductChildOf(IGFX_PVC)) {
    bool abortOnSpills =
        IGC_GET_FLAG_VALUE(AllowSIMD16DropForXE2Plus) && (CGCtx->platform.isCoreXE2() || CGCtx->platform.isCoreXE3());
    auto FG = FGA ? FGA->getGroup(F) : nullptr;
    bool hasStackCall = (FG && FG->hasStackCall()) || (F && F->hasFnAttribute("visaStackCall"));
    bool isIndirectGroup = FG && FGA->isIndirectCallGroup(FG);
    bool hasSubroutine = FG && !FG->isSingle() && !hasStackCall && !isIndirectGroup;
    if (abortOnSpills || hasSubroutine) {
      return SIMDMode::SIMD16;
    }
    return SIMDMode::SIMD32;
  }

  return SIMDMode::SIMD8;
}

ValueSet IGCLivenessAnalysisBase::getDefs(llvm::BasicBlock &BB) {

  ValueSet &BBIn = In[&BB];
  ValueSet &BBOut = Out[&BB];

  ValueSet Difference;
  for (auto *elem : BBOut)
    if (!BBIn.count(elem))
      Difference.insert(elem);

  return Difference;
}

// we scan through all successors and merge their INSETs as our OUTSET
void IGCLivenessAnalysisBase::combineOut(llvm::BasicBlock *BB) {
  ValueSet *OutSet = &Out[BB];
  for (llvm::succ_iterator SI = llvm::succ_begin(BB), SE = llvm::succ_end(BB); SI != SE; ++SI) {
    // for every successor we take all of it's IN values
    // and the PHI values that are coming from our BB
    OutSet->insert(In[*SI].begin(), In[*SI].end());
  }
}

unsigned int IGCLivenessAnalysisBase::getSizeForType(Value *V, const DataLayout &DL) {

  if (V->getType()->isPointerTy())
    return CGCtx->getRegisterPointerSizeInBits(V->getType()->getPointerAddressSpace());

  return (unsigned int)DL.getTypeSizeInBits(V->getType());
}

PressurePair IGCLivenessAnalysisBase::computeSizeInBytes(Value *V, unsigned int SIMD, WIAnalysisRunner *WI,
                                                         const DataLayout &DL) {
  // when we check size of operands, this check is redundant
  // but allows for a nicer code
  bool NoRetVal = V->getType()->isVoidTy();
  if (NoRetVal)
    return {};

  unsigned int TypeSizeInBits = getSizeForType(V, DL);

  if (WI && WI->isUniform(V))
    return {TypeSizeInBits / 8, 0};

  return {0, TypeSizeInBits * SIMD / 8};
}

PressurePair IGCLivenessAnalysisBase::addOperandsToSet(llvm::Instruction *Inst, ValueSet &Set, unsigned int SIMD,
                                                       WIAnalysisRunner *WI, const DataLayout &DL) {

  // we do not process PHI's there
  auto Phi = llvm::dyn_cast<llvm::PHINode>(Inst);
  if (Phi)
    return {};

  // do not process debug instructions and lifetimehints in any way
  if (Inst->isDebugOrPseudoInst() || Inst->isLifetimeStartOrEnd())
    return {};

  PressurePair ResultSizeInBytes = {};

  for (auto &Op : Inst->operands()) {
    llvm::Value *V = Op.get();
    // We are counting only instructions right now
    // potetntially we should also count globals, but
    // we defintely shouldn't count:
    // br label %bb1 (basic block names)
    // call %functionName (function names)
    // add %a, 1 (constants)
    if (!(llvm::isa<llvm::Instruction>(V) || llvm::isa<llvm::Argument>(V)))
      continue;

    PressurePair SizeOfOperand = computeSizeInBytes(V, SIMD, WI, DL);
    if (!Set.count(V))
      ResultSizeInBytes += SizeOfOperand;

    Set.insert(V);
  }

  return ResultSizeInBytes;
}

void IGCLivenessAnalysisBase::addNonLocalOperandsToSet(llvm::Instruction *Inst, ValueSet &Set) {

  // do not process debug instructions and lifetimehints in any way
  if (Inst->isDebugOrPseudoInst() || Inst->isLifetimeStartOrEnd())
    return;

  for (auto &Op : Inst->operands()) {
    llvm::Value *V = Op.get();
    // We are counting only instructions right now
    // potetntially we should also count globals, but
    // we defintely shouldn't count:
    // br label %bb1 (basic block names)
    // call %functionName (function names)
    // add %a, 1 (constants)
    Instruction *I = dyn_cast<Instruction>(V);
    bool IsInstruction = I != nullptr;
    bool OperandInDifferentBB = IsInstruction && (I->getParent() != Inst->getParent());
    bool IsArgument = !IsInstruction && llvm::isa<llvm::Argument>(V);
    if (OperandInDifferentBB || IsArgument) {
      Set.insert(V);
    }
  }
}

// idea is simple, each predecessor block that converges into our block
// has its own set of PHI values, that it has to deliver
// so we take values that are coming from each block
// and add them to their OUT set directly
void IGCLivenessAnalysisBase::addToPhiSet(llvm::PHINode *Phi, PhiSet *InPhiSet) {
  for (auto BB : Phi->blocks()) {
    auto ValueFromOurBlock = Phi->getIncomingValueForBlock(BB);
    auto *OutSet = &Out[BB];
    if (llvm::isa<llvm::Constant>(ValueFromOurBlock))
      continue;
    OutSet->insert(ValueFromOurBlock);
    // for visualization purposes also add them to INPHI set
    // this way it will be easier to scan them and print
    (*InPhiSet)[BB].insert(ValueFromOurBlock);
  }
}

// scan through block in reversed order and add each operand
// into IN block while deleting defined values
void IGCLivenessAnalysisBase::processBlock(llvm::BasicBlock *BB, ValueSet &Set, PhiSet *PhiSet) {
  for (auto RI = BB->rbegin(), RE = BB->rend(); RI != RE; ++RI) {
    llvm::Instruction *Inst = &(*RI);
    Set.erase(Inst);

    auto Phi = llvm::dyn_cast<llvm::PHINode>(Inst);
    if (Phi) {
      addToPhiSet(Phi, PhiSet);
      continue;
    }
    addNonLocalOperandsToSet(Inst, Set);
  }
}

void IGCLivenessAnalysisBase::livenessAnalysis(llvm::Function &F, BBSet *StartBBs) {
  std::queue<llvm::BasicBlock *> Worklist;

  if (StartBBs != nullptr) {
    // If StartBBs are provided we know that only these BBs could be changed
    // Add only them to the initial Worklist
    for (BasicBlock *BB : *StartBBs)
      Worklist.push(BB);
  } else {
    // Start with adding all BBs to the Worklist
    // to make sure In set is populated for every BB
    for (auto BBIt = IGCLLVM::rbegin(&F); BBIt != IGCLLVM::rend(&F); ++BBIt) {
      Worklist.push(&*BBIt);
    }
  }

  while (!Worklist.empty()) {

    llvm::BasicBlock *BB = Worklist.front();
    Worklist.pop();

    ValueSet *InSet = &In[BB];
    ValueSet *OutSet = &Out[BB];
    PhiSet *InPhiSet = &InPhi[BB];

    combineOut(BB);

    unsigned int SizeBefore = InSet->size();
    unsigned int SizeBeforePhi = InPhiSet->size();

    // this should be a copy, everything that should go OUT
    // minus what is defined and what was intermediate, should go IN
    // so we are feeding copy of the OUT as IN, and delete everything
    // that was defined in the block while processing it
    *InSet = *OutSet;
    processBlock(BB, *InSet, InPhiSet);

    // sets can only grow monotonically, updated set is a bigger set
    bool IsSetChanged = InSet->size() != SizeBefore;
    bool IsPhiSetChanged = InPhiSet->size() != SizeBeforePhi;
    if (IsSetChanged || IsPhiSetChanged)
      for (auto *Pred : predecessors(BB))
        Worklist.push(Pred);
  }
}

PressurePair IGCLivenessAnalysisBase::estimateSizeInBytes(const llvm::SmallSetVector<llvm::Instruction *, 16> &Set,
                                                          Function &F, unsigned int SIMD, WIAnalysisRunner *WI) {
  const DataLayout &DL = F.getParent()->getDataLayout();

  PressurePair Result = {};
  for (auto El : Set) {
    PressurePair SizeInBytes = computeSizeInBytes(El, SIMD, WI, DL);
    Result += SizeInBytes;
  }
  return Result;
}

PressurePair IGCLivenessAnalysisBase::estimateSizeInBytes(ValueSet &Set, Function &F, unsigned int SIMD,
                                                          WIAnalysisRunner *WI) {
  const DataLayout &DL = F.getParent()->getDataLayout();

  PressurePair Result = {};
  for (auto El : Set) {
    PressurePair SizeInBytes = computeSizeInBytes(El, SIMD, WI, DL);
    Result += SizeInBytes;
  }
  return Result;
}

void IGCLivenessAnalysisBase::collectPressureForBB(llvm::BasicBlock &BB, InsideBlockPressureMap &BBListing,
                                                   unsigned int SIMD, WIAnalysisRunner *WI) {

  const DataLayout &DL = BB.getParent()->getParent()->getDataLayout();
  ValueSet &BBOut = Out[&BB];
  // this should be a copy
  ValueSet BBSet = BBOut;

  PressurePair SizeInBytes = estimateSizeInBytes(BBSet, *BB.getParent(), SIMD, WI);

  for (auto RI = BB.rbegin(), RE = BB.rend(); RI != RE; ++RI) {

    llvm::Instruction *Inst = &(*RI);

    PressurePair SizeUpdate = {};
    SizeUpdate = addOperandsToSet(Inst, BBSet, SIMD, WI, DL);

    BBListing[Inst] = SizeInBytes;
    SizeInBytes += SizeUpdate;

    if (!BBSet.count(Inst))
      continue;
    PressurePair InstSizeInBytes = computeSizeInBytes(Inst, SIMD, WI, DL);
    SizeInBytes -= InstSizeInBytes;
    BBSet.erase(Inst);
  }
}

bool IGCLivenessAnalysis::runOnFunction(llvm::Function &F) {
  auto *FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>();
  auto *MDUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
  auto *CGCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  LivenessRunner = IGCLivenessAnalysisRunner(CGCtx, MDUtils, FGA);
  LivenessRunner.livenessAnalysis(F, nullptr);
  return false;
}

char IGCRegisterPressurePrinter::ID = 0;
// Register pass to igc-opt
#define PASS_FLAG1 "igc-pressure-printer"
#define PASS_DESCRIPTION1 "prints register pressure estimation"
#define PASS_CFG_ONLY1 false
#define PASS_ANALYSIS1 false
IGC_INITIALIZE_PASS_BEGIN(IGCRegisterPressurePrinter, PASS_FLAG2, PASS_DESCRIPTION2, PASS_CFG_ONLY2, PASS_ANALYSIS2)
IGC_INITIALIZE_PASS_DEPENDENCY(WIAnalysis)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(IGCFunctionExternalRegPressureAnalysis)
IGC_INITIALIZE_PASS_DEPENDENCY(IGCLivenessAnalysis)
IGC_INITIALIZE_PASS_END(IGCRegisterPressurePrinter, PASS_FLAG1, PASS_DESCRIPTION1, PASS_CFG_ONLY1, PASS_ANALYSIS1)
#define PRINT(str) llvm::errs() << str

IGCRegisterPressurePrinter::IGCRegisterPressurePrinter() : FunctionPass(ID) {
  initializeIGCRegisterPressurePrinterPass(*PassRegistry::getPassRegistry());
};

IGCRegisterPressurePrinter::IGCRegisterPressurePrinter(const std::string &FileName) : FunctionPass(ID) {
  initializeIGCRegisterPressurePrinterPass(*PassRegistry::getPassRegistry());
  DumpFileName = FileName;
  DumpToFile = true;
};

void IGCRegisterPressurePrinter::printInstruction(llvm::Instruction *Inst, std::string &Str) {
  llvm::raw_string_ostream rso(Str);
  Inst->print(rso, false);
  rso << "\n";
}

void IGCRegisterPressurePrinter::printIntraBlock(llvm::BasicBlock &BB, std::string &Output,
                                                 InsideBlockPressureMap &BBListing) {
  for (auto &I : BB) {
    llvm::Instruction *Inst = &I;
    if (WI->isUniform(Inst)) {
      Output += "U: ";
    } else {
      Output += "N: ";
    }
    PressurePair SizeInBytes = BBListing[Inst];
    // to prevent overcounting by one
    unsigned int AmountOfRegistersRoundUp = RPE->bytesToRegisters((unsigned int)SizeInBytes);
    PressurePair PairInRegisters = RPE->bytesToRegisters(SizeInBytes);
    MaxPressureInFunction = std::max(MaxPressureInFunction, AmountOfRegistersRoundUp);
    MaxPressurePair = std::max(MaxPressurePair, PairInRegisters);

    if (PrinterType > 5)
      Output += "UP: " + std::to_string(PairInRegisters.Uniform) +
                " NP: " + std::to_string(PairInRegisters.NonUniform) + " \t";
    Output +=
        std::to_string((unsigned int)SizeInBytes) + " \t (" + std::to_string(AmountOfRegistersRoundUp) + ")" + "    \t";
    printInstruction(Inst, Output);
  }
}

void IGCRegisterPressurePrinter::intraBlock(llvm::BasicBlock &BB, std::string &Output, unsigned int SIMD) {
  InsideBlockPressureMap BBListing;
  RPE->collectPressureForBB(BB, BBListing, SIMD, &WI->Runner);
  printIntraBlock(BB, Output, BBListing);
}

void IGCRegisterPressurePrinter::dumpRegPressure(llvm::Function &F, unsigned int SIMD) {
  const char *Filter = IGC_GET_REGKEYSTRING(DumpRegPressureEstimateFilter);
  if (Filter && *Filter != '\0' && !std::regex_search(F.getName().str(), std::regex(Filter)))
    return;

  std::string Output;
  Output.reserve(32768);

  IGC::Debug::DumpLock();
  {
    const int MaxFunctionLength = 100;

    std::stringstream ss;
    std::string FunctionName = F.getName().str();

    std::replace_if(FunctionName.begin(), FunctionName.end(), [](char c) { return !std::isalnum(c) && c != '_'; }, '_');

    if (FunctionName.length() > MaxFunctionLength) {
      FunctionName = FunctionName.substr(0, MaxFunctionLength);
    }

    ss << FunctionName << "_" << DumpFileName << "_RegEst";
    auto Name = Debug::DumpName(IGC::Debug::GetShaderOutputName())
                    .Hash(CGCtx->hash)
                    .Type(CGCtx->type)
                    .Retry(CGCtx->m_retryManager->GetRetryId())
                    .Pass(ss.str().c_str())
                    .Extension("ll");

    int Counter = 1;
    std::string originalName = Name.str();
    std::ifstream FileToCheck(Name.str());
    while (FileToCheck.good()) {
      FileToCheck.close();
      Name = Debug::DumpName(IGC::Debug::GetShaderOutputName())
                 .Hash(CGCtx->hash)
                 .Type(CGCtx->type)
                 .Retry(CGCtx->m_retryManager->GetRetryId())
                 .Pass((ss.str() + "_" + std::to_string(Counter)).c_str())
                 .Extension("ll");
      FileToCheck.open(Name.str());
      Counter++;
    }
    FileToCheck.close();

    std::ofstream OutputFile(Name.str());
    OutputFile << "SIMD: " << SIMD << ", external pressure: " << ExternalPressure << "\n";

    for (BasicBlock &BB : F) {
      // prints information for one BB
      printSets(&BB, Output, SIMD);
      if (OutputFile.is_open())
        OutputFile << Output;
      Output.clear();
    }

    OutputFile << "==============================================" << "\n";
    OutputFile << "MaxPressure In Function: " << MaxPressureInFunction + ExternalPressure << "\n";
    OutputFile << "MaxUniformPressure In Function Isolated: " << MaxPressurePair.Uniform << "\n";
    OutputFile << "MaxNonUniformPressure In Function Isolated: " << MaxPressurePair.NonUniform << "\n";
    OutputFile << "MaxPressure In Function Isolated: " << MaxPressureInFunction << "\n";

    OutputFile.close();
  }
  IGC::Debug::DumpUnlock();
}

void IGCRegisterPressurePrinter::printDefs(const ValueSet &In, const ValueSet &Out, std::string &Output) {
  ValueSet Difference, Common, Kills;

  for (auto *elem : Out)
    if (!In.count(elem))
      Difference.insert(elem);

  for (auto *elem : Out)
    if (In.count(elem))
      Common.insert(elem);

  for (auto *elem : In)
    if (!Common.count(elem))
      Kills.insert(elem);

  if (!Kills.empty()) {
    Output += "KILL:\t[\t" + std::to_string(Kills.size()) + "\t] ";
    if (PrinterType > 2)
      printNames(Kills, Output);
    Output += "\n";
  }

  if (!Difference.empty()) {
    Output += "DEF:\t[\t" + std::to_string(Difference.size()) + "\t] ";
    if (PrinterType > 2)
      printDefNames(Difference, Output);
    Output += "\n";
  }
}

void IGCRegisterPressurePrinter::printName(llvm::Value *Val, std::string &String) {
  llvm::raw_string_ostream Rso(String);
  // false means don't print type
  Val->printAsOperand(Rso, false);
}

void IGCRegisterPressurePrinter::printSets(llvm::BasicBlock *BB, std::string &Output, unsigned int SIMD) {
  if (PrinterType <= 0)
    return;

  Output += "block: " + BB->getName().str() + " ";
  Output += "function: " + BB->getParent()->getName().str() + "\n";

  ValueSet &PtrSetIn = RPE->getInSet()[BB];
  ValueSet &PtrSetOut = RPE->getOutSet()[BB];
  PhiSet &PtrSetPhi = RPE->getInPhiSet()[BB];

  printPhi(PtrSetPhi, Output);

  Output += "IN: \t[\t" + std::to_string(PtrSetIn.size()) + "\t] ";
  if (PrinterType > 2)
    printNames(PtrSetIn, Output);
  Output += "\n";

  printDefs(PtrSetIn, PtrSetOut, Output);

  Output += "OUT:\t[\t" + std::to_string(PtrSetOut.size()) + "\t] ";
  if (PrinterType > 2)
    printNames(PtrSetOut, Output);
  Output += "\n";

  intraBlock(*BB, Output, SIMD);
}

void IGCRegisterPressurePrinter::printNames(const ValueSet &Set, std::string &Name) {
  for (auto *val : Set) {
    llvm::raw_string_ostream Rso(Name);
    // false means don't print type
    val->printAsOperand(Rso, false);
    Rso << ", ";
  }
}

int getAmountOfBB(llvm::Value *Val) {
  llvm::SmallPtrSet<llvm::BasicBlock *, 16> BBSet;
  for (auto U : Val->users()) {

    auto UserInst = llvm::dyn_cast<Instruction>(U);
    if (!UserInst)
      continue;
    auto UserBB = UserInst->getParent();
    BBSet.insert(UserBB);
  }
  return BBSet.size();
}

void IGCRegisterPressurePrinter::printDefNames(const ValueSet &Set, std::string &Name) {

  for (auto *Val : Set) {
    llvm::raw_string_ostream Rso(Name);
    // false means don't print type
    Val->printAsOperand(Rso, false);
    int HowManyBB = getAmountOfBB(Val);
    Rso << "(" << std::to_string(Val->getNumUses()) << ")[" + std::to_string(HowManyBB) << "], ";
  }
}

void IGCRegisterPressurePrinter::printPhi(const PhiSet &Set, std::string &Output) {

  if (Set.empty())
    return;

  Output += "PHIS:\t[\t" + std::to_string(Set.size()) + "\t] ";
  for (const auto &Elem : Set) {
    auto &OneBBSet = Elem.second;
    unsigned int Size = OneBBSet.size();
    if (PrinterType > 1) {
      Output += "{ ";
      printNames(OneBBSet, Output);
      Output += "} ";
      Output += "(" + std::to_string(Size) + ")";
    }
  }
  Output += "\n";
}

bool IGCRegisterPressurePrinter::runOnFunction(llvm::Function &F) {

  ExternalPressure = getAnalysis<IGCFunctionExternalRegPressureAnalysis>().getExternalPressureForFunction(&F);
  RPE = &getAnalysis<IGCLivenessAnalysis>().getLivenessRunner();
  WI = &getAnalysis<WIAnalysis>();
  CGCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  MaxPressureInFunction = 0;
  MaxPressurePair = {};
  FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>();

  unsigned int SIMD = numLanes(RPE->bestGuessSIMDSize(&F, FGA));

  if (DumpToFile) {
    dumpRegPressure(F, SIMD);
  } else {
    // basically only for LIT testing
    std::string Output;
    // no particular reason behind this, just big enough power of 2
    // helps to reduce printing time, by preemptively allocating memory
    Output.reserve(32768);
    Output += "SIMD: " + std::to_string(SIMD) + ", external pressure: " + std::to_string(ExternalPressure) + "\n";
    for (BasicBlock &BB : F) {
      printSets(&BB, Output, SIMD);
    }
    Output += "\n";
    Output += "==============================================\n";
    Output += "MaxPressure In Function: " + F.getName().str() + " --> " + std::to_string(MaxPressureInFunction) + "\n";
    Output += "Max Uniform Pressure In Function: " + F.getName().str() + " --> " +
              std::to_string(MaxPressurePair.Uniform) + "\n";
    Output += "Max NonUniform Pressure In Function: " + F.getName().str() + " --> " +
              std::to_string(MaxPressurePair.NonUniform) + "\n";
    PRINT(Output);
    Output.clear();
  }

  return false;
}

char IGCRegisterPressurePublisher::ID = 0;
// Register pass to igc-opt
#define PASS_FLAG3 "igc-pressure-publisher"
#define PASS_DESCRIPTION3 "puts metadata with register pressure estimation"
#define PASS_CFG_ONLY1 false
#define PASS_ANALYSIS1 false
IGC_INITIALIZE_PASS_BEGIN(IGCRegisterPressurePublisher, PASS_FLAG3, PASS_DESCRIPTION3, PASS_CFG_ONLY1, PASS_ANALYSIS1)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(IGCFunctionExternalRegPressureAnalysis)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
IGC_INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
IGC_INITIALIZE_PASS_END(IGCRegisterPressurePublisher, PASS_FLAG3, PASS_DESCRIPTION3, PASS_CFG_ONLY1, PASS_ANALYSIS1)

IGCRegisterPressurePublisher::IGCRegisterPressurePublisher() : llvm::ModulePass(ID) {
  initializeIGCRegisterPressurePublisherPass(*PassRegistry::getPassRegistry());
};

bool IGCRegisterPressurePublisher::runOnModule(llvm::Module &M) {
  if (skipModule(M)) {
    return false;
  }

  auto *CGCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
  auto *MDUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
  auto *ModMD = CGCtx->getModuleMetaData();
  auto *FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>();

  for (auto &F : M) {
    if (F.isDeclaration() || F.isIntrinsic() || F.hasOptNone()) {
      continue;
    }

    IGCLivenessAnalysisRunner RPE(CGCtx, MDUtils, FGA);

    unsigned int SimdSize = numLanes(RPE.bestGuessSIMDSize(&F, FGA));
    unsigned int Existing = RPE.checkPublishRegPressureMetadata(F);
    bool LegacyPublished = (Existing != 0);
    bool PerSIMDPublished = false;
    if (MDUtils->findFunctionsInfoItem(&F) != MDUtils->end_FunctionsInfo()) {
      auto md = MDUtils->getFunctionsInfoItem(&F);
      PerSIMDPublished = (md->getMaxRegPressureForSIMDSize(numLanes(SIMDMode::SIMD16))->getMaxPressure() != 0) &&
                         (md->getMaxRegPressureForSIMDSize(numLanes(SIMDMode::SIMD32))->getMaxPressure() != 0);
    }

    // Compile-time guard: passes between CodeLoopSinking and this point are
    // assumed to only reduce RP. If the early estimate is already below the
    // lowest possible early-retry threshold, the refined value can never trip
    // any threshold -- skip the WIAnalysis + IGCLiveness re-run entirely.
    unsigned int MinThreshold =
        std::min(IGC_GET_FLAG_VALUE(EarlyRetryLargeGRFThreshold), IGC_GET_FLAG_VALUE(EarlyRetryDefaultGRFThreshold));
    bool RepublishLate = IGC_IS_FLAG_ENABLED(EnableLateRPRepublish) && LegacyPublished && Existing >= MinThreshold;

    if (LegacyPublished && !RepublishLate && PerSIMDPublished) {
      continue;
    }

    RPE.livenessAnalysis(F, nullptr);
    PressurePair ExternalPressurePair =
        getAnalysis<IGCFunctionExternalRegPressureAnalysis>().getExternalPressurePairForFunction(&F);

    auto *DT = &getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
    auto *PDT = &getAnalysis<PostDominatorTreeWrapperPass>(F).getPostDomTree();
    auto *LI = &getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();

    TranslationTable TT;
    TT.run(F);
    WIAnalysisRunner WI(&F, LI, DT, PDT, MDUtils, CGCtx, ModMD, &TT, false);
    WI.run();

    if (!LegacyPublished || RepublishLate) {
      unsigned int MaxPressureInFunction = RPE.getMaxRegCountForFunction(F, SimdSize, &WI);
      RPE.publishRegPressureMetadata(F, MaxPressureInFunction + RPE.bytesToRegisters(ExternalPressurePair));
    }

    if (!PerSIMDPublished) {
      // Publish per-SIMD estimates so later SIMD-selection heuristics
      // (e.g. shouldDropToSIMD16 on VRT platforms) can compare SIMD16 vs SIMD32
      // pressure instead of a single measurement.
      PressurePair ExternalPressurePair16 = {
          ExternalPressurePair.Uniform, 16 * (uint32_t)llvm::divideCeil(ExternalPressurePair.NonUniform, SimdSize)};
      PressurePair Pair16Pressure =
          RPE.getMaxPressurePairForFunction(F, numLanes(SIMDMode::SIMD16), &WI) + ExternalPressurePair16;
      unsigned int Simd16Pressure = RPE.bytesToRegisters(Pair16Pressure);
      unsigned int Simd32Pressure = RPE.bytesToRegisters({Pair16Pressure.Uniform, 2 * Pair16Pressure.NonUniform});
      RPE.publishRegPressureMetadataForSIMD(F, numLanes(SIMDMode::SIMD16), Simd16Pressure);
      RPE.publishRegPressureMetadataForSIMD(F, numLanes(SIMDMode::SIMD32), Simd32Pressure);
    }
  }

  return false;
}
