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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#ifndef __REMAT_H__
#define __REMAT_H__

#include "FlowGraph.h"
#include "GraphColor.h"
#include "RPE.h"
#include <list>
#include <map>

namespace vISA
{
// Remat will trigger only for vars that have less than following uses
#define MAX_USES_REMAT 6

// Minimum def-use distance for remat to trigger
#define MIN_DEF_USE_DISTANCE 20

// Distance in instructions to reuse rematted value in BB
#define MAX_LOCAL_REMAT_REUSE_DISTANCE 40

    typedef std::pair<G4_INST*, G4_BB*> Reference;
    class References
    {
    public:
        std::vector<Reference> def;
        //std::vector<Reference> uses;

        unsigned int numUses = 0;

        // lastUseLexId provides a quick and dirty way
        // to determine end of liveness for a variable.
        // This is not always accurate due to holes
        // in live-ranges but should be fine most times.
        unsigned int lastUseLexId = 0;

        // Store number of times this var has been used in
        // a remat'd operation. This forms part of heuristic
        // that decides if it is profitable to remat an
        // operation inside a loop.
        unsigned int numRemats = 0;

        // Store set of rows of this variable ever used.
        // This is useful for samplers.
        std::unordered_set<unsigned int> rowsUsed;
    };

    class Rematerialization
    {
    private:
        G4_Kernel& kernel;
        const LivenessAnalysis& liveness;
        GraphColor& coloring;
        GlobalRA& gra;
        G4_Declare* samplerHeader = nullptr;
        unsigned int numRematsInLoop = 0;
        bool IRChanged = false;
        bool samplerHeaderMapPopulated = false;
        unsigned int loopInstsBeforeRemat = 0;
        unsigned int totalInstsBeforeRemat = 0;
        RPE& rpe;

        static const unsigned int cRematLoopRegPressure128GRF = 85;
        static const unsigned int cRematRegPressure128GRF = 120;

        unsigned int rematLoopRegPressure = 0;
        unsigned int rematRegPressure = 0;

        std::vector<G4_Declare*> preDefinedVars;
        std::vector<G4_Declare*> spills;
        // For each top dcl, this map holds all defs
        std::unordered_map<G4_Declare*, References> operations;
        // This vector contains declares that could potentially save spill
        // if remat'd.
        std::vector<bool> rematCandidates;

        // Map each sampler instruction with instruction initializing
        // samplerHeader instruction. This is required when inserting
        // remat'd samplers.
        std::unordered_map<G4_INST*, G4_INST*> samplerHeaderMap;
        std::unordered_set<G4_BB*> deLVNedBBs;
        // Map BB->subroutine it belongs to
        // BBs not present are assumed to belong to main kernel
        std::unordered_map<G4_BB*, const FuncInfo*> BBPerSubroutine;
        bool cr0DefBB = false;

        void populateRefs();
        void populateSamplerHeaderMap();
        void deLVNSamplers(G4_BB*);
        bool usesNoMaskWA(const Reference* uniqueDef);
        bool canRematerialize(G4_SrcRegRegion*, G4_BB*, const Reference*&, INST_LIST_ITER instIter);
        G4_SrcRegRegion* rematerialize(G4_SrcRegRegion*, G4_BB*, const Reference*, std::list<G4_INST*>&, G4_INST*&);
        G4_SrcRegRegion* createSrcRgn(G4_SrcRegRegion*, G4_DstRegRegion*, G4_Declare*);
        const Reference* findUniqueDef(References&, G4_SrcRegRegion*);
        bool areInSameLoop(G4_BB*, G4_BB*, bool&);
        bool isRangeSpilled(G4_Declare*);
        bool areAllDefsInBB(G4_Declare*, G4_BB*, unsigned int);
        unsigned int getLastUseLexId(G4_Declare*);
        bool checkLocalWAR(G4_INST*, G4_BB*, INST_LIST_ITER);
        void updateSplitInfo(G4_INST *dstInst, int srcNum);
        void reduceNumUses(G4_Declare* dcl)
        {
            auto opIt = operations.find(dcl);
            if (opIt != operations.end())
            {
                auto numUses = (*opIt).second.numUses;
                if (numUses > 0)
                    opIt->second.numUses = numUses - 1;

                if (numUses == 1)
                {
                    for (auto ref : opIt->second.def)
                    {
                        ref.second->remove(ref.first);
                    }
                    opIt->second.def.clear();
                }
            }


        }

