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

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

/*!
 * @file A tool that detects uninitialized memory accesses in kernels
 *
 *       The tool supports the following Read-Only and Read-Modify-Write accesses to global memory and SLM
 * 
 *       SLM:     any RO and RMW access
 *       Global memory: 
 *         - Any RO and RMW access done with A64 addressing model
 *         - Only RO accesses done with any other addressing model
 */

#include <cstring>
#include <map>
#include <algorithm>

#include "memory_check.h"
#include "gen_send_decoder.h"

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

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

/* ============================================================================================= */
// Class MemAccess
/* ============================================================================================= */
/// Information about memory access instruction
struct MemAccess
{
    /// Constructor. If memory access is unsupported, the reason can be queried by Error()
    explicit MemAccess(const IGtIns &ins);

    bool         IsValid()      const { return _isValid; }            ///< @return true for a supported memory access
    InsId        Id()           const { return _insId; }              ///< @return Instruction ID
    GtAccessType AccessType()   const { return _accessType; }         ///< @return Access type: read, write or read-write
    GtRegNum     FirstDstReg()  const { return _firstDstReg; }        ///< @return First register in the destination payload
    GtRegNum     FirstAddrReg() const { return _firstAddrReg; }       ///< @return First register in the address payload
    uint32_t     NumAccesses()  const { return _numAccesses; }        ///< @return Number of elements in the address payload
    uint32_t     DataSize()     const { return _dataSize; }           ///< @return Data size, in bytes, per each address
    bool         IsSlm()        const { return _isSlm; }              ///< @return true for SLM access
    const GtMemoryAddrModel& AddrModel() const { return _addrModel; } ///< @return Address model of the memory access
    const std::string&       Error()     const { return _errMsg; }    ///< @return Error message on unsupported memory access

    /*!
     * @return Arguments of the HLI function that detects OOB violations in this memory access
     * @note This object owns the MemoryCheckArgs structure, but its content is controlled externally
     */
    const MemoryCheckArgs& GetMemoryCheckArgs() const { return _mcArgs; }
    MemoryCheckArgs&       GetMemoryCheckArgs()       { return _mcArgs; }

private:
    bool                _isValid = false;       ///< True, if this structure represents supported memory access
    InsId               _insId;                 ///< ID of the memory access instruction
    GtAccessType        _accessType;            ///< Access type: read-only, write-only or read-write
    GtMemoryAddrModel   _addrModel;             ///< Address model of the memory access
    GtRegNum            _firstDstReg;           ///< First register in the address payload of the instruction
    GtRegNum            _firstAddrReg;          ///< First register in the address payload of the instruction
    uint32_t            _numAccesses = 0;       ///< Number of elements in the address payload of the instruction
    uint32_t            _dataSize = 0;          ///< Size, in bytes, of the memory range referenced by a single address
    MemoryCheckArgs     _mcArgs;                ///< Common arguments of memory check functions

    std::string         _errMsg;                ///< Error message on unsupported memory access
    bool                _isSlm = false;         ///< True if the memory access is to SLM
};

/* ============================================================================================= */
// Struct ProfileResults
/* ============================================================================================= */
/*!
 * Profile results per kernel dispatch / per insruction
 */
struct ProfileResults
{
    ProfileResults(const IGtKernelDispatch& dispatcher, const MemAccess& memAccess);

    uint64_t          dispatchId;     ///< Unique ID of the kernel dispatch assigned by GTPin
    GtKernelExecDesc  kernelExecDesc; ///< Kernel execution descriptor
    InsId             insId;          ///< ID of the memory access instruction
    uint32_t          umrCount;       ///< Counter of UMR violations in memory accesses performed by the instruction
};

/* ============================================================================================= */
// Class KernelProfile
/* ============================================================================================= */
/// Static properties of the kernel, and its profile data updated on each kernel run
class KernelProfile
{
public:
    using MemAccessMap = std::map<InsId, MemAccess>;              ///< Information about memory accesses by kernel instructions

public:
    KernelProfile(const IGtKernel& kernel, const IGtCfg& cfg);    ///< Constructor

