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

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

/*!
 * @file Helper class that facilitates trace collection in the profile buffer
 */

#ifndef GT_PROFILE_TRACE_H_
#define GT_PROFILE_TRACE_H_

#include "ged.h"
#include "gt_basic_defs.h"
#include "igt_ins.h"
#include "igt_ins_factory.h"
#include "igt_vreg_factory.h"
#include "igt_profile_buffer.h"
#include "igt_gen_coder.h"
#include "igt_gen_procedure.h"

namespace gtpin
{
#pragma pack(push, 8)

/*! @addtogroup GTPIN_PROFILE_BUFFER 
@{
*/

/* ===================================================================================================== */
// Class GtProfileTrace
/* ===================================================================================================== */
/*!
 * @brief Class that represents a trace - sequence of data records chronologically collected and stored
 * in the profile buffer.
  * 
 * @par Availability:
 * - IGtTool::OnKernelBuild
 * - IGtTool::OnKernelRun
 * - IGtTool::OnKernelComplete
 */
class GtProfileTrace
{
public:
    /// Default constructor
    inline GtProfileTrace();

    /*!
     * @brief Constructor
     * @param[in]   maxTraceSize        Maximum size of the trace, in bytes
     * @param[in]   maxRecordSize       Maximum size of the trace record, in bytes. The record size is variable,
     *                                  not exceeding maxRecordSize bytes
     * @param[in]   maxThreads          Maximum number of threads concurrently accessing the trace
     * @param[in]   alignment           Alignment of the trace in the profile buffer; must be power of two
     */
    inline GtProfileTrace(uint32_t maxTraceSize, uint32_t maxRecordSize, uint32_t maxThreads = 0x1000, uint32_t alignment = 64);

    /*!
     * @brief Reserve space required to store the trace in the profile buffer
     * @param[in]   allocator           Allocator to be used to reserve space
     * @return true - success, false - failure
     *
     * @pre !IsAllocated()
     */
    inline bool Allocate(IGtProfileBufferAllocator& allocator);

    /*!
     * @brief Get Size of the trace collected in the profile buffer
     * @pre IsAllocated()
     * @return Actual size of the trace collected in the specified profile buffer
     */
    inline uint32_t Size(const IGtProfileBuffer& profileBuffer) const;

    /// Return true, if the trace collected in the specified profile buffer has been trimmed to avoid buffer overflow
    inline bool IsTruncated(const IGtProfileBuffer& profileBuffer) const;

    /*!
      * @brief Read data from the trace located in the specified profile buffer
      * @param[in]   profileBuffer       Profile buffer in which this trace is allocated
      * @param[in]   data                Pointer to an array of bytes that receives data from the trace.
      *                                  There are should be at least 'size' bytes in this array
      * @param[in]   offset              Offset within the trace of the first byte to read from
      * @param[in]   size                Number of bytes to read from the trace into 'data' array
      * @return true - success, false - failure
      *
      * @pre IsAllocated() && (offset + size <= MaxTraceSize())
      */
    inline bool Read(const IGtProfileBuffer& profileBuffer, void* data, uint32_t offset, uint32_t size) const;

    /*!
      * @brief Write data into the trace located in the specified profile buffer
      * @param[in]   profileBuffer       Profile buffer in which this trace is allocated
      * @param[in]   data                Pointer to an array of bytes that holds data to be written.
      *                                  There are should be at least 'size' bytes in this array
      * @param[in]   offset              Offset within the trace of the first byte to write into
      * @param[in]   size                Number of bytes to write from 'data' array into the trace 
      * @return true - success, false - failure
      *
      * @pre IsAllocated() && (offset + size <= MaxTraceSize())
      */
    inline bool Write(IGtProfileBuffer& profileBuffer, const void* data, uint32_t offset, uint32_t size) const;

    /*!
     * @brief Initialize the trace header in the specified profile buffer and, optionally, zero-initialize the
     * memory area allocated for trace records
     * @param[in]   profileBuffer       Profile buffer in which this trace is allocated
     * @param[in]   zeroInitRecords     true  - zero-initialize the memory area allocated for trace records
     *                                  false - leave trace records uninitialized
     * @return true - success, false - failure
     *
     * @pre IsAllocated()
     */
    inline bool Initialize(IGtProfileBuffer& profileBuffer, bool zeroInitRecords = true) const;