        unsigned int getNumUses(G4_Declare* dcl) const
        {
            auto opIt = operations.find(dcl);
            if (opIt != operations.end())
                return opIt->second.numUses;

            return 0;
        }

        void incNumRemat(G4_Declare* dcl)
        {
            auto opIt = operations.find(dcl);
            if (opIt != operations.end())
                opIt->second.numRemats += 1;
        }

        unsigned int getNumRemats(G4_Declare* dcl) const
        {
            auto opIt = operations.find(dcl);
            if (opIt != operations.end())
                return opIt->second.numRemats;

            return 0;
        }

        bool isRematCandidateOp(G4_INST* inst) const
        {
            if (inst->isFlowControl() || inst->isWait() ||
                (inst->isSend() && inst->asSendInst()->isFence()) ||
                inst->isLifeTimeEnd() || inst->isAccDstInst() || inst->isAccSrcInst() ||
                inst->getImplAccDst() || inst->getImplAccSrc() || inst->isRelocationMov())
            {
                return false;
            }

            if (inst->isCallerRestore() || inst->isCallerSave() || inst->isCalleeRestore() || inst->isCalleeSave())
            {
                return false;
            }

            G4_Declare* dcl = nullptr;
            if (inst->getDst() && inst->getDst()->getTopDcl())
                dcl = inst->getDst()->getTopDcl();

            if (kernel.fg.builder->isPreDefArg(dcl) || kernel.fg.builder->isPreDefRet(dcl) ||
                kernel.fg.builder->isPreDefFEStackVar(dcl))
                return false;

            return true;
        }

        void cleanRedundantSamplerHeaders();

        unsigned int getNumRematsInLoop() const { return numRematsInLoop; }
        void incNumRematsInLoop() { numRematsInLoop++; }
        bool inSameSubroutine(G4_BB*, G4_BB*);

        bool isPartGRFBusyInput(G4_Declare* inputDcl, unsigned int atLexId);

    public:
        Rematerialization(G4_Kernel& k, const LivenessAnalysis& l, GraphColor& c, RPE& r, GlobalRA& g) :
            kernel(k), liveness(l), coloring(c), gra(g), rpe(r)
        {
            unsigned numGRFs = k.getNumRegTotal();
            auto scale = [=](unsigned threshold) -> unsigned {
                float ratio = 1.0f - (128 - threshold) / 128.0f;
                return static_cast<unsigned>(numGRFs * ratio);
            };
            rematLoopRegPressure = scale(cRematLoopRegPressure128GRF);
            rematRegPressure = scale(cRematRegPressure128GRF);

            rematCandidates.resize(l.getNumSelectedVar(), false);

            for (auto&& lr : coloring.getSpilledLiveRanges())
            {
                auto dcl = lr->getDcl()->getRootDeclare();
                if (!dcl->isSpilled())
                {
                    spills.push_back(dcl);
                    dcl->setSpillFlag();
                }

                // Mark all simultaneously live variables as remat candidates
                unsigned int spillId = dcl->getRegVar()->getId();
                const auto& intfVec = coloring.getIntf()->getSparseIntfForVar(spillId);

                for (auto intfId : intfVec)
                {
                    rematCandidates[intfId] = true;
                }
            }

            std::set<G4_BB*> bbsInLoop;
            for (auto&& be : kernel.fg.backEdges)
            {
                auto loopIt = kernel.fg.naturalLoops.find(be);

                if (loopIt != kernel.fg.naturalLoops.end())
                {
                    bbsInLoop.insert(loopIt->second.begin(), loopIt->second.end());
                }
            }

            for (auto& bb : kernel.fg)
            {
                bool bbInLoop = (bbsInLoop.find(bb) != bbsInLoop.end());
                if (bbInLoop)
                {
                    for (auto& inst : *bb)
                    {
                        if (!inst->isLabel() && !inst->isPseudoKill())
                        {
                            loopInstsBeforeRemat++;
                        }
                    }
                }
            }

            // Map BBs in subroutines
            for (auto curFuncInfo : kernel.fg.funcInfoTable)
            {
                const auto& bbList = curFuncInfo->getBBList();
                for (auto bb : bbList)
                {
                    BBPerSubroutine.insert(std::make_pair(bb, curFuncInfo));
                }
            }
        }

        ~Rematerialization()
        {
            for (auto&& dcl : spills)
            {
                dcl->resetSpillFlag();
            }
        }

        bool getChangesMade() const { return IRChanged; }

        void run();
    };
}
#endif
