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

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

/*!
 * @file Itrace tool definitions
 */

#ifndef ITRACE_H_
#define ITRACE_H_

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

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

using namespace gtpin;

 /* ============================================================================================= */
 // Struct ItraceRecord
 /* ============================================================================================= */
 /*!
  * Structure of the trace record header.
  * The header details architectural state during execution of a BBL.
  */
struct ItraceRecord
{
    uint16_t bblId;     ///< BBL identifier
    uint16_t sr0;       ///< LSB-16 of the State register sr0.0:ud
    uint32_t tileId;    ///< TileId
};

/* ============================================================================================= */
// Class ItraceDispatch
/* ============================================================================================= */
/*!
 * Class that holds trace collected during a single kernel dispatch
 */
class ItraceDispatch
{
public:
    /// Construct a ItraceDispatch object with the empty trace
    explicit ItraceDispatch(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 ItraceKernel
/* ============================================================================================= */
/*!
 * Class that contains
 *  - Static information about basic block offsets in the kernel
 *  - Collection of instruction traces recorded by kernel dispatches
 */
class ItraceKernel
{
public:
    ItraceKernel() = default;

    /// Construct a ItraceKernel object intended to hold traces of the specified kernel
    explicit ItraceKernel(const IGtKernelInstrument& kernelInstrument, uint32_t numTiles);

    /*!
     * Read a trace recorded by the specified kernel dispatch. Create and add the corresponding ItraceDispatch
     * instance to this object
     */
    ItraceDispatch& AddItrace(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 Control-Flow Graph edges
    typedef std::pair<BblId, BblId> Edge;
    typedef std::set<Edge>          Edges;
    const Edges& GetEdges() const { return _edges; }

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

    typedef std::pair<ImgOffset, ImgOffset> BblBounds;    ///< Basic block's head and tail offsets
    typedef std::map<BblId, BblBounds>      BblBoundsMap; ///< Basic Block Id to Offsets map
    /// @return Basic Block to Offset map
    const BblBoundsMap& GetBblBounds() const { return _bblBoundsMap; }

    /// @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
    Edges             _edges;             ///< Kernel's set of control-flow graph's edges
    GtProfileTrace    _traceAccessor;     ///< Trace accessor
    Traces            _traces;            ///< Traces collected in kernel's dispatches
    BblBoundsMap      _bblBoundsMap;      ///< Basic Block to BBlBounds map
    uint32_t          _numTiles;          ///< The number of supported tiles
};

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

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

public:
    static void OnFini();                        ///< Callback function registered with atexit()

    static Itrace* Instance();                   ///< @return Single instance of this class

private:
    Itrace() {}
    Itrace(const Itrace&) = delete;
    Itrace& operator = (const Itrace&) = delete;
    ~Itrace() = default;

    /*!
     * 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] ItraceKernel      Object that holds information about basic blocks in the kernel
     * @return success/failure status
     */
    bool InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const ItraceKernel& ItraceKernel);

    /*!
     * Generate procedure that allocates space for the new trace record in the trace and stores the record
     * for the specified basic block. The procedure sets _offsetReg equal to the offset of the location within the
     * profile buffer immediately following the record header.
     * If new record can not be allocated due to the trace capacity limitations, the procedure zeroes _offsetReg.
     *
     * @param[in, out]  proc            Procedure, the generated code is appended to
     * @param[in]       coder           GEN code generator
     * @param[in]       bbl             Basic block being instrumented
     * @param[in]       ItraceKernel    Object that holds information about basic blocks in the kernel
     * @param[in]       recordSize      Size of the new record, in bytes
     */
    void StoreRecord(GtGenProcedure& proc, const IGtGenCoder& coder, const IGtBbl& bbl,
        const ItraceKernel& ItraceKernel, uint32_t recordSize);

private:
    std::map<GtKernelId, ItraceKernel> _kernels;   ///< Collection of traces per kernel

    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 _offsetReg;    ///< Virtual register that holds offset within the trace
    GtReg _tileIdReg;    ///< Virtual register that holds tile ID
};

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

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

private:
    /// KernelWeight interface overrides (@see description of KernelWeight functions)
    uint32_t GetBblWeight(IGtKernelInstrument& kernelInstrument, const IGtBbl& bbl) const;
    void AggregateDispatchCounters(KernelWeightCounters& kc, KernelWeightCounters dc) const;

private:
    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
};

/* ============================================================================================= */
// Class ItracePostProcessor
/* ============================================================================================= */
/*!
 * Function object that processes kernel traces - stores them in files within the profile directory:
 *
 *    kernel_name
 *    |
 *        |- kernel_dispatch_1
 *           |- itrace_compressed.bin
 *           |- itrace_total.cfg
 *        |- kernel_dispatch_2
 *           |- itrace_compressed.bin
 *           |- itrace_total.cfg
 * The .bin trace files can be uncompressed by the uncompress_itrace.exe utility.
 *
 * Format of .bin trace files:
 * - Static information:
 *     - Number of BBLs in the kernel
 *     - For each BBL:
 *        - BBL ID
 *        - Offset:
 *          value = 0xFFFFXXXX for BBLs which do not end with control flow instruction (XXXX == offset of the BBL within the original binary)
 *          value = 0xYYYYXXXX for BBLs which end with control flow instruction (XXXX == offset of the BBL within the original binary, YYYY == the offset of the control flow instruction)
 * - 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
 *
 * The format of *.cfg files:
 * Each line related to a separate edge of the control-flow graph:
 *     - source BBL ID, destination BBL ID, frequency of visiting the edge between source and destination BBLs.
 */
class ItracePostProcessor
{
public:
    /// Construct a ItracePostProcessor object for the specified collection of kernel traces
    ItracePostProcessor(const IGtCore& gtpinCore, const ItraceKernel& ItraceKernel);

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

private:

    /// Process the specified trace
    void ProcessTrace(const ItraceDispatch& trace);

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

    /// Store static information about BBL offsets in the kernel in the specified file stream
    void StoreBblBoundsInfo(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)); }

    /// Store kernel's Control Flow Graph
    void StoreCfg(std::ofstream& fs);

private:

    using TraceRecord = std::pair<uint32_t, uint32_t>;         ///< Trace record
    using TraceRecordList  = std::list<TraceRecord>;          ///< List of ItraceRecords per single HW thread
    using PerTileTraceRecords = std::vector<TraceRecordList>; ///< Per tile Itrace records
    typedef std::map<ItraceKernel::Edge, uint32_t> ItraceCfg;   ///< Weightened Control-flow graph
 
private:
    const ItraceKernel*              _kernel;            ///< Kernel&traces to be processed
    std::string                      _kernelDir;         ///< Directory to store kernel's trace files
    std::vector<PerTileTraceRecords> _itraceRecords;     ///< Map of tile ID to list of BBL trace with their frequencies, indexed by the thread ID
    std::vector<uint32_t>            _numProfiledThreads;///< Number of profiled (active) threads per tile
    ItraceCfg                        _cfg;               ///< Control-flow graph
    static const char*               _traceFileName;     ///< Name of the file to store trace in
};

#endif