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

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

/*!
 * @file FunLatency Tool definitions
 */

#ifndef FUNLATENCY_H_
#define FUNLATENCY_H_

#include <map>
#include <set>
#include <string>

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

using namespace gtpin;

/// Identifier of the instrumentation site : function ID
using InstrumentSiteId = FunctionId;

/// Array of instrumentation site IDs, indexed by ordinal numbers of corresponding records in the profile buffer
using InstrumentSites = std::vector<InstrumentSiteId>;

/* ============================================================================================= */
// Struct FunLatencyRecord
/* ============================================================================================= */
/*!
 * Layout of records collected in profile buffer by the Latency tool
 */
struct alignas(uint64_t) FunLatencyRecord
{
    uint64_t cycles;    ///< Total number of cycles
    uint32_t freq;      ///< Total number of executions
};

/* ============================================================================================= */
// Class FunLatencyDispatchProfile
/* ============================================================================================= */
/*!
 * Profiling data collected during a single kernel dispatch
 */
struct  FunLatencyDispatchProfile
{
    FunLatencyDispatchProfile(const IGtKernelDispatch& kernelDispatch, InstrumentSiteId siteId);
    void Accumulate(const FunLatencyRecord& record);

    InstrumentSiteId  siteId;         ///< Identifier of the instrumentation site: FuntionID
    GtKernelExecDesc  kernelExecDesc; ///< Kernel execution descriptor
    uint64_t          cycles;         ///< Total number of cycles
    uint64_t          freq;           ///< Total number of executions
};

/* ============================================================================================= */
// struct FunctionDescriptor
/* ============================================================================================= */
/*!
 * Static information about the function
 */
struct FunctionDescriptor
{
    FunctionDescriptor(const FunctionId fId, const std::string fName,
        const bool fIsMain, const InsId firstIns, const InsId lastIns, const uint32_t recNum);

    FunctionId      id;         ///< Identifier of the function
    std::string     name;       ///< function name
    bool            isMain;     ///< true for main kernel function
    InsId           firstInsId; ///< first Instruction in the function 
    InsId           lastInsId;  ///< last Instruction in the function 
    uint32_t        recordNum;  ///< record number in profile array
    GtReg           freqReg;    ///< Allocated VReg for freq count
    GtReg           cycleReg;    ///< Allocated VReg for cycles count
};

/* ============================================================================================= */
// Class FunLatencyKernelProfile
/* ============================================================================================= */
/*!
 * Aggregated profile of all instrumented kernel dispatches
 */
class FunLatencyKernelProfile
{
public: 
    typedef std::map<FunctionId, FunctionDescriptor>  FunctionMap;
    FunLatencyKernelProfile(const IGtKernel& kernel, const IGtCfg& cfg,
        const GtProfileArray& profileArray, InstrumentSites&& instrumentSites, FunctionMap&& functions);
    
    /// Add the specified profile of a kernel dispatch
    FunLatencyDispatchProfile& AddDispatchProfile(const FunLatencyDispatchProfile& dispatchProfile);

    std::string             ToString()        const;                          ///< @return Text representation of the profile data
    void                    DumpAsm()         const;                          ///< Dump kernel's assembly text to file
    std::string             GetName()         const { return _name; }         ///< @return  Kernel's name
    const GtProfileArray&   GetProfileArray() const { return _profileArray; } ///< @return Profile buffer accessor
    
    /// Given a record number in the profile buffer, return corresponding ID of the instrumentation site   
    InstrumentSiteId        GetSiteId(uint32_t recordNum) const;

