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

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

/*!
 * @file Pertrace tool definitions
 */

#ifndef PERFTRACE_H_
#define PERFTRACE_H_

#include <list>
#include <map>
#include <vector>

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

using namespace gtpin;

#pragma pack(push, 1)

/* ============================================================================================= */
// Struct PerfTraceRecord
/* ============================================================================================= */
/*!
 * Structure of the trace records.
 */
struct alignas(16) PerfTraceFuntimeRecord
{
    uint32_t sr0;       ///< State register sr0.0:ud
    uint32_t tileId;    ///< Tile ID
    uint64_t timeStart; ///< Time stamp counter before the first instruction of the kernel
    uint64_t timeEnd;   ///< Time stamp counter before the last (EOT) instruction of the kernel
};

struct alignas(16) PerfTraceMemoryLatencyRecord
{
    uint32_t sr0;       ///< State register sr0.0:ud
    uint32_t tileId;    ///< Tile ID
    uint32_t cycles;    ///< Accumulated cycles
};

/* ============================================================================================= */
// Class PerfTraceDispatch
/* ============================================================================================= */
/*!
 * Class that holds memory trace collected during a single kernel dispatch
 */
class PerfTraceDispatch
{
public:
    /// Construct a PerfTraceDispatch object with the empty trace
    explicit PerfTraceDispatch(const IGtKernelDispatch& dispatch) : _isTrimmed(false) { dispatch.GetExecDescriptor(_kernelExecDesc); }
    
    /// Read the entire trace from the specified profile buffer into this object
    bool ReadTrace(const GtProfileTrace& traceAccessor, const IGtProfileBuffer& profileBuffer);

    const GtKernelExecDesc& KernelExecDesc() const { return _kernelExecDesc; } ///< @return Descriptor of this kernel dispatch
    uint32_t        Size()      const { return (uint32_t)_rawTrace.size(); }   ///< @return Trace size in bytes
    const uint8_t*  Data()      const { return _rawTrace.data(); }             ///< @return Trace data collected in this dispatch
    uint8_t*        Data()            { return _rawTrace.data(); }             ///< @return Trace data collected in this dispatch
    bool            IsEmpty()   const;                                         ///< @return true if the trace is empty
    bool            IsTrimmed() const { return _isTrimmed; }                   ///< @return true if the trace has been trimmed

private:
    GtKernelExecDesc        _kernelExecDesc; ///< Kernel execution descriptor
    std::vector<uint8_t>    _rawTrace;       ///< Trace data collected in this kernel dispatch
    bool                    _isTrimmed;      ///< true if the trace has been trimmed to avoid buffer overflow
};

/* ============================================================================================= */
// Class PerfTraceKernel
/* ============================================================================================= */
class PerfTraceKernel
{
public:
    PerfTraceKernel() = default;

    explicit PerfTraceKernel(const IGtKernelInstrument& kernelInstrument, const uint32_t recordSize, uint32_t numTiles);

    /*!
     * Read a trace recorded by the specified kernel dispatch. Create and add the corresponding PerfTraceDispatch
     * instance to this object
     */
    PerfTraceDispatch& AddPerfTrace(IGtKernelDispatch& kernelDispatch);

    std::string           Name()            const { return _name; }               ///< @return Kernel's name
    std::string           ExtendedName()    const { return _extName; }            ///< @return Kernel's extended name
    std::string           UniqueName()      const { return _uniqueName; }            ///< @return Kernel's unique name
    const GtGpuPlatform   Platform()        const { return _platform; }           ///< @return Kernel's platform
    const IGtGenModel&    GenModel()        const { return GetGenModel(_genId); } ///< @return Kernel's GEN model
    const GtProfileTrace& TraceAccessor()   const { return _traceAccessor; }      ///< @return Trace accessor
    void  DumpAsm()                         const;                                ///< Dump kernel's assembly text to file

     /// @return true, if tracing of this kernel is enabled
    uint32_t IsEnabled() const { return (_traceAccessor.MaxTraceSize() != 0); }

