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

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

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

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

#include "funtime.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// 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");
Knob<bool> knobSkipZeroResults("skip_zero_results", false, "Skip zero results in the Funtime output");

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

void Funtime::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel&            kernel          = instrumentor.Kernel();
    const IGtCfg&               cfg             = instrumentor.Cfg();
    const IGtGenCoder&          coder           = instrumentor.Coder();
    const IGtGenArch&           genArch         = GTPin_GetCore()->GenArch();
    const IGtGenModel&          genModel        = kernel.GenModel();
    IGtProfileBufferAllocator&  allocator       = instrumentor.ProfileBufferAllocator();
    IGtVregFactory&             vregs           = coder.VregFactory();
    bool                        is64BitCounter  = Use64BitCounters(coder);

    // Allocate the profile buffer. It will hold single FuntimeRecord per each thread bucket
    uint32_t numThreadBuckets = (knobNumThreadBuckets == 0) ? genModel.MaxThreadBuckets() : knobNumThreadBuckets;
    uint32_t numTiles         = (knobPerTileProfiling && coder.IsTileIdSupported()) ? genArch.MaxTiles(kernel.GpuPlatform()) : 1;
    GtProfileArray profileArray(sizeof(FuntimeRecord), numTiles, numThreadBuckets);
    profileArray.Allocate(allocator);

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

    // Generate code that starts/stops timer at entry/exit of the kernel
    GtGenProcedure preCode;  GeneratePreCode(preCode, coder);
    GtGenProcedure postCode; GeneratePostCode(postCode, coder, profileArray);

    // Instrument kernel entries
    instrumentor.InstrumentEntries(preCode);

    // Instrument kernel exits
    for (auto bblPtr : cfg.ExitBbls())
    {
        const IGtIns& ins = bblPtr->LastIns(); GTPIN_ASSERT(ins.IsEot());
        GtGenProcedure fakeConsumers;
        coder.GenerateFakeSrcConsumers(fakeConsumers, ins);
        instrumentor.InstrumentInstruction(ins, GtIpoint::Before(), fakeConsumers);
        instrumentor.InstrumentInstruction(ins, GtIpoint::Before(), postCode);
    }

    _kernels.emplace(kernel.Id(), FuntimeKernelProfile(kernel, profileArray));
}

void Funtime::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);
            FuntimeKernelProfile&   kernelProfile   = it->second;
            const GtProfileArray&   profileArray    = kernelProfile.GetProfileArray();
            if (profileArray.Initialize(*buffer))
            {
                isProfileEnabled = true;
            }
            else
            {
                GTPIN_ERROR_MSG(string("FUNTIME : ") + string(kernel.Name()) + " : Failed to write into memory buffer");
            }
        }
    }
    dispatcher.SetProfilingMode(isProfileEnabled);
}

void Funtime::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);
        FuntimeKernelProfile&   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
        for (uint32_t tileId = 0; tileId < numTiles; tileId++)
        {
            FuntimeDispatchProfile& dispatchProfile = kernelProfile.AddKernelDispatch(dispatcher, tileId);

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

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

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

string Funtime::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(15) << "Freq." << setw(15) << "Total-Cycle" << setw(15) << "Avg-Cycles" << setw(15) << "Skipped";
    ostr << setw(20) << "Platform";
    if (knobPerTileProfiling)
    {
        ostr << setw(10) << "Tile";
    }
    ostr << " " << setw(35) << "Execution descriptor";
    ostr << endl;
    for (const auto& kernelEntry : _kernels)
    {
        ostr << kernelEntry.second.ToString();
    }
    return ostr.str();
}

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

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

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

    // Generate code that computes elapsed time, and sets flagReg in case of timer overflow
    coder.StopTimerExt(proc, _timeReg);

    // _addrReg =  address of the current thread's FuntimeRecord in the profile buffer
    if (profileArray.NumRecords() > 1)
    {
        // The record number = tile ID
        GtReg& offsetReg = _tmpReg32;
        coder.LoadTileId(proc, offsetReg);
        proc += insF.MakeMul(offsetReg, offsetReg, sizeof(FuntimeRecord));
        profileArray.ComputeAddress(coder, proc, _addrReg, offsetReg);
    }
    else
    {
        profileArray.ComputeAddress(coder, proc, _addrReg);
    }

    int32_t base = 0;
    int32_t offset;

    // cycles += _timeReg
    offset = offsetof(FuntimeRecord, cycles) - base;
    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offset); base += offset;
    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);
    }
    proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg, (is64BitCounter? GED_DATA_TYPE_uq : GED_DATA_TYPE_ud));

    // freq++
    offset = offsetof(FuntimeRecord, freq) - base;
    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offset); base += offset;
    proc += insF.MakeAtomicInc(NullReg(), _addrReg, GED_DATA_TYPE_ud);

    // if (flagReg) skipped++
    offset = offsetof(FuntimeRecord, skipped) - base;
    profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offset); base += offset;
    proc += insF.MakeAtomicInc(NullReg(), _addrReg, GED_DATA_TYPE_ud).SetPredicate(flagReg);

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

/* ============================================================================================= */
// FuntimeDispatchProfile implementation
/* ============================================================================================= */
FuntimeDispatchProfile::FuntimeDispatchProfile(const IGtKernelDispatch& kernelDispatch, uint32_t tile) :
    tileId(tile), cycles(0), freq(0), skipped(0)
{
    kernelDispatch.GetExecDescriptor(kernelExecDesc);
}

void FuntimeDispatchProfile::Accumulate(const FuntimeRecord& record)
{
    cycles  += record.cycles;
    freq    += record.freq;
    skipped += record.skipped;
}

/* ============================================================================================= */
// FuntimeKernelProfile implementation
/* ============================================================================================= */
FuntimeKernelProfile::FuntimeKernelProfile(const IGtKernel& kernel, const GtProfileArray& profileArray) :
    _name(GlueString(kernel.Name())), _type(kernel.Type()), _platform(kernel.GpuPlatform()), _hashId(kernel.HashId()),
    _simd(kernel.SimdWidth()), _profileArray(profileArray) {}

FuntimeDispatchProfile& FuntimeKernelProfile::AddKernelDispatch(const IGtKernelDispatch& kernelDispatch, const uint32_t tile)
{
    _dispatchProfiles.emplace_back(kernelDispatch, tile);
    return _dispatchProfiles.back();
}

string FuntimeKernelProfile::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
            }

            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(15) << dp.freq << setw(15) << dp.cycles << setw(15) << avgCycles << setw(15) << dp.skipped;
            ostr << setw(20) << _platform.ToString();
            if (knobPerTileProfiling)
            {
                ostr << setw(10) << dp.tileId;
            }
            ostr << " " << 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(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);
    Funtime::Instance()->Register();
    atexit(Funtime::OnFini);
}
