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

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

/*!
 * @file Implementation of the SIMD operation counting tool
 */

#include <algorithm>
#include <vector>
#include <map>
#include <string>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <assert.h>

#include "simdprof.h"

using namespace gtpin;
using namespace std;

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int>  knobNumThreadBuckets("num_thread_buckets", 32, "Number of thread buckets. 0 - maximum thread buckets");

/* ============================================================================================= */
// SimdProfArgs implementation
/* ============================================================================================= */

bool SimdProfArgs::operator < (const SimdProfArgs& other) const
{
    return std::make_tuple(maskCtrl, execMask, predicate, isSendIns) <
           std::make_tuple(other.maskCtrl, other.execMask, other.predicate, other.isSendIns);
}

/* ============================================================================================= */
// SimdProfSection implementation
/* ============================================================================================= */

void SimdProfSection::AddInstruction(const IGtIns& ins)
{
    uint32_t    execMask    = ins.ExecMask().Bits();
    GtPredicate predicate   = ins.Predicate();
    bool        maskCtrl    = !ins.IsWriteMaskEnabled();
    bool        isSendIns   = ins.IsSendMessage();

    auto it = groups.emplace(SimdProfArgs(maskCtrl, execMask, predicate, isSendIns), 0).first;
    ++(it->second);
}

/* ============================================================================================= */
// SimdprofKernelProfile implementation
/* ============================================================================================= */

SimdProfKernelProfile::SimdProfKernelProfile(const IGtKernel& kernel) :
    _name(GlueString(kernel.Name())), _extName(ExtendedKernelName(kernel)), _type(kernel.Type()), _platform(kernel.GpuPlatform()),
    _hashId(kernel.HashId()), _simd(kernel.SimdWidth()), _binarySignature(kernel.BinarySignature()),
    _totalOpCount(0) {}

void SimdProfKernelProfile::Instrument(IGtKernelInstrument& instrumentor)
{
    const IGtGenCoder&  coder           = instrumentor.Coder();
    const IGtKernel&    kernel          = instrumentor.Kernel();
    const IGtCfg&       cfg             = instrumentor.Cfg();
    IGtVregFactory&     vregs           = coder.VregFactory();
    bool                is64BitCounter  = Use64BitCounters(coder);

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

    // Identify SimdProf sections and #groups in the kernel
    std::vector<SimdProfSection> sections;      // All SimdProf sections in the kernel
    uint32_t                     numGroups = 0; // Number of SimdProf groups in the kernel

    for (auto bblPtr : cfg.Bbls())
    {
        bool isSectionBegin = true;

        // Iterate through sections within the current BBL
        for (auto insPtr : bblPtr->Instructions())
        {
            const IGtIns& ins = *insPtr;

            if (ins.Id() < (uint32_t)knobMinInstrumentIns || ins.Id() > (uint32_t)knobMaxInstrumentIns)
            {
                continue;
            }

            if (isSectionBegin)
            {
                sections.emplace_back(ins);
                isSectionBegin = false;
            }

            SimdProfSection& section = sections.back();
            section.AddInstruction(ins);

            if (ins.IsFlagModifier() || (ins.Id() == bblPtr->LastIns().Id())) //section end
            {
                numGroups += (uint32_t)section.groups.size();
                isSectionBegin = true;
            }
        }

        if (isSectionBegin == false)
        {
            numGroups += (uint32_t)sections.back().groups.size();
        }
    }

    // Allocate the profile buffer. It will hold single SimdProfRecord per each group in each thread bucket
    uint32_t numThreadBuckets = (knobNumThreadBuckets == 0) ? kernel.GenModel().MaxThreadBuckets() : knobNumThreadBuckets;
    _profileArray = GtProfileArray(sizeof(SimdProfRecord), numGroups, numThreadBuckets);
    _profileArray.Allocate(instrumentor.ProfileBufferAllocator());

    // Instrument SimdProf sections and initialize the _profileData array
    for (auto& section : sections) { InstrumentSection(instrumentor, section); }

    // Save BBL information for the post processing phase
    for (auto bblPtr : cfg.Bbls())
    {
        _bblInsInfo.emplace(bblPtr->Id(), std::make_pair(bblPtr->FirstIns().Id(), bblPtr->LastIns().Id()));
    }
}