    /// @return Traces collected in kernel's dispatches
    typedef std::list<PerfTraceDispatch>  Traces;
    const Traces& GetTraces() const { return _traces; }

    /// @return Number of tiles
    uint32_t NumTiles() const { return _numTiles; }

private:
    std::string         _name;              ///< Kernel's name
    std::string         _uniqueName;        ///< Kernel's unique name
    std::string         _extName;           ///< Kernel's extended name
    GtGpuPlatform       _platform;          ///< Kernel's platform
    GtGenModelId        _genId;             ///< Identifier of the GEN model, the kernel is compiled for
    std::string         _asmText;           ///< Kernel's assembly text
    GtProfileTrace      _traceAccessor;     ///< Trace accessor
    Traces              _traces;            ///< Traces collected in kernel's dispatches
    uint32_t            _numTiles;          ///< The number of supported tiles
};

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

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

public:
    virtual uint32_t RecordSize() const = 0; ///< @return Record size aligned to OWORD

protected:
    /*!
     * Generate instrumentation for kernel
     * @param[in] instrumentor      Interface of the kernel being instrumented
     * @param[in] perfTraceKernel   Object that holds information about the kernel
     */
    virtual void Instrument(IGtKernelInstrument& instrumentor, const PerfTraceKernel& perfTraceKernel) = 0;

protected:
    PerfTrace() = default;
    PerfTrace(const PerfTrace&) = delete;
    PerfTrace& operator = (const PerfTrace&) = delete;
    ~PerfTrace() = default;

protected:
    std::map<GtKernelId, PerfTraceKernel>   _kernels;  ///< Collection of traces per kernel
};

/* ============================================================================================= */
// Class PerfTraceFuntime
/* ============================================================================================= */
class PerfTraceFuntime : public PerfTrace
{
public:
    static PerfTraceFuntime* Instance();
    static void OnFini(void);

    inline uint32_t RecordSize() const { return sizeof(PerfTraceFuntimeRecord); }

private:

    /*!
     * Generate instrumentation for kernel
     * @param[in] instrumentor      Interface of the kernel being instrumented
     * @param[in] perfTraceKernel   Object that holds information about the kernel
     */
    void Instrument(IGtKernelInstrument& instrumentor, const PerfTraceKernel& perfTraceKernel);

    /*!
     * Generate instrumentation for the specified basic block
     * @param[in] instrumentor      Interface of the kernel being instrumented
     * @param[in] bbl               Basic block to be instrumented
     * @param[in] perfTraceKernel   Object that holds information about the kernel
     */
    void InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const PerfTraceKernel& perfTraceKernel);

    void GeneratePreCode(GtGenProcedure& proc, const IGtGenCoder& coder);
    void GeneratePostCode(GtGenProcedure& proc, const IGtGenCoder& coder, const PerfTraceKernel& perfTraceKernel);

private:
    PerfTraceFuntime() : PerfTrace() {}
    PerfTraceFuntime(const PerfTraceFuntime&) = delete;
    PerfTraceFuntime& operator = (const PerfTraceFuntime&) = delete;
    ~PerfTraceFuntime() {}

private:

    GtReg _timeReg; ///< Virtual timer register
};

/* ============================================================================================= */
// Class PerfTraceMemoryLatency
/* ============================================================================================= */
class PerfTraceMemoryLatency : public PerfTrace
{
public:
    static PerfTraceMemoryLatency* Instance();

    inline uint32_t RecordSize() const { return sizeof(PerfTraceMemoryLatencyRecord); }

    static void OnFini(void);

private:

    /*!
     * Generate instrumentation for kernel
     * @param[in] instrumentor      Interface of the kernel being instrumented
     * @param[in] perfTraceKernel   Object that holds information about the kernel
     */
    void Instrument(IGtKernelInstrument& instrumentor, const PerfTraceKernel& perfTraceKernel);

