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

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

/*!
* @file Interface of the host-to-device memory mapper in GTpin
*/

#ifndef IGT_MEMORY_MAPPER_H_
#define IGT_MEMORY_MAPPER_H_

#include <vector>

#include "gt_basic_defs.h"
#include "igt_profile_buffer.h"

namespace gtpin
{
#pragma pack(push, 8)

/*! @addtogroup GTPIN_MEMORY_MAPPER
@{
*/

/* ===================================================================================================== */
// Class GtMmapMethod
/* ===================================================================================================== */
/*!
 * @brief Method of the host-to-device memory mapping and sharing
 * 
 * The host-to-device memory mapping is a GTPin service that allows tools to specify host pointers in HLI
 * function arguments, and instruct GTPin to share the referenced data between host and device address spaces.\n
 * Depending on the sharing method, GTPin synchronizes content of mapped memory regions before kernel run,
 * and/or after the kernel completion.
 *
 * GTPin supports the following methods of memory mapping and sharing:\n
 * GT_MMAP_SHARE       - The content of the mapped region is synchronized at the start and completion of the kernel.
 *                       Can be used for sharing in&out objects passed by reference\n
 * GT_MMAP_SHARE_CONST - The content of the mapped region is synchronized only at the start of the kernel.
 *                       Can be used for sharing input objects passed by the reference\n
 * GT_MMAP_SHARE_RET   - The content of the mapped region is synchronized only at the completion of the kernel.
 *                       Can be used for sharing output objects passed by the reference\n
 * GT_MMAP_NO_SHARE    - Content of the mapped region is not shared at all; only address range in the host space
 *                       is mapped to the corresponding range in the device address space. The @ref IGtMemoryMapper
 *                       class provides access to such regions using addresses in the host space\n
 */
enum GT_MMAP_METHOD :uint32_t
{
    GT_MMAP_INVALID,
    GT_MMAP_SHARE,       ///< Map host-to-device address range; share data at the start and completion of the kernel
    GT_MMAP_SHARE_CONST, ///< Map host-to-device address range; share data at the start of the kernel
    GT_MMAP_SHARE_RET,   ///< Map host-to-device address range; share data at the completion of the kernel
    GT_MMAP_NO_SHARE     ///< Map host-to-device address range; no data sharing 
};

/*!
 * @brief Method of the host-to-device memory mapping and sharing
 */
class GtMmapMethod
{
public:
    inline constexpr GtMmapMethod(GT_MMAP_METHOD v = GT_MMAP_INVALID) : _value(v) {}    ///< Conversion constructor
    inline constexpr operator GT_MMAP_METHOD() const { return _value; }                 ///< Conversion operator

    inline constexpr bool IsValid()      const { return _value != GT_MMAP_INVALID; }    ///< Return true for a valid mapping method
    inline constexpr bool IsShare()      const { return _value == GT_MMAP_SHARE; }      ///< Return true for 'in&out' object sharing
    inline constexpr bool IsShareConst() const { return _value == GT_MMAP_SHARE_CONST; }///< Return true for 'input' object sharing
    inline constexpr bool IsShareRet()   const { return _value == GT_MMAP_SHARE_RET; }  ///< Return true for 'output' object sharing
    inline const char*    ToString()     const;                                         ///< Return string representation

private:
    GT_MMAP_METHOD _value;
};
GTPIN_ASSERT_STANDARD_LAYOUT(GtMmapMethod);


/* ===================================================================================================== */
// Class IGtMemoryMapper
/* ===================================================================================================== */
/*!
 * @brief Class that provides the host-to-device memory mapping service. This feature enables transparent
 * data sharing between host and device address spaces.
 * 
 * This mechanism is especially useful in mapping arguments of HLI functions. These functions are executed
 * by a GPU device, so all pointer parameters of the function should refer to objects in the profile buffer.
 * However, for tools that define instrumentation functions in the host (CPU) space, addresses in the
 * device space are unknown.\n
 * The host-to-device memory mapping interface allows tools to specify host pointers in function arguments,
 * and instruct GTPin to copy the referenced objects to/from the profile buffer in the device address space.
 */
class IGtMemoryMapper
{
public:
    /*!
     * @brief Reserve a memory block of the specified size in the profile buffer.
     *
     * Establish the specified sharing relation between [base, size] memory range in the host address space
     * and the reserved memory block in the device address space.
     *   
     * @param[in]   base        Base address of the host memory range to be mapped
     * @param[in]   size        Size of the memory range to be mapped, in bytes
     * @param[in]   mapMethod   The method of the memory mapping and sharing
     * @param[in]   alignment   Minimal alignment of the memory block in the profile buffer.
     *                          Must be power of two, not exceeding page size
     * @return Offset of the block allocated in the profile buffer, or -1 if mapping failed.
     *         In case of error, the IGtCore::LastError() function provides the specific failure reason
     * 
     * @par Availability:
     * - IGtTool::OnKernelBuild
     */
    virtual GtProfileBufferOffset Map(const void* base, uint32_t size,
                                      GtMmapMethod mapMethod = GT_MMAP_SHARE, uint32_t alignment = 8) = 0;

