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

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

/*!
 * @file A tool that detects out-of-bounds (OOB) memory accesses in kernels
 */

#include <cstring>
#include <map>
#include <vector>
#include <set>

#include "bounds_check.h"

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

/* ============================================================================================= */
// Configuration
/* ============================================================================================= */
Knob<int>  KNOB_MAX_KERNEL_BUFFERS("max_kernel_buffers", 256, "Maximum number of global buffers available to the kernel");
Knob<bool> KNOB_CHECK_ACCESS_TYPE("check_access_type", false, "Check access type compatibility (readonly vs. write)");
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, const DcSendMsg &msg);

    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     FirstAddrReg() const { return _firstAddrReg; }       ///< @return First register in the address payload
    uint32_t     NumAddresses() const { return _numAddresses; }       ///< @return Number of elements in the address payload
    uint32_t     DataSize()     const { return _dataSize; }           ///< @return Data size, in bytes, per each address
    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 BoundsCheckArgs structure, but its content is controlled externally
     */
    const BoundsCheckArgs& GetBoundsCheckArgs() const { return _bcArgs; }
    BoundsCheckArgs&       GetBoundsCheckArgs()       { return _bcArgs; }

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            _firstAddrReg;          ///< First register in the address payload of the instruction
    uint32_t            _numAddresses = 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
    BoundsCheckArgs     _bcArgs;                ///< Common arguments of bounds check functions
    std::string         _errMsg;                ///< Error message on unsupported memory access
};

/* ============================================================================================= */
// Struct ProfileResults
/* ============================================================================================= */
/*!
 * Profile results per kernel dispatch / per instuction
 */
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          surfaceSize;         ///< Size, in bytes, of the available memory space for a surface access
    uint32_t          oobCount;            ///< Counter of OOB 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         BoundsCheckResults() 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

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

    void RecordUnsupportedInstruction(const IGtIns& ins, const std::string& errMsg); ///< Record unsupported instruction
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::vector<std::pair<InsId, std::string>> _unsupportedInstructions;  ///< Instructions that could not be instrumented

    std::map<InsId, MemAccess> _memAccessMap;   ///< Map: Instruction ID to memory access information
    std::list<ProfileResults>  _profileResults; ///< Profile results per kernel dispatch / per instuction 
};

/* ============================================================================================= */
// Class BoundsCheck
/* ============================================================================================= */
/*!
 * A tool that detects out-of-bounds (OOB) memory accesses in kernels
 */
class BoundsCheck : public GtTool
{
public:
    // Implementation of the IGtTool interface
    const char* Name()          const                   override { return "Bounds check"; }
    void        OnKernelBuild(IGtKernelInstrument&)     override;
    void        OnKernelRun(IGtKernelDispatch&)         override;
    void        OnKernelComplete(IGtKernelDispatch&)    override;

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

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

    /*!
     * Insert a call to HLI function that detects out-of-bounds (OOB) 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 InsertBoundsCheck(const IGtIns &ins, const MemAccess& memAccess, IGtKernelInstrument& instrumentor);

    /*!
     * Given BTI index, return memory buffer referenced by the corresponding stateful argument.
     * Return empty buffer, if BTI index cannot be dereferenced
     */
    GtMemoryBuffer GetBtiBuffer(uint32_t bti, const GtMemoryBufferConstSpan& allocatedBuffers,
                                const IGtKernelDispatch& dispatcher);

    /// @return Pointer to a buffer in the 'buffers' sequence that contains 'ptr', or nullptr if buffer is not found
    const GtMemoryBuffer* FindBuffer(const GtMemoryBufferConstSpan& buffers, uintptr_t ptr);

private:
    // Bounds check functions
    CheckA64AccessFunc                  _checkA64AccessFunc;
    CheckA32AccessFunc                  _checkA32AccessFunc;
    CheckSurfaceAccessFunc              _checkSurfaceAccessFunc;
    CheckA64ScatteredAccessFunc         _checkA64ScatteredAccessFunc;
    CheckA32ScatteredAccessFunc         _checkA32ScatteredAccessFunc;
    CheckSurfaceScatteredAccessFunc     _checkSurfaceScatteredAccessFunc;
    CheckBlock2DAccessFunc              _checkBlock2DAccessFunc;

