/*========================== begin_copyright_notice ============================
Copyright (C) 2024-2025 Intel Corporation

SPDX-License-Identifier: MIT
============================= end_copyright_notice ===========================*/

/*!
 * @file Implementation of the SIMD operation counting tool
 */

#include <algorithm>
#include <array>
#include <vector>
#include <map>
#include <string>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <assert.h>

#include "gtpin_api.h"
#include "gtpin_tool_utils.h"

#include "simdprof_hlif.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Struct SimdProfArgs
/* ============================================================================================= */
/*!
 * SimdProf instrumentation arguments (instruction properties).
 * Each unique combination of these arguments requires a separate instrumentation procedure
 * to be generated for each group of instructions with these properties
 */
struct SimdProfArgs
{
    SimdProfArgs(bool ctrl, uint32_t mask, GtPredicate pred, bool isSend = false) :
                 maskCtrl(ctrl), execMask(mask), predicate(pred), isSendIns(isSend) {}

    inline bool operator< (const SimdProfArgs& other) const;

    bool               maskCtrl;  ///< 'MaskCtrl' flag of instrumented instructions
    uint32_t           execMask;  ///< Execution mask of instrumented instructions
    GtPredicate        predicate; ///< Predicate of instrumented instructions
    bool               isSendIns; ///< true if instrumented instructions are SEND instructions
};

/* ============================================================================================= */
// Struct SimdProfGroup
/* ============================================================================================= */
/*!
 * Structure that holds information and profiling results for a group of instructions being
 * instrumented by a single instrumentation routine.
 * @note All instructions within a group have exactly the same SimdProfArgs.
 * @note In order to provide separate channel counters per instruction category (e.g. integer, FP, etc.),
 *       replace the {insCount, opCount} pair with an array of counter pairs per category.
 */
struct SimdProfGroup
{
    SimdProfGroup(uint32_t bbl, uint32_t count) : bblId(bbl), insCount(count), totalOpCount(0), simdProfCountArgs() { memset(opCounts.data(), 0, 32 * sizeof(uint64_t)); }

    BblId                    bblId;            ///< Identifier of a BBL that contains this group of instructions
    uint32_t                 insCount;         ///< Number of instructions in the group
    uint64_t                 totalOpCount;
    std::array<uint64_t, 32> opCounts;         ///< Number of SIMD operations (effective channels) executed by each instruction in the group
    SimdProfCountArgs        simdProfCountArgs;
};

/* ============================================================================================= */
// Struct SimdProfSection
/* ============================================================================================= */
/*!
 * Structure that holds information on a SimdProf section - sequence of instructions for which
 * instrumentation routines can be inserted at the same point.
 * @note All instructions within a section are executed with the same value of the flag register -
 * single dynamic parameter of the SIMD operation calulator
 */
struct SimdProfSection
{
    SimdProfSection(BblId bid, const IGtIns& headIns) :  bblId(bid), headInsId(headIns.Id()) {}

    /// Add a new instruction to the section. Update the corresponding SimdProf group within this section
    void AddInstruction(const IGtIns& ins);

    BblId                            bblId;     ///< Id of the corresponding basic block
    InsId                            headInsId; ///< First intruction of the section - common 
                                                ///< instrumentation point for all groups in the section
    std::map<SimdProfArgs, uint32_t> groups;    ///< SimdProf groups along with the number of instructions
    std::map<SimdProfArgs, InsId>    groupHeadIds;
};

/* ============================================================================================= */
// Class SimdProfKernelProfile
/* ============================================================================================= */
/*!
 * Class that represents a kernel profiled by the SimdProf instrumentation
 */
class SimdProfKernelProfile
{
public:
    SimdProfKernelProfile(const IGtKernel& kernel);

    /*!
     * Instrument the kernel.
     * The function is called by the OnKernelBuild handler
     * @return success/failure status
     */
    void Instrument(IGtKernelInstrument& instrumentor);

    void WriteArgumentsToDevice(IGtMemoryMapper& memMapper);

