/*========================== begin_copyright_notice ============================
Copyright (C) 2016-2026 Intel Corporation

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

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

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

#include "latency.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int>  knobMode("mode", 0, "Latency measurment mode: 0 - basic block latency, 1 - memory read latency");
Knob<int>  knobNumThreadBuckets("num_thread_buckets", 0, "Number of thread buckets. 0 (default) - maximum thread buckets");
Knob<bool> knobUseRegInstrument("no_reg_instrument", true, "Disable register-bound instrumentation");
Knob<int>  knobMaxFreq("max_freq", 100, "Maximum amount of instrument executions for bbl per kernel execution. 100 (default) - Applied only when use_conditional_instrument is on");
Knob<bool> knobUseConditionalInstrument("use_conditional_instrument", false, "Enable conditional instrumentation");
Knob<bool> knobSkipZeroResults("skip_zero_results", false, "Skip zero results in the Latency output");

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

void Latency::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    /*
     * Register allocation policy
     * ======================================================================================================
     * To ensure accuracy of latency measurments, the instrumentation overhead should be as low as possible.
     * It is believed that register-bound procedures incur lower overhead than memory-bound procedures.
     * Assuming that, we will try to use register-bound instrumentation for as many sites as possible, and
     * apply memory-bound instrumentation to the rest of the sites.
     * 
     * To check and guarantee availability of free physical registers we pre-allocate registers
     * in the register-bound procedures created by the GeneratePostCodeRegBound() function.
     * For this purpose we use the IGtRegAllocator::ReserveVregOperands() interface.
     * 
     * For memory-bound procedures, pre-allocation of registers is not needed, and even can be harmful, as
     * explained in the IGtRegAllocator description. On the other hand, excessive pre-allocation of registers
     * in the register-bound procedures may cause multiple spills/fills in memory-bound procedures. To avoid
     * this situation, we do the following:
     *      A) RESERVE registers in memory-bound procedures BEFORE generating register-bound procedures using
     *         IGtRegAllocator::Reserve()
     *      B) RELEASE registers in memory-bound procedures AFTER generating register-bound procedures using
     *         IGtRegAllocator::ReleaseReserved()
     */

    const IGtGenCoder& coder          = instrumentor.Coder();
    IGtVregFactory&    vregs          = coder.VregFactory();
    IGtRegAllocator&   ra             = coder.RegAllocator();
    bool               is64BitCounter = Use64BitCounters(coder);

    // Initialize virtual registers
    _timeReg = vregs.Make(VREG_TYPE_DWORD);
    _addrReg = vregs.MakeMsgAddrScratch();
    _dataReg = vregs.MakeMsgDataScratch(is64BitCounter? VREG_TYPE_QWORD : VREG_TYPE_DWORD);

    // A) Reserve registers used in the memory-bound procedures
    ra.Reserve(_dataReg.VregNumber());
    ra.Reserve(_addrReg.VregNumber());
    ra.Reserve(vregs.GetProfileBufferAddrVreg().Num());

    switch (knobMode)
    {
    case 0:     InstrumentBasicBlocks(instrumentor); break;
    case 1:     InstrumentMemoryInstructions(instrumentor); break;
    default:    GTPIN_ERROR_MSG("LATENCY : Invalid value of the 'mode' option");
    }

    // B) Release registers used in the memory-bound procedures
    ra.ReleaseReserved(_dataReg.VregNumber());
    ra.ReleaseReserved(_addrReg.VregNumber());
    ra.ReleaseReserved(vregs.GetProfileBufferAddrVreg().Num());
}

