Tensor.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
#include "kompute/Memory.hpp"
#include "logger/Logger.hpp"
#include <memory>
#include <string>
 
namespace kp {
 
// Forward-declare the Image class
class Image;
class Tensor : public Memory
{
  public:
    Tensor(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           void* data,
           uint32_t elementTotalCount,
           uint32_t elementMemorySize,
           const DataTypes& dataType,
           const MemoryTypes& tensorType = MemoryTypes::eDevice);
 
    Tensor(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           uint32_t elementTotalCount,
           uint32_t elementMemorySize,
           const DataTypes& dataType,
           const MemoryTypes& memoryType = MemoryTypes::eDevice);
 
    Tensor(const Tensor&) = delete;
    Tensor(const Tensor&&) = delete;
    Tensor& operator=(const Tensor&) = delete;
    Tensor& operator=(const Tensor&&) = delete;
 
    virtual ~Tensor();
 
    void destroy() override;
 
    bool isInit() override;
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Tensor> copyFromTensor) override;
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Image> copyFromImage) override;
 
    void recordCopyFromStagingToDevice(
      const vk::CommandBuffer& commandBuffer) override;
 
    void recordCopyFromDeviceToStaging(
      const vk::CommandBuffer& commandBuffer) override;
 
    void recordPrimaryMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask) override;
    void recordStagingMemoryBarrier(
      const vk::CommandBuffer& commandBuffer,
      vk::AccessFlagBits srcAccessMask,
      vk::AccessFlagBits dstAccessMask,
      vk::PipelineStageFlagBits srcStageMask,
      vk::PipelineStageFlagBits dstStageMask) override;
 
    vk::WriteDescriptorSet constructDescriptorSet(
      vk::DescriptorSet descriptorSet,
      uint32_t binding) override;
 
    std::shared_ptr<vk::Buffer> getPrimaryBuffer();
 
    Type type() override { return Type::eTensor; }
 
  protected:
    // -------------- ALWAYS OWNED RESOURCES
    vk::DescriptorBufferInfo mDescriptorBufferInfo;
 
  private:
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::Buffer> mPrimaryBuffer;
    bool mFreePrimaryBuffer = false;
    std::shared_ptr<vk::Buffer> mStagingBuffer;
    bool mFreeStagingBuffer = 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 recordCopyBufferFromImage(const vk::CommandBuffer& commandBuffer,
                                   std::shared_ptr<vk::Image> imageFrom,
                                   std::shared_ptr<vk::Buffer> bufferTo,
                                   vk::ImageLayout fromLayout,
                                   vk::DeviceSize /*bufferSize*/,
                                   vk::BufferImageCopy 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::BufferUsageFlags getStagingBufferUsageFlags();
 
    vk::DescriptorBufferInfo constructDescriptorBufferInfo();
 
    void reserve();
};
 
template<typename T>
class TensorT : public Tensor
{
  public:
    TensorT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
            std::shared_ptr<vk::Device> device,
            const size_t size,
            const MemoryTypes& tensorType = MemoryTypes::eDevice)
      : Tensor(physicalDevice,
               device,
               size,
               sizeof(T),
               Memory::dataType<T>(),
               tensorType)
    {
        KP_LOG_DEBUG("Kompute TensorT constructor with data size {}", size);
    }
 
    TensorT(
      std::shared_ptr<vk::PhysicalDevice> physicalDevice,
      std::shared_ptr<vk::Device> device,
      const std::vector<T>& data,
      const Memory::MemoryTypes& tensorType = Memory::MemoryTypes::eDevice)
      : Tensor(physicalDevice,
               device,
               (void*)data.data(),
               static_cast<uint32_t>(data.size()),
               sizeof(T),
               Memory::dataType<T>(),
               tensorType)
    {
        KP_LOG_DEBUG("Kompute TensorT filling constructor with data size {}",
                     data.size());
    }
 
    TensorT(const TensorT&) = delete;
    TensorT(const TensorT&&) = delete;
    TensorT& operator=(const TensorT&) = delete;
    TensorT& operator=(const TensorT&&) = delete;
 
 
    ~TensorT() { KP_LOG_DEBUG("Kompute TensorT destructor"); }
 
    DataTypes dataType() { return Memory::dataType<T>(); }
    std::vector<T> vector() { return Memory::vector<T>(); }
    T* data() { return Memory::data<T>(); }
};
 
} // End namespace kp