    /// Global memory buffers available to the kernel; shared and passed as an argument to bounds check functions
    MemBounds*                          _allocatedBuffers = nullptr; ///< All allocated buffers
    MemBounds*                          _writableBuffers  = nullptr; ///< Allocated writable buffers

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

/* ============================================================================================= */
// MemAccess implementation
/* ============================================================================================= */
MemAccess::MemAccess(const IGtIns &ins, const DcSendMsg &msg) : _insId(ins.Id())
{


    // Get and check data port (SFID)
    GtSfid sfid = ins.Sfid();
    bool isHdc = (sfid == GED_SFID_DP_DC0) || (sfid == GED_SFID_DP_DC1);
    if ((sfid != GED_SFID_UGM) && (sfid != GED_SFID_SLM) && !isHdc)
    {
        _errMsg = "Unsupported data port " + std::string(sfid.ToString());
        return;
    }

    // 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          isLoadLsc  = (opcode == GED_DP_OPCODE_LOAD) || (opcode == GED_DP_OPCODE_LOAD_2D_BLOCK);
    bool          isStoreLsc = (opcode == GED_DP_OPCODE_STORE) || (opcode == GED_DP_OPCODE_STORE_2D_BLOCK);
    bool          isAtomic   = ins.IsAtomic();
    bool          isLoad     = isLoadLsc || (isHdc && !isAtomic && ins.HasDstOperand() && ins.DstRegFile().IsGrf());
    bool          isStore    = isStoreLsc || (isHdc && !isAtomic && ins.DstOperand().Reg().IsNullReg());

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

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

    // Finally, initialize the rest of data members...
    GTPIN_ASSERT(ins.SrcRegFile(0).IsGrf());
    _firstAddrReg = ins.SrcRegOperand(0).Reg().RegNum(); 
    _numAddresses = ins.NumAccesses();              GTPIN_ASSERT(_numAddresses != 0);
    _dataSize     = msg.ElementSize() * msg.NumElements();            GTPIN_ASSERT(_dataSize != 0);
    _accessType   = (isLoad ? GT_ACCESS_READ : (isStore ? GT_ACCESS_WRITE : GT_ACCESS_READ_WRITE));

    _isValid = true;
}

/* ============================================================================================= */
// ProfileResults implementation
/* ============================================================================================= */
ProfileResults::ProfileResults(const IGtKernelDispatch& dispatcher, const MemAccess& memAccess) :
    dispatchId(dispatcher.DispatchId()), insId(memAccess.Id()),
    surfaceSize(memAccess.GetBoundsCheckArgs().in.surfaceSize),
    oobCount(memAccess.GetBoundsCheckArgs().out.oobCount)
{
    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;
            const DcSendMsg msg = DcSendMsg::Decode(ins.GetGedIns());
            // Exclude EOT and fence instructions; fence is used for synchronization, not direct memory access.
            if (ins.IsMemAccess() && !ins.IsEot() && !msg.IsMemFence())
            {
                MemAccess memAccess(ins, msg);
                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; }

std::string KernelProfile::BoundsCheckResults() 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;

    // Print suspicious instructions explanation (zero surface size), if needed
    bool printSuspicious = false;
    for (const auto& res : _profileResults)
    {
        if (res.oobCount != 0 && res.surfaceSize == 0)
        {
            os << std::endl << "Note: Z marks out-of-bounds accesses to zero-sized surfaces (stateful/BTS, bindless/SS, or SLM)." << std::endl << std::endl;
            printSuspicious = true;
            break;
        }
    }

    uint32_t oobTotal = 0;
    for (const auto& res : _profileResults)
    {
        if (res.oobCount != 0)
        {
            if (oobTotal == 0)
            {
                os << std::setw(20) << "Dispatch ID" << std::setw(20) << "Instruction ID" << std::setw(20) << "#OOB violations";
                os << " " << std::setw(45) << "Execution descriptor";
                if (printSuspicious)
                {
                    os << std::setw(20) << "0-sized Surface";
                }
                os << std::endl;
                os << std::string(printSuspicious ? 140 : 120, '-') << std::endl;
            }
            os << std::setw(20) << res.dispatchId << std::setw(20) << res.insId << std::setw(20) << res.oobCount;
            os << " " << std::setw(45) << res.kernelExecDesc.ToString(_platform, ExecDescAlignedFormat());
            if (printSuspicious && res.surfaceSize == 0)
            {
                os << std::setw(20) << "Z";
            }
            else
            {
                os << std::setw(20) << " ";
            }
            os << std::endl;
            oobTotal += res.oobCount;
        }
    }
    if (oobTotal == 0)
    {
        os << "No OOB accesses detected" << std::endl;
    }

    // Print unsupported instructions
    if (!_unsupportedInstructions.empty())
    {
        os << std::endl;
        os << " ---------------------------" << std::endl;
        os << " Unsupported memory accesses:" << std::endl;
        os << " ---------------------------" << std::endl;
        for (const auto& entry : _unsupportedInstructions)
        {
            os << entry.second << ": Instruction ID: [" << std::setw(3) << entry.first << "]" << std::endl;
        }
        os << std::string(120, '-') << std::endl;
    }

    return os.str();
}

void KernelProfile::RecordBoundsCheckResults(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())
    {
        _unsupportedInstructions.emplace_back(ins.Id(), errMsg);
    }
}