    /*!
     * @brief Given a memory block [dst, size] mapped into profile buffer, copy the corresponding buffer
     * content to the host block.
     * 
     * This function provides convenient access to mapped but not shared memory blocks (@ref GT_MMAP_NO_SHARE).
     * All other sharing methods synchronize mapped blocks automatically.\n
     * The specified address range must be entirely contained in a single host memory block mapped by the
     * @ref Map function.
     * 
     * @param[in]   dst     Destination address in the host memory
     * @param[in]   size    Number of bytes to copy. If this parameter is omitted (equal to zero), the
     *                      entire mapped memory block that contains 'dst' is copied
     * @return true - success, false - failure
     *         In case of error, the IGtCore::LastError() function provides the specific failure reason
     * 
     * @par Availability:
     * - IGtTool::OnKernelRun
     * - IGtTool::OnKernelComplete
     */
    virtual bool Read(void* dst, uint32_t size = 0) const = 0;

    /*!
     * @brief Given a memory block [src, size] mapped into profile buffer, copy its content to the corresponding
     * location within profile buffer.
     * 
     * This function provides convenient access to mapped but not shared memory blocks (@ref GT_MMAP_NO_SHARE).
     * All other sharing methods synchronize mapped blocks automatically.\n
     * The specified address range must be entirely contained in a single host memory block mapped by the
     * @ref Map function.
     * 
     * @param[in]   src     Source address in the host memory
     * @param[in]   size    Number of bytes to copy. If this parameter is omitted (equal to zero), the
     *                      entire mapped memory block that contains 'src' is copied
     * @return true - success, false - failure
     *         In case of error, the IGtCore::LastError() function provides the specific failure reason
     * 
     * @par Availability:
     * - IGtTool::OnKernelRun
     * - IGtTool::OnKernelComplete
     */
    virtual bool Write(const void* src, uint32_t size = 0) const = 0;

    /*!
     * @brief Find a location in the profile buffer, the specified address is mapped to.
     * 
     * @param[in]   ptr     An address in the host memory
     * @return Offset in the profile buffer, the specified address is mapped to,
     *         or invalid GtProfileBufferOffset value if 'ptr' does not belong in any mapped memory block
     * 
     * @note The function returns correct offset even if 'ptr' is not a base address of a mapped block but
     *       belongs in some mapped block
     * 
     * @par Availability:
     * - Any IGtTool callback
     */
    virtual GtProfileBufferOffset GetBufferOffset(const void* ptr) const = 0;

    #if (_GTPIN_CPP_STD >= 11)
    /*!
     * @brief Template variant of the virtual Map function
     *   
     * @param       T           Trivially copyable type
     * @param[in]   var         Variable to be mapped
     * @param[in]   mapMethod   The method of the memory mapping and sharing
     * @return Offset of the block allocated in the profile buffer, or -1 if mapping failed.
     * 
     * @par Availability:
     * - Any callback that provides IGtMemoryMapper
     */
    template <typename T> inline GtProfileBufferOffset Map(const T& var, GtMmapMethod mapMethod = GT_MMAP_SHARE);

    /// Specialization of Map<T> method for T = Span<E>
    template <typename E> inline GtProfileBufferOffset Map(const Span<E>& var, GtMmapMethod mapMethod = GT_MMAP_SHARE);