    /*!
     * Generate instrumentation for the specified basic block
     * @param[in] instrumentor      Interface of the kernel being instrumented
     * @param[in] bbl               Basic block to be instrumented
     * @param[in] perfTraceKernel   Object that holds information about the kernel
     */
    void InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const PerfTraceKernel& perfTraceKernel);

    void GeneratePreCode(GtGenProcedure& proc, const IGtGenCoder& coder);
    void GeneratePostCode(GtGenProcedure& proc, const IGtGenCoder& coder);
    void GenerateFiniCode(GtGenProcedure& proc, const IGtGenCoder& coder, const PerfTraceKernel& perfTraceKernel);

private:
    PerfTraceMemoryLatency() : PerfTrace() {}
    PerfTraceMemoryLatency(const PerfTraceMemoryLatency&) = delete;
    PerfTraceMemoryLatency& operator = (const PerfTraceMemoryLatency&) = delete;
    ~PerfTraceMemoryLatency() {}

private:

    GtReg _timeReg;         ///< Virtual register to read time register
    GtReg _cyclesAccumReg;  ///< Virtual register to accumulate cycles values
};


/* ============================================================================================= */
// Class PerfTracePreProcessor
/* ============================================================================================= */
/*!
 * Class that computes per-kernel trace sizes in the preprocessing phase, and provides access to
 * this data in the trace gathering phase
 */
class PerfTracePreProcessor : public KernelWeight
{
public:
    uint64_t TraceSize(const std::string& extKernelName) const; ///< Given extended kernel name, return the trace size in bytes
    static void OnFini();                                       ///< Callback function registered with atexit()

    static PerfTracePreProcessor* Instance();

    /*!
     * @return Weight of the specified basic block in the kernel.
     * By default this function returns the bbl.NumIns() value, which means that the tool counts the number of executed
     * instructions.
     * Derived classes may give a different interpretation of the basic block's "weight" by overriding this function.
     * For example, in order to count the number of executed basic blocks. the override function may return 1 (one).
     */
    uint32_t GetBblWeight(IGtKernelInstrument& ki, const IGtBbl& bbl) const;

private:
    PerfTracePreProcessor();
    PerfTracePreProcessor(const PerfTracePreProcessor&) = delete;
    PerfTracePreProcessor& operator = (const PerfTracePreProcessor&) = delete;

private:
    void AggregateDispatchCounters(KernelWeightCounters& kc, KernelWeightCounters dc) const;

protected:
    KernelWeightProfileData _kernelCounters; ///< Per-kernel counters of required trace records; collected in preprocessing phase

    static const char* _kernelPreProcessFileName;   ///< Name of the file that contains preprocessing data per kernel
    static const char* _dispatchPreProcessFileName; ///< Name of the file that contains preprocessing data per kernel dispatch

    PerfTrace*   _perfTrace;      ///< Corresponding PerfTrace object
};

/* ============================================================================================= */
// Class PerfTracePostProcessor
/* ============================================================================================= */
/*!
 * Function object that processes kernel traces - stores them in files within the profile directory:
 *
 *    kernel_name
 *    |
 *        |- kernel_dispatch_1
 *           |- memorytrace_compressed.bin
 *        |- kernel_dispatch_2
 *           |- memorytrace_compressed.bin
 * The .bin trace files can be uncompressed by the uncompress_memtrace.exe utility.
 *
 * Format of .bin trace files:
 * - Static information:
 *     - Number of BBLs that access memory
 *     - For each BBL that accesses memory:
 *        - BBL ID
 *        - Number of SEND instructions in this basic block
 *        - For each SEND instruction:
 *            - Decoded information including offset, address model, address type, address payload size, etc
 * 
 * - Dynamic trace data:
 *      - Number of HW threads in which the trace was collected
 *      - For each HW thread:
 *          - HW Thread ID (in the format of sr0.0)
 *          - Number of records collected for this HW thread
 *          - All the records collected for this HW thread
 */
