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

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

/*!
 * @file Implementation of the kernel weight measuring tool
 */


#include <fstream>
#include <sstream>
#include <iomanip>
#include <algorithm>
#include "kernel_weight.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
#if defined(KERNEL_WEIGHT_STANDALONE)
Knob<int> knobNumThreadBuckets("num_thread_buckets", 32, "Number of thread buckets. Default - 32, zero - maximum thread buckets");
#else
const int knobNumThreadBuckets = 32;
#endif

/* ============================================================================================= */
// KernelWeight implementation
/* ============================================================================================= */
void KernelWeight::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();

    // Allocate the profile buffer. It will hold single KernelWeightRecord per each thread bucket
    uint32_t numThreadBuckets = (knobNumThreadBuckets == 0) ? genModel.MaxThreadBuckets() : knobNumThreadBuckets;
    GtProfileArray profileArray(sizeof(KernelWeightRecord), 1, numThreadBuckets);
    profileArray.Allocate(allocator);

    // Initialize virtual registers
    _weightReg  = vregs.MakeCounter();
    _addrReg    = vregs.MakeMsgAddrScratch();
    _dataReg    = vregs.MakeMsgDataScratch();

    // Instrument basic blocks
    for (auto bblPtr : cfg.Bbls())
    {
        InstrumentBbl(instrumentor, *bblPtr, profileArray);
    }

    _kernelProfiles.emplace(ExtendedKernelName(kernel), KernelWeightProfile(kernel, profileArray));
}

void KernelWeight::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 = _kernelProfiles.find(ExtendedKernelName(kernel));

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

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

    const IGtKernel& kernel = dispatcher.Kernel();
    auto it = _kernelProfiles.find(ExtendedKernelName(kernel));

    if (it != _kernelProfiles.end())
    {
        const IGtProfileBuffer* buffer          = dispatcher.GetProfileBuffer(); GTPIN_ASSERT(buffer);
        KernelWeightProfile&    kernelProfile   = it->second;
        const GtProfileArray&   profileArray    = kernelProfile.ProfileArray();
        DispatchWeightProfile&  dispatchProfile = kernelProfile.AddKernelDispatch(dispatcher);

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

void KernelWeight::InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const GtProfileArray& profileArray)
{
    const IGtGenCoder&  coder   = instrumentor.Coder();
    IGtInsFactory&      insF    = coder.InstructionFactory();

    uint32_t weight = GetBblWeight(instrumentor, bbl);
    if (weight != 0)
    {
        GtGenProcedure proc;
        proc += insF.MakeAdd(_weightReg, _weightReg, weight);
        proc.front()->AppendAnnotation("KernelWeight compute");
        instrumentor.InstrumentBbl(bbl, GtIpoint::Before(), proc);
    }

    if (bbl.IsEot())
    {
        GtGenProcedure proc;

        // _addrReg =  address of KernelWeightRecord in the current thread's bucket
        profileArray.ComputeAddress(coder, proc, _addrReg);
        int32_t base = 0;
        int32_t offset;

        // _addrReg[weight] += _weightReg
        offset = offsetof(KernelWeightRecord, weight) - base;
        profileArray.ComputeRelAddress(coder, proc, _addrReg, _addrReg, offset); base += offset;
        proc += insF.MakeMov(_dataReg, _weightReg);
        proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg);

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

        proc.front()->AppendAnnotation("KernelWeight store");
        instrumentor.InstrumentBbl(bbl, GtIpoint::After(), proc);
    }
}

bool KernelWeight::DumpKernelProfiles(const std::string& filePath) const
{

    KernelWeightProfileData kernelProfileData;
    for (const auto& kernelEntry : GetKernelProfiles())
    {
        const auto& kernelProfile = kernelEntry.second;
        KernelWeightCounters counters {0, 0};
        for (const auto& dispatchProfile : kernelProfile.GetDispatchProfiles())
        {
            AggregateDispatchCounters(counters, dispatchProfile.data);
        }
        kernelProfileData[kernelEntry.first] = counters;
    }

    std::stringstream outStream;
    outStream << kernelProfileData;
    
    LockableFile file(filePath);
    bool res = file.Write(outStream.str(), true);

    return res;
}

