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

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

/*!
 * @file Simple 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 HlifBcount
/* ============================================================================================= */
/*!
 * Simple tool that generates dynamic execution counters for each basic block in each kernel
 */
class HlifBcount : 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 HlifBcount* Instance();                      ///< Return single instance of this class
    static void        OnFini();                        ///< Termination handler registered with atexit()

private:
    HlifBcount() : _atomicIncFunc("AtomicInc"), _hliModule(nullptr) {} ///< Constructor

    GtHliFunction<uint64_t, uint64_t*>  _atomicIncFunc; ///< Instrumentation function: uint64_t AtomicInc(uint64_t* ptr);
    IGtHliModuleHandle                  _hliModule;     ///< Module of HLI functions

    using KernelProfile = std::vector<uint64_t>;        ///< Vector of execution counters indexed by basic block IDs
    std::map<GtKernelId, KernelProfile> _profiles;      ///< Collection of kernel profiles
    std::string                         _profileStr;    ///< Profiling results of all kernel runs, in text format
};

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

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

void HlifBcount::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    // Create profile for this kernel
    auto result = _profiles.emplace(instrumentor.Kernel().Id(), instrumentor.Cfg().NumBbls());
    GTPIN_ASSERT(result.second);
    KernelProfile& profile = result.first->second;

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

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

    // Instrument basic blocks
    for (auto bblPtr : instrumentor.Cfg().Bbls())
    {
        if (!bblPtr->IsEmpty())
        {
            GTPIN_ASSERT(bblPtr->Id() < profile.size());
            _atomicIncFunc.InsertCallAtBbl(instrumentor, *bblPtr, GtIpoint::Before(),
                                           NullReg(),                       // Unused return value
                                           profile.data() + bblPtr->Id());  // arg[0]: Counter to be incremented
        }
    }
}

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

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

void HlifBcount::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) << "BBL" << setw(15) << "Counter" << std::endl << std::endl;
    for (uint32_t bblId = 0; bblId != profile.size(); ++bblId)
    {
        auto count = profile[bblId];
        if (count != 0)
        {
            os << setw(15) << bblId << 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
    HlifBcount::Instance()->Register();

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

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