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

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

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

#include <fstream>

#include "itrace.h"
#include "gtpin_tool_utils.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int>  gKnobMaxTraceBufferInMB("max_buffer_mb", 3072, "itrace - the max allowed size of the trace buffer per kernel in MB");
Knob<int>  gKnobPhase("phase", 0, "tracing tool - processing phase\n { 1 - pre-processing, 2 - processing - trace gathering}");
Knob<bool> gKnobCfgOnly("cfg-only", false, "indicates that the collected trace should not be saved - only resulting cfg file");

/* ============================================================================================= */
// ItraceDispatch implementation
/* ============================================================================================= */
bool ItraceDispatch::ReadTrace(const GtProfileTrace& traceAccessor, const IGtProfileBuffer& profileBuffer)
{
    uint32_t traceSize = traceAccessor.Size(profileBuffer);
    _rawTrace.resize(traceSize);
    _isTrimmed = traceAccessor.IsTruncated(profileBuffer);
    return traceAccessor.Read(profileBuffer, _rawTrace.data(), 0, traceSize);
}

bool ItraceDispatch::IsEmpty() const
{
    return _rawTrace.size() < sizeof(ItraceRecord);
}

/* ============================================================================================= */
// ItraceKernel implementation
/* ============================================================================================= */
ItraceKernel::ItraceKernel(const IGtKernelInstrument& kernelInstrument, uint32_t numTiles) : _numTiles(numTiles)
{
    const IGtKernel& kernel = kernelInstrument.Kernel();
    const IGtCfg& cfg = kernelInstrument.Cfg();

    _name       = GlueString(kernel.Name());
    _extName    = ExtendedKernelName(kernel);
    _platform   = kernel.GpuPlatform();
    _genId      = kernel.GenModel().Id();
    _asmText    = CfgAsmText(cfg);
    _uniqueName = kernel.UniqueName();

    // Initialize trace accessor. The trace capacity is expected to be computed during the preprocessing phase.
    uint64_t traceCapacity = ItracePreProcessor::Instance()->TraceSize(_extName);
    if (traceCapacity == 0)
    {
        // Unknown trace capacity
        GTPIN_WARNING("ITRACE: unknown trace capacity for kernel " + _name + ". Assuming the kernel is filtered out. "
                      "Allocating a buffer of 8KB size. If the kernel is supposed to run, expect buffer overflow. "
                      "In this case, please re-run phase 1 and make sure the kernel is not filtered out.");
        traceCapacity = 0x2000;
    }
    else
    {
        traceCapacity += 0x2000; // Add some space to account for possible fluctuation of trace sizes between phases
        if (traceCapacity > UINT32_MAX)
        {
            GTPIN_WARNING("ITRACE: The kernel " + _name + " exceedeed maximum trace capacity.");
            traceCapacity = UINT32_MAX;
        }
    }
    if (traceCapacity > (uint64_t(gKnobMaxTraceBufferInMB) * 0x100000))
    {
        GTPIN_WARNING("ITRACE: required capacity (" + DecStr(traceCapacity) + ") for kernel " + _name + " is too big - cut to " + DecStr(gKnobMaxTraceBufferInMB) + "MB. "
                      "Expect the final trace to contain partial data.");
        traceCapacity = uint64_t(gKnobMaxTraceBufferInMB) * 0x100000;
    }
    uint32_t maxRecordSize = sizeof(ItraceRecord);
    _traceAccessor = GtProfileTrace((uint32_t)traceCapacity, maxRecordSize);
    _traceAccessor.Allocate(kernelInstrument.ProfileBufferAllocator());
    // Fill basic block offsets info
    for (auto bblPtr : cfg.Bbls())
    {
        const IGtBbl&     bbl = *bblPtr;
        BblId           bblId = bbl.Id();
        const IGtIns& insHead = bbl.FirstIns();
        const IGtIns& insTail = bbl.LastIns();
        uint32_t offsetHead   = cfg.GetInstructionOffset(insHead);
        uint32_t offsetTail   = insTail.IsChangingIP() ? uint32_t(cfg.GetInstructionOffset(insTail)) : 0xFFFFFFFF;
        _bblBoundsMap[bblId]  = BblBounds(offsetHead, offsetTail);
        const EdgeSpan& outgoingEdges = bbl.OutgoingEdges();
        for (auto outEdge : outgoingEdges)
        {
            const IGtBbl& dstBbl = outEdge->DstBbl();
            uint32_t dstBblId = dstBbl.Id();
            _edges.emplace(bblId, dstBblId);
        }
    }
}

