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

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

/*!
 * @file Implementation of the MyGTPin tool as demonstrated in GTPin manual
 */

#include <fstream>
#include <iomanip>
#include <map>

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

#include <ostream>

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int> numOrigDispatches("num_orig_dispatches", 0, "Number of original kernel dispatches");

/* ============================================================================================= */
// Struct MyRecord
/* ============================================================================================= */
// Layout of records collected in the profile buffer by MyGTPinTool
struct MyRecord
{
    uint32_t cycleCount; // Number of cycles
    uint32_t exeCount;   // Number of kernel executions in HW threads
};

/* ============================================================================================= */
// Class MyEventHandler
/* ============================================================================================= */
// Handler of events reported by GTPin 
class MyEventHandler : public IGtEventHandler
{
    GtEventReaction HandleEvent(const IGtEvent& event) const
    {
        std::cout << event.ToString() << std::endl;
        return EVENT_REACTION_DEFAULT;
    }

    GtEventReaction OnEvent(const IGtEvent& event)         { return HandleEvent(event); }
    GtEventReaction OnEvent(const IGtKernelEvent& event)   { return HandleEvent(event); }
    GtEventReaction OnEvent(const IGtInternalEvent& event) { return HandleEvent(event); }
};

/* ============================================================================================= */
// Class MyGTPinTool
/* ============================================================================================= */
// Singleton that implements tool's callbacks
class MyGTPinTool : public IGtTool
{
public:
    // Implementation of the IGtTool interface
    const char* Name()          const                   override { return "My GTPin Tool"; }
    uint32_t    ApiVersion()    const                   override { return GTPIN_API_VERSION; }
    void        OnKernelBuild(IGtKernelInstrument&)     override; 
    void        OnKernelRun(IGtKernelDispatch&)         override;
    void        OnKernelComplete(IGtKernelDispatch&)    override;

    // Register the MyGTPinTool instance with the GTPin core
    bool Register()
    {
        IGtCore* gtpinCore = GTPin_GetCore();
        if (!gtpinCore->RegisterTool(*this))
        {
            GTPIN_ERROR_MSG(std::string(Name()) + ": Failed registration with GTPin core.");
            return false;
        }
        gtpinCore->RegisterEventHandler(_eventHandler);
        return true;
    }

    static MyGTPinTool* Instance();     // Return single instance of this class
    static void OnFini();               // Termination handler registered with atexit()

private:
    void SetProfileArray(GtKernelId kernelId, const GtProfileArray& profileArray);  // Associate the specified profile array with the specified kernel ID
    GtProfileArray& GetProfileArray(GtKernelId kernelId);                           // Get profile array associated with the specified kernel ID
    std::map<GtKernelId, GtProfileArray> _profileArrays;                            // Profile arrays associated with kernels

private:
    uint64_t _totalCycleCount = 0;  // Total number of cycles executed by GPU kernels
    uint64_t _totalExeCount   = 0;  // Total number of kernel executions in HW threads

    MyEventHandler _eventHandler;   // Event handler
};

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

void MyGTPinTool::OnFini()
{
    MyGTPinTool& me = *Instance();
    std::cout << std::endl << "MyGTPinTool resuts:" << std::endl;
    std::cout << "totalCycleCount = " << me._totalCycleCount << ", totalExeCount = " << me._totalExeCount << std::endl;
}

void MyGTPinTool::SetProfileArray(GtKernelId kernelId, const GtProfileArray& profileArray)
{
    _profileArrays[kernelId] = profileArray;
}

GtProfileArray& MyGTPinTool::GetProfileArray(GtKernelId kernelId)
{
    return _profileArrays.at(kernelId);
}

void MyGTPinTool::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel&            kernel    = instrumentor.Kernel();                  // The kernel being instrumented
    const IGtGenCoder&          coder     = instrumentor.Coder();                   // GEN code generator
    IGtProfileBufferAllocator&  allocator = instrumentor.ProfileBufferAllocator();  // Profile buffer allocator
    IGtVregFactory&             vregs     = coder.VregFactory();                    // Factory of virtual registers
    IGtInsFactory&              insF      = coder.InstructionFactory();             // GEN instruction factory

    // Create virtual registers
    GtReg timeReg = vregs.Make(VREG_TYPE_DWORD);
    GtReg addrReg = vregs.MakeMsgAddrScratch();
    GtReg dataReg = vregs.MakeMsgDataScratch();

    // There will be one MyRecord per each HW thread in the profile buffer. Initialize and store the corresponding profile array
    GtProfileArray profileArray(sizeof(MyRecord), 1, kernel.GenModel().MaxThreadBuckets());
    profileArray.Allocate(allocator);
    SetProfileArray(kernel.Id(), profileArray);

    // Generate the 'preCode' procedure that starts timer at kernel’s entry
    GtGenProcedure preCode;
    coder.StartTimer(preCode, timeReg);                 // timeReg = tm0.0

    // Generate the 'postCode' procedure that stops timer at kernel's exits
    GtGenProcedure postCode;

    // buffer[TID].cycleCount += (tm0.0 – timeReg)
    coder.StopTimerExt(postCode, timeReg);              // timeReg = (tm0.0 – timeReg)
    postCode += insF.MakeMov(dataReg, timeReg);
    profileArray.ComputeAddress(coder, postCode, addrReg, 0, offsetof(MyRecord, cycleCount));
    postCode += insF.MakeAtomicAdd(NullReg(), addrReg, dataReg, GED_DATA_TYPE_ud);

    // buffer[TID].exeCount += 1
    profileArray.ComputeAddress(coder, postCode, addrReg, 0, offsetof(MyRecord, exeCount));
    postCode += insF.MakeAtomicInc(NullReg(), addrReg, GED_DATA_TYPE_ud);

    // Insert preCode at kernel entries
    instrumentor.InstrumentEntries(preCode);

    // Insert postCode at kernel exits
    instrumentor.InstrumentExits(postCode);
}

void MyGTPinTool::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    // The number of kernel dispatches so far
    static int dispatchCount = 0;

    // Skip first numOrigDispatches dispatches
    if (++dispatchCount >= numOrigDispatches)
    {
        // Create and initialize memory buffer for this kernel dispatch
        IGtProfileBuffer* buffer       = dispatcher.CreateProfileBuffer();
        GtProfileArray&   profileArray = GetProfileArray(dispatcher.Kernel().Id());
        if (profileArray.Initialize(*buffer))
        {
            // Tell GTPin to run instrumented code
            dispatcher.SetProfilingMode(true);
        }
    }
}

void MyGTPinTool::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    const IGtProfileBuffer& buffer = *dispatcher.GetProfileBuffer();

    // Read all records stored in the profile buffer, and accumulate their data in _totalCycleCount and _totalExeCount
    GtProfileArray& profileArray = GetProfileArray(dispatcher.Kernel().Id());
    for (uint32_t tBucket = 0; tBucket < profileArray.NumThreadBuckets(); ++tBucket)
    {
        MyRecord record;
        bool ret = profileArray.Read(buffer, &record, 0, 1, tBucket);
        if (ret)
        {
            _totalCycleCount += record.cycleCount;
            _totalExeCount += record.exeCount;
        }
    }
}

/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    // Parse command line and configure GTPin
    ConfigureGTPin(argc, argv);

    // Register the tool (callbacks) with the GTPin core
    MyGTPinTool::Instance()->Register();

    // Register the termination function
    atexit(MyGTPinTool::OnFini);
}