    /*!
     * @brief Generate code that allocates space for the new record in the trace, and computes offset/address
     * within the record.
     *
     * The offset is relative to the profile buffer base. If new record can not be allocated due to the trace
     * capacity limitations, the generated procedure sets f0.0:uw bits on. Otherwise, these bit are cleared off.
     * 
     * @param[in]   coder               GEN code generator
     * @param[in]   proc                Procedure, the generated code is appended to
     * @param[in]   recordSize          Size of the new record, in bytes
     * @param[in]   offsetReg           Register that receives the calculated offset within the profile buffer
     * @param[in]   addrReg             Register that receives the calculated address
     * @param[in]   offsetInRecord      Offset within the record
     *
     * @pre IsAllocated() && (recordSize <= MaxRecordSize())
     */
    inline void ComputeNewRecordOffset(const IGtGenCoder& coder, IGtGenProcedure& proc, uint32_t recordSize,
                                       const GtReg& offsetReg, uint32_t offsetInRecord = 0) const;

    /*! @copydoc ComputeNewRecordOffset */
    inline void ComputeNewRecordAddress(const IGtGenCoder& coder, IGtGenProcedure& proc, uint32_t recordSize,
                                        const GtReg& addrReg, uint32_t offsetInRecord = 0) const;

    // Accessors
    // =========
    inline uint32_t MaxTraceSize()  const { return _maxTraceSize; }     ///< Return maximum size of the trace, in bytes
    inline uint32_t MaxRecordSize() const { return _maxRecordSize; }    ///< Return maximum size of the trace record, in bytes
    inline bool     IsAllocated()   const { return _offset.IsValid(); } ///< Return true if this trace has been allocated

public:
        inline void ResetRegs() { _addrReg = GtReg(); _dataReg = GtReg(); }

private:

    struct Header                  ///< The header of the trace
    {
        uint32_t nextRecordOffset; ///< Offset of the next trace record, relative to the start of the allocated block
    };

    ///< @return Aligned size of the header, in bytes
    inline uint32_t AlignedHeaderSize() const { return AlignUp(uint32_t(sizeof(Header)), _alignment); }

private:
    uint32_t              _maxTraceSize;  ///< Maximum size of the trace, in bytes
    uint32_t              _maxRecordSize; ///< Maximum size of the trace record, in bytes
    uint32_t              _alignment;     ///< Alignment of the trace in the profile buffer
    GtProfileBufferOffset _offset;        ///< Offset of the memory block allocated in the profile buffer, or -1