    inline GtKernelId             Id()                 const;        ///< @return Unique identifier of the kernel
    inline GtGpuPlatform          Platform()           const;        ///< @return Kernel's platform
    inline const std::string&     Name()               const;        ///< @return Name of the kernel
    inline const std::string&     UniqueName()         const;        ///< @return Unique name of the kernel
    inline std::string            MemoryCheckResults() const;        ///< @return Profiling results of kernel runs, in text format
    inline void                   DumpAsm()            const;        ///< Store kernel's assembly text in the file
    inline const MemAccessMap&    GetMemAccessMap()    const;        ///< @return Information about memory accesses in the kernel
    inline MemAccessMap&          GetMemAccessMap();                 ///< @return Information about memory accesses in the kernel
    inline const SlmInitArg&      GetSlmInitArg()      const;        ///< @return argument for SLM initialization function
    inline SlmInitArg&            GetSlmInitArg();                   ///< @return argument for SLM initialization function

    void RecordMemoryCheckResults(IGtKernelDispatch& dispatcher); ///< Update profile data with the latest memory check results

    void RecordUnsupportedInstruction(const IGtIns& ins, const std::string& errMsg); ///< Record unsupported instruction
    void RecordUnhandledAccess(InsId insId, const std::string& errMsg);              ///< Record unhandled memory access

private:
    GtKernelId      _id;                        ///< Unique identifier of the kernel
    GtGpuPlatform   _platform;                  ///< Kernel's platform
    std::string     _name;                      ///< Name of the kernel
    std::string     _uniqueName;                ///< Unique name of the kernel
    std::string     _asmText;                   ///< Assembly text of the kernel
    std::string     _unhandledAccesses;         ///< List of unhadled memory accesses, in text format

    SlmInitArg                 _slmInitMemArg;  ///< Argument for SLM initialization function
    std::map<InsId, MemAccess> _memAccessMap;   ///< Map: Instruction ID to memory access information
    std::list<ProfileResults>  _profileResults; ///< Profile results per kernel dispatch / per insruction 
};

/* ============================================================================================= */
// Class MemoryCheck
/* ============================================================================================= */
/*!
 * A tool that detects uininitialized memory read (UMR) accesses in kernels
 */
class MemoryCheck : public GtTool
{
public:
    // Implementation of the IGtTool interface
    const char* Name()          const                   override { return "Memory check"; }
    void        OnKernelBuild(IGtKernelInstrument&)     override;
    void        OnKernelRun(IGtKernelDispatch&)         override;
    void        OnKernelComplete(IGtKernelDispatch&)    override;

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

private:
    
    MemoryCheck();                                          ///< Default constructor
    MemoryCheck(const MemoryCheck&) = delete;               ///< Disabled copy constructor
    MemoryCheck& operator = (const MemoryCheck&) = delete;  ///< Disabled assignment operator
    ~MemoryCheck();                                         ///< Destructor
    void Fini();                                            ///< Post process and dump profiling data

    /*!
     * Insert a call to HLI function that detects uininitialized memory read (UMR) memory accesses in the specified instruction
     * @param[in]   ins             The memory access instruction
     * @param[in]   memAccess       Information about memory access
     * @param[in]   instrumentor    Instrumentation interface
     * @return true - success, false - the instruction or memory operation is not supported
     */
    bool InsertMemoryCheck(const IGtIns &ins, const MemAccess& memAccess, IGtKernelInstrument& instrumentor);
    bool InsertSlmInit(const IGtIns& ins, const SlmInitArg& slmInitArg, IGtKernelInstrument& instrumentor);

private:
    // Memory check functions
    CheckReadAccessFunc                 _checkReadAccessFunc;
    CheckSlmReadAccessFunc              _checkSlmReadAccessFunc;
    CheckUgmA64ReadAccessFunc           _checkUgmA64ReadAccessFunc;
    InitSlmFunc                         _initSlmFunc;

    IGtHliModuleHandle                  _hliModule = nullptr;        ///< Module of HLI functions
    std::map<GtKernelId, KernelProfile> _kernels;                    ///< Collection of kernel profiles
};

