Image.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
#include "kompute/Memory.hpp"
#include "kompute/Tensor.hpp"
#include "logger/Logger.hpp"
#include <memory>
#include <string>
 
namespace kp {
 
class Image : public Memory
{
  public:
    Image(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
          std::shared_ptr<vk::Device> device,
          void* data,
          size_t dataSize,
          uint32_t x,
          uint32_t y,
          uint32_t numChannels,
          const DataTypes& dataType,
          vk::ImageTiling tiling,
          const MemoryTypes& memoryType = MemoryTypes::eDevice)
      : Memory(physicalDevice, device, dataType, memoryType, x, y)
    {
        if (dataType == DataTypes::eCustom) {
            throw std::runtime_error(
              "Custom data types are not supported for Kompute Images");
        }
 
        init(data, dataSize, numChannels, tiling);
    }
 
    Image(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
          std::shared_ptr<vk::Device> device,
          uint32_t x,
          uint32_t y,
          uint32_t numChannels,
          const DataTypes& dataType,
          vk::ImageTiling tiling,
          const MemoryTypes& memoryType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              nullptr,
              0,
              x,
              y,
              numChannels,
              dataType,
              tiling,
              memoryType)
    {
    }
 
    Image(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
          std::shared_ptr<vk::Device> device,
          void* data,
          size_t dataSize,
          uint32_t x,
          uint32_t y,
          uint32_t numChannels,
          const DataTypes& dataType,
          const MemoryTypes& memoryType = MemoryTypes::eDevice)
      : Memory(physicalDevice, device, dataType, memoryType, x, y)
    {
        vk::ImageTiling tiling;
 
        if (dataType == DataTypes::eCustom) {
            throw std::runtime_error(
              "Custom data types are not supported for Kompute Images");
        }
 
        if (memoryType == MemoryTypes::eHost ||
            memoryType == MemoryTypes::eDeviceAndHost) {
            // Host-accessible memory must be linear-tiled.
            tiling = vk::ImageTiling::eLinear;
        } else if (memoryType == MemoryTypes::eDevice ||
                   memoryType == MemoryTypes::eStorage) {
            tiling = vk::ImageTiling::eOptimal;
        } else {
            throw std::runtime_error("Kompute Image unsupported memory type");
        }
 
        init(data, dataSize, numChannels, tiling);
    }
 
    Image(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
          std::shared_ptr<vk::Device> device,
          uint32_t x,
          uint32_t y,
          uint32_t numChannels,
          const DataTypes& dataType,
          const MemoryTypes& memoryType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              nullptr,
              0,
              x,
              y,
              numChannels,
              dataType,
              memoryType)
    {
    }
 
 
    Image(const Image&) = delete;
    Image(const Image&&) = delete;
    Image& operator=(const Image&) = delete;
    Image& operator=(const Image&&) = delete;
 
    virtual ~Image();
 
    void destroy() override;
 
    bool isInit() override;
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Image> copyFromImage) override;
 
    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                        std::shared_ptr<Tensor> copyFromTensor) 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;
 
    void recordPrimaryImageBarrier(const vk::CommandBuffer& commandBuffer,
                                   vk::AccessFlagBits srcAccessMask,
                                   vk::AccessFlagBits dstAccessMask,
                                   vk::PipelineStageFlagBits srcStageMask,
                                   vk::PipelineStageFlagBits dstStageMask,
                                   vk::ImageLayout dstLayout);
 
    vk::WriteDescriptorSet constructDescriptorSet(
      vk::DescriptorSet descriptorSet,
      uint32_t binding) override;
 
    std::shared_ptr<vk::Image> getPrimaryImage();
    vk::ImageLayout getPrimaryImageLayout();
 
    /***
     * Retreive the number of channels in the image
     *
     * @return Number of channels in the image
     */
    uint32_t getNumChannels();
 
    Type type() override { return Type::eImage; }
 
  protected:
    // -------------- ALWAYS OWNED RESOURCES
    uint32_t mNumChannels;
    vk::DescriptorImageInfo mDescriptorImageInfo;
    vk::ImageLayout mPrimaryImageLayout = vk::ImageLayout::eUndefined;
    vk::ImageLayout mStagingImageLayout = vk::ImageLayout::eUndefined;
    std::shared_ptr<vk::ImageView> mImageView = nullptr;
    vk::ImageTiling mTiling = vk::ImageTiling::eOptimal;
 
  private:
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::Image> mPrimaryImage;
    bool mFreePrimaryImage = false;
    std::shared_ptr<vk::Image> mStagingImage;
    bool mFreeStagingImage = false;
 
    void allocateMemoryCreateGPUResources(); // Creates the vulkan image
    void createImage(std::shared_ptr<vk::Image> image,
                     vk::ImageUsageFlags imageUsageFlags,
                     vk::ImageTiling imageTiling);
    void allocateBindMemory(std::shared_ptr<vk::Image> image,
                            std::shared_ptr<vk::DeviceMemory> memory,
                            vk::MemoryPropertyFlags memoryPropertyFlags);
    void recordCopyImage(const vk::CommandBuffer& commandBuffer,
                         std::shared_ptr<vk::Image> srcImage,
                         std::shared_ptr<vk::Image> dstImage,
                         vk::ImageLayout srcLayout,
                         vk::ImageLayout dstLayout,
                         vk::ImageCopy copyRegion);
    void recordCopyImageFromTensor(const vk::CommandBuffer& commandBuffer,
                                   std::shared_ptr<vk::Buffer> srcBuffer,
                                   std::shared_ptr<vk::Image> dstImage,
                                   vk::ImageLayout dstLayout,
                                   vk::BufferImageCopy copyRegion);
 
    void recordStagingImageBarrier(const vk::CommandBuffer& commandBuffer,
                                   vk::AccessFlagBits srcAccessMask,
                                   vk::AccessFlagBits dstAccessMask,
                                   vk::PipelineStageFlagBits srcStageMask,
                                   vk::PipelineStageFlagBits dstStageMask,
                                   vk::ImageLayout dstLayout);
 
    void recordImageMemoryBarrier(const vk::CommandBuffer& commandBuffer,
                                  const vk::Image& image,
                                  vk::AccessFlagBits srcAccessMask,
                                  vk::AccessFlagBits dstAccessMask,
                                  vk::PipelineStageFlagBits srcStageMask,
                                  vk::PipelineStageFlagBits dstStageMask,
                                  vk::ImageLayout oldLayout,
                                  vk::ImageLayout newLayout);
 
    // Private util functions
    vk::ImageUsageFlags getPrimaryImageUsageFlags();
    vk::ImageUsageFlags getStagingImageUsageFlags();
 
    vk::Format getFormat();
 
    vk::DescriptorImageInfo constructDescriptorImageInfo();
 
    void init(void* data,
              size_t dataSize,
              uint32_t numChannels,
              vk::ImageTiling tiling);
    void reserve();
};
 
