/*========================== 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 "GenISAIntrinsics/GenIntrinsicInst.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) {
    switch (IGC_GET_FLAG_VALUE(ForceOCLSIMDWidth)) {
    case 0:
        break;
    case 8:
        return SIMDMode::SIMD8;
    case 16:
        return SIMDMode::SIMD16;
    case 32:
        return SIMDMode::SIMD32;
    }

    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 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(SimdSize)
            return lanesToSIMDMode(SimdSize);
    }

    // rule for pvc
    if(CGCtx->platform.isProductChildOf(IGFX_PVC)) {
        if(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 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 0;

    auto Type = V->getType();
    unsigned int TypeSizeInBits = (unsigned int) DL.getTypeSizeInBits(Type);
    unsigned int Multiplier = SIMD;
    if (WI && WI->isUniform(V))
        Multiplier = 1;
    unsigned int SizeInBytes = TypeSizeInBits * Multiplier / 8;
    return SizeInBytes;
}

unsigned int 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 0;

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

    unsigned int ResultSizeInBytes = 0;

    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;

        unsigned int 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);
    }
}

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

    unsigned int Result = 0;
    for (auto El : Set) {
        unsigned int 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;

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

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

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

        unsigned int SizeUpdate = 0;
        SizeUpdate = addOperandsToSet(Inst, BBSet, SIMD, WI, DL);

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

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

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

    FGA = getAnalysisIfAvailable<GenXFunctionGroupAnalysis>();
    MDUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
    CGCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
    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: ";
        }
        unsigned int SizeInBytes = BBListing[Inst];
        unsigned int AmountOfRegistersRoundUp = RPE->bytesToRegisters(SizeInBytes);
        MaxPressureInFunction = std::max(MaxPressureInFunction, AmountOfRegistersRoundUp);
        Output += std::to_string(SizeInBytes) + " (" +
                  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;

    // force print all
    PrinterType = 2;

    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 << "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>();
    WI = &getAnalysis<WIAnalysis>();
    CGCtx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
    MaxPressureInFunction = 0;

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

    if (DumpToFile) {
        PrinterType = 1;
        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";
        PRINT(Output);
        Output.clear();
    }

    return false;
}