ItraceDispatch& ItraceKernel::AddItrace(IGtKernelDispatch& kernelDispatch)
{
    // Create a new ItraceDispatch object and store the entire trace within this object
    _traces.emplace_back(kernelDispatch);
    ItraceDispatch& ItraceDispatch = _traces.back();
    if (!ItraceDispatch.ReadTrace(_traceAccessor, *kernelDispatch.GetProfileBuffer()))
    {
        GTPIN_ERROR_MSG("ITRACE: Failed to read profile buffer for kernel " + _name);
    }
    return ItraceDispatch;
}

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

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

void Itrace::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel& kernel = instrumentor.Kernel();
    uint32_t numTiles = (instrumentor.Coder().IsTileIdSupported()) ? GTPin_GetCore()->GenArch().MaxTiles(kernel.GpuPlatform()) : 1;

    // Create new KernelData object and add it to the data base
    auto ret = _kernels.emplace(piecewise_construct, forward_as_tuple(kernel.Id()), forward_as_tuple(instrumentor, numTiles));
    if (ret.second)
    {
        ItraceKernel& ItraceKernel = (*ret.first).second;
        if (!ItraceKernel.IsEnabled())
        {
            GTPIN_WARNING("ITRACE: The trace won't be generated for kernel " + ItraceKernel.Name());
            return;
        }

        const IGtCfg& cfg = instrumentor.Cfg();
        IGtVregFactory& vregs = instrumentor.Coder().VregFactory();

        // Initialize virtual registers
        _addrReg = vregs.MakeMsgAddrScratch();
        _dataReg = vregs.MakeMsgDataScratch();
        _offsetReg = vregs.MakeScratch(VREG_TYPE_DWORD);
        _tileIdReg = vregs.Make(VREG_TYPE_DWORD);

        GtGenProcedure preCode;
        instrumentor.Coder().LoadTileId(preCode, _tileIdReg);

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

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

void Itrace::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())
        {
            const ItraceKernel& ItraceKernel = it->second;
            if (ItraceKernel.IsEnabled())
            {
                IGtProfileBuffer* buffer = dispatcher.CreateProfileBuffer(); GTPIN_ASSERT(buffer);
                const GtProfileTrace& traceAccessor = ItraceKernel.TraceAccessor();
                if (traceAccessor.Initialize(*buffer))
                {
                    isProfileEnabled = true;
                }
                else
                {
                    GTPIN_ERROR_MSG("ITRACE: Failed to write into memory buffer for kernel " + string(kernel.Name()));
                }
            }
        }
    }
    dispatcher.SetProfilingMode(isProfileEnabled);
}

void Itrace::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())
    {
        // Read the trace from the profile buffer
        ItraceKernel& ItraceKernel = it->second;
        ItraceKernel.AddItrace(dispatcher);
    }
}

bool Itrace::InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const ItraceKernel& ItraceKernel)
{
    const IGtGenCoder& coder = instrumentor.Coder();
    const IGtCfg& cfg = instrumentor.Cfg();

    // Generate code that allocates space for the new record in the trace and stores the trace record.
    // Insert this procedure before the first instruction in the basic block.
    GtGenProcedure headerProc;
    auto firstInsIt = bbl.Instructions().begin();
    const IGtIns& firstIns = cfg.GetInstruction((*firstInsIt)->Id());
    StoreRecord(headerProc, coder, bbl, ItraceKernel, sizeof(ItraceRecord));
    instrumentor.InstrumentInstruction(firstIns, GtIpoint::Before(), headerProc);

    return true;
}