template<typename T>
class ImageT : public Image
{
 
  public:
    ImageT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           const std::vector<T>& data,
           uint32_t x,
           uint32_t y,
           uint32_t numChannels,
           vk::ImageTiling tiling,
           const MemoryTypes& imageType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              (void*)data.data(),
              data.size(),
              x,
              y,
              numChannels,
              Memory::dataType<T>(),
              tiling,
              imageType)
    {
        // Images cannot be created with custom types
        static_assert(Memory::dataType<T>() != DataTypes::eCustom,
                      "Custom data types are not supported for Kompute Images");
 
        KP_LOG_DEBUG("Kompute imageT constructor with data size {}, x {}, "
                     "y {}, and num channels {}",
                     data.size(),
                     x,
                     y,
                     numChannels);
    }
 
    ImageT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           const std::vector<T>& data,
           uint32_t x,
           uint32_t y,
           uint32_t numChannels,
           const MemoryTypes& imageType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              (void*)data.data(),
              data.size(),
              x,
              y,
              numChannels,
              Memory::dataType<T>(),
              imageType)
    {
        // Images cannot be created with custom types
        static_assert(Memory::dataType<T>() != DataTypes::eCustom,
                      "Custom data types are not supported for Kompute Images");
 
        KP_LOG_DEBUG("Kompute imageT constructor with data size {}, x {}, "
                     "y {}, and num channels {}",
                     data.size(),
                     x,
                     y,
                     numChannels);
    }
 
    ImageT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           uint32_t x,
           uint32_t y,
           uint32_t numChannels,
           vk::ImageTiling tiling,
           const MemoryTypes& imageType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              x,
              y,
              numChannels,
              Memory::dataType<T>(),
              tiling,
              imageType)
    {
        // Images cannot be created with custom types
        static_assert(Memory::dataType<T>() != DataTypes::eCustom,
                      "Custom data types are not supported for Kompute Images");
 
        KP_LOG_DEBUG("Kompute imageT constructor with no data, x {}, "
                     "y {}, and num channels {}",
                     x,
                     y,
                     numChannels);
    }
 
    ImageT(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
           std::shared_ptr<vk::Device> device,
           uint32_t x,
           uint32_t y,
           uint32_t numChannels,
           const MemoryTypes& imageType = MemoryTypes::eDevice)
      : Image(physicalDevice,
              device,
              x,
              y,
              numChannels,
              Memory::dataType<T>(),
              imageType)
    {
        // Images cannot be created with custom types
        static_assert(Memory::dataType<T>() != DataTypes::eCustom,
                      "Custom data types are not supported for Kompute Images");
 
        KP_LOG_DEBUG("Kompute imageT constructor with no data, x {}, "
                     "y {}, and num channels {}",
                     x,
                     y,
                     numChannels);
    }
 
    ~ImageT() { KP_LOG_DEBUG("Kompute imageT destructor"); }
 
    std::vector<T> vector() { return Memory::vector<T>(); }
 
    T& operator[](int index) { return *(Memory::data<T>() + index); }
};
 
} // End namespace kp