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

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

/*!
 * @file Implementation of the edge_count tool
 */

#include <fstream>
#include <sstream>
#include <iomanip>
#include <algorithm>
#include <numeric>

#include "edge_count.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Utils
/* ============================================================================================= */
static uint32_t GetEdgeRecordNum(const IGtEdge* edge, int numOfBbls)
{
    /*  
    *    Calculation of an edge's corresponding record number:
    *    Since there are numOfBbls^2 possible edges (edge between every two basic blocks), there are numOfBbls^2 records.
    *    For a given edge = (srcBbl, dstBbl), the record number is in cell [srcBbl, dstBbl] and therefore the calculation is:
    *        record number = (srcBbl * number_of_bbls) + dstBbl
    */
    return (uint32_t)edge->SrcBbl().Id() * (uint32_t)numOfBbls + (uint32_t)edge->DstBbl().Id();
}

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int>  knobNumThreadBuckets("num_thread_buckets", 32, "Number of thread buckets. Default - 32, zero - maximum thread buckets");
Knob<bool> knobPerTileProfiling("per_tile_profiling", false, "Enable per-tile (subdevice) profiling");

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

void EdgeCount::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel&            kernel          = instrumentor.Kernel();
    const IGtCfg&               cfg             = instrumentor.Cfg();
    const IGtGenCoder&          coder           = instrumentor.Coder();
    const IGtGenModel&          genModel        = kernel.GenModel();
    IGtProfileBufferAllocator&  allocator       = instrumentor.ProfileBufferAllocator();
    IGtVregFactory&             vregs           = coder.VregFactory();
    int                         numOfBbls       = (int)cfg.NumBbls();

    // Allocate the profile buffer. It will hold single EdgeCountRecord per each instrumentation site in each thread bucket
    uint32_t numThreadBuckets = (knobNumThreadBuckets == 0) ? genModel.MaxThreadBuckets() : knobNumThreadBuckets;
    uint32_t numRecords       = (uint32_t)cfg.Bbls().size()* (uint32_t)cfg.Bbls().size();
    GtProfileArray profileArray(sizeof(EdgeCountRecord), numRecords, numThreadBuckets);
    profileArray.Allocate(allocator);

    // Initialize virtual registers
    _addrReg   = vregs.MakeMsgAddrScratch();

    // Instrument kernel edges
    // fill vec with values 0 to (numOfBbls * numOfBbls) - 1
    InstrumentSites sites(numOfBbls * numOfBbls);
    std::iota(sites.begin(), sites.end(), 0);
    std::vector<bool> edgeAdded(numOfBbls * numOfBbls, false);
    uint32_t recordNum;

    for (auto bbl : cfg.Bbls())
    {
        for (auto edge : bbl->IncomingEdges())
        {
            if (!edgeAdded[GetEdgeRecordNum(edge, numOfBbls)])
            {
                // Generate code that increaments the counter
                GtGenProcedure edgeCode;
                recordNum = GetEdgeRecordNum(edge,  numOfBbls);
                GenerateEdgeCode(recordNum, edgeCode, coder, profileArray);
                InstrumentEdge(instrumentor,*edge, edgeCode);
                edgeAdded[GetEdgeRecordNum(edge, numOfBbls)] = true;
            }
        }
    }

    // Create EdgeCountKernelProfile object that represents profile of this kernel
    _kernels.emplace(kernel.Id(), EdgeCountKernelProfile(kernel, cfg, profileArray, std::move(sites)));
}

void EdgeCount::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    bool isProfileEnabled = false;

    const IGtKernel& kernel = dispatcher.Kernel();
    GtKernelExecDesc execDesc; dispatcher.GetExecDescriptor(execDesc);
    if (kernel.IsInstrumented() && IsKernelExecProfileEnabled(execDesc, kernel.GpuPlatform(), kernel.Name().Get()))
    {
        auto it = _kernels.find(kernel.Id());

        if (it != _kernels.end())
        {
            IGtProfileBuffer*       buffer          = dispatcher.CreateProfileBuffer(); GTPIN_ASSERT(buffer);
            EdgeCountKernelProfile&    kernelProfile   = it->second;
            const GtProfileArray&   profileArray    = kernelProfile.GetProfileArray();
            if (profileArray.Initialize(*buffer))
            {
                isProfileEnabled = true;
            }
            else
            {
                GTPIN_ERROR_MSG(string("EdgeCount : ") + string(kernel.Name()) + " : Failed to write into memory buffer");
            }
        }
    }
    dispatcher.SetProfilingMode(isProfileEnabled);
}

void EdgeCount::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())
    {
        const IGtProfileBuffer* buffer          = dispatcher.GetProfileBuffer(); GTPIN_ASSERT(buffer);
        EdgeCountKernelProfile&    kernelProfile   = it->second;
        const GtProfileArray&   profileArray    = kernelProfile.GetProfileArray();
        uint32_t                numTiles        = profileArray.NumRecords(); // There is a single record for each tile in each thread bucket

        // accumulate results and add dispatch
        for (uint32_t recordNum = 0; recordNum < numTiles; recordNum++)
        {
            EdgeCountDispatchProfile dispatchProfile(dispatcher, kernelProfile.GetSiteId(recordNum));
            for (uint32_t threadBucket = 0; threadBucket < profileArray.NumThreadBuckets(); ++threadBucket)
            {
                EdgeCountRecord record;
                if (!profileArray.Read(*buffer, &record, recordNum, 1, threadBucket))
                {
                    GTPIN_ERROR_MSG(string("EdgeCount : ") + string(kernel.Name()) + " : Failed to read from memory buffer");
                }
                else
                {
                    dispatchProfile.Accumulate(record);
                }
            }
            kernelProfile.AddDispatchProfile(dispatchProfile);
        }
    }
}