void Latency::InstrumentBasicBlocks(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();

    // Select basic blocks to be instrumented. Exclude BBLs whose only instruction is a branch
    vector<const IGtBbl*> bbls;
    for (auto bblPtr : cfg.Bbls())
    {
        if (!bblPtr->IsEmpty() && !bblPtr->FirstIns().IsChangingIP())
        {
            bbls.push_back(bblPtr);
        }
    }

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


    // Instrument selected basic blocks
    GtGenProcedure finiCode; // Procedure that stores results of register-bound measurments in the profile buffer
    bool           tryRegInstrument  = knobUseRegInstrument; // true - try to apply register-bound instrumentation
    for (uint32_t idx = 0; idx != bbls.size(); ++idx)
    {
        const IGtBbl& bbl = *bbls[idx];

        int32_t bblId = bbl.Id();

        if (bblId < knobMinInstrumentBbl || bblId > knobMaxInstrumentBbl)
        {
            continue;
        }

        // Insert code that starts timer at BBL's entry
        GtGenProcedure preCode;
        GeneratePreCode(preCode, coder);
        InstrumentBbl(instrumentor, bbl, GtIpoint::Before(), preCode);

        // Insert code that stops timer at BBL's exit and stores result in the profile buffer
        GtGenProcedure postCode;
        if (bbl.IsEot())
        {
            // Insert fake consumers of source registers to expose hidden latency of EOT instructions
            const IGtIns& eotIns = bbl.LastIns();
            coder.GenerateFakeSrcConsumers(postCode, eotIns);
            GeneratePostCodeMemBound(postCode, coder, profileArray, idx);
            InstrumentInstruction(instrumentor, eotIns, GtIpoint::Before(), postCode);
        }
        else
        {
            if (tryRegInstrument)
            {
                tryRegInstrument = GeneratePostCodeRegBound(postCode, finiCode, coder, profileArray, idx);
            }
            if (!tryRegInstrument)
            {
                GeneratePostCodeMemBound(postCode, coder, profileArray, idx);
            }
            InstrumentBbl(instrumentor, bbl, GtIpoint::After(), postCode);
        }
    }

    // Insert 'finiCode' at all exits of the kernel - before each EOT instruction
    instrumentor.InstrumentExits(finiCode);

    // Transform 'bbls' into array of instrumentation sites
    InstrumentSites sites(bbls.size());
    std::transform(bbls.begin(), bbls.end(), sites.begin(), [](const IGtBbl* b) -> InsId { return b->FirstIns().Id(); });

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

void Latency::InstrumentMemoryInstructions(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();

    // Select memory instructions to be instrumentated
    InstrumentSites sites;
    for (auto bblPtr : cfg.Bbls())
    {
        for (auto insPtr : bblPtr->Instructions())
        {
            if (insPtr->IsMemRead() && !insPtr->IsEot())
            {
                sites.push_back(insPtr->Id());
            }
        }
    }

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

    // Instrument each instruction whose ID appears in 'sites':
    //  - Before instruction, insert fake consumers of source registers + procedure that starts timer
    //  - After instruction, insert fake consumers of destination registers + procedure that stops timer
    // Use memory-bound instrumentation because, in case of a single instruction latency, the instrumentation overhead does
    // not compromise accuracy, while the register allocation flexibility is important.
    for (uint32_t idx = 0; idx != sites.size(); ++idx)
    {
        const IGtIns& ins = cfg.GetInstruction(sites[idx]);

        GtGenProcedure preCode;
        coder.GenerateFakeSrcConsumers(preCode, ins);
        GeneratePreCode(preCode, coder);
        InstrumentInstruction(instrumentor, ins, GtIpoint::Before(), preCode);

        GtGenProcedure postCode;
        coder.GenerateFakeDstConsumers(postCode, ins);
        GeneratePostCodeMemBound(postCode, coder, profileArray, idx);
        InstrumentInstruction(instrumentor, ins, GtIpoint::After(), postCode);
    }

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

void Latency::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);
            LatencyKernelProfile&   kernelProfile   = it->second;
            const GtProfileArray&   profileArray    = kernelProfile.GetProfileArray();

            if (profileArray.Size() != 0)
            {
                if (profileArray.Initialize(*buffer))
                {
                    isProfileEnabled = true;
                }
                else
                {
                    GTPIN_ERROR_MSG(string("LATENCY : ") + string(kernel.Name()) + " : Failed to write into memory buffer");
                }
            }
        }
    }
    dispatcher.SetProfilingMode(isProfileEnabled);
}