/* ============================================================================================= */
// MemAccess implementation
/* ============================================================================================= */
MemAccess::MemAccess(const IGtIns &ins) : _insId(ins.Id())
{
    // Get and check data port (SFID)
    GtSfid sfid = ins.Sfid();
    if ((sfid != GED_SFID_UGM) && (sfid != GED_SFID_SLM) && (sfid != GED_SFID_DP_DC0) && (sfid != GED_SFID_DP_DC1))
    {
        _errMsg = "Unsupported data port " + std::string(sfid.ToString());
        return;
    }

    bool isHdc = (sfid == GED_SFID_DP_DC0) || (sfid == GED_SFID_DP_DC1);

    _isSlm = (sfid == GED_SFID_SLM);

    // Retrieve message descriptor
    if (!ins.MsgDescRegFile().IsImm())
    {
        _errMsg = "SEND message descriptor is not immediate";
        return;
    }

    // Check opcode of the memory operation
    GED_DP_OPCODE opcode = ins.DPOpCode();
    bool          isAtomic   = ins.IsAtomic();
    bool          isLoad = (isHdc && !isAtomic && ins.HasDstOperand() && ins.DstRegFile().IsGrf()) || (opcode == GED_DP_OPCODE_LOAD);
    bool          isStore = (isHdc && !isAtomic && !ins.HasDstOperand() && ins.DstOperand().Reg().IsNullReg()) || (opcode == GED_DP_OPCODE_STORE);
    _accessType = (isLoad ? GT_ACCESS_READ : (isStore ? GT_ACCESS_WRITE : GT_ACCESS_READ_WRITE));

    if (!isLoad && !isAtomic)
    {
        _errMsg = "Unsupported SEND operation (not load/atomic)";
        return;
    }

    // Initialize address model
    _addrModel = ins.MemAddrModel();
    if (!_addrModel.IsValid())
    {
        _errMsg = "Unsupported/unknown address model"; // @fixme Support BSS model
        return;
    }

    // Finally, initialize the rest of data members...
    GTPIN_ASSERT(ins.SrcRegFile(0).IsGrf());
    _firstDstReg  = ins.HasDstOperand() ? ins.DstOperand().Reg().RegNum() : GtRegNum();
    _firstAddrReg = ins.SrcRegOperand(0).Reg().RegNum(); 
    _numAccesses  = ins.NumAccesses(); GTPIN_ASSERT(_numAccesses != 0);

    DcSendMsg msg = DcSendMsg::Decode(ins.GetGedIns());
    _dataSize = msg.ElementSize() * msg.NumElements();

    if (_dataSize == 0)
    {
        _errMsg = "Unsupported SEND operation";
        return;
    }

    if (_accessType == GT_ACCESS_READ_WRITE)
    {
        if ((_dataSize != 4 && _dataSize != 8) || (!_addrModel.IsSlm() && !_addrModel.IsA64()))
        {
            _errMsg = "Unsupported Read-Modify-Write instruction: only Dword and Qword data sizes are supported for SLM and A64 addressing modes";
            return;
        }
    }

    if (_dataSize & 0x3)
    {
        _errMsg = "Unsupported data size: " + DecStr(_dataSize) + "(should be multiplication of dword)";
        return;
    }

    _isValid = true;
}

/* ============================================================================================= */
// ProfileResults implementation
/* ============================================================================================= */
ProfileResults::ProfileResults(const IGtKernelDispatch& dispatcher, const MemAccess& memAccess) :
    dispatchId(dispatcher.DispatchId()), insId(memAccess.Id()), umrCount(memAccess.GetMemoryCheckArgs().out.umrCount)
{
    dispatcher.GetExecDescriptor(kernelExecDesc);
}

/* ============================================================================================= */
// KernelProfile implementation
/* ============================================================================================= */
KernelProfile::KernelProfile(const IGtKernel& kernel, const IGtCfg& cfg) :
    _id(kernel.Id()), _platform(kernel.GpuPlatform()), _name(GlueString(kernel.Name())), _uniqueName(kernel.UniqueName()),
    _asmText(CfgAsmText(cfg))
{
    // Populate this object with the information about memory accesses
    for (auto bblPtr : cfg.Bbls())
    {
        for (auto insPtr : bblPtr->Instructions())
        {
            const IGtIns& ins = *insPtr;
            if (ins.IsMemAccess() && !ins.IsEot())
            {
                MemAccess memAccess(ins);

                if (!memAccess.IsValid() && memAccess.AccessType().IsWriteOnly())
                {
                    continue;
                }

                if (memAccess.IsValid())
                {
                    _memAccessMap.emplace(ins.Id(), memAccess);
                }
                else
                {
                    RecordUnsupportedInstruction(ins, memAccess.Error());
                }
            }
        }
    }
}