void SimdProfKernelProfile::InstrumentSection(IGtKernelInstrument& instrumentor, const SimdProfSection& section)
{
    const IGtGenCoder&  coder           = instrumentor.Coder();
    IGtInsFactory&      insF            = coder.InstructionFactory();
    const IGtCfg&       cfg             = instrumentor.Cfg();
    bool                is64BitCounter  = Use64BitCounters(coder);
    GtReg               dataRegL        = {_dataReg, sizeof(uint32_t), 0};  // Low 32-bits of the data payload register

    // Instrument each SimdProf group:
    //  - If a group is associated with a non-SEND instructions, compute the SIMD count by aplying CBIT to the SIMD mask.
    //  - Otherwise, if a group is created for SEND instructions, increment the SIMD count for each SEND whose SIMD mask
    //    is nonzero. From the EU perspective, SEND instruction is 1 operation, unless the SIMD mask is zero
    // Insert each per-group instrumentation procedure at the beginning of the corresponding section

    //Insert SimdProf instrumentaion at the beginning of the current section
    const IGtIns& ins = cfg.GetInstruction(section.headInsId);
    const IGtBbl& bbl = cfg.GetBbl(ins);

    for (auto& group : section.groups)
    {
        GtGenProcedure      proc;
        const SimdProfArgs& args = group.first;

        if (is64BitCounter)
        {
            // Clear the high 32-bits of the data payload register
            GtReg dataRegH = {_dataReg, sizeof(uint32_t), 1};
            proc += insF.MakeMov(dataRegH, 0);
        }
		
        // dataRegL = SIMD mask
        coder.ComputeSimdMask(proc, dataRegL, args.maskCtrl, args.execMask, args.predicate);

        // dataRegL = number SIMD operations executed
        if (!args.isSendIns)
        {
            proc += insF.MakeCbit(dataRegL, dataRegL);
        }
        else
        {
            proc += insF.MakeSel(dataRegL, dataRegL, 1).SetCondModifier(GED_COND_MODIFIER_l); // dataRegL = min(dataRegL, 1)
        }

        // Generate code that updates the SIMD operation counter in the corresponding SimdProfRecord
        uint32_t recordNum = (uint32_t)_profileData.size();
        _profileArray.ComputeAddress(coder, proc, _addrReg, recordNum);

        proc += insF.MakeAtomicAdd(NullReg(), _addrReg, _dataReg, (is64BitCounter? GED_DATA_TYPE_uq : GED_DATA_TYPE_ud));

        // Insert a new instrumentation routine and append the new group to _profileData
        proc.front()->AppendAnnotation(__func__);
        SimdProf::Instance()->InstrumentInstruction(instrumentor, ins, GtIpoint::Before(), proc);
        _profileData.emplace_back(bbl.Id(), group.second);
    }
}

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

void SimdProfKernelProfile::ReadProfileData(const IGtProfileBuffer* buffer)
{
    GTPIN_ASSERT(_profileData.size() == _profileArray.NumRecords());
    uint32_t recordNum = 0;

    // Iterate through all SimdProf groups and read counters of executed operations (channels).
    for (auto& group : _profileData)
    {
        // Accumulate counters for all threads in which this group of instructions was executed
        for (uint32_t threadBucket = 0; threadBucket < _profileArray.NumThreadBuckets(); ++threadBucket)
        {
            SimdProfRecord record;
            if (!_profileArray.Read(*buffer, &record, recordNum, 1, threadBucket))
            {
                GTPIN_ERROR_MSG(string("SIMDPROF : ") + _name + " : Failed to read from memory buffer");
            }
            else
            {
                // Update counters of executed operations
                uint64_t opCount = record.opCount * group.insCount;
                group.opCount += opCount;
                UpdateBblOpCounter(group.bblId, opCount);
                _totalOpCount += opCount;
            }
        }
        recordNum++;
    }
}

void SimdProfKernelProfile::UpdateBblOpCounter(BblId bblId, uint64_t incValue)
{
    auto it = _bblOpCounts.emplace(bblId, 0).first;
    it->second += incValue;
}

std::string SimdProfKernelProfile::ToString() const
{
    ostringstream ostr;
    ostr << ExtendedName() << endl;

    if (!_bblOpCounts.empty())
    {
        ostr << setw(10) << "BBL" << setw(15) << "Head Ins ID" << setw(15) << "Tail Ins ID" << setw(20) << "Channels" << endl;
        for (const auto& bc : _bblOpCounts)
        {
            ostr << setw(10) << bc.first << setw(15) << _bblInsInfo.at(bc.first).first << setw(15) << _bblInsInfo.at(bc.first).second << setw(20) << bc.second << endl;
        }
        ostr << setw(10) << "Total" << setw(15) << _totalOpCount << endl;
    }
    else
    {
        ostr << "No channels executed" << endl;
    }

    return ostr.str();
}

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

void SimdProf::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel& kernel = instrumentor.Kernel();
    auto it = _kernels.emplace(kernel.Id(), kernel).first;
    it->second.Instrument(instrumentor);
}

void SimdProf::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    bool isProfileEnabled = false;

    const IGtKernel& kernel = dispatcher.Kernel();
    GtKernelExecDesc execDesc; dispatcher.GetExecDescriptor(execDesc);
    if (kernel.IsInstrumented() && IsKernelExecProfileEnabled(execDesc, kernel.GpuPlatform()))
    {
        auto it = _kernels.find(kernel.Id());

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

void SimdProf::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())
    {
        const IGtProfileBuffer* buffer = dispatcher.GetProfileBuffer(); GTPIN_ASSERT(buffer);
        SimdProfKernelProfile& kernelProfile = it->second;
        kernelProfile.ReadProfileData(buffer);
    }
}

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

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

string SimdProf::ToString() const
{
    ostringstream ostr;
    ostr << "Channels (SIMD operations) executed by kernels/BBLs" << endl;
    ostr << "===================================================" << endl;
    
    uint64_t totalOpCount = 0;
    for (const auto& k : _kernels)
    {
        ostr << string(100, '-') << endl;
        ostr << k.second.ToString() << endl;
        totalOpCount += k.second.GetTotalOpCounter();
    }
    ostr << "Total number of kernels:                    " << _kernels.size() << std::endl;
    ostr << "Total number of channels (SIMD operations): " << totalOpCount << std::endl;

    return ostr.str();
}

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