    /*!
     * Read profiling results which are assumed to be collected and stored in the buffer
     * associated with the kernel.
     * The function is called by the OnKernelComplete handler
     */
     void ReadProfileData();

     uint64_t GetTotalOpCounter() const { return _totalOpCount; }

    std::string           ToString()        const; ///< @return Text representation of the profile data

private:
    /*!
     * Generate instrumentation procedures for all SimdProf groups of the specified SimdProf section.
     * Insert instrumentation at the beginning of the section.
     * Initialize the _profileData array
     * @param[in]      instrumentor     Instrumentor of the GEN kernel
     * @param[in]      section          SimdProf section to be instrumented
     */
    void InstrumentSection(IGtKernelInstrument& instrumentor, const SimdProfSection& section);

    /// Increment counter of SIMD operations for the specified BBL by 'incValue'
    void UpdateBblOpCounter(BblId bblId, uint64_t* perChannelValues);

    /// @return Extended kernel name
    std::string ExtendedName() const { return _extName; }

private:
    /// Kernel descriptor
    std::string                         _name;              ///< Kernel's name
    std::string                         _extName;           ///< Kernel's extended name
    GtKernelType                        _type;              ///< Kernel's type
    GtGpuPlatform                       _platform;          ///< Kernel's platform
    uint64_t                            _hashId;            ///< Kernel's hash identifier
    GtSimdWidth                         _simd;              ///< Kernel's SIMD width
    uint64_t                            _binarySignature;   ///< Kernel's binary signature
    uint64_t                            _totalOpCount = 0;
    std::list<SimdProfGroup>            _profileData;       ///< Profiling data for instrumented SimdProf groups

    std::map<BblId, std::pair<InsId, InsId> >         _bblInsInfo;   ///< Head and tail instructions per BBL

    typedef struct OpCounts {
        uint64_t counts[32] = {};
        uint64_t total = 0;
    } OpCounts;

    std::map<BblId, OpCounts>                         _bblOpCounts;  ///< Number of executed SIMD operations per BBL

    CountEnableChannelsFunc                           _countEnableChannelsFunc;
    CountEnableChannelsForSendFunc                    _countEnableChannelsForSendFunc;
};

/* ============================================================================================= */
// Class SimdProfHlif
/* ============================================================================================= */
/*!
 * Implementation of the IGtTool interface for the SimdProfHlif tool
 */
class SimdProfHlif : public GtTool
{
public:
    /// Implementation of the IGtTool interface
    const char* Name() const { return "simdprof_hlif"; }

    void OnKernelBuild(IGtKernelInstrument& instrumentor);
    void OnKernelRun(IGtKernelDispatch& dispatcher);
    void OnKernelComplete(IGtKernelDispatch& dispatcher);

public:
    std::string ToString() const;                ///< @return Text representation of the profile data
    static SimdProfHlif* Instance();             ///< @return Single instance of this class
    static void OnFini() { Instance()->Fini(); } ///< Callback function registered with atexit()
    void   LoadHliLibrary();                     ///< Compile and load library of HLI functions

protected:
    SimdProfHlif() = default;
    SimdProfHlif(const SimdProfHlif&) = delete;
    SimdProfHlif& operator = (const SimdProfHlif&) = delete;
    ~SimdProfHlif() = default;

    void Fini();                                ///< Post process and dump profiling data

private:

    /// Collection of kernel profiles
    using KernelProfilers = std::map<GtKernelId, SimdProfKernelProfile>;

private:

    KernelProfilers             _kernels;
    IGtHliModuleHandle          _hliModule = nullptr;        ///< Module of HLI functions
};


/* ============================================================================================= */
// SimdProfArgs implementation
/* ============================================================================================= */

bool SimdProfArgs::operator < (const SimdProfArgs& other) const
{
    return std::make_tuple(maskCtrl, execMask, predicate, isSendIns) <
           std::make_tuple(other.maskCtrl, other.execMask, other.predicate, other.isSendIns);
}

/* ============================================================================================= */
// SimdProfSection implementation
/* ============================================================================================= */

