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

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

/*!
 * @file Memorytrace tool definitions
 */

#ifndef MEMORYTRACE_H_
#define MEMORYTRACE_H_

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

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

using namespace gtpin;

/*
 * The trace of memory accesses consists of records, each of which holds data pertaining to a single
 * execution of a basic block that accesses memory.
 * Each trace record comprises two data portions:
 *  - Header that details architectural state during this BBL execution. The MemTraceRecordHeader
 *    structure defines layout of the data elements in the header
 *  - Array of address payload register values in SEND instructions of this BBL. The values are
 *    stored in the order of the corresponding instructions in the basic block. This portion of the
 *    record is of a variable length - different BBLs may access different number of data elements.
 * Both portions (header and address payloads) are aligned with the GRF register size.
 */

/* ============================================================================================= */
// Struct MemTraceRecordHeader
/* ============================================================================================= */
/*!
 * Structure of the trace record header.
 * The header details architectural state during execution of a BBL that accesses memory. In the
 * trace record, the header is followed by the array of address payload registers.
 */
#pragma pack(push, 1)
struct MemTraceRecordHeader 
{
    uint16_t bblId;     ///< BBL identifier
    uint16_t sr0;       ///< LSB-16 of the State register sr0.0:ud
    uint32_t tileId;    ///< Tile ID
    uint32_t ce;        ///< Channel Enable register ce:ud
    uint32_t dm;        ///< Dispatch Mask register dm:ud
    uint32_t res;       ///< Reserved
    uint32_t tm00;      ///< tm0.0:ud
    uint32_t tm01;      ///< tm0.1:ud

    /// @return Size of this structure aligned with the GRF register size in the specified GEN model 
    static uint32_t AlignedSize(const IGtGenModel& genModel)
    {
        return (uint32_t)AlignUp(sizeof(MemTraceRecordHeader), genModel.GrfRegSize());
    }
};
#pragma pack(pop)

/* ============================================================================================= */
// Struct MemIns
/* ============================================================================================= */
/// SEND instruction descriptor
struct MemIns
{
    InsId       id;         ///< Instruction ID
    uint32_t    offset;     ///< Offset ot the instruction within the kernel
    DcSendMsg   msg;        ///< Decoded SEND message instruction
};
using MemInsList = std::list<MemIns>;

/* ============================================================================================= */
// Class BblMemAccessInfo
/* ============================================================================================= */
/*!
 * Class that retrieves information about memory accesses in the basic block, and provides an
 * interface for querying this information
 */
class BblMemAccessInfo
{
public:
    BblMemAccessInfo() : _recordSize(0) {}

    /// Populate this object with the information about memory accesses in the specified basic block
    BblMemAccessInfo(const IGtKernelInstrument& kernelInstrument, const IGtBbl& bbl) : BblMemAccessInfo() { Build(kernelInstrument, bbl); }
    BblMemAccessInfo& Build(const IGtKernelInstrument& kernelInstrument, const IGtBbl& bbl);

    /// @return List of memory instructions in the basic block
    const MemInsList& MemInstructions() const { return _memInstructions; }

    /// @return Size, in bytes, of the trace record in the basic block
    uint32_t RecordSize() const { return _recordSize; }

    /// @return true if basic block does not contain any memory instruction of interest
    bool IsEmpty() const { return (_recordSize == 0); }

private:
    MemInsList  _memInstructions;   ///< List of memory instructions in the basic block
    uint32_t    _recordSize;        ///< Size, in bytes, of the trace record in the basic block
};

/* ============================================================================================= */
// Class KernelMemAccessInfo
/* ============================================================================================= */
/*!
 * Class that retrieves information about memory accesses in the kernel, and provides an
 * interface for querying this information
 */
class KernelMemAccessInfo
{
public:
    KernelMemAccessInfo() : _maxRecordSize(0) {}

    /// Populate this object with the information about memory accesses in the specified kernel
    explicit KernelMemAccessInfo(const IGtKernelInstrument& kernelInstrument) { Build(kernelInstrument); }
    KernelMemAccessInfo& Build(const IGtKernelInstrument& kernelInstrument);

    /// @return Information about memory accesses in the kernel, indexed by BBL identifiers
    using MemAccessMap = std::map<BblId, BblMemAccessInfo>;
    const MemAccessMap& GetMemAccessMap() const { return _memAccessMap; }

    /// @return Information about memory accesses in the specified BBL, or NULL if requested information is not found
    const BblMemAccessInfo* GetBblInfo(BblId bblId) const;

    /// @return Number of basic blocks that access memory
    uint32_t NumMemBbls() const { return (uint32_t)_memAccessMap.size(); }

