Memory.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
#include "logger/Logger.hpp"
#include <memory>
#include <string>
 
namespace kp {
 
// Forward declare the Tensor and Image classes
class Tensor;
class Image;
 
class Memory
{
    // This is the base class for Tensors and Images.
    // It's required so that algorithms and sequences can mix tensors and
    // images.
  public:
    enum class MemoryTypes
    {
        eDevice = 0,  
        eHost = 1,    
        eStorage = 2, 
        eDeviceAndHost =
          3, 
    };
 
    enum class DataTypes
    {
        eBool = 0,
        eInt = 1,
        eUnsignedInt = 2,
        eFloat = 3,
        eDouble = 4,
        eCustom = 5,
        eShort = 6,
        eUnsignedShort = 7,
        eChar = 8,
        eUnsignedChar = 9
    };
 
    enum class Type
    {
        eTensor = 0,
        eImage = 1
    };
 
    static std::string toString(MemoryTypes dt);
    static std::string toString(DataTypes dt);
 
    Memory(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           const DataTypes& dataType,
           const MemoryTypes& memoryType,
           uint32_t x,
           uint32_t y);
 
 
    Memory(const Memory&) = delete;
    Memory(const Memory&&) = delete;
    Memory& operator=(const Memory&) = delete;
    Memory& operator=(const Memory&&) = delete;
 
    virtual ~Memory(){};
 
    virtual void destroy();
 
    MemoryTypes memoryType();
 
    template<typename T>
    static constexpr DataTypes dataType();
 
    DataTypes dataType();
 
    virtual bool isInit() = 0;
 
    virtual void recordCopyFromStagingToDevice(
      const vk::CommandBuffer& commandBuffer) = 0;
 
    virtual void recordCopyFromDeviceToStaging(
      const vk::CommandBuffer& commandBuffer) = 0;
    virtual void recordPrimaryMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask) = 0;
    virtual void recordStagingMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask) = 0;
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Memory> copyFromMemory);
 
    virtual vk::WriteDescriptorSet constructDescriptorSet(
      vk::DescriptorSet descriptorSet,
      uint32_t binding) = 0;
 
    uint32_t size();
 
    uint32_t dataTypeMemorySize();
 
    static uint32_t dataTypeMemorySize(DataTypes dataType);
 
    uint32_t memorySize();
 
    vk::DescriptorType getDescriptorType() { return mDescriptorType; }
 
    void* rawData();
 
    void setData(const void* data, size_t size);
 
    template<typename T>
    void setData(const std::vector<T>& data)
    {
        KP_LOG_DEBUG("Kompute Memory setting data with data size {}",
                     data.size() * sizeof(T));
 
        this->setData(data.data(), data.size() * sizeof(T));
    }
 
    template<typename T>
    T* data()
    {
        if (this->mRawData == nullptr) {
            this->mapRawData();
        }
 
        return (T*)this->mRawData;
    }
 
    template<typename T>
    std::vector<T> vector()
    {
        if (this->mRawData == nullptr) {
            this->mapRawData();
        }
 
        return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
    }
 
    /***
     * Retreive the size of the x-dimension of the memory
     *
     * @return Size of the x-dimension of the memory
     */
    uint32_t getX() { return this->mX; }
 
    /***
     * Retreive the size of the y-dimension of the memory
     *
     * @return Size of the y-dimension of the memory
     */
    uint32_t getY() { return this->mY; };
 
    virtual Type type() = 0;
 
  protected:
    // -------------- ALWAYS OWNED RESOURCES
    MemoryTypes mMemoryType;
    DataTypes mDataType;
    uint32_t mSize;
    uint32_t mDataTypeMemorySize;
    void* mRawData = nullptr;
    vk::DescriptorType mDescriptorType;
    bool mUnmapMemory = false;
    uint32_t mX;
    uint32_t mY;
 
    // -------------- NEVER OWNED RESOURCES
    std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice;
    std::shared_ptr<vk::Device> mDevice;
 
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::DeviceMemory> mPrimaryMemory;
    bool mFreePrimaryMemory = false;
    std::shared_ptr<vk::DeviceMemory> mStagingMemory;
    bool mFreeStagingMemory = false;
 
    // Private util functions
    void mapRawData();
    void unmapRawData();
    void updateRawData(void* data);
    vk::MemoryPropertyFlags getPrimaryMemoryPropertyFlags();
    vk::MemoryPropertyFlags getStagingMemoryPropertyFlags();
 
    virtual void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                                std::shared_ptr<Tensor> copyFromMemory) = 0;
    virtual void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                                std::shared_ptr<Image> copyFromMemory) = 0;
};
 
template<>
constexpr Memory::DataTypes
Memory::dataType<bool>()
{
    return DataTypes::eBool;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<int8_t>()
{
    return DataTypes::eChar;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<uint8_t>()
{
    return DataTypes::eUnsignedChar;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<int16_t>()
{
    return DataTypes::eShort;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<uint16_t>()
{
    return DataTypes::eUnsignedShort;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<int32_t>()
{
    return DataTypes::eInt;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<uint32_t>()
{
    return DataTypes::eUnsignedInt;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<float>()
{
    return DataTypes::eFloat;
}
 
template<>
constexpr Memory::DataTypes
Memory::dataType<double>()
{
    return DataTypes::eDouble;
}
 
} // End namespace kp