/* ============================================================================================= */
// BoundsCheck implementation
/* ============================================================================================= */
BoundsCheck::BoundsCheck() : _checkA64AccessFunc("CheckA64Access"),
                             _checkA32AccessFunc("CheckA32Access"),
                             _checkSurfaceAccessFunc("CheckSurfaceAccess"),
                             _checkA64ScatteredAccessFunc("CheckA64ScatteredAccess"),
                             _checkA32ScatteredAccessFunc("CheckA32ScatteredAccess"),
                             _checkSurfaceScatteredAccessFunc("CheckSurfaceScatteredAccess"),
                             _checkBlock2DAccessFunc("CheckBlock2DAccess")
{
    size_t size = MemBounds::SizeOf(KNOB_MAX_KERNEL_BUFFERS);
    #if (_GTPIN_CPP_STD >= 17)
        constexpr std::align_val_t alignment { alignof(MemBounds) };
        _allocatedBuffers = reinterpret_cast<MemBounds*>(operator new [](size, alignment));
        _writableBuffers  = reinterpret_cast<MemBounds*>(operator new [](size, alignment));
    #else
        _allocatedBuffers = reinterpret_cast<MemBounds*>(operator new [](size));
        _writableBuffers  = reinterpret_cast<MemBounds*>(operator new [](size));
    #endif
}

BoundsCheck::~BoundsCheck()
{
    #if (_GTPIN_CPP_STD >= 17)
        constexpr std::align_val_t alignment { alignof(MemBounds) };
        operator delete [](_allocatedBuffers, alignment);
        operator delete [](_writableBuffers, alignment);
    #else
        operator delete [](_allocatedBuffers);
        operator delete [](_writableBuffers);
    #endif
}

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

void BoundsCheck::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;

    // Handle common instrumentation knobs
    HandleCommonInstrumnetationKnobs(instrumentor);
    // 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);
        InsertBoundsCheck(ins, memAccess, instrumentor);

        // 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.GetBoundsCheckArgs(), GT_MMAP_SHARE);
    }

    // Map the array of global buffers available to kernel. This is an HLI argument common to all global memory accesses.
    // The array will be populated at the start of the kernel.
    uint32_t sizeOfMemBounds = (uint32_t)MemBounds::SizeOf(KNOB_MAX_KERNEL_BUFFERS);
    memMapper.Map(_allocatedBuffers, sizeOfMemBounds, GT_MMAP_NO_SHARE, alignof(MemBounds));
    if (KNOB_CHECK_ACCESS_TYPE)
    {
        memMapper.Map(_writableBuffers, sizeOfMemBounds, GT_MMAP_NO_SHARE, alignof(MemBounds));
    }

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

