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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "Compiler/CISACodeGen/RegisterEstimator.hpp"
#include "Compiler/CISACodeGen/ShaderCodeGen.hpp"
#include "Compiler/CISACodeGen/LiveVars.hpp"
#include "Compiler/CodeGenPublic.h"
#include "Compiler/IGCPassSupport.h"
#include "common/debug/Debug.hpp"
#include "common/igc_regkeys.hpp"
#include "common/LLVMWarningsPush.hpp"
#include "llvm/ADT/SparseBitVector.h"
#include <llvm/IR/CFG.h>
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/SetVector.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/IR/DataLayout.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/Support/MathExtras.h>
#include <llvmWrapper/IR/DerivedTypes.h>
#include "common/LLVMWarningsPop.hpp"
#include "Probe/Assertion.h"

using namespace llvm;
using namespace IGC;


char RegisterEstimator::ID = 0;

// Register pass to igc-opt
#define PASS_FLAG "igc-registerestimator"
#define PASS_DESCRIPTION "Use LivenessAnalysis to calculate register pressure"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS true
IGC_INITIALIZE_PASS_BEGIN(RegisterEstimator, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(LivenessAnalysis)
IGC_INITIALIZE_PASS_END(RegisterEstimator, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)


uint32_t RegisterEstimator::getNumLiveGRFAtInst(Instruction* I, uint16_t simdsize)
{
    InstToRegUsageMap::iterator MI = m_LiveVirtRegs.find(I);
    if (MI != m_LiveVirtRegs.end())
    {
        RegUse& ruse = MI->second.allUses[REGISTER_CLASS_GRF];
        return getNumRegs(ruse, simdsize);
    }

    BasicBlock* BB = I->getParent();
    RegUsage nCurrLiveIns = m_BBLiveInVirtRegs[BB];

    ValueToIntMap& ValueIds = m_LVA->ValueIds;
    ValueToValueVecMap& KillInfo = m_LVA->KillInsts;

    // Calculate the number of lives for each instruction of this BB
    for (BasicBlock::iterator II = BB->begin(), IE = BB->end();
        II != IE; ++II)
    {
        Instruction* Inst = &*II;
        if (Inst == I)
        {
            RegUse& ruse = nCurrLiveIns.allUses[REGISTER_CLASS_GRF];
            return getNumRegs(ruse, simdsize);
        }

        // Kills
        ValueToValueVecMap::iterator II0 = KillInfo.find(Inst);
        if (II0 != KillInfo.end())
        {
            ValueVec& VS = II0->second;
            for (int i = 0, e = (int)VS.size(); i < e; ++i)
            {
                Value* killVal = VS[i];
                if (const RegUse * pregs = getRegUse(killVal))
                {
                    RegUse& ruse_curr = nCurrLiveIns.allUses[pregs->rClass];
                    ruse_curr -= (*pregs);
                }
            }
        }

        // Defs
        ValueToIntMap::iterator II1 = ValueIds.find(Inst);
        if (II1 != ValueIds.end())
        {
            RegUse regs = estimateNumOfRegs(Inst);
            RegUse& ruse_curr = nCurrLiveIns.allUses[regs.rClass];
            ruse_curr += regs;
        }
    }
    IGC_ASSERT_MESSAGE(0, "Instruction is not in its BB, something wrong!");
    return 0;
}

// Add register usage of all values in BV into RUsage:
// RUsage:  both in and out.
void RegisterEstimator::addRegUsage(RegUsage& RUsage, SBitVector& BV)
{
    for (SBitVector::iterator I = BV.begin(), E = BV.end();
        I != E; ++I)
    {
        int id = *I;
        if (const RegUse * pregs = getRegUse(id))
        {
            RUsage.allUses[pregs->rClass] += (*pregs);
        }
    }
}

RegUse RegisterEstimator::estimateNumOfRegs(Value* V) const
{
    // Assume no sharing GRF among Values.
    RegUse regs;
    if (getValueRegClass(V) == REGISTER_CLASS_FLAG)
    {
        // Flag register
        regs.rClass = (RegClass)REGISTER_CLASS_FLAG;
        regs.nregs_simd16 = 1;
        return regs;
    }

    // GRF
    Type* Ty = V->getType();
    if (!Ty->isVoidTy())
    {
        IGCLLVM::FixedVectorType* VTy = dyn_cast<IGCLLVM::FixedVectorType>(Ty);
        Type* eltTy = VTy ? VTy->getElementType() : Ty;
        uint32_t nelts = VTy ? int_cast<uint32_t>(VTy->getNumElements()) : 1;
        uint32_t eltBits = (uint32_t)m_DL->getTypeSizeInBits(eltTy);
        uint32_t nBytes = nelts * ((eltBits + 7) / 8);

        regs.rClass = (RegClass)REGISTER_CLASS_GRF;
        if (m_LVA->isUniform(V))
        {
            // round up to DW
            uint16_t sz = (nBytes + DWORD_SIZE_IN_BYTE - 1) / DWORD_SIZE_IN_BYTE;
            regs.uniformInBytes += sz;
        }
        else
        {
            uint16_t sz16 = (nBytes * 16 + GRF_SIZE_IN_BYTE - 1) / GRF_SIZE_IN_BYTE; // SIMD16
            regs.nregs_simd16 = sz16;
        }
    }
    return regs;
}

void RegisterEstimator::calculate(bool doRPEPerInst)
{
    if (m_BBMaxLiveVirtRegs.size() > 0)
    {
        // Already computed, just return.
        return;
    }
    if (!doRPEPerInst && IGC_GET_FLAG_VALUE(ForceRPE) == FLAG_LEVEL_2)
    {
        doRPEPerInst = true;
    }

    m_LVA->calculate(m_F);

    // Pre-allocate maps
    // We will allocate a bit more (~10%) in case it needs expansion.
    // Also, as llvm::DenseMap will resize if the Map's capacity is 75% full,
    // allocate even more to avoid such automatic resizing.
    if (doRPEPerInst)
    {
        // Use IdValues's size as the number of values to be considered.
        size_t nVals = m_LVA->IdValues.size();

        m_LiveVirtRegs.clear();
        uint32_t mapCap1 = int_cast<uint32_t>((size_t)(nVals * 1.40f));
        m_LiveVirtRegs.reserve(mapCap1);
    }

    uint32_t mapCap2 = int_cast<uint32_t>((size_t)(m_F->size() * 1.40f));
    m_BBMaxLiveVirtRegs.clear();
    m_BBLiveInVirtRegs.clear();
    m_BBMaxLiveVirtRegs.reserve(mapCap2);
    m_BBLiveInVirtRegs.reserve(mapCap2);

    BBLiveInMap& BBLiveIns = m_LVA->BBLiveIns;
    ValueToIntMap& ValueIds = m_LVA->ValueIds;
    ValueToValueVecMap& KillInfo = m_LVA->KillInsts;
    for (Function::iterator BI = m_F->begin(), BE = m_F->end();
        BI != BE; ++BI)
    {
        BasicBlock* BB = &*BI;
        SBitVector& BitVec = BBLiveIns[BB];
        RegUsage nCurrLiveIns;

        // Calculate the number of live-ins at entry to BB
        for (SBitVector::iterator I = BitVec.begin(), E = BitVec.end();
            I != E; ++I)
        {
            int id = *I;
            if (const RegUse * pregs = getRegUse(id))
            {
                nCurrLiveIns.allUses[pregs->rClass] += (*pregs);
            }
        }

        m_BBLiveInVirtRegs.insert(std::make_pair(BB, nCurrLiveIns));

        RegUsage bbMaxRegs = nCurrLiveIns;

        // Calculate the number of lives for each instruction of this BB.
        // For simplicity, the number of lives are the one that is at exit
        // of the instruction, not at entry (nor the largest of the both).
        for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
        {
            Instruction* Inst = &*I;
            ValueToIntMap::iterator IDef = ValueIds.find(Inst);
            if (IDef != ValueIds.end())
            {
                if (const RegUse * pregs = getRegUse(Inst))
                {
                    RegUse& ruse_curr = nCurrLiveIns.allUses[pregs->rClass];
                    ruse_curr += (*pregs);
                }
            }

            // Kills
            ValueToValueVecMap::iterator IKill = KillInfo.find(Inst);
            if (IKill != KillInfo.end())
            {
                ValueVec& VS = IKill->second;
                for (int i = 0, e = (int)VS.size(); i < e; ++i)
                {
                    Value* killVal = VS[i];
                    if (const RegUse * pregs = getRegUse(killVal))
                    {
                        RegUse& ruse_curr = nCurrLiveIns.allUses[pregs->rClass];
                        ruse_curr -= (*pregs);
                    }
                }
            }

            for (int i = 0; i < REGISTER_CLASS_TOTAL; ++i)
            {
                RegClass RC = (RegClass)i;
                RegUse& ruse1 = bbMaxRegs.allUses[RC];
                RegUse& ruse2 = nCurrLiveIns.allUses[RC];

                if (ruse1 < ruse2)
                {
                    ruse1 = ruse2;
                }
            }
            if (doRPEPerInst)
            {
                m_LiveVirtRegs.insert(std::make_pair(Inst, nCurrLiveIns));
            }
        }

        // save it in map
        m_BBMaxLiveVirtRegs.insert(std::make_pair(BB, bbMaxRegs));

        for (int i = 0; i < REGISTER_CLASS_TOTAL; ++i)
        {
            RegUse& ruse_max = m_MaxRegs.allUses[(RegClass)i];
            RegUse& ruse_curr = bbMaxRegs.allUses[(RegClass)i];
            if (ruse_max < ruse_curr)
            {
                ruse_max = ruse_curr;
            }
        }
    }

#if 0
    // Sort LiveVirtRegs in decreasing order. As Map cannot be sorted, a list
    // is used for sorting it.
    for (ValueToIntMap::iterator I = m_LiveVirtRegs.begin(),
        E = m_LiveVirtRegs.end();
        I != E; ++I)
    {
        Value* V = I->first;
        m_AllValues.push_back(V);
    }
    std::sort(m_AllValues.begin(), m_AllValues.end(), isNRegGreater(m_LiveVirtRegs));
#endif
    int dumpLevel = IGC_GET_FLAG_VALUE(RPEDumpLevel);
    if (dumpLevel > FLAG_LEVEL_0)
    {
        print(errs(), dumpLevel);
    }
}

int RegisterEstimator::getNUsesInBB(Value* V, BasicBlock* BB)
{
    int nUses = 0;
    for (auto UI = V->user_begin(), UE = V->user_end(); UI != UE; ++UI)
    {
        Instruction* I = cast<Instruction>(*UI);
        if (I->getParent() == BB)
        {
            ++nUses;
        }
    }
    return nUses;
}

bool RegisterEstimator::runOnFunction(Function& F)
{
    m_F = &F;
    m_DL = &F.getParent()->getDataLayout();
    m_LVA = &getAnalysis<LivenessAnalysis>();

    m_WIA = getAnalysisIfAvailable<WIAnalysis>();

    uint32_t nVals = (uint32_t)m_LVA->IdValues.size();
    uint32_t Caps = (uint32_t)m_LVA->IdValues.capacity();
    m_ValueRegUses.reserve(Caps);
    m_ValueRegUses.resize(nVals);

    // 1. Pre-compute register needed for each values.
    // 2. The max possible registers needed, assume all values are live,
    RegUsage estNumRegs;  // RPE assuming all values are live
    ValueToIntMap::iterator VI, VE;
    for (VI = m_LVA->ValueIds.begin(), VE = m_LVA->ValueIds.end(); VI != VE; ++VI)
    {
        Value* Val = VI->first;
        uint32_t valId = VI->second;
        RegUse regs = estimateNumOfRegs(Val);
        RegUse& ruse = estNumRegs.allUses[regs.rClass];
        ruse += regs;
        IGC_ASSERT_MESSAGE(valId < nVals, "ValueIds does not match IdValues!");
        m_ValueRegUses[valId] = regs;
    }

    m_noGRFPressure = isGRFPressureLow(16, estNumRegs);

    // Note that runOnFunction does not do RPE calculation unless ForceRPE
    // is enabled (for debugging).  The RPE is calcualted for users to call
    // calculate() explicitly.
    if (IGC_IS_FLAG_ENABLED(ForceRPE))
    {
        // Calculate register pressure for each BB
        calculate();
    }
    return false;
}

void RegisterEstimator::print(raw_ostream& OS, BasicBlock* BB, int dumpLevel)
{
    uint32_t grf_livein_simd8 = getNumLiveInGRFAtBB(BB, 8);
    uint32_t grf_livein_simd16 = getNumLiveInGRFAtBB(BB, 16);
    uint32_t grf_livein_simd32 = getNumLiveInGRFAtBB(BB, 32);
    uint32_t grf_maxlive_simd8 = getMaxLiveGRFAtBB(BB, 8);
    uint32_t grf_maxlive_simd16 = getMaxLiveGRFAtBB(BB, 16);
    uint32_t grf_maxlive_simd32 = getMaxLiveGRFAtBB(BB, 32);

    if (BB->hasName())
        OS << "  " << BB->getName() << "  ";
    OS << "\n  ";
    OS << "livein simd8|16|32=<"
        << grf_livein_simd8 << ", "
        << grf_livein_simd16 << ", "
        << grf_livein_simd32 << ">  ";
    OS << "maxLive simd8|16|32=<"
        << grf_maxlive_simd8 << ", "
        << grf_maxlive_simd16 << ", "
        << grf_maxlive_simd32 << ">\n";
    if (dumpLevel < 2)
    {
        return;
    }

    // dumpLevel > 1
    IntToValueVector& IdValues = m_LVA->IdValues;
    //ValueToValueSetMap& KillInfo = m_LVA.KillInsts;
    BBLiveInMap& BBLiveIns = m_LVA->BBLiveIns;

    SBitVector& BitVec = BBLiveIns[BB];
    int nVals = 0;
    for (SBitVector::iterator I = BitVec.begin(), E = BitVec.end();
        I != E; ++I)
    {
        int id = *I;
        Value* V = IdValues[id];
        IGC_ASSERT_MESSAGE(nullptr != V, "Value should be in Value Map!");

        const RegUse* pregs = getRegUse(id);
        if (pregs == nullptr) {
            IGC_ASSERT_MESSAGE(0, "Missing RegUse!");
            continue;
        }

        if (nVals == 0) {
            OS << "    ";
        }
        if (V->hasName())
        {
            OS << V->getName() << " ";
        }
        else
        {
            OS << LivenessAnalysis::getllvmValueName(V) << " ";
        }
        OS << "{"
            << getNumRegs(*pregs, 8) << ", "
            << getNumRegs(*pregs, 16) << ", "
            << getNumRegs(*pregs, 32) << "}, ";
        ++nVals;
        if (nVals == 6)
        {
            // 6 values per line
            OS << "\n";
            nVals = 0;
        }
    }
    OS << (nVals != 0 ? "\n\n" : "\n");

    for (BasicBlock::iterator it = BB->begin(), ie = BB->end(); it != ie; ++it)
    {
        Instruction* I = &*it;
        if (m_LiveVirtRegs.count(I) == 0) {
            OS << "  { not computed } ";
        }
        else
        {
            RegUsage& ruse = m_LiveVirtRegs[I];
            RegUse& grfuse = ruse.allUses[REGISTER_CLASS_GRF];
            OS << "  {"
                << getNumRegs(grfuse, 8) << ", "
                << getNumRegs(grfuse, 16) << ", "
                << getNumRegs(grfuse, 32) << "} ";
        }
        OS << *I;
        if (!m_pBB2ID.empty()) {
            if (BranchInst * BrI = dyn_cast<BranchInst>(I))
            {
                OS << "  (";
                for (int i = 0, e = (int)BrI->getNumSuccessors(); i < e; ++i)
                {
                    BasicBlock* b = BrI->getSuccessor(i);
                    OS << " BB[" << m_pBB2ID[b] << "]";
                }
                OS << " )";
            }
        }
        OS << "\n";
    }
}

void RegisterEstimator::print(raw_ostream& OS, int dumpLevel)
{
    if (!m_F) return;

    std::stringstream ss;
    ss << "RegisterEstimator: " << m_F->getName().str();
    Debug::Banner(OS, ss.str());

    int bid = 0;
    m_pBB2ID.clear();
    for (Function::iterator I = m_F->begin(), E = m_F->end(); I != E; ++I) {
        BasicBlock* BB = &*I;
        m_pBB2ID[BB] = bid;
        ++bid;
    }
    for (Function::iterator I = m_F->begin(), E = m_F->end(); I != E; ++I) {
        BasicBlock* BB = &*I;
        if (dumpLevel > FLAG_LEVEL_1)
        {
            OS << "\n";
        }
        OS << "BB[" << m_pBB2ID[BB] << "]:";
        print(OS, BB, dumpLevel);
        ++bid;
    }
    OS << "\n\n";
    m_pBB2ID.clear();
}

RegPressureTracker::RegPressureTracker(RegisterEstimator* RPE) :
    m_BB(nullptr),
    m_pRPE(RPE)
{
    m_TrackRegPressure = !m_pRPE->hasNoGRFPressure();
    if (!m_TrackRegPressure)
    {
        return;
    }

    m_pRPE->calculate();

#if 0
    // TODO: need per-BB possible Max estimate
    if (m_pRPE->isGRFPressureLow(32))
    {
        m_TrackRegPressure = false;
        return;
    }
#endif

    // Pre-allocate DenseMap
    size_t nVals = m_pRPE->getNumValues();
    uint32_t mapCap = int_cast<uint32_t>((size_t)(nVals * 1.40f));
    m_DeadValueNumUses.grow(mapCap);
}

void RegPressureTracker::init(BasicBlock* BB, bool doMaxRegInBB)
{
    m_BB = BB;
    if (!m_TrackRegPressure)
    {
        return;
    }

    m_LiveOutSet.clear();
    m_DeadValueNumUses.clear();

    LivenessAnalysis* LVA = m_pRPE->getLivenessAnalysis();

    BBLiveInMap& BBLiveIns = LVA->BBLiveIns;
    SBitVector& BitVec = BBLiveIns[BB];
    for (succ_iterator SI = succ_begin(BB), E = succ_end(BB); SI != E; ++SI)
    {
        BasicBlock* SuccBB = *SI;
        SBitVector& succBitVec = BBLiveIns[SuccBB];
        m_LiveOutSet |= succBitVec;
    }

    for (SBitVector::iterator I = BitVec.begin(), E = BitVec.end();
        I != E; ++I)
    {
        int id = *I;
        if (m_LiveOutSet.test(id))
        {
            continue;
        }
        // This Value has the last use in this BB, add it into the map.
        Value* V = LVA->IdValues[id];
        int nUses = m_pRPE->getNUsesInBB(V, BB);
        m_DeadValueNumUses[V] = nUses;
    }

    for (BasicBlock::iterator BI = m_BB->begin(), BE = m_BB->end();
        BI != BE; ++BI)
    {
        Instruction* Inst = &*BI;
        Value* V = Inst;
        ValueToIntMap::iterator VI = LVA->ValueIds.find(V);
        if (VI != LVA->ValueIds.end())
        {
            int id = VI->second;
            if (m_LiveOutSet.test(id))
            {
                continue;
            }
            int nUses = m_pRPE->getNUsesInBB(V, BB);
            m_DeadValueNumUses[V] = nUses;
        }
    }

    m_pRPE->getMaxLiveinRegsAtBB(m_RUsage, BB);
}

void RegPressureTracker::advance(llvm::Instruction* I)
{
    if (!m_TrackRegPressure)
    {
        return;
    }

    Value* V = I;
    LivenessAnalysis* LVA = m_pRPE->getLivenessAnalysis();
    auto VI = LVA->ValueIds.find(V);
    if (VI != LVA->ValueIds.end())
    {
        int id = VI->second;
        if (const RegUse * pRegs = m_pRPE->getRegUse(id))
        {
            m_RUsage.allUses[pRegs->rClass] += *pRegs;
        }
    }

    // Update current register pressure
    for (auto OI = I->op_begin(), OE = I->op_end(); OI != OE; ++OI)
    {
        Value* Opr = *OI;
        if (isa<Constant>(Opr))
        {
            continue;
        }
        auto DVI = m_DeadValueNumUses.find(Opr);
        if (DVI != m_DeadValueNumUses.end())
        {
            int& numUses = DVI->second;
            --numUses;
            if (numUses == 0)
            {
                if (const RegUse * pRegs = m_pRPE->getRegUse(Opr))
                {
                    m_RUsage.allUses[pRegs->rClass] -= *pRegs;
                }
            }
        }
    }
}


#if defined( _DEBUG )

void RegisterEstimator::dump()
{
    dump(1);
}

void RegisterEstimator::dump(int dumpLevel)
{
    print(dbgs(), dumpLevel);

}

void RegisterEstimator::dump(BasicBlock* BB)
{
    dump(BB, 1);
}

void RegisterEstimator::dump(BasicBlock* BB, int dumpLevel)
{
    print(dbgs(), BB, dumpLevel);
}

#endif