bool KernelWeight::DumpDispatchProfiles(const std::string& filePath) const
{
    std::stringstream outStream;

    for (const auto& kernelEntry : GetKernelProfiles())
    {
        const auto& kernelProfile = kernelEntry.second;
        auto        platform      = kernelProfile.Platform();
        for (const auto& dispatchProfile : kernelProfile.GetDispatchProfiles())
        {
            outStream << kernelEntry.first << " " << dispatchProfile.data.weight << " " << dispatchProfile.data.freq << " " <<
                platform.ToString() << " " << dispatchProfile.kernelExecDesc.ToString(platform) << std::endl;
        }
    }
    
    LockableFile file(filePath);
    bool res = file.Write(outStream.str(), true);

    return res;
}

/* ============================================================================================= */
// DispatchWeightProfile implementation
/* ============================================================================================= */
DispatchWeightProfile::DispatchWeightProfile(const IGtKernelDispatch& kernelDispatch) : data({0, 0})
{
    kernelDispatch.GetExecDescriptor(kernelExecDesc);
}

/* ============================================================================================= */
// KernelWeightProfile implementation
/* ============================================================================================= */
KernelWeightProfile::KernelWeightProfile(const IGtKernel& kernel, const GtProfileArray& profileArray) :
    _name(GlueString(kernel.Name())), _platform(kernel.GpuPlatform()), _profileArray(profileArray) {}

DispatchWeightProfile& KernelWeightProfile::AddKernelDispatch(const IGtKernelDispatch& kernelDispatch)
{
    _dispatchProfiles.emplace_back(kernelDispatch);
    return _dispatchProfiles.back();
}

/* ============================================================================================= */
// KernelWeightProfileData operators
/* ============================================================================================= */
std::ostream& operator<<(std::ostream& os, KernelWeightProfileData& data)
{
    for (const auto& kernelEntry : data)
    {
        const std::string&          extKernelName = kernelEntry.first;
        const KernelWeightCounters& counters      = kernelEntry.second;
        os << extKernelName << " " << counters.weight << " " << counters.freq << std::endl;
    }
    return os;
}

std::istream& operator>>(std::istream& is, KernelWeightProfileData& data)
{
    auto flags = is.setf(std::ios_base::skipws);
    while (is && !is.eof())
    {
        std::string extKernelName;
        KernelWeightCounters counters;

        is >> extKernelName >> counters.weight >> counters.freq >> std::ws;
        if (data.find(extKernelName) != data.end())
        {
            counters.weight = std::max(counters.weight, data[extKernelName].weight);
            counters.freq = std::max(counters.freq, data[extKernelName].freq);
        }
        data[extKernelName] = counters;
    }
    is.setf(flags);
    return is;
}

#if defined(KERNEL_WEIGHT_STANDALONE)
/* ============================================================================================= */
// Standalone (command line) kernel_weight tool
/* ============================================================================================= */
class KernelWeightStandalone : public KernelWeight
{
public:
    using KernelWeight::KernelWeight;

    /// @return Single instance of this class
    static KernelWeightStandalone* Instance()
    {
        static KernelWeightStandalone instance;
        return &instance;
    }

    /// Callback function registered with atexit()
    static void OnFini()
    {
        KernelWeightStandalone& tool            = *Instance();
        std::string             profileDir      = GTPin_GetCore()->ProfileDir();
        std::string             kernelFilePath  = JoinPath(profileDir, "kernel_weight.txt");
        std::string             dispathFilePath = JoinPath(profileDir, "dispatch_weight.txt");

        tool.DumpKernelProfiles(kernelFilePath);
        tool.DumpDispatchProfiles(dispathFilePath);
    }
};

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