void Latency::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    const IGtKernel& kernel = dispatcher.Kernel();
    GtKernelExecDesc execDesc; dispatcher.GetExecDescriptor(execDesc);
    bool isProfilingEnabled = dispatcher.IsProfilingEnabled();
    if (!isProfilingEnabled || !IsKernelExecProfileEnabled(execDesc, kernel.GpuPlatform(), kernel.Name().Get()))
    {
        return; // Do nothing with unprofiled kernel dispatches
    }

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

    if (it != _kernels.end())
    {
        const IGtProfileBuffer* buffer          = dispatcher.GetProfileBuffer(); GTPIN_ASSERT(buffer);
        LatencyKernelProfile&   kernelProfile   = it->second;
        const GtProfileArray&   profileArray    = kernelProfile.GetProfileArray();

        for (uint32_t recordNum = 0; recordNum != profileArray.NumRecords(); ++recordNum)
        {
            LatencyDispatchProfile dispatchProfile(dispatcher, kernelProfile.GetSiteId(recordNum));

            for (uint32_t threadBucket = 0; threadBucket < profileArray.NumThreadBuckets(); ++threadBucket)
            {
                LatencyRecord record;
                if (!profileArray.Read(*buffer, &record, recordNum, 1, threadBucket))
                {
                    GTPIN_ERROR_MSG(string("LATENCY : ") + string(kernel.Name()) + " : Failed to read from memory buffer");
                }
                else
                {
                    dispatchProfile.Accumulate(record);
                }
            }
            kernelProfile.AddDispatchProfile(dispatchProfile);
        }
    }
}

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

void Latency::GeneratePreCode(GtGenProcedure& proc, const IGtGenCoder& coder)
{
    coder.StartTimer(proc, _timeReg);
    if (!proc.empty()) { proc.front()->AppendAnnotation(__func__); }
}

void Latency::GeneratePostCodeMemBound(GtGenProcedure& proc, const IGtGenCoder& coder,
                                       const GtProfileArray& profileArray, uint32_t recordNum)
{
    IGtInsFactory&  insF           = coder.InstructionFactory();
    bool            is64BitCounter = Use64BitCounters(coder);
    GtReg           dataRegL       = {_dataReg, sizeof(uint32_t), 0};  // Low 32-bits of the data payload register
    GtPredicate     pred(FlagReg(0));

    // Generate code that computes elapsed time
    coder.StopTimerExt(proc, _timeReg);

    // cycles += _timeReg
    proc += insF.MakeMov(dataRegL, _timeReg);   // Move timer value to the low 32-bits of the data register
    if (is64BitCounter)
    {
        // Clear the high 32-bits of the data payload register
        GtReg dataRegH = {_dataReg, sizeof(uint32_t), 1};
        proc += insF.MakeMov(dataRegH, 0);
    }

    // _addrReg =  address of the current thread's CILatencyRecord in the profile buffer
    profileArray.ComputeAddress(coder, proc, _addrReg, recordNum);
    proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg, (is64BitCounter? GED_DATA_TYPE_uq : GED_DATA_TYPE_ud));

    // freq++
    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offsetof(LatencyRecord, freq));
    proc += insF.MakeAtomicInc(NullReg(), _addrReg, GED_DATA_TYPE_ud).SetPredicate(!pred);;

    if (knobUseConditionalInstrument)
    {
        /*
            Since we don't have a way to atomically read-modify-write a memory location with a condition,
            print a warning that the conditional instrumentation is not supported in memory-bound mode.
        */
        GTPIN_WARNING("LATENCY : Skipping Conditional instrumentation for record number: " + DecStr(recordNum) + ", since it is not supported in memory-bound mode.");
    }

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