    /// @return Collection of profiles per kernel dispatch
    typedef std::list<FunLatencyDispatchProfile>  DispatchProfiles;
    const DispatchProfiles& GetDispatchProfiles() const { return _dispatchProfiles; }

private:
    std::string        _name;              ///< Kernel's name
    std::string        _uniqueName;        ///< Kernel's unique name (IGC style)
    GtKernelType       _type;              ///< Kernel's type
    GtGpuPlatform      _platform;          ///< Kernel's platform
    uint64_t           _hashId;            ///< Kernel's hash identifier
    GtSimdWidth        _simd;              ///< Kernel's SIMD width
    std::string        _asmText;           ///< Kernel's assembly text
    GtProfileArray     _profileArray;      ///< Profile buffer accessor
    uint64_t           _totalKernelCycles; ///< Total cycles of the kernel
    InstrumentSites    _instrumentSites;   ///< Array of instrumentation site IDs
    FunctionMap        _functionMap;       ///< Map each function ID to it's Function descriptor
    DispatchProfiles   _dispatchProfiles;  ///< Profiles per kernel dispatch
};

/* ============================================================================================= */
// Class FunLatency
/* ============================================================================================= */
/*!
 * Implementation of the IGtTool interface for the FunLatency tool
 */
class FunLatency : public GtTool
{
public:
    /// Implementation of the IGtTool interface
    const char* Name() const { return "FunLatency"; }

    void OnKernelBuild(IGtKernelInstrument& instrumentor);
    void OnKernelRun(IGtKernelDispatch& dispatcher);
    void OnKernelComplete(IGtKernelDispatch& dispatcher);

public:
    std::string ToString() const;                ///< @return Text representation of the profile data
    void        DumpAsm()  const;                ///< Dump assembly text of profiled kernels to files

    static FunLatency* Instance();               ///< @return Single instance of this class
    static void OnFini() { Instance()->Fini(); } ///< Callback function registered with atexit()
protected:
    FunLatency() = default;
    FunLatency(const FunLatency&) = delete;
    FunLatency& operator = (const FunLatency&) = delete;
    ~FunLatency() = default;

    /// @return Collection of kernel profiles
    typedef std::map<GtKernelId, FunLatencyKernelProfile>  KernelProfiles;
    const KernelProfiles& GetKernelProfiles() const { return _kernels; }

    /// Post process and dump profiling data
    void Fini();
private:
    /// Generate instrumentation for functions
    void InstrumentFunctions(IGtKernelInstrument& instrumentor);
	
    /// Generate instrumentation to be inserted before the code fragment whose latency is measured
    void GeneratePreCode(GtGenProcedure& proc, const IGtGenCoder& coder);

    /*!
     * Generate instrumentation that computes latency and stores the result in the profile buffer.
     *  - GeneratePostCodeMemBound  - Generate memory-bound instrumentation
     *  - GeneratePostCodeRegBound  - Generate register-bound instrumentation. Return success/failure status
     *
     * @param[in, out] postProc         Procedure to be inserted after the code fragment whose latency is measured
     * @param[in, out] finiProc         Procedure to be inserted before each EOT
     * @param[in]      coder            GEN code generator
     * @param[in]      profileArray     Array of 'FunLatencyRecord's in the profile buffer
     * @param[in]      funcDesc         pointer to the decriptor of the function to be instrumented
     * @param[in]      isFuncEntry      true if the BBL is the first BBL in the Function
     */
    bool GeneratePostCodeRegBound(GtGenProcedure& postProc, GtGenProcedure& finiProc, const IGtGenCoder& coder,
        const GtProfileArray& profileArray, FunctionDescriptor* funcDesc, bool isFuncEntry, bool updateFrequency, uint32_t simdMask);
    void GeneratePostCodeMemBound(GtGenProcedure& proc, const IGtGenCoder& coder,
        const GtProfileArray& profileArray, FunctionDescriptor* funcDesc, bool updateFrequency, uint32_t simdMask);

    /// @return true/false - use 64-bit/32-bit integer for the cycle counter
    static bool Use64BitCounters(const IGtGenCoder& coder);

private:
    KernelProfiles  _kernels;   ///< Collection of kernel profiles

    GtReg _addrReg;             ///< Virtual register that holds address within profile buffer
    GtReg _dataReg;             ///< Virtual register that holds data to be read from/written to profile buffer 
    GtReg _timeReg;             ///< Virtual timer register
};

#endif