void EdgeCount::OnFini()
{
    EdgeCount&    me          = *Instance();
    string     profileDir  = GTPin_GetCore()->ProfileDir();
    string     filePath    = JoinPath(profileDir, "edge_count.txt");

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

string EdgeCount::ToString() const
{
    ostringstream ostr;
    ostr << "### Kernel/Shader execution-time profile generated by GTPin ###" << endl << endl;
    ostr << "Legend:" << endl;
    ostr << "NA - kernel was not instrumented." << endl << endl;
    ostr << setw(30) << "Name" << setw(20) << "HashID";
    ostr << setw(15) << "SrcBbl" << setw(15) << "DstBbl" << setw(15) << "Counter";
    ostr << setw(20) << "Platform";
    ostr << endl;
    for (const auto& kernelEntry : _kernels)
    {
        ostr << kernelEntry.second.ToString();
    }
    return ostr.str();
}

bool EdgeCount::Use64BitCounters(const IGtGenCoder& coder)
{
    return coder.InstructionFactory().CanAccessAtomically(GED_DATA_TYPE_uq);
}

void EdgeCount::GenerateEdgeCode(uint32_t recordNum,  GtGenProcedure& proc, const IGtGenCoder& coder, GtProfileArray profileArray)
{
    IGtInsFactory& insF = coder.InstructionFactory();

    // addrReg =  address of the current thread's EdgeCountRecord in the profile buffer
    profileArray.ComputeAddress(coder, proc, _addrReg, recordNum);

    // counter++
    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offsetof(EdgeCountRecord, counter));
    proc += insF.MakeAtomicInc(NullReg(), _addrReg, GED_DATA_TYPE_ud);

    if (!proc.empty()) { proc.front()->AppendAnnotation(__func__); }
}

/* ============================================================================================= */
// EdgeCountDispatchProfile implementation
/* ============================================================================================= */
EdgeCountDispatchProfile::EdgeCountDispatchProfile(const IGtKernelDispatch& kernelDispatch, InstrumentSiteId siteId) :
    siteId(siteId), counter(0)
{
    kernelDispatch.GetExecDescriptor(kernelExecDesc);
}

void EdgeCountDispatchProfile::Accumulate(const EdgeCountRecord& record)
{
    counter += record.counter;
}


InstrumentSiteId EdgeCountKernelProfile::GetSiteId(uint32_t recordNum) const
{
    GTPIN_ASSERT(recordNum < _instrumentSites.size());
    return _instrumentSites[recordNum];
}

/* ============================================================================================= */
// EdgeCountKernelProfile implementation
/* ============================================================================================= */
EdgeCountKernelProfile::EdgeCountKernelProfile(const IGtKernel& kernel,
                                           const IGtCfg& cfg,
                                           const GtProfileArray& profileArray,
                                           InstrumentSites&& instrumentSites) :
    _name(GlueString(kernel.Name())), _uniqueName(kernel.UniqueName()), _type(kernel.Type()), _platform(kernel.GpuPlatform()), _hashId(kernel.HashId()),
    _simd(kernel.SimdWidth()), _asmText(CfgAsmText(cfg)), _profileArray(profileArray),
    _instrumentSites(std::move(instrumentSites)), _numOfBbls(cfg.Bbls().size()) {}

EdgeCountDispatchProfile& EdgeCountKernelProfile::AddDispatchProfile(const EdgeCountDispatchProfile& dispathProfile)
{
    _dispatchProfiles.push_back(dispathProfile);
    string execDesStr = dispathProfile.kernelExecDesc.ToString(_platform);
    return _dispatchProfiles.back();
}

string EdgeCountKernelProfile::ToString() const
{
    ostringstream ostr;
    if (!_dispatchProfiles.empty())
    {
        for (const auto& dp: _dispatchProfiles)
        {
            if (dp.counter != 0)
            {
                auto srcBblId = dp.siteId / this->_numOfBbls;
                auto dstBblId = (int)dp.siteId % _numOfBbls;
                ostr << setw(30) << _name << setw(20) << hex << _hashId << dec;
                ostr << setw(15) << srcBblId << setw(15) << dstBblId << setw(15) << dp.counter;
                ostr << setw(20) << _platform.ToString();
                ostr << endl;
            }
        }
    }
    else
    {
        ostr << setw(30) << _name << setw(20) << hex << _hashId << dec << setw(10) << _simd << setw(10) << _type.ToString();
        ostr << setw(15) << "NA" << setw(15) << "NA" << setw(15) << "NA" << setw(15) << "NA" << setw(20) << "NA";
        if (knobPerTileProfiling)
        {
            ostr << setw(10) << "NA";
        }
        ostr << " " << setw(35) << "NA";
    }
    ostr << endl;
    return ostr.str();
}

/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    ConfigureGTPin(argc, argv);
    EdgeCount::Instance()->Register();
    atexit(EdgeCount::OnFini);
}




