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

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

/*!
 * @file Simple edge (BBL->BBL) counting tool that demonstrates HLIF interface
 */

#include <string>
#include <vector>
#include <map>
#include <utility>

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

using namespace gtpin;
using namespace std;

Knob<bool> KNOB_NO_COUT("no_cout", false, "Do not send profiling results to standard output device");

/* ============================================================================================= */
// Class HlifEcount
/* ============================================================================================= */
/*!
 * Simple tool that generates per-kernel dynamic counters of control transfers (edges) between
 * basic blocks
 */
class HlifEcount : public GtTool
{
public:
    // Implementation of the IGtTool interface
    const char* Name()          const                   override { return "HLIF bcount"; }
    void        OnKernelBuild(IGtKernelInstrument&)     override; 
    void        OnKernelRun(IGtKernelDispatch&)         override;
    void        OnKernelComplete(IGtKernelDispatch&)    override;

    void               LoadHliLibrary();                ///< Compile and load library of HLI functions
    static HlifEcount* Instance();                      ///< Return single instance of this class
    static void        OnFini();                        ///< Termination handler registered with atexit()

private:
    /// Static properties of the kernel, and its profile data updated on each kernel run
    struct KernelProfile
    {
        uint32_t                numBbls;      ///< Number of BBLs in the kernel's CFG
        BblId                   entryBbl;     ///< ID of the entry BBL
        std::vector<uint64_t>   edgeCounters; ///< Matrix of dynamic edge counters indexed by pair of BBL IDs: [toId][fromId]
        /// Constructor
        KernelProfile(uint32_t numKernelBbls, BblId entryBblId) :
            numBbls(numKernelBbls), entryBbl(entryBblId), edgeCounters(numKernelBbls * numKernelBbls) {}
    };
private:
    /// Constructor
    HlifEcount() : _ecountFunc("AtomicElementInc"), _hliModule(nullptr) {}

    /*!
     * Instrumentation function that atomically increments value of elements[idx]
     * uint32_t AtomicElementInc(uint64_t* elements, uint32_t idx, uint32_t retVal);
     * @return  The 'retVal' value
     */
    GtHliFunction<uint32_t, uint64_t*, uint32_t, uint32_t>  _ecountFunc;

    IGtHliModuleHandle                  _hliModule;     ///< Module of HLI functions
    std::map<GtKernelId, KernelProfile> _profiles;      ///< Collection of kernel profiles
    std::string                         _profileStr;    ///< Profiling results of all kernel runs, in text format
};

/* ============================================================================================= */
// HlifEcount implementation
/* ============================================================================================= */
HlifEcount* HlifEcount::Instance()
{
    static HlifEcount instance;
    return &instance;
}

void HlifEcount::LoadHliLibrary()
{
    std::string installDir = GetKnobValue<std::string>("installDir");
    _hliModule = GTPin_GetCore()->HliLibrary().CompileModuleFromFile(JoinPath(installDir, "Examples", "hli_sample.cl").c_str());
    GTPIN_ASSERT(_hliModule != nullptr);
}

void HlifEcount::OnFini()
{
    // Dump profiling results of all kernel runs
    ofstream fs(JoinPath(GTPin_GetCore()->ProfileDir(), "hlif_ecount.txt"));
    GTPIN_ASSERT(fs.is_open());
    const std::string& profileStr = Instance()->_profileStr;
    fs << profileStr;
    if (!KNOB_NO_COUT)
    {
        std::cout << profileStr;
    }
}

void HlifEcount::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    // Create profile for this kernel
    uint32_t numBbls   = instrumentor.Cfg().NumBbls();
    BblSpan  entryBbls = instrumentor.Cfg().EntryBbls(); GTPIN_ASSERT(entryBbls.size() == 1);
    auto     result    = _profiles.emplace(std::piecewise_construct,
                                            std::forward_as_tuple(instrumentor.Kernel().Id()),
                                            std::forward_as_tuple(numBbls, entryBbls[0]->Id()));
    GTPIN_ASSERT(result.second);
    KernelProfile& profile = result.first->second;

    // Share profile data between host and device memory
    instrumentor.MemoryMapper().Map(profile.edgeCounters);

    // Link the kernel with the library of HLI functions
    instrumentor.LinkHliModule(_hliModule);

    // Initialize virtual register that will convey BBL ID from predecessor to successor blocks
    IGtVregFactory& vregs = instrumentor.Coder().VregFactory();
    GtReg bblIdReg = vregs.MakeCounter();

    // Instrument basic blocks
    for (auto bblPtr : instrumentor.Cfg().Bbls())
    {
        if (!bblPtr->IsEmpty())
        {
            uint32_t currentBblId = bblPtr->Id();
            GTPIN_ASSERT(currentBblId < numBbls);
            auto rowPtr = profile.edgeCounters.data() + (bblPtr->Id() * numBbls);
            _ecountFunc.InsertCallAtBbl(instrumentor, *bblPtr, GtIpoint::Before(),
                                        bblIdReg,        // Return value: ID of the current block
                                        rowPtr,          // arg[0]: Row of counters for incoming edges of the current block
                                        bblIdReg,        // arg[1]: ID of the predecessor block
                                        currentBblId);   // arg[2]: ID of the current block
        }
    }
}

void HlifEcount::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    // Zero-initialize profile counters before this kernel dispatch
    KernelProfile& profile = _profiles.at(dispatcher.Kernel().Id());
    std::fill(profile.edgeCounters.begin(), profile.edgeCounters.end(), 0);

    // Tell GTPin to run instrumented code
    dispatcher.SetProfilingMode(true);
}

void HlifEcount::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    // Translate profile data of this kernel dispatch into text
    KernelProfile& profile = _profiles.at(dispatcher.Kernel().Id());

    std::ostringstream os;
    os << std::string(100, '-') << std::endl;
    os << setw(4) << dispatcher.DispatchId() << ":    " << ExtendedKernelName(dispatcher.Kernel()) << std::endl;
    os << std::string(100, '-') << std::endl;
    os << setw(15) << "From BBL" << setw(15) << "To BBL" << setw(15) << "Counter" << std::endl << std::endl;
    for (uint32_t fromId = 0; fromId != profile.numBbls; ++fromId)
    {
        for (uint32_t toId = 0; toId != profile.numBbls; ++toId)
        {
            auto count = profile.edgeCounters[(toId * profile.numBbls) + fromId];
            if (count != 0)
            {
                std::string fromStr = (toId == profile.entryBbl) ? "Entry" : ToString(fromId);
                os << setw(15) << fromStr << setw(15) << toId << setw(15) << count << std::endl;
            }
        }
    }
    os << std::endl;

    // Add profile data of this kernel dispatch to _profileStr
    _profileStr += os.str();
}

/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    // Parse command line and configure GTPin
    ConfigureGTPin(argc, argv);

    // Register the tool (callbacks) with the GTPin core
    HlifEcount::Instance()->Register();

    // Compile and load library of HLI functions
    HlifEcount::Instance()->LoadHliLibrary();

    // Register the termination function
    atexit(HlifEcount::OnFini);
}