    /// @return Maximum size, in bytes, of the trace record in the kernel
    uint32_t MaxRecordSize() const { return _maxRecordSize; }

private:
    MemAccessMap _memAccessMap;  ///< Information about memory accesses, indexed by BBL identifiers
    uint32_t     _maxRecordSize; ///< Max. size, in bytes, of the trace record in the kernel
};

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

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

    /*!
     * Read a trace recorded by the specified kernel dispatch. Create and add the corresponding MemTraceDispatch
     * instance to this object
     */
    MemTraceDispatch& AddMemTrace(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 Static information about kernel's memory accesses
    const KernelMemAccessInfo& GetMemAccessInfo() const { return _memAccessInfo; }

    /// @return Traces collected in kernel's dispatches
    typedef std::list<MemTraceDispatch>  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
    KernelMemAccessInfo _memAccessInfo;     ///< Static information about kernel's memory accesses
    GtProfileTrace      _traceAccessor;     ///< Trace accessor
    Traces              _traces;            ///< Traces collected in kernel's dispatches
    uint32_t            _numTiles;          ///< The number of supported tiles
};

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

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

public:
    static void      OnFini();      ///< Callback function registered with atexit()
    static MemTrace* Instance();    ///< @return Single instance of this class

private:
    MemTrace() = default;
    MemTrace(const MemTrace&) = delete;
    MemTrace& operator = (const MemTrace&) = delete;
    ~MemTrace() = 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] memTraceKernel    Object that holds information about memory accesses in the kernel
     * @return success/failure status
     */
    bool InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const MemTraceKernel& memTraceKernel);

    /*!
     * Generate procedure that allocates space for the new trace record in the trace and stores the record header
     * 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]       memTraceKernel  Object that holds information about memory accesses in the kernel
     * @param[in]       recordSize      Size of the new record, in bytes
     * @param[in]       firstSendInBbl  Indicates whether this is the first Send instruction in BBL
     */
    void StoreRecordHeader(GtGenProcedure& proc, const IGtGenCoder& coder, const IGtBbl& bbl,
                           const MemTraceKernel& memTraceKernel, uint32_t recordSize, bool firstSendInBbl = true);

    /*!
     * Generate procedure that stores the specified range of GRF registers in the trace.
     * On entry, the procedure assumes that _offsetReg holds offset of the location within the profile buffer
     * at which address payload is about to be stored.
     * On exit, the procedure advances _offsetReg to the location just after the stored register range.
     * If input value of _offsetReg is zero. no registers are stored, and _offsetReg retains value zero.
     * 
     * @param[in, out]  proc            Procedure, the generated code is appended to
     * @param[in]       coder           GEN code generator
     * @param[in]       firstRegNum     First GRF register to be stored
     * @param[in]       numRegs         Number of registers to be stored
     */
    void StoreRegRange(GtGenProcedure& proc, const IGtGenCoder& coder, uint32_t firstRegNum, uint32_t numRegs);

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

private:
    MemoryTracePreProcessor();
    MemoryTracePreProcessor(const MemoryTracePreProcessor&) = delete;
    MemoryTracePreProcessor& operator = (const MemoryTracePreProcessor&) = 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 MemoryTracePostProcessor
/* ============================================================================================= */
/*!
 * 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 MemoryTracePostProcessor
{
public:
    /// Construct a MemoryTracePostProcessor object for the specified collection of kernel traces
    MemoryTracePostProcessor(const IGtCore& gtpinCore, const MemTraceKernel&  memTraceKernel);

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

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

    /// Store static information about memory accesses in the kernel
    void StoreMemAccessInfo(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)); }

private:
    struct TraceRecord                              ///< Reference to the trace record
    {
        const MemTraceRecordHeader* header;         ///< Pointer to the header of 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
    ///The SEND instruction description stored in the trace
    struct MemInsInfo
    {
        explicit MemInsInfo(const MemIns& memIns);    ///< Construct packed variant of the specified MemIns structure

        uint32_t    offset;             ///< Offset of the instruction within the kernel
        uint32_t    isWrite;            ///< Is write (true) or read (false) operation
        uint32_t    isBlock2D;          ///< Is block 2D load or store access
        uint32_t    isScatter;          ///< Is scatter
        uint32_t    isBTS;              ///< Is Binding Table State address model
        uint32_t    isSLM;              ///< Is Shared Local Memory access
        uint32_t    isScratch;          ///< Is access to scratch block
        uint32_t    isAtomic;           ///< Is atomic operation
        uint32_t    isFence;            ///< Is fence
        uint32_t    addressWidth;       ///< Address width (32bit or 64bit)
        uint32_t    simdWidth;          ///< SIMD width
        uint32_t    bti;                ///< BTI value
        uint32_t    addrPayloadLength;  ///< Address payload length
        uint32_t    dataPort;           ///< Data port ID: 0=DP0, 1=DP1, 2=UGM, 3=UGML, 4=TGM, 5=SLM, 6-all others
        uint32_t    isEOT;              ///< Is SEND.EOT
        uint32_t    isMedia;            ///< Is media block access
        uint32_t    elementSize;        ///< Data element size in bytes
        uint32_t    numElements;        ///< Number of elements
        uint32_t    execSize;           ///< Instruction execution size
        uint32_t    channelOffset;      ///< Channel mask offset
    };

private:
    const MemTraceKernel*            _kernel;            ///< Kernel&traces to be processed
    const KernelMemAccessInfo*       _memAccessInfo;     ///< Static information about memory accesses in the kernel
    std::string                      _kernelDir;         ///< Directory to store kernel's trace files
    std::vector<PerTileTraceRecords> _threadTraceRecords;///< Lists of trace records, indexed by the thread ID

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

#endif