GtKernelId         KernelProfile::Id()          const { return _id; }
GtGpuPlatform      KernelProfile::Platform()    const { return _platform; }
const std::string& KernelProfile::Name()        const { return _name; }
const std::string& KernelProfile::UniqueName()  const { return _uniqueName; }
void               KernelProfile::DumpAsm()     const { DumpKernelAsmText(_name, _uniqueName, _asmText); }
const KernelProfile::MemAccessMap& KernelProfile::GetMemAccessMap() const { return _memAccessMap; }
KernelProfile::MemAccessMap&       KernelProfile::GetMemAccessMap()       { return _memAccessMap; }
const SlmInitArg& KernelProfile::GetSlmInitArg() const { return _slmInitMemArg; }
SlmInitArg&       KernelProfile::GetSlmInitArg()       { return _slmInitMemArg; }

std::string KernelProfile::MemoryCheckResults() const
{
    std::ostringstream os;

    os << std::string(120, '-') << std::endl;
    os << std::setw(4) << _id << ":    " << _name << "___" << _uniqueName << std::endl;
    os << std::string(120, '-') << std::endl;

    if (!_unhandledAccesses.empty())
    {
        os << " Unhandled memory accesses:" << std::endl;
        os << " --------------------------" << std::endl;
        os << _unhandledAccesses << std::endl;
        os << std::string(120, '-') << std::endl;
    }

    uint64_t umrTotal = 0;
    for (const auto& res : _profileResults)
    {
        if (res.umrCount != 0)
        {
            if (umrTotal == 0)
            {
                os << std::setw(20) << "Dispatch ID" << std::setw(20) << "Instruction ID" << std::setw(20) << "#UMR violations";
                os << " " << std::setw(45) << "Execution descriptor" << std::endl;
                os << std::string(120, '-') << std::endl;
            }
            os << std::setw(20) << res.dispatchId << std::setw(20) << res.insId << std::setw(20) << res.umrCount;
            os << " " << std::setw(45) << res.kernelExecDesc.ToString(_platform, ExecDescAlignedFormat()) << std::endl;

            umrTotal += res.umrCount;
        }
    }
    if (umrTotal == 0)
    {
        os << "No UMR accesses detected" << std::endl;
    }
    return os.str();
}

void KernelProfile::RecordMemoryCheckResults(IGtKernelDispatch& dispatcher)
{
    for (auto& entry: _memAccessMap)
    {
        MemAccess& memAccess = entry.second;
        _profileResults.emplace_back(dispatcher, memAccess);
    }
}

void KernelProfile::RecordUnsupportedInstruction(const IGtIns& ins, const std::string& errMsg)
{
    if (!errMsg.empty())
    {
        std::ostringstream os;
        os << errMsg << ": [" << std::setw(3) << ins.Id() << "] " << ins.ToString() << std::endl;
        _unhandledAccesses.append(os.str());
    }
}

void KernelProfile::RecordUnhandledAccess(InsId insId, const std::string& errMsg)
{
    if (!errMsg.empty())
    {
        std::ostringstream os;
        os << errMsg << ": Instruction ID: [" << std::setw(3) << insId << "] " << std::endl;
        _unhandledAccesses.append(os.str());
    }
}

/* ============================================================================================= */
// MemoryCheck implementation
/* ============================================================================================= */
MemoryCheck::MemoryCheck() :
    _checkReadAccessFunc("CheckReadAccess"),
    _checkSlmReadAccessFunc("CheckSlmReadAccess"),
    _checkUgmA64ReadAccessFunc("CheckUgmA64ReadAccess"),
    _initSlmFunc("InitSlm") {}
MemoryCheck::~MemoryCheck() {}

MemoryCheck* MemoryCheck::Instance()
{
    static MemoryCheck instance;
    return &instance;
}