bool Latency::GeneratePostCodeRegBound(GtGenProcedure& postProc, GtGenProcedure& finiProc, const IGtGenCoder& coder,
                                       const GtProfileArray& profileArray, uint32_t recordNum)
{
    IGtInsFactory&   insF   = coder.InstructionFactory();
    IGtVregFactory&  vregs  = coder.VregFactory();
    IGtRegAllocator& ra     = coder.RegAllocator();
    bool            is64BitCounter = Use64BitCounters(coder);

    GtReg flagReg   = FlagReg(0);
    GtReg freqReg   = vregs.MakeCounter(VREG_TYPE_DWORD);                                    // Frequency counter
    GtReg cycleReg  = vregs.MakeCounter(is64BitCounter ? VREG_TYPE_QWORD : VREG_TYPE_DWORD); // Cycle counter
    GtReg cycleRegL = {cycleReg, sizeof(uint32_t), 0};                                       // Low 32-bits of cycle counter
    GtReg dataRegL  = {_dataReg, sizeof(uint32_t), 0};                                       // Low 32-bits of '_dataReg'

    // Generate procedure that computes and aggregates cycles and frequency
    GtGenProcedure proc;
    coder.StopTimerExt(proc, _timeReg);

    // Generate procedure that updates registers
    GtGenProcedure updateResProc;
    updateResProc += insF.MakeAdd(cycleRegL, cycleRegL, _timeReg); // Add elapsed time to lower 32-bits of of 'cycleReg'
    if (is64BitCounter)
    {
        // If cycleRegL overflowed, increment the high 32-bits of 'cycleReg'
        GtReg cycleRegH = {cycleReg, sizeof(uint32_t), 1};
        updateResProc += insF.MakeCmp(GED_COND_MODIFIER_l, flagReg, cycleRegL, _timeReg);
        updateResProc += insF.MakeAdd(cycleRegH, cycleRegH, 1).SetPredicate(flagReg);
    }
    updateResProc += insF.MakeAdd(freqReg, freqReg, 1);  // freq++

    if (!ra.ReserveVregOperands(updateResProc))
    {
        return false; // No more free registers
    }
    updateResProc.front()->AppendAnnotation(__func__);

    // Convert to conditional instrumentation if knob is turned on. Otherwise - append original procedure.
    if (knobUseConditionalInstrument)
    {
        /*
        *   This code demonstrates a use of the conditional instrumentation feature.
        *   It converts the updating results procedure into conditional, its execution will be determined during runtime
        *   when the condition is met. In this case - if the value of freqReg is smaller than max_freq for BBL.
        *
        *   It has two steps:
        *   1. Creates procedure that calculates the condition's value.
        *   2. Converts original procedure into conditional based of the condition's value stored in a virtual register.
        *
        */

        // Generate condition procedure for register bound basic block - check if freqReg is lower than knobMaxFreq
        GtGenProcedure condPostCode;
        coder.GenerateConditionalProcedure(updateResProc, condPostCode, GED_COND_MODIFIER_l, freqReg, knobMaxFreq.GetValue());

        proc += condPostCode;
    }
    else
    {
        proc += updateResProc;
    }

    postProc.MoveAfter(std::move(proc));

    // Generate 'finiProc' procedure that stores aggregated cycles and frequency counters in the profile buffer
    profileArray.ComputeAddress(coder, proc, _addrReg, recordNum);
    proc += insF.MakeRegMov(_dataReg, cycleReg);
    proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg, (is64BitCounter? GED_DATA_TYPE_uq : GED_DATA_TYPE_ud));

    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offsetof(LatencyRecord, freq));
    proc += insF.MakeRegMov(dataRegL, freqReg);
    proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg, GED_DATA_TYPE_ud);

    proc.front()->AppendAnnotation("GenerateFiniCode");
    finiProc.MoveAfter(std::move(proc));

    return true;
}