void BoundsCheck::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    const IGtKernel& kernel        = dispatcher.Kernel();
    if (_kernels.find(kernel.Id()) == _kernels.end())
    {
        return;
    }

    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 final dispatch stage.
            // So, on the initial dispatch stage, we only enable instrumentation, and request GTPin to invoke BoundsCheck::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 initial dispatch stage.
            if (dispatcher.ReportFinalDispatchStage())
            {
                return;
            }
        }
        else
        {
            dispatcher.SetProfilingMode(false); // Disable instrumentation
            return;
        }
    }

    IGtMemoryMapper& memMapper = dispatcher.MemoryMapper();

    // Populate and share array of global buffers available to the kernel
    GtMemoryBufferConstSpan allocatedBuffers = dispatcher.GetAllocatedBuffers();
    GTPIN_ASSERT_MSG(allocatedBuffers.size() <= uint32_t(KNOB_MAX_KERNEL_BUFFERS),
                     "Number of global buffers " + ToString(allocatedBuffers.size()) +
                     " exceeded " + ToString(KNOB_MAX_KERNEL_BUFFERS.GetValue()) + ". Increase max_kernel_buffers value");

    _allocatedBuffers->numRanges = 0;
    for (const auto& buffer : allocatedBuffers)
    {
        _allocatedBuffers->ranges[_allocatedBuffers->numRanges++] = buffer.Range();
    }
    memMapper.Write(_allocatedBuffers, (uint32_t)_allocatedBuffers->SizeOf());

    if (KNOB_CHECK_ACCESS_TYPE)
    {
        _writableBuffers->numRanges = 0;
        for (const auto& buffer : allocatedBuffers)
        {
            if (buffer.Access().IsWrite())
            {
                _writableBuffers->ranges[_writableBuffers->numRanges++] = buffer.Range();
            }
        }
        memMapper.Write(_writableBuffers, (uint32_t)_writableBuffers->SizeOf());
    }

    // 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;
        GtMemoryAddrModel addrModel = memAccess.AddrModel();
        BoundsCheckArgs&  bcArgs    = memAccess.GetBoundsCheckArgs();

        bcArgs.out.oobCount    = 0;
        bcArgs.in.dataSize     = memAccess.DataSize();
        bcArgs.in.numAddresses = memAccess.NumAddresses();

        // Compute size of the available memory space for BTS, SLM and scratch accesses
        bcArgs.in.surfaceSize = UINT32_MAX;
        if (addrModel.IsSlm())
        {
            bcArgs.in.surfaceSize = dispatcher.SlmSize();
        }
        else if (addrModel.IsSurfaceState())
        {
            bcArgs.in.surfaceSize = dispatcher.ScratchSpaceSize(); // Assuming SS address model used in scratch access only
        }
        else if (addrModel.IsBts())
        {
            GtMemoryBuffer buffer = GetBtiBuffer(addrModel.Bti(), allocatedBuffers, dispatcher);
            if (!buffer.IsEmpty() && (!KNOB_CHECK_ACCESS_TYPE || buffer.Access().Includes(memAccess.AccessType())))
            {
                bcArgs.in.surfaceSize = (uint32_t)buffer.Size(); GTPIN_ASSERT(buffer.Size() <= UINT32_MAX);
            }
        }

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

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

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

    uint32_t numAddresses = memAccess.NumAddresses();
    if (numAddresses == 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                    firstReg             = GrfReg(memAccess.FirstAddrReg(), 0, regSize);
    uint32_t                 numRegs              = RoundUp(memAccess.NumAddresses() * addrSize, regSize) / regSize;
    BoundsCheckArgs*         checkArgs            = const_cast<BoundsCheckArgs*>(&memAccess.GetBoundsCheckArgs());
    bool                     checkWritableBuffers = (KNOB_CHECK_ACCESS_TYPE && memAccess.AccessType().IsWrite());
    MemBounds*               globalBuffers        = (checkWritableBuffers ? _writableBuffers : _allocatedBuffers);
    IargConstGrfRange        addrPayload(firstReg.RegNum(), numRegs);
    IargInsOpMask            accessMask(ins);

    if (addrModel.IsA64())
    {
        GTPIN_ASSERT(addrSize == sizeof(uint64_t));

        if(ins.IsBlock2DAccess())
        {
            _checkBlock2DAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                globalBuffers,      // arg[0]: Available global buffers
                                addrPayload,        // arg[1]: Address payload
                                accessMask,         // arg[2]: Per-channel mask of memory accesses
                                checkArgs           // arg[3]: Bounds check arguments
                                );
        }
        else if (numAddresses == 1)
        {
            _checkA64AccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                globalBuffers,      // arg[0]: Available global buffers
                                firstReg,           // arg[1]: Base address of the accessed memory range
                                accessMask,         // arg[2]: Per-channel mask of memory accesses
                                checkArgs           // arg[3]: Bounds check arguments
                                );
        }
        else
        {
            _checkA64ScatteredAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                globalBuffers,      // arg[0]: Available global buffers
                                addrPayload,        // arg[1]: Base addresses of accessed memory ranges
                                accessMask,         // arg[2]: Per-channel mask of memory accesses
                                checkArgs           // arg[3]: Bounds check arguments
                                );
        }
    }
    else if (addrModel.IsA32())
    {
        GTPIN_ASSERT(addrSize == sizeof(uint32_t));

        if (numAddresses == 1)
        {
            _checkA32AccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                globalBuffers,      // arg[0]: Available global buffers
                                firstReg,           // arg[1]: Base address of the accessed memory range
                                accessMask,         // arg[2]: Per-channel mask of memory accesses
                                checkArgs           // arg[3]: Bounds check arguments
                                );
        }
        else
        {
            _checkA32ScatteredAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                globalBuffers,      // arg[0]: Available global buffers
                                addrPayload,        // arg[1]: Base addresses of accessed memory ranges
                                accessMask,         // arg[2]: Per-channel mask of memory accesses
                                checkArgs           // arg[3]: Bounds check arguments
                                );
        }
    }
    else
    {
        GTPIN_ASSERT(addrSize == sizeof(uint32_t));
        GTPIN_ASSERT_MSG(addrModel.IsBts() || addrModel.IsSlm() || addrModel.IsBss() || addrModel.IsSurfaceState(),
                         "Unknown address model: " + addrModel.ToString());

        if (numAddresses == 1)
        {
            _checkSurfaceAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                firstReg,           // arg[0]: Base address of the accessed memory range
                                accessMask,         // arg[1]: Per-channel mask of memory accesses
                                checkArgs           // arg[2]: Bounds check arguments
                                );
        }
        else
        {
            _checkSurfaceScatteredAccessFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                NullReg(),          // Unused return value
                                addrPayload,        // arg[0]: Base addresses of accessed memory ranges
                                accessMask,         // arg[1]: Per-channel mask of memory accesses
                                checkArgs           // arg[2]: Bounds check arguments
                                );
        }
    }
    return true;
}