void SimdProfSection::AddInstruction(const IGtIns& ins)
{
    uint32_t    execMask    = ins.ExecMask().Bits();
    GtPredicate predicate   = ins.Predicate();
    bool        maskCtrl    = !ins.IsWriteMaskEnabled();
    bool        isSendIns   = ins.IsSendMessage();

    SimdProfArgs simdProfArgs(maskCtrl, execMask, predicate, isSendIns);

    auto it = groups.emplace(simdProfArgs, 0).first;

    if (it->second == 0)
    {
        groupHeadIds.emplace(simdProfArgs, ins.Id());
    }

    ++(it->second);
}

/* ============================================================================================= */
// SimdprofKernelProfile implementation
/* ============================================================================================= */

SimdProfKernelProfile::SimdProfKernelProfile(const IGtKernel& kernel) :
    _name(GlueString(kernel.Name())), _extName(ExtendedKernelName(kernel)), _type(kernel.Type()), _platform(kernel.GpuPlatform()),
    _hashId(kernel.HashId()), _simd(kernel.SimdWidth()), _binarySignature(kernel.BinarySignature()),
    _countEnableChannelsFunc("CountEnableChannels"), _countEnableChannelsForSendFunc("CountEnableChannelsForSend") {}

void SimdProfKernelProfile::Instrument(IGtKernelInstrument& instrumentor)
{
    const IGtCfg&    cfg       = instrumentor.Cfg();
    IGtMemoryMapper& memMapper = instrumentor.MemoryMapper();

    std::vector<SimdProfSection> sections;      // All SimdProf sections in the kernel

    // Identify SimdProf sections and #groups in the kernel
    for (auto bblPtr : cfg.Bbls())
    {
        bool isSectionBegin = true;

        // Iterate through sections within the current BBL
        for (auto insPtr : bblPtr->Instructions())
        {
            const IGtIns& ins = *insPtr;

            if (insPtr->Id() < uint32_t(knobMinInstrumentIns))
            {
                continue;
            }

            if (insPtr->Id() > uint32_t(knobMaxInstrumentIns))
            {
                continue;
            }

            if (isSectionBegin)
            {
                sections.emplace_back(bblPtr->Id(), ins);
                isSectionBegin = false;
            }

            SimdProfSection& section = sections.back();
            section.AddInstruction(ins);


            if (ins.IsFlagModifier() || (ins.Id() == bblPtr->LastIns().Id()) || (insPtr->Id() == uint32_t(knobMaxInstrumentIns))) //section end
            {
                isSectionBegin = true;
            }
        }
    }

    // Instrument SimdProf sections and initialize the _profileData array
    for (auto& section : sections) { InstrumentSection(instrumentor, section); }

    std::map<uint32_t, uint32_t> bblGroupNumber;

    for (auto& group : _profileData)
    {
        bblGroupNumber[group.bblId]++;
        memMapper.Map(group.simdProfCountArgs, GT_MMAP_SHARE);
    }

    // Save BBL information for the post processing phase
    for (auto bblPtr : cfg.Bbls())
    {
        _bblInsInfo.emplace(bblPtr->Id(), std::make_pair(bblPtr->FirstIns().Id(), bblPtr->LastIns().Id()));
    }
}

void SimdProfKernelProfile::WriteArgumentsToDevice(IGtMemoryMapper& memMapper)
{
    for (auto& group : _profileData)
    {
        void* ptr = (void*)&(group.simdProfCountArgs);
        memset(ptr, 0, sizeof(SimdProfCountArgs));
        memMapper.Write(ptr, sizeof(SimdProfCountArgs));
    }
}