void Itrace::StoreRecord(GtGenProcedure& proc, const IGtGenCoder& coder, const IGtBbl& bbl,
    const ItraceKernel& ItraceKernel, uint32_t recordSize)
{
    IGtInsFactory& insF = coder.InstructionFactory();

    GtPredicate    predicate(FlagReg(0));

    // Set values of ItraceRecord fields in _dataReg
    proc += insF.MakeShl(_dataReg, StateReg(0), 16);                    // idFieldReg[16:31] = sr0.0
    proc += insF.MakeAdd(_dataReg, _dataReg, GtImmU32(bbl.Id()));       // idFieldReg[0:15]  = bbl.Id()

    // Allocate new record in the trace.
    // Set _offsetReg = offset of the allocated record in the profile buffer, _addrReg = address of the allocated record
    ItraceKernel.TraceAccessor().ComputeNewRecordOffset(coder, proc, recordSize, _offsetReg);
    coder.ComputeAddress(proc, _addrReg, _offsetReg);

    // Zero _offsetReg if the trace buffer is overflowed (predicate == true)
    proc += insF.MakeMov(_offsetReg, 0).SetPredicate(predicate);

    // Store Sr0.0 and BBL ID
    //if (!predicate) { STORE buffer[_addrReg] = _dataReg;
    proc += insF.MakeAtomicStore(_addrReg, _dataReg).SetPredicate(!predicate);

    // Store tile ID
    //if (!predicate) { STORE buffer[_addrReg] = _dataReg;
    proc += insF.MakeMov(_dataReg, _tileIdReg);
    coder.ComputeRelAddress(proc, _addrReg, _addrReg, offsetof(ItraceRecord, tileId));
    proc += insF.MakeAtomicStore(_addrReg, _dataReg).SetPredicate(!predicate);

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

void Itrace::OnFini()
{
    Itrace& me = *Instance();
    IGtCore* gtpinCore = GTPin_GetCore();
    for (auto& ref : me._kernels)
    {
        const ItraceKernel& ItraceKernel = ref.second;
        ItracePostProcessor(*gtpinCore, ItraceKernel)();
        ItraceKernel.DumpAsm();
    }
}

/* ============================================================================================= */
// ItracePreProcessor implementation
/* ============================================================================================= */
const char* ItracePreProcessor::_kernelPreProcessFileName = "itrace_pre_process.txt";
const char* ItracePreProcessor::_dispatchPreProcessFileName = "itrace_pre_process_dispatch.txt";

ItracePreProcessor::ItracePreProcessor()
{
    if (gKnobPhase == 2)
    {
        // Read the data collected during the preprocessing phase
        std::ifstream is(_kernelPreProcessFileName);
        GTPIN_ASSERT_MSG(is, string("File ") + _kernelPreProcessFileName + " does not exist. The trace won't be generated");
        is >> _kernelCounters;
    }
    else if (gKnobPhase == 1)
    {
        // Create pre_process files or remove old pre_process files's content if they exist
        CreateCleanFile(_kernelPreProcessFileName);
        CreateCleanFile(_dispatchPreProcessFileName);
    }
}

ItracePreProcessor* ItracePreProcessor::Instance()
{
    static ItracePreProcessor instance;
    return &instance;
}

void ItracePreProcessor::OnFini()
{
    ItracePreProcessor& tool = *Instance();
    tool.DumpKernelProfiles(_kernelPreProcessFileName);
    tool.DumpDispatchProfiles(_dispatchPreProcessFileName);
}

uint64_t ItracePreProcessor::TraceSize(const string& extKernelName) const
{
    auto it = _kernelCounters.find(extKernelName);
    return ((it == _kernelCounters.end()) ? 0 : it->second.weight);
}

uint32_t ItracePreProcessor::GetBblWeight(IGtKernelInstrument&, const IGtBbl&) const
{
    return sizeof(ItraceRecord);
}

void ItracePreProcessor::AggregateDispatchCounters(KernelWeightCounters& kc, KernelWeightCounters dc) const
{
    kc.weight = std::max(kc.weight, dc.weight);
    kc.freq += dc.freq;
}

/* ============================================================================================= */
// ItracePostProcessor implementation
/* ============================================================================================= */
const char* ItracePostProcessor::_traceFileName = "itrace_compressed.bin";

ItracePostProcessor::ItracePostProcessor(const IGtCore& gtpinCore, const ItraceKernel& ItraceKernel) :
    _kernel(&ItraceKernel),
    _kernelDir(JoinPath(string(gtpinCore.ProfileDir()), ItraceKernel.UniqueName())) {}

bool ItracePostProcessor::operator()()
{
    if (!MakeDirectory(_kernelDir))
    {
        GTPIN_WARNING("ITRACE: Could not create directory " + _kernelDir);
        return false;
    }

    // Process traces recorded in kernel dispatches
    for (const ItraceDispatch& trace : _kernel->GetTraces())
    {
        if (!trace.IsEmpty())
        {
            if (trace.IsTrimmed())
            {
                GTPIN_WARNING("ITRACE: Detected trace buffer overflow in kernel " + _kernel->Name());
            }

            ProcessTrace(trace);

            string subdir = trace.KernelExecDesc().ToString(_kernel->Platform(), ExecDescFileNameFormat());
            string dir = MakeSubDirectory(_kernelDir, subdir);

            if (!gKnobCfgOnly)
            {
                string filePath = JoinPath(dir, _traceFileName);
                ofstream fs(filePath, std::ios::binary);
                if (!fs)
                {
                    GTPIN_WARNING("ITRACE: Could not create file " + filePath);
                    continue;
                }
                StoreTrace(fs);
            }

            string cfgFilePath = JoinPath(dir, "itrace_total.cfg");
            ofstream cfgfs(cfgFilePath);
            if (!cfgfs)
            {
                GTPIN_WARNING("ITRACE: Could not create file " + cfgFilePath);
                continue;
            }
            StoreCfg(cfgfs);
        }
    }

    return true;
}

void ItracePostProcessor::ProcessTrace(const ItraceDispatch& trace)
{
    const uint8_t* traceData = trace.Data();
    uint32_t       traceSize = trace.Size();

    // Associate trace records with threads - populate _threadTraceRecords array
    const GtStateRegAccessor& sra = _kernel->GenModel().StateRegAccessor();
    uint32_t maxThreads = _kernel->GenModel().MaxThreads(); // Max number of HW threads

    // Build control-flow graph
    _cfg.clear();
    ItraceKernel::Edges edges = _kernel->GetEdges();
    for (auto& edge : edges)
    {
        _cfg.emplace(edge, 0);
    }

    // Reference to the trace record
    struct Record
    {
        const ItraceRecord* header; ///< Pointer to the header of the record
        uint32_t              size; ///< Size of the record in bytes, including header
    };

    using RecordList = std::list<Record>;           ///< List of Records
    using PerTileRecords = std::vector<RecordList>; ///< Per tile Records
    std::vector<PerTileRecords> threadRecords;      ///< Vector of per tile records. Tile ID is an index to this vector

    threadRecords.resize(_kernel->NumTiles());
    _itraceRecords.resize(_kernel->NumTiles());
    _numProfiledThreads.resize(_kernel->NumTiles(), 0);

    for (uint32_t tile = 0; tile < _kernel->NumTiles(); tile++)
    {
        threadRecords[tile].resize(maxThreads);

        _itraceRecords[tile].clear();
        _itraceRecords[tile].resize(maxThreads);
    }

    for (uint32_t recordOffset = 0; recordOffset + sizeof(ItraceRecord) <= traceSize;)
    {
        // Retrive thread ID and BBL ID from the record
        const ItraceRecord* record = (const ItraceRecord*)(traceData + recordOffset);
        uint32_t tid = sra.GetGlobalTid(record->sr0);
        uint32_t tileId = record->tileId; GTPIN_ASSERT(tileId < _kernel->NumTiles());
        uint32_t recordSize = sizeof(ItraceRecord);
        if (recordOffset + recordSize > traceSize)
        {
            break; // end of trace
        }

        auto& tileRecords = threadRecords[tileId];
        auto& records = tileRecords[tid];

        // Add a new trace record reference to _threadTraceRecords
        if (records.empty()) { ++_numProfiledThreads[tileId]; } // Increment thread count on the first relevant record
        records.emplace_back(Record{ record, recordSize });

        recordOffset += recordSize;
    }

    for (uint32_t tileId = 0; tileId < threadRecords.size(); tileId++)
    {
        if (_numProfiledThreads[tileId] == 0)
        {
            continue;
        }

        const auto& tileRecords = threadRecords[tileId];
        auto& tileItraceRecords = _itraceRecords[tileId];

        // Store per-thread traces
        for (uint32_t tid = 0; tid < maxThreads; tid++)
        {
            const auto& threadRecordList = tileRecords[tid];
            auto& records = tileItraceRecords[tid];

            if (threadRecordList.empty())
            {
                continue;
            }

            // Store trace records
            BblId prevBblId;
            for (const auto& record : threadRecordList)
            {
                const auto& header = *(record.header);
                BblId       bblId  = header.bblId;

                if (prevBblId != bblId)
                {
                    records.emplace_back(bblId, 1);
                }
                else
                {
                    auto& r = records.back();
                    r.second++;
                }

                if (prevBblId.IsValid())
                {
                    auto it = _cfg.find(ItraceKernel::Edge(prevBblId, bblId));
                    if (it != _cfg.end())
                    {
                        it->second += 1;
                    }
                }

                prevBblId = bblId;
            }
        }
    }
}

void ItracePostProcessor::StoreTrace(std::ofstream& fs)
{
    StoreBblBoundsInfo(fs);

    // Compute and store the number of involved tiles
    uint32_t numOfTiles = 0;
    for (uint32_t i = 0; i < _numProfiledThreads.size(); i++)
    {
        numOfTiles += (_numProfiledThreads[i] == 0) ? 0 : 1;
    }
    Store(numOfTiles, fs);

    for (uint32_t tileId = 0; tileId < _itraceRecords.size(); tileId++)
    {
        if (_numProfiledThreads[tileId] == 0) { continue; }

        const auto& tileRecords = _itraceRecords[tileId];

        Store(tileId, fs);

        // Store the number of profiled threads
        Store(_numProfiledThreads[tileId], fs);

        // Store per-thread traces
        for (uint32_t tid = 0; tid < tileRecords.size(); tid++)
        {
            const auto& records = tileRecords[tid];

            if (records.empty())
            {
                continue;
            }

            StoreGlobalTid(tid, fs);    // Store Global Thread Identifier

            uint32_t numRecords = (uint32_t)records.size();
            Store(numRecords, fs);      // Store #records collected in the thread

            // Store trace records
            for (const auto& record : records)
            {
                uint32_t bblId = record.first;
                uint32_t loopCount = record.second;
                Store(bblId, fs);       // Store BBL ID
                Store(loopCount, fs);   // Store loopCount
            }
        }
    }
}

void ItracePostProcessor::StoreBblBoundsInfo(std::ofstream& fs)
{
    // Store static information about memory accesses in BBLs
    ItraceKernel::BblBoundsMap bblBoundsMap = _kernel->GetBblBounds();
    uint32_t numBbls = (uint32_t)bblBoundsMap.size();
    Store(numBbls, fs);                                 // Store the number of BBLs that access memory

    for (uint32_t bblId = 0; bblId < numBbls; bblId++)
    {
        auto bounds = bblBoundsMap[bblId];
        Store(bblId, fs);                               // Store BBL ID
        uint32_t val = bounds.first;
        Store(val, fs);
        val = bounds.second;
        Store(val, fs);
    }
}

void ItracePostProcessor::StoreGlobalTid(uint32_t gtid, std::ofstream& fs)
{
    const GtStateRegAccessor& sra = _kernel->GenModel().StateRegAccessor();
    uint32_t sr0 = sra.SetGlobalTid(0, gtid);

    auto storeSr0Field = [&](const ScatteredBitFieldU32& sbf)
    {
        uint32_t val = (sbf.IsEmpty() ? UINT32_MAX : sbf.GetValue(sr0));
        Store(val, fs);
    };

    storeSr0Field(sra.SliceIdField());
    storeSr0Field(sra.DualSubSliceIdField());
    storeSr0Field(sra.SubSliceIdField());
    storeSr0Field(sra.EuIdField());
    storeSr0Field(sra.ThreadSlotField());
}

void ItracePostProcessor::StoreCfg(std::ofstream& fs)
{
    ostringstream ostr;
    ostr << "srcBBL, dstBBL, Frequency" << std::endl;
    ostr << "=========================" << std::endl;
    for (const auto& it : _cfg)
    {
        ItraceKernel::Edge edge = it.first;
        uint32_t frequency = it.second;

        // print srcBBL, dstBBL, frequency
        ostr << std::dec << edge.first << "," << edge.second << "," << frequency << std::endl;
    }
    fs << ostr.str();
    fs.close();
}

/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char* argv[])
{
    ConfigureGTPin(argc, argv);
    if (gKnobPhase == 1)
    {
        ItracePreProcessor::Instance()->Register();
        atexit(ItracePreProcessor::OnFini);
    }
    else
    {
        GTPIN_ASSERT_MSG((gKnobPhase == 2), "Itrace: Invalid phase value. Should be 1 or 2, provided " + std::to_string(gKnobPhase));
        Itrace::Instance()->Register();
        atexit(Itrace::OnFini);
    }
}