GtMemoryBuffer BoundsCheck::GetBtiBuffer(uint32_t bti, const GtMemoryBufferConstSpan& allocatedBuffers,
                                         const IGtKernelDispatch& dispatcher)
{
    Uint32Validatable<>  argIndex;
    UintPtrValidatable<> argPtr;

    for (const auto& arg : dispatcher.Kernel().PayloadArguments())
    {
        if (arg.IsMemoryPointer() && arg.IsExplicit() && arg.addrModel.IsValid() && (arg.addrModel.Bti() == bti))
        {
            ConstByteSpan argValue = dispatcher.GetPayloadArgumentValue(arg.index);
            if (argValue.size() == sizeof(uintptr_t))
            {
                argPtr   = *reinterpret_cast<const uintptr_t*>(argValue.data());
                argIndex = arg.index;
            }
        }
    }
    if (!argPtr.IsValid())
    {
        return {};
    }

    const GtMemoryBuffer* buffer = FindBuffer(allocatedBuffers, argPtr);
    if (!buffer)
    {
        return {};
    }

    // IF buffer_offset is specified THEN
    //    BTI_base == allocation_base_of(arg_value), AND BTI_size = allocation_size_of(arg_value)
    // ELSE
    //    BTI_base == arg_value, AND BTI_size = allocation_size_of(arg_value) - (arg_value - allocation_base_of(arg_value))

    for (const auto& arg : dispatcher.Kernel().PayloadArguments())
    {
        if ((arg.index == argIndex) && (std::strcmp(arg.type, "buffer_offset") == 0))
        {
            return *buffer;
        }
    }

    size_t offset = argPtr - buffer->Base(); GTPIN_ASSERT(offset <= buffer->Size());
    return GtMemoryBuffer(argPtr, buffer->Size() - offset, buffer->Access());
}

const GtMemoryBuffer* BoundsCheck::FindBuffer(const GtMemoryBufferConstSpan& buffers, uintptr_t ptr)
{
    for (const auto& buffer : buffers)
    {
        if (buffer.Range().Contains(ptr))
        {
            return &buffer;
        }
    }
    return nullptr;
}

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

void BoundsCheck::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.BoundsCheckResults();
        kernelProfile.DumpAsm();
    }

    std::ofstream fs(JoinPath(GTPin_GetCore()->ProfileDir(), "bounds_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[])
{
    // Parse command line and configure GTPin
    ConfigureGTPin(argc, argv);

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

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

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