    /// Specialization of Map<T> method for T = std::vector<E>
    template <typename E> inline GtProfileBufferOffset Map(const std::vector<E>& var, GtMmapMethod mapMethod = GT_MMAP_SHARE);

#endif

protected:
    /// @internal Protected destructor
    virtual ~IGtMemoryMapper() = default;
};
GTPIN_ASSERT_PURE_ABSTRACT(IGtMemoryMapper);

#if (_GTPIN_CPP_STD >= 11)
/*!
 * @brief Template container of host-to-device mapped variables
 * @param   MM      The method of the memory mapping and sharing
 * @param   VARS    Types of variables to be mapped
 * 
   @verbatim
   Usage:
      auto mySharedVars      = GtMakeShared(myVar1, myVar2);        // Declare myVar... variables as host<->device synchronized
      auto mySharedConstVars = GtMakeSharedConst(myCvar1, myCvar2); // Declare myCvar... variables as host->device synchronized
      auto mySharedRetVars   = GtMakeSharedRet(myRvar1, myRvar2);   // Declare myRvar... variables as host<-device synchronized
      
      MyTool::OnKernelBuild(IGtKernelInstrument& instrumentor)
      {
          auto& mapper = instrumentor.MemoryMapper();
          mySharedVars.Map(mapper);        // Share myVar... variables using GT_MMAP_SHARE method
          mySharedConstVars.Map(mapper);   // Share myCvar... variables using GT_MMAP_SHARE_CONST method
          mySharedRetVars.Map(mapper);     // Share myRvar... variables using GT_MMAP_SHARE_RET method
      }
   @endverbatim
 */
template <GT_MMAP_METHOD MM, typename... VARS> class GtMapped
{
public:
    /*!
     * @brief Constructor
     * @param[in]   vars    List of variables to be mapped
     */
    GtMapped(const VARS&... vars) : _vars({&vars... }) {}

    /*!
     * @brief Map all variables in this container using the specified memory mapper
     * @param[in]   mapper      Object that provides the host-to-device memory mapping service
     * 
     * @par Availability:
     * - IGtTool::OnKernelBuild
     */
    void Map(IGtMemoryMapper& mapper) const { MapImpl<0, VARS...>(mapper); }

private:
    /// Recursive implementation of the Map() function
    template <uint32_t IDX, typename V, typename... VL> void MapImpl(IGtMemoryMapper& mapper) const
    {
        mapper.Map(*static_cast<const V*>(_vars[IDX]), MM, alignof(V));
        MapImpl<IDX + 1, VL...>(mapper);
    }
    template <uint32_t IDX> void MapImpl(IGtMemoryMapper&) const {}

private:
    Array<const void*, sizeof... (VARS)> _vars; ///< Array of pointers to mapped variables
};

/// Create a GtMapped object, deducing the template types from the types of arguments
template <GT_MMAP_METHOD MM, typename... VARS> GtMapped<MM, VARS...> GtMakeMapped(const VARS&... vars)
{
    return GtMapped<MM, VARS...>(vars...);
}

/// Alias of the GtMakeMapped<GT_MMAP_SHARE> function
template <typename... VARS> GtMapped<GT_MMAP_SHARE, VARS...> GtMakeShared(const VARS&... vars)
{
    return GtMapped<GT_MMAP_SHARE, VARS...>(vars...);
}

/// Alias of the GtMakeMapped<GT_MMAP_SHARE_CONST> function
template <typename... VARS> GtMapped<GT_MMAP_SHARE_CONST, VARS...> GtMakeSharedConst(const VARS&... vars)
{
    return GtMapped<GT_MMAP_SHARE_CONST, VARS...>(vars...);
}

/// Alias of the GtMakeMapped<GT_MMAP_SHARE_RET> function
template <typename... VARS> GtMapped<GT_MMAP_SHARE_RET, VARS...> GtMakeSharedRet(const VARS&... vars)
{
    return GtMapped<GT_MMAP_SHARE_RET, VARS...>(vars...);
}

#endif
/*! @} */

#if (_GTPIN_CPP_STD >= 11)
/* ===================================================================================================== */
// IGtMemoryMapper inline implementation
/* ===================================================================================================== */
template <typename T> GtProfileBufferOffset IGtMemoryMapper::Map(const T& var, GtMmapMethod mapMethod)
{
    GTPIN_ASSERT_TRIVIALLY_COPYABLE(T);
    return Map(&var, uint32_t(sizeof(T)), mapMethod, alignof(T));
}

template <typename E> GtProfileBufferOffset IGtMemoryMapper::Map(const Span<E>& var, GtMmapMethod mapMethod)
{
    GTPIN_ASSERT_TRIVIALLY_COPYABLE(E);
    return Map(var.data(), uint32_t(var.size() * sizeof(E)), mapMethod, alignof(E));
}

template <typename E> GtProfileBufferOffset IGtMemoryMapper::Map(const std::vector<E>& var, GtMmapMethod mapMethod)
{
    GTPIN_ASSERT_TRIVIALLY_COPYABLE(E);
    return Map(var.data(), uint32_t(var.size() * sizeof(E)), mapMethod, alignof(E));
}
#endif

#pragma pack(pop)
}
#endif
