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

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

/*!
 * @file edge_count tool definitions
 */

#ifndef EDGE_COUNT_H_
#define EDGE_COUNT_H_

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

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

using namespace gtpin;

/* ============================================================================================= */
// Struct EdgeCountRecord
/* ============================================================================================= */
/*!
 * Layout of records collected in profile buffer by the edge_count tool
 */
struct EdgeCountRecord
{
    uint32_t srcBbl;    ///< srcBbl id
    uint32_t dstBbl;    ///< dstBbl id
    uint32_t counter;   ///< Total number of edge passages
};

/*!
 * Identifiers of the instrumentation edge
 *  - ID of the src basic block
 *  - ID of the dstrc basic block
 */
using InstrumentSiteId = uint32_t;

/// Array of instrumentation site IDs.
using InstrumentSites = std::vector<uint32_t>;

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

    GtKernelExecDesc  kernelExecDesc; ///< Kernel execution descriptor
    InstrumentSiteId  siteId;         ///< Identifier of the instrumentation site: SrcBbL ID and DstBbl ID
    uint32_t          counter;        ///< Total number of edge passages
};

/* ============================================================================================= */
// Class EdgeCountKernelProfile
/* ============================================================================================= */
/*!
 * Aggregated profile of all instrumented kernel dispatches
 */
class EdgeCountKernelProfile
{
public:
    EdgeCountKernelProfile(const IGtKernel& kernel, const IGtCfg& cfg, const GtProfileArray& profileArray,
                         InstrumentSites&& instrumentSites);

    /// Add the specified profile of a kernel dispatch
    EdgeCountDispatchProfile& AddDispatchProfile(const EdgeCountDispatchProfile& dispathProfile);

    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<EdgeCountDispatchProfile>  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
    InstrumentSites                     _instrumentSites;   ///< Array of instrumentation site IDs
    DispatchProfiles                    _dispatchProfiles;  ///< Profiles per kernel dispatch
    std::map<std::string, uint64_t>     _execTotalCycles;   ///< Total number of cycles per execution descriptor
    uint64_t                            _numOfBbls;         ///< Number of original basic blocks
};

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

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

public:
    static void OnFini();                        ///< Callback function registered with atexit()
    std::string ToString() const;                ///< @return Text representation of the profile data

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

private:
    EdgeCount() = default;
    EdgeCount(const EdgeCount&) = delete;
    EdgeCount& operator = (const EdgeCount&) = delete;
    ~EdgeCount() = default;

    /// Generate code at the edges of the kernel
    void GenerateEdgeCode(uint32_t recordNum, GtGenProcedure& proc, const IGtGenCoder& coder, GtProfileArray profileArray);

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

private:
    std::map<GtKernelId, EdgeCountKernelProfile> _kernels;  ///< Collection of kernel profiles

    GtReg _addrReg;     ///< Virtual register that holds address within profile buffer
};

#endif