void SimdProfKernelProfile::InstrumentSection(IGtKernelInstrument& instrumentor, const SimdProfSection& section)
{
    const IGtCfg&       cfg       = instrumentor.Cfg();

    // Instrument each SimdProf group:
    //  - If a group is associated with a non-SEND instructions, compute the SIMD count by applying CBIT to the SIMD mask.
    //  - Otherwise, if a group is created for SEND instructions, increment the SIMD count for each SEND whose SIMD mask
    //    is nonzero. From the EU perspective, SEND instruction is 1 operation, unless the SIMD mask is zero
    // Insert each per-group instrumentation procedure at the beginning of the corresponding section

    //Insert SimdProf instrumentaion at the beginning of the current section
    const IGtIns& ins = cfg.GetInstruction(section.headInsId);
    const IGtBbl& bbl = cfg.GetBbl(ins);

    for (auto& group : section.groups)
    {
    
        auto it = section.groupHeadIds.find(group.first); GTPIN_ASSERT(it != section.groupHeadIds.end());
        const IGtIns& groupHeadIns = cfg.GetInstruction(it->second);
        
        _profileData.emplace_back(bbl.Id(), group.second);

        SimdProfGroup&      newProfGroup = _profileData.back();
        SimdProfCountArgs*  countArgs = const_cast<SimdProfCountArgs*>(&newProfGroup.simdProfCountArgs);
        IargInsOpMask       accessMask(groupHeadIns);

        if (group.first.isSendIns)
        {
            _countEnableChannelsForSendFunc.InsertCallAtInstruction(instrumentor, groupHeadIns, GtIpoint::Before(),
                NullReg(),          // Unused return value
                accessMask,         // arg[1]: Per-channel mask of memory accesses
                countArgs           // arg[2]: Counter arguments
            );
        }
        else
        {
            _countEnableChannelsFunc.InsertCallAtInstruction(instrumentor, groupHeadIns, GtIpoint::Before(),
                NullReg(),          // Unused return value
                accessMask,         // arg[1]: Per-channel mask of memory accesses
                countArgs           // arg[2]: Counter arguments
            );
        }
    }
}

void SimdProfKernelProfile::ReadProfileData()
{
    // Iterate through all SimdProf groups and read counters of executed operations (channels).
    for (auto& group : _profileData)
    {
        uint64_t groupCounts[32] = {};

        SimdProfCountArgs& args = group.simdProfCountArgs;
        uint32_t numOfInsts = group.insCount;

        for (uint32_t i = 0; i < 32; i++)
        {
            groupCounts[i] = args.out.channelCounts[i] * numOfInsts;
            _totalOpCount += groupCounts[i];
        }
        // Update counters of executed operations
        UpdateBblOpCounter(group.bblId._value, groupCounts);
    }
}

void SimdProfKernelProfile::UpdateBblOpCounter(BblId bblId, uint64_t* perChannelValues)
{
    auto it = _bblOpCounts.emplace(bblId, OpCounts()).first;

    auto& opCounts = it->second;

    for (uint32_t i = 0; i < 32; ++i)
    {
        opCounts.counts[i] += perChannelValues[i];
        opCounts.total += perChannelValues[i];
    }
}

std::string SimdProfKernelProfile::ToString() const
{
    ostringstream ostr;
    ostr << ExtendedName() << endl;

    if (!_bblOpCounts.empty())
    {
        ostr << setw(10) << "BBL" << setw(15) << "Head Ins ID" << setw(15) << "Tail Ins ID" << setw(20) << "Channels" << endl;
        for (const auto& bc : _bblOpCounts)
        {
            auto bblId = bc.first;
            uint32_t firstIns = _bblInsInfo.at(bc.first).first;
            uint32_t lastIns = _bblInsInfo.at(bc.first).second;
            ostr << setw(10) << bblId << setw(15) << firstIns << setw(15) << lastIns << setw(20) << bc.second.total << endl;
            if (bc.second.total)
            {
                for (uint32_t i = 0; i < 32; ++i)
                {
                    if (bc.second.counts[i])
                    {
                        ostr << setw(80) << "chan[" << setw(2) << i << "]: " << bc.second.counts[i] << endl;
                    }
                }
            }
        }
        ostr << setw(10) << "Total" << setw(15) << _totalOpCount << endl;
    }
    else
    {
        ostr << "No channels executed" << endl;
    }

    return ostr.str();
}

