Tensor.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
#include "logger/Logger.hpp"
#include <memory>
#include <string>
 
namespace kp {
 
class Tensor
{
  public:
    enum class TensorTypes
    {
        eDevice = 0,  
        eHost = 1,    
        eStorage = 2, 
    };
    enum class TensorDataTypes
    {
        eBool = 0,
        eInt = 1,
        eUnsignedInt = 2,
        eFloat = 3,
        eDouble = 4,
    };
 
    static std::string toString(TensorDataTypes dt);
    static std::string toString(TensorTypes dt);
 
    Tensor(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           void* data,
           uint32_t elementTotalCount,
           uint32_t elementMemorySize,
           const TensorDataTypes& dataType,
           const TensorTypes& tensorType = TensorTypes::eDevice);
 
    virtual ~Tensor();
 
    void rebuild(void* data,
                 uint32_t elementTotalCount,
                 uint32_t elementMemorySize);
 
    void destroy();
 
    bool isInit();
 
    TensorTypes tensorType();
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Tensor> copyFromTensor);
 
    void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer);
 
    void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer);
 
    void recordPrimaryBufferMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask);
    void recordStagingBufferMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask);
 
    vk::DescriptorBufferInfo constructDescriptorBufferInfo();
 
    uint32_t size();
 
    uint32_t dataTypeMemorySize();
 
    uint32_t memorySize();
 
    TensorDataTypes dataType();
 
    void* rawData();
 
    void setRawData(const void* data);
 
    template<typename T>
    T* data()
    {
        return (T*)this->mRawData;
    }
 
    template<typename T>
    std::vector<T> vector()
    {
        return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
    }
 
  protected:
    // -------------- ALWAYS OWNED RESOURCES
    TensorTypes mTensorType;
    TensorDataTypes mDataType;
    uint32_t mSize;
    uint32_t mDataTypeMemorySize;
    void* mRawData;
 
  private:
    // -------------- NEVER OWNED RESOURCES
    std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice;
    std::shared_ptr<vk::Device> mDevice;
 
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::Buffer> mPrimaryBuffer;
    bool mFreePrimaryBuffer = false;
    std::shared_ptr<vk::Buffer> mStagingBuffer;
    bool mFreeStagingBuffer = false;
    std::shared_ptr<vk::DeviceMemory> mPrimaryMemory;
    bool mFreePrimaryMemory = false;
    std::shared_ptr<vk::DeviceMemory> mStagingMemory;
    bool mFreeStagingMemory = false;
 
    void allocateMemoryCreateGPUResources(); // Creates the vulkan buffer
    void createBuffer(std::shared_ptr<vk::Buffer> buffer,
                      vk::BufferUsageFlags bufferUsageFlags);
    void allocateBindMemory(std::shared_ptr<vk::Buffer> buffer,
                            std::shared_ptr<vk::DeviceMemory> memory,
                            vk::MemoryPropertyFlags memoryPropertyFlags);
    void recordCopyBuffer(const vk::CommandBuffer& commandBuffer,
                          std::shared_ptr<vk::Buffer> bufferFrom,
                          std::shared_ptr<vk::Buffer> bufferTo,
                          vk::DeviceSize bufferSize,
                          vk::BufferCopy copyRegion);
    void recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
                                   const vk::Buffer& buffer,
                                   vk::AccessFlagBits srcAccessMask,
                                   vk::AccessFlagBits dstAccessMask,
                                   vk::PipelineStageFlagBits srcStageMask,
                                   vk::PipelineStageFlagBits dstStageMask);
 
    // Private util functions
    vk::BufferUsageFlags getPrimaryBufferUsageFlags();
    vk::MemoryPropertyFlags getPrimaryMemoryPropertyFlags();
    vk::BufferUsageFlags getStagingBufferUsageFlags();
    vk::MemoryPropertyFlags getStagingMemoryPropertyFlags();
 
    void mapRawData();
    void unmapRawData();
};
 
template<typename T>
class TensorT : public Tensor
{
 
  public:
    TensorT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
            std::shared_ptr<vk::Device> device,
            const std::vector<T>& data,
            const TensorTypes& tensorType = TensorTypes::eDevice)
      : Tensor(physicalDevice,
               device,
               (void*)data.data(),
               data.size(),
               sizeof(T),
               this->dataType(),
               tensorType)
    {
        KP_LOG_DEBUG("Kompute TensorT constructor with data size {}",
                     data.size());
    }
 
    ~TensorT() { KP_LOG_DEBUG("Kompute TensorT destructor"); }
 
    T* data() { return (T*)this->mRawData; }
 
    std::vector<T> vector()
    {
        return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
    }
 
    T& operator[](int index) { return *(((T*)this->mRawData) + index); }
 
    void setData(const std::vector<T>& data)
    {
 
        KP_LOG_DEBUG("Kompute TensorT setting data with data size {}",
                     data.size());
 
        if (data.size() != this->mSize) {
            throw std::runtime_error(
              "Kompute TensorT Cannot set data of different sizes");
        }
 
        Tensor::setRawData(data.data());
    }
 
    TensorDataTypes dataType();
};
 
} // End namespace kp