void MemoryCheck::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    const IGtKernel& kernel     = instrumentor.Kernel();
    const IGtCfg&    cfg        = instrumentor.Cfg();
    IGtMemoryMapper& memMapper  = instrumentor.MemoryMapper();

    // Create profile for this kernel
    auto result = _kernels.emplace(std::piecewise_construct,
                                   std::forward_as_tuple(instrumentor.Kernel().Id()),
                                   std::forward_as_tuple(kernel, cfg));
    KernelProfile& kernelProfile = result.first->second;

    bool hasSlm = false;

    // Instrument memory accesses and share per-access arguments with HLI functions
    for (const auto& entry : kernelProfile.GetMemAccessMap())
    {
        const auto& memAccess = entry.second;
        auto        insId     = entry.first;

        if (int32_t(insId) < knobMinInstrumentIns || knobMaxInstrumentIns < int32_t(insId))
        {
            continue;
        }
        const IGtIns& ins = cfg.GetInstruction(insId);

        hasSlm |= memAccess.IsSlm();

        auto accessType = memAccess.AccessType();
        if (accessType == GT_ACCESS_READ || memAccess.IsSlm() || memAccess.AddrModel().IsA64())
        {
            InsertMemoryCheck(ins, memAccess, instrumentor);
        }
        else
        {
            GTPIN_ERROR();
        }

        // Share per-access HLI arguments.
        // They will be initialized at the start of the kernel, and copied back to the host memory at completion of the kernel
        memMapper.Map(memAccess.GetMemoryCheckArgs(), GT_MMAP_SHARE);
    }

    if (hasSlm)
    {
        for (auto& bblPtr : cfg.EntryBbls())
        {
            const IGtIns& ins = bblPtr->FirstIns();
            InsertSlmInit(ins, kernelProfile.GetSlmInitArg(), instrumentor);
            memMapper.Map(kernelProfile.GetSlmInitArg(), GT_MMAP_SHARE);
        }
    }

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

void MemoryCheck::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    const IGtKernel& kernel        = dispatcher.Kernel();
    KernelProfile&   kernelProfile = _kernels.at(kernel.Id());

    if (dispatcher.ExecStage().IsDispatch())
    {
        GtKernelExecDesc execDesc; dispatcher.GetExecDescriptor(execDesc);
        if (kernel.IsInstrumented() && IsKernelExecProfileEnabled(execDesc, kernel.GpuPlatform(), kernel.Name().Get()))
        {
            dispatcher.SetProfilingMode(true);  // Enable instrumentation

            //// This tool needs an accurate information about memory allocations, which is available on the the final dispatch stage.
            //// So, on the initial dispatch stage, we only enable instrumentation, and request GTPin to invoke MemoryCheck::OnKernelRun
            //// one more time, on the final dispatch stage. If this request is accepted, the initialization of the profile buffer will
            //// be done on the final dispatch stage, otherwise - on the intial dispatch stage.
            //if (dispatcher.ReportFinalDispatchStage())
            //{
            //    return;
            //}
        }
        else
        {
            dispatcher.SetProfilingMode(false); // Disable instrumentation
            return;
        }
    }

    IGtMemoryMapper& memMapper = dispatcher.MemoryMapper();

    bool hasSlm = false;

    // Initialize per-access arguments of HLI functions
    for (auto& entry: kernelProfile.GetMemAccessMap())
    {
        auto        insId = entry.first;

        if (int32_t(insId) < knobMinInstrumentIns || knobMaxInstrumentIns < int32_t(insId))
        {
            continue;
        }

        MemAccess&        memAccess = entry.second;
        MemoryCheckArgs&  mcArgs    = memAccess.GetMemoryCheckArgs();

        hasSlm |= memAccess.IsSlm();

        mcArgs.out.umrCount    = 0;
        mcArgs.in.dataSize     = memAccess.DataSize();
        mcArgs.in.numAccesses  = memAccess.NumAccesses();

        memMapper.Write(&mcArgs, sizeof(mcArgs));
    }

    if (hasSlm)
    {
        uint32_t slmSize = dispatcher.SlmSize(); GTPIN_ASSERT(slmSize);

        SlmInitArg& slmInitMemArgs = kernelProfile.GetSlmInitArg();
        slmInitMemArgs.in.size = slmSize;

        memMapper.Write(&slmInitMemArgs, sizeof(SlmInitArg));
    }
}

void MemoryCheck::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    if (dispatcher.IsProfilingEnabled())
    {
        KernelProfile& kernelProfile = _kernels.at(dispatcher.Kernel().Id());
        kernelProfile.RecordMemoryCheckResults(dispatcher);
    }
}