/* ============================================================================================= */
// SimdProfHlif implementation
/* ============================================================================================= */
SimdProfHlif* SimdProfHlif::Instance()
{
    static SimdProfHlif instance;
    return &instance;
}

void SimdProfHlif::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel& kernel = instrumentor.Kernel();
    auto it = _kernels.emplace(kernel.Id(), kernel).first;
    it->second.Instrument(instrumentor);
    
    // Link the kernel with the library of HLI functions
    instrumentor.LinkHliModule(_hliModule);
}

void SimdProfHlif::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    const IGtKernel& kernel = dispatcher.Kernel();
    GtKernelExecDesc execDesc; dispatcher.GetExecDescriptor(execDesc);
    if (kernel.IsInstrumented() && IsKernelExecProfileEnabled(execDesc, kernel.GpuPlatform(), kernel.Name().Get()))
    {
        dispatcher.SetProfilingMode(true);  // Enable instrumentation

        // This tool needs an accurate information about memory allocations, which is available on the the final dispatch stage.
        // So, on the initial dispatch stage, we only enable instrumentation, and request GTPin to invoke OnKernelRun
        // one more time, on the final dispatch stage. If this request is accepted, the initialization of the profile buffer will
        // be done on the final dispatch stage, otherwise - on the intial dispatch stage.
        if (dispatcher.ReportFinalDispatchStage())
        {
            return;
        }

        auto it = _kernels.find(kernel.Id());

        if (it != _kernels.end())
        {
            SimdProfKernelProfile& kernelProfile = it->second;
            kernelProfile.WriteArgumentsToDevice(dispatcher.MemoryMapper());
        }
    }
    else
    {
        dispatcher.SetProfilingMode(false);
    }
}

void SimdProfHlif::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    if (!dispatcher.IsProfilingEnabled())
    {
        return; // Do nothing with unprofiled kernel dispatches
    }

    const IGtKernel& kernel = dispatcher.Kernel();
    auto it = _kernels.find(kernel.Id());

    if (it != _kernels.end())
    {
        SimdProfKernelProfile& kernelProfile = it->second;
        kernelProfile.ReadProfileData();
    }
}

void SimdProfHlif::LoadHliLibrary()
{
    std::string modulePath = JoinPath(GetKnobValue<std::string>("installDir"), "Examples", "simdprof_hlif.cl");
    _hliModule = GTPin_GetCore()->HliLibrary().CompileModuleFromFile(modulePath.c_str());
    GTPIN_ASSERT_MSG(_hliModule != nullptr, "Could not load HLI module " + modulePath);
}

void SimdProfHlif::Fini()
{
    string profileDir = GTPin_GetCore()->ProfileDir();
    string filePath = JoinPath(profileDir, "simdprof_hlif.txt");

    ofstream fs(filePath);
    if (fs.is_open())
    {
        fs << ToString();
        fs.close();
    }
    else
    {
        GTPIN_WARNING("SIMDPROF_HLIF : could not create file: " + filePath);
    }
}

string SimdProfHlif::ToString() const
{
    ostringstream ostr;
    ostr << "Channels (SIMD operations) executed by kernels/BBLs" << endl;
    ostr << "===================================================" << endl;
    
    uint64_t totalOpCount = 0;
    for (const auto& k : _kernels)
    {
        ostr << string(100, '-') << endl;
        ostr << k.second.ToString() << endl;
        totalOpCount += k.second.GetTotalOpCounter();
    }
    ostr << "Total number of kernels:                    " << _kernels.size() << std::endl;
    ostr << "Total number of channels (SIMD operations): " << totalOpCount << std::endl;

    return ostr.str();
}

// Define DETACHED_SIMDPROF to use SimdProf functionality in a different tool
#if !defined (DETACHED_SIMDPROF)
/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char* argv[])
{
    ConfigureGTPin(argc, argv);
    SimdProfHlif::Instance()->Register();
    
    // Compile and load library of HLI functions
    SimdProfHlif::Instance()->LoadHliLibrary();

    atexit(SimdProfHlif::OnFini);
}
#endif