    mutable GtReg         _addrReg;       ///< SendMessage address payload vreg
    mutable GtReg         _dataReg;       ///< SendMessage data payload vreg
};
/*! @} */

/* ===================================================================================================== */
// GtProfileTrace implementation
/* ===================================================================================================== */
GtProfileTrace::GtProfileTrace() : _maxTraceSize(0), _maxRecordSize(0), _alignment(0) {}

GtProfileTrace::GtProfileTrace(uint32_t maxTraceSize, uint32_t maxRecordSize, uint32_t maxThreads, uint32_t alignment) :
    _maxRecordSize(maxRecordSize), _alignment(std::max(alignment, uint32_t(alignof(Header))))
{
    // The trace memory block consists of trace header followed by trace records.
    // The upper limit of the block size is (4Gb - redZone). The redZone prevents wrapping of Header::nextRecordOffset
    // value around 32-bit boundary when multiple threads simultaneously allocate records in the trace approaching 4Gb size.
    // Note, the procedure generated by ComputeNewRecordOffset() never writes beyond trace bounds - it does not
    // touch memory within redZone.

    uint64_t redZone = (int64_t)maxThreads * maxRecordSize;
    uint32_t upLimit = ((redZone > UINT32_MAX) ? 0 : 0 - (uint32_t)redZone);
    if (upLimit <= AlignedHeaderSize())
    {
        _maxTraceSize = 0;
    }
    else
    {
        _maxTraceSize = std::min(upLimit - AlignedHeaderSize(), maxTraceSize);
    }
}

bool GtProfileTrace::Allocate(IGtProfileBufferAllocator& allocator)
{
    if (IsAllocated())
    {
        return true;
    }
    if (_maxTraceSize > 0)
    {
        // Allocated block = {heasder, trace, one overflowed record}
        _offset = allocator.Allocate(AlignedHeaderSize() + _maxTraceSize  + _maxRecordSize, _alignment);
    }
    return _offset.IsValid();
}

bool GtProfileTrace::Initialize(IGtProfileBuffer& profileBuffer, bool zeroInitRecords) const
{
    GTPIN_ASSERT(IsAllocated());

    bool status;
    Header header;
    header.nextRecordOffset = AlignedHeaderSize();
    if ((status = profileBuffer.Write(_offset, &header, sizeof(header))))
    {
        if (zeroInitRecords)
        {
            status = profileBuffer.Fill(_offset + AlignedHeaderSize(), _maxTraceSize, 0);
        }
    }
    return status;
}

uint32_t GtProfileTrace::Size(const IGtProfileBuffer& profileBuffer) const
{
    GTPIN_ASSERT(IsAllocated());

    Header header;
    bool status = profileBuffer.Read(_offset, &header, sizeof(header)); GTPIN_ASSERT(status);
    GTPIN_ASSERT(header.nextRecordOffset >= AlignedHeaderSize());
    return std::min(header.nextRecordOffset - AlignedHeaderSize(), _maxTraceSize);
}

bool GtProfileTrace::IsTruncated(const IGtProfileBuffer& profileBuffer) const
{
    GTPIN_ASSERT(IsAllocated());

    Header header;
    bool status = profileBuffer.Read(_offset, &header, sizeof(header)); GTPIN_ASSERT(status);
    return (header.nextRecordOffset > (AlignedHeaderSize() + _maxTraceSize));
}

bool GtProfileTrace::Read(const IGtProfileBuffer& profileBuffer, void* data, uint32_t offset, uint32_t size) const
{
    GTPIN_ASSERT(IsAllocated());
    GTPIN_ASSERT(offset + size <= _maxTraceSize);
    return profileBuffer.Read(_offset + AlignedHeaderSize() + offset, data, size);
}

bool GtProfileTrace::Write(IGtProfileBuffer& profileBuffer, const void* data, uint32_t offset, uint32_t size) const
{
    GTPIN_ASSERT(IsAllocated());
    GTPIN_ASSERT(offset + size <= _maxTraceSize);
    return profileBuffer.Write(_offset + AlignedHeaderSize() + offset, data, size);
}

void GtProfileTrace::ComputeNewRecordOffset(const IGtGenCoder& coder, IGtGenProcedure& proc, uint32_t recordSize,
                                            const GtReg& offsetReg, uint32_t offsetInRecord) const
{
    GTPIN_ASSERT(IsAllocated());
    GTPIN_ASSERT(recordSize <= _maxRecordSize);

    IGtVregFactory& vregs = coder.VregFactory();
    IGtInsFactory&  insF  = coder.InstructionFactory();

    if (!_addrReg.IsValid()) { _addrReg = vregs.MakeMsgAddrScratch(); }
    if (!_dataReg.IsValid()) { _dataReg = vregs.MakeMsgDataScratch(); }

    GtReg flagReg = FlagReg(0);
    GtPredicate predicate(flagReg);

    // _addrReg = &Header::nextRecordOffset
    coder.ComputeAddress(proc, _addrReg, _offset + uint32_t(offsetof(Header, nextRecordOffset)));

    // Allocate the next trace record. Set _dataReg =  offset of the new record relative to the start of the allocated block:
    // _dataReg = [_addrReg]; [_addrReg] += recordSize
    proc += insF.MakeMov(_dataReg, recordSize);
    proc += insF.MakeAtomicAdd(_dataReg, _addrReg, _dataReg, GED_DATA_TYPE_ud);

    // If the new record can not be allocated, set the f0.0:uw bits on. Otherwise, clear these bits off.
    proc += insF.MakeCmp(GED_COND_MODIFIER_g, flagReg, _dataReg, GtImmU32(AlignedHeaderSize() + _maxTraceSize), {16});

    // If allocation failed, restore the Header::nextRecordOffset value: [_addrReg] -= RecordSize()
    proc += insF.MakeMov(_dataReg, recordSize).SetPredicate(predicate);
    proc += insF.MakeAtomicSub(NullReg(), _addrReg, _dataReg, GED_DATA_TYPE_ud).SetPredicate(predicate);

    // Set offsetReg = offset of the requested location within the new record, relative to the base of the profile buffer
    uint32_t offset = _offset + offsetInRecord;
    if ((offset != 0) || (offsetReg != _dataReg))
    {
        proc += insF.MakeAdd(offsetReg, _dataReg, GtImmU32(offset)).SetPredicate(!predicate);
    }
}

void GtProfileTrace::ComputeNewRecordAddress(const IGtGenCoder& coder, IGtGenProcedure& proc, uint32_t recordSize,
                                             const GtReg& addrReg, uint32_t offsetInRecord) const
{
    IGtVregFactory& vregs = coder.VregFactory();
    if (!_dataReg.IsValid()) { _dataReg = vregs.MakeMsgDataScratch(); }

    ComputeNewRecordOffset(coder, proc, recordSize, _dataReg, offsetInRecord);
    coder.ComputeAddress(proc, addrReg, _dataReg);
}

#pragma pack(pop)
}
#endif