bool MemoryCheck::InsertMemoryCheck(const IGtIns &ins, const MemAccess& memAccess, IGtKernelInstrument& instrumentor)
{
    GTPIN_ASSERT(memAccess.IsValid() && (memAccess.Id() == ins.Id()));

    uint32_t numAccesses = memAccess.NumAccesses();
    if (numAccesses == 0)
    {
        return false; // Nothing to check
    }

    const IGtKernel&         kernel               = instrumentor.Kernel();
    const IGtGenModel&       genModel             = kernel.GenModel();
    uint32_t                 regSize              = genModel.GrfRegSize();
    const GtMemoryAddrModel& addrModel            = memAccess.AddrModel();
    uint32_t                 addrSize             = addrModel.PtrSize();
    GtReg                    firstAddrReg         = GrfReg(memAccess.FirstAddrReg(), 0, regSize);
    GtReg                    firstDstReg          = GrfReg(memAccess.FirstDstReg(), 0, regSize);
    uint32_t                 dataSize             = memAccess.DataSize();
    uint32_t                 numOfElements        = memAccess.NumAccesses();
    uint32_t                 numDstRegs           = RoundUp(numOfElements * dataSize, regSize) / regSize;
    uint32_t                 numAddrRegs          = RoundUp(numOfElements * addrSize, regSize) / regSize;
    MemoryCheckArgs*         checkArgs            = const_cast<MemoryCheckArgs*>(&memAccess.GetMemoryCheckArgs());
    IargConstGrfRange        dstPayload(firstDstReg.RegNum(), numDstRegs);
    IargConstGrfRange        addrPayload(firstAddrReg.RegNum(), numAddrRegs);
    IargInsOpMask            accessMask(ins);
    IargSlmPtr               slmPtr;

    if (memAccess.IsSlm())
    {
        _checkSlmReadAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
            NullReg(),          // Unused return value
            addrPayload,        // arg[1]: Base address of the accessed memory range
            accessMask,         // arg[2]: Per-channel mask of memory accesses
            checkArgs,          // arg[3]: Memory check arguments
            slmPtr              // arg[4]: SLM pointer
        );
    }
    else if (memAccess.AddrModel().IsA64())
    {
        _checkUgmA64ReadAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
            NullReg(),          // Unused return value
            addrPayload,        // arg[1]: Base address of the accessed memory range
            accessMask,         // arg[2]: Per-channel mask of memory accesses
            checkArgs           // arg[3]: Memory check arguments
        );
    }
    else
    {
        _checkReadAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::After(),
            NullReg(),          // Unused return value
            dstPayload,         // arg[1]: Base address of the destination
            accessMask,         // arg[2]: Per-channel mask of memory accesses
            checkArgs           // arg[3]: Memory check arguments
        );
    }
    return true;
}

bool MemoryCheck::InsertSlmInit(const IGtIns& ins, const SlmInitArg& slmInitArg, IGtKernelInstrument& instrumentor)
{
    IargSlmPtr slmPtr;

    _initSlmFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
        NullReg(),          // Unused return value
        &slmInitArg,        // arg[1]: Slm initialization argument
        slmPtr              // arg[2]: SLM pointer
    );

    return true;
}

void MemoryCheck::LoadHliLibrary()
{
    std::string modulePath = JoinPath(GetKnobValue<std::string>("installDir"), "Examples", "memory_check.cl");
    _hliModule = GTPin_GetCore()->HliLibrary().CompileModuleFromFile(modulePath.c_str());
    GTPIN_ASSERT_MSG(_hliModule != nullptr, "Could not load HLI module " + modulePath);
}

void MemoryCheck::Fini()
{
    std::string str;

    // Dump profiling results and assembly code of all kernels
    for (const auto& entry : _kernels)
    {
        const auto& kernelProfile = entry.second;
        str += kernelProfile.MemoryCheckResults();
        kernelProfile.DumpAsm();
    }

    std::ofstream fs(JoinPath(GTPin_GetCore()->ProfileDir(), "memory_check.txt"));
    GTPIN_ASSERT(fs.is_open());
    fs << str;

    if (!KNOB_NO_COUT)
    {
        std::cout << str;
    }
}

/* ============================================================================================= */
// GTPin_Entry
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    SetKnobValue<bool>(true, "uninitialized_buffers_check_on");
    ConfigureGTPin(argc, argv);

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

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

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