string Latency::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" << setw(10) << "SIMD" << setw(10) << "Type";
    ostr << setw(10) << "Ins. ID";
    ostr << setw(15) << "Freq." << setw(20) << "Total-latency(%)" << setw(20) << "Total-latency" << setw(20) << "Avg-Cycles";
    ostr << setw(20) << "Platform" << " " << setw(35) << "Execution descriptor";
    ostr << endl;
    for (const auto& kernelEntry : _kernels)
    {
        ostr << kernelEntry.second.ToString();
    }
    return ostr.str();
}

void Latency::DumpAsm() const
{
    for (const auto& kernelEntry : _kernels)
    {
        kernelEntry.second.DumpAsm();
    }
}

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

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

/* ============================================================================================= */
// LatencyDispatchProfile implementation
/* ============================================================================================= */
LatencyDispatchProfile::LatencyDispatchProfile(const IGtKernelDispatch& kernelDispatch, InstrumentSiteId id) :
    siteId(id), cycles(0), freq(0)
{
    kernelDispatch.GetExecDescriptor(kernelExecDesc);
}

void LatencyDispatchProfile::Accumulate(const LatencyRecord& record)
{
    cycles  += record.cycles;
    freq    += record.freq;
}

/* ============================================================================================= */
// LatencyKernelProfile implementation
/* ============================================================================================= */
LatencyKernelProfile::LatencyKernelProfile(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)) {}

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

LatencyDispatchProfile& LatencyKernelProfile::AddDispatchProfile(const LatencyDispatchProfile& dispathProfile)
{
    _dispatchProfiles.push_back(dispathProfile);
    string execDesStr = dispathProfile.kernelExecDesc.ToString(_platform);
    _execTotalCycles[execDesStr] += dispathProfile.cycles;
    return _dispatchProfiles.back();
}

string LatencyKernelProfile::ToString() const
{
    ostringstream ostr;
    if (!_dispatchProfiles.empty())
    {
        for (const auto& dp: _dispatchProfiles)
        {
            if (knobSkipZeroResults && (dp.freq == 0))
            {
                continue; // Skip zero results if the knob is set
            }

            string      execDesStr          = dp.kernelExecDesc.ToString(_platform);
            uint64_t    execTotalCycles     = _execTotalCycles.at(execDesStr);
            double      dispatchToExecTotal = (execTotalCycles ? ((100.0 * (double)dp.cycles) / execTotalCycles) : 0.0);
            uint64_t    avgCycles           = (dp.freq ? (dp.cycles / dp.freq) : 0);


            ostr << setw(30) << _name << setw(20) << hex << _hashId << dec << setw(10) << _simd << setw(10) << _type.ToString();
            ostr << setw(10) << uint32_t(dp.siteId) << setw(15) << dp.freq;
            ostr << fixed << setprecision(2) << setw(20) << dispatchToExecTotal << setw(20) << dp.cycles << setw(20) << avgCycles;
            ostr << setw(20) << _platform.ToString() << " " << setw(35) << dp.kernelExecDesc.ToString(_platform, ExecDescAlignedFormat());
            ostr << endl;
        }
    }
    else
    {
        ostr << setw(30) << _name << setw(20) << hex << _hashId << dec << setw(10) << _simd << setw(10) << _type.ToString();
        ostr << setw(10) << "NA" << setw(15) << "NA";
        ostr << setw(20) << "NA" << setw(20) << "NA" << setw(20) << "NA";
        ostr << setw(20) << "NA" << " " << setw(35) << "NA";
    }
    ostr << endl;
    return ostr.str();
}

void LatencyKernelProfile::DumpAsm() const
{
    DumpKernelAsmText(_name, _uniqueName, _asmText);
}

// Define DETACHED_LATENCY to use Latency functionality in a different tool
#if !defined (DETACHED_LATENCY)
/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    ConfigureGTPin(argc, argv);
    Latency::Instance()->Register();
    atexit(Latency::OnFini);
}
#endif