class PerfTracePostProcessor
{
public:
    /// Construct a PerfTracePostProcessor object for the specified collection of kernel traces
    PerfTracePostProcessor(const IGtCore& gtpinCore, const PerfTraceKernel&  perfTraceKernel, const PerfTrace* perfTrace);

    /// Process all kernel traces associated with this object - store them in files within the profile directory
    bool operator()();

protected:
    /// Derives global tid from the trace record
    virtual uint32_t GetGlobalTid(const uint8_t* traceRecord) const = 0;

    /// Derives global tid from the trace record
    virtual uint32_t GetTileId(const uint8_t* traceRecord) const = 0;

    /// Stores data from the trace record
    virtual void StoreTraceRecordData(const uint8_t* traceRecord, std::ofstream& fs) const = 0;

    /// Store the specified trace in the specified file stream
    void StoreTrace(const PerfTraceDispatch& trace, std::ofstream& fs);

    /// Store Global Thread Identifier
    void StoreGlobalTid(uint32_t gtid, std::ofstream& fs);

    /// Store the specified value in the specified file stream in the binary format
    template <typename T> void Store(const T& val, std::ofstream& fs) { fs.write((const char*)&val, sizeof(val)); }

protected:
    struct TraceRecord       ///< Reference to the trace record
    {
        const void* record;  ///< Pointer to the record
        uint32_t    size;    ///< Size of the record in bytes, including header
    };
    using TraceRecordList     = std::list<TraceRecord>;       ///< List of references to trace records
    using PerTileTraceRecords = std::vector<TraceRecordList>; ///< Per tile trace records

protected:
    const PerfTraceKernel*           _kernel;            ///< Kernel&traces to be processed
    std::string                      _kernelDir;         ///< Directory to store kernel's trace files
    std::vector<PerTileTraceRecords> _threadTraceRecords;///< Map of tile ID to Lists of trace records, indexed by the thread ID

    static const char*               _traceFileName;     ///< Name of the file to store trace in

    const PerfTrace*                 _perfTrace;         ///< Pointer to a corresponding instance of PerfTrace
};

/* ============================================================================================= */
// Class PerfTracePostProcessorFuntime
/* ============================================================================================= */
class PerfTracePostProcessorFuntime : public PerfTracePostProcessor
{
public:
    /// Construct a PerfTracePostProcessorFuntime object for the specified collection of kernel traces
    PerfTracePostProcessorFuntime(const IGtCore& gtpinCore, const PerfTraceKernel&  perfTraceKernel) :
        PerfTracePostProcessor(gtpinCore, perfTraceKernel, PerfTraceFuntime::Instance()) {}

private:

    /// Derives global tid from the trace record
    uint32_t GetGlobalTid(const uint8_t* traceRecord) const;

    /// Derives global tid from the trace record
    uint32_t GetTileId(const uint8_t* traceRecord) const;

    /// Stores data from the trace record
    void StoreTraceRecordData(const uint8_t* traceRecord, std::ofstream& fs) const;
};

/* ============================================================================================= */
// Class PerfTracePostProcessorMemoryLatency
/* ============================================================================================= */
class PerfTracePostProcessorMemoryLatency : public PerfTracePostProcessor
{
public:
    /// Construct a PerfTracePostProcessorMemoryLatency object for the specified collection of kernel traces
    PerfTracePostProcessorMemoryLatency(const IGtCore& gtpinCore, const PerfTraceKernel&  perfTraceKernel) :
        PerfTracePostProcessor(gtpinCore, perfTraceKernel, PerfTraceMemoryLatency::Instance()) {}

private:

    /// Derives global tid from the trace record
    uint32_t GetGlobalTid(const uint8_t* traceRecord) const;

    /// Derives global tid from the trace record
    uint32_t GetTileId(const uint8_t* traceRecord) const;

    /// Stores data from the trace record
    void StoreTraceRecordData(const uint8_t* traceRecord, std::ofstream& fs) const;
};

#pragma pack(pop)

#endif
