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

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

/*!
 * @file Kernel weight measuring tool
 */

#ifndef KERNEL_WEIGHT_H_
#define KERNEL_WEIGHT_H_

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

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


using namespace gtpin;

/* ============================================================================================= */
// Struct KernelWeightCounters, KernelWeightRecord
/* ============================================================================================= */
/*!
 * Counters collected by the KernelWeight tool
 * @arg T     Intergral type of the counter
 */
template <typename T> struct KweightCounters
{
    T weight;    ///< Weight of the kernel
    T freq;      ///< Number of kernel/thread executions

    template <typename U> void Accumulate(const U& other) { weight += other.weight; freq += other.freq; }
};
using KernelWeightRecord     = KweightCounters<uint32_t>; ///< Data record in the KernelWeight profile array
using KernelWeightCounters   = KweightCounters<uint64_t>; ///< Structure that holds accumulated profile counters

/* ============================================================================================= */
// Class DispatchWeightProfile
/* ============================================================================================= */
/*!
 * Profiling data collected during a single kernel dispatch
 */
struct DispatchWeightProfile
{
    explicit DispatchWeightProfile(const IGtKernelDispatch& kernelDispatch);

    GtKernelExecDesc        kernelExecDesc; ///< Kernel execution descriptor
    KernelWeightCounters    data;           ///< Profile data collected in this kernel dispatch
};

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

    /// Add new dispatched kernel instance,  and return reference to its (empty) profile
    DispatchWeightProfile& AddKernelDispatch(const IGtKernelDispatch& kernelDispatch);

    std::string           Name()         const { return _name; }          ///< @return Kernel's name
    const GtGpuPlatform   Platform()     const { return _platform; }      ///< @return Kernel's platform
    const GtProfileArray& ProfileArray() const { return _profileArray; }  ///< @return Profile buffer accessor

    /// @return Profile data collected in kernel's dispatches
    typedef std::list<DispatchWeightProfile>  DispatchProfiles;
    const DispatchProfiles& GetDispatchProfiles() const { return _dispatchProfiles; }

private:
    std::string         _name;              ///< Kernel's name
    GtGpuPlatform       _platform;          ///< Kernel's platform
    GtProfileArray      _profileArray;      ///< Profile buffer accessor
    DispatchProfiles    _dispatchProfiles;  ///< Profiles per kernel dispatch
};

/* ============================================================================================= */
// Class KernelWeight
/* ============================================================================================= */
/*!
 * Implementation of the IGtTool interface for the kernel_weight tool
 * @note This class can be inherited to provide customized implementation of the GetBblWeight() method
 */
class KernelWeight : public GtTool
{
public:
    /// Implementation of the IGtTool interface
    virtual const char* Name()       const { return "kernel_weight"; }

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

public:
    /// @return Collection of kernel profiles indexed by kernel's extended names
    typedef std::map<std::string, KernelWeightProfile>  KernelProfiles;
    const KernelProfiles& GetKernelProfiles() const { return _kernelProfiles; }

    bool DumpKernelProfiles(const std::string& filePath)    const; ///< Store kernel profiles in the specified file
    bool DumpDispatchProfiles(const std::string& filePath)  const; ///< Store kernel dispatch profiles in the specified file

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

    /*!
     * @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).
     */
    virtual uint32_t GetBblWeight(IGtKernelInstrument& ki, const IGtBbl& bbl) const { (void)ki; (void)bbl; return 1; }

    /*!
     * Aggregate 'dc' counters collected by a single kernel dispatch into accumulated 'kc' counters of the kernel.
     * By default this function uses KernelWeightCounters::Accumulate() method. Derived classes can specify a different
     * method by overriding this function.
     */
    virtual void AggregateDispatchCounters(KernelWeightCounters& kc, KernelWeightCounters dc) const { kc.Accumulate(dc); }

private:
    /// Instrument the specified basic block
    void InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, const GtProfileArray& profileArray);

private:
    KernelProfiles  _kernelProfiles;    ///< 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           _weightReg;         ///< Virtual register that holds accumulated weight of executed basic blocks
};

/* ============================================================================================= */
// Class KernelWeightProfileData
/* ============================================================================================= */
/*!
 * Helper class for streaming profile data; implemented as an associative container that holds
 * kernel profile data indexed by extended kernel names
 */
using KernelWeightProfileData = std::map<std::string, KernelWeightCounters>;

std::ostream& operator<<(std::ostream& os, KernelWeightProfileData& data);
std::istream& operator>>(std::istream& is, KernelWeightProfileData& data);


#endif