Manager.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
 
#include "kompute/Image.hpp"
#include "kompute/Sequence.hpp"
#include "logger/Logger.hpp"
 
#define KP_DEFAULT_SESSION "DEFAULT"
 
namespace kp {
 
class Manager
{
  public:
    Manager();
 
    Manager(uint32_t physicalDeviceIndex,
            const std::vector<uint32_t>& familyQueueIndices = {},
            const std::vector<std::string>& desiredExtensions = {});
 
    Manager(std::shared_ptr<vk::Instance> instance,
            std::shared_ptr<vk::PhysicalDevice> physicalDevice,
            std::shared_ptr<vk::Device> device);
 
 
    Manager(const Manager&) = delete;
    Manager(const Manager&&) = delete;
    Manager& operator=(const Manager&) = delete;
    Manager& operator=(const Manager&&) = delete;
 
    ~Manager();
 
    std::shared_ptr<Sequence> sequence(uint32_t queueIndex = 0,
                                       uint32_t totalTimestamps = 0);
 
    template<typename T>
    std::shared_ptr<TensorT<T>> tensorT(
      const std::vector<T>& data,
      Memory::MemoryTypes tensorType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager tensor creation triggered");
 
        std::shared_ptr<TensorT<T>> tensor{ new kp::TensorT<T>(
          this->mPhysicalDevice, this->mDevice, data, tensorType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(tensor);
        }
 
        return tensor;
    }
 
    template<typename T>
    std::shared_ptr<TensorT<T>> tensorT(
      size_t size,
      Memory::MemoryTypes tensorType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager tensor creation triggered");
 
        std::shared_ptr<TensorT<T>> tensor{ new kp::TensorT<T>(
          this->mPhysicalDevice, this->mDevice, size, tensorType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(tensor);
        }
 
        return tensor;
    }
    std::shared_ptr<TensorT<float>> tensor(
      const std::vector<float>& data,
      Memory::MemoryTypes tensorType = Memory::MemoryTypes::eDevice)
    {
        return this->tensorT<float>(data, tensorType);
    }
 
    std::shared_ptr<Tensor> tensor(
      void* data,
      uint32_t elementTotalCount,
      uint32_t elementMemorySize,
      const Memory::DataTypes& dataType,
      Memory::MemoryTypes tensorType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Tensor> tensor{ new kp::Tensor(this->mPhysicalDevice,
                                                       this->mDevice,
                                                       data,
                                                       elementTotalCount,
                                                       elementMemorySize,
                                                       dataType,
                                                       tensorType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(tensor);
        }
 
        return tensor;
    }
 
    std::shared_ptr<Tensor> tensor(
      uint32_t elementTotalCount,
      uint32_t elementMemorySize,
      const Memory::DataTypes& dataType,
      Memory::MemoryTypes tensorType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Tensor> tensor{ new kp::Tensor(this->mPhysicalDevice,
                                                       this->mDevice,
                                                       elementTotalCount,
                                                       elementMemorySize,
                                                       dataType,
                                                       tensorType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(tensor);
        }
 
        return tensor;
    }
 
    template<typename T>
    std::shared_ptr<ImageT<T>> imageT(
      const std::vector<T>& data,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager image creation triggered");
 
        std::shared_ptr<ImageT<T>> image{ new kp::ImageT<T>(
          this->mPhysicalDevice,
          this->mDevice,
          data,
          width,
          height,
          numChannels,
          tiling,
          imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    template<typename T>
    std::shared_ptr<ImageT<T>> imageT(
      const std::vector<T>& data,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager image creation triggered");
 
        std::shared_ptr<ImageT<T>> image{ new kp::ImageT<T>(
          this->mPhysicalDevice,
          this->mDevice,
          data,
          width,
          height,
          numChannels,
          imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    template<typename T>
    std::shared_ptr<ImageT<T>> imageT(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager image creation triggered");
 
        std::shared_ptr<ImageT<T>> image{ new kp::ImageT<T>(
          this->mPhysicalDevice,
          this->mDevice,
          width,
          height,
          numChannels,
          tiling,
          imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    template<typename T>
    std::shared_ptr<ImageT<T>> imageT(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        KP_LOG_DEBUG("Kompute Manager image creation triggered");
 
        std::shared_ptr<ImageT<T>> image{ new kp::ImageT<T>(
          this->mPhysicalDevice,
          this->mDevice,
          width,
          height,
          numChannels,
          imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
    std::shared_ptr<ImageT<float>> image(
      const std::vector<float>& data,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        return this->imageT<float>(
          data, width, height, numChannels, tiling, imageType);
    }
 
    std::shared_ptr<ImageT<float>> image(
      const std::vector<float>& data,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        return this->imageT<float>(data, width, height, numChannels, imageType);
    }
 
    std::shared_ptr<ImageT<float>> image(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        return this->imageT<float>(
          width, height, numChannels, tiling, imageType);
    }
 
    std::shared_ptr<ImageT<float>> image(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        return this->imageT<float>(width, height, numChannels, imageType);
    }
 
    std::shared_ptr<Image> image(
      void* data,
      size_t dataSize,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      const Image::DataTypes& dataType,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Image> image{ new kp::Image(this->mPhysicalDevice,
                                                    this->mDevice,
                                                    data,
                                                    dataSize,
                                                    width,
                                                    height,
                                                    numChannels,
                                                    dataType,
                                                    tiling,
                                                    imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    std::shared_ptr<Image> image(
      void* data,
      size_t dataSize,
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      const Memory::DataTypes& dataType,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Image> image{ new kp::Image(this->mPhysicalDevice,
                                                    this->mDevice,
                                                    data,
                                                    dataSize,
                                                    width,
                                                    height,
                                                    numChannels,
                                                    dataType,
                                                    imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    std::shared_ptr<Image> image(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      const Memory::DataTypes& dataType,
      vk::ImageTiling tiling,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Image> image{ new kp::Image(this->mPhysicalDevice,
                                                    this->mDevice,
                                                    width,
                                                    height,
                                                    numChannels,
                                                    dataType,
                                                    tiling,
                                                    imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    std::shared_ptr<Image> image(
      uint32_t width,
      uint32_t height,
      uint32_t numChannels,
      const Memory::DataTypes& dataType,
      Memory::MemoryTypes imageType = Memory::MemoryTypes::eDevice)
    {
        std::shared_ptr<Image> image{ new kp::Image(this->mPhysicalDevice,
                                                    this->mDevice,
                                                    width,
                                                    height,
                                                    numChannels,
                                                    dataType,
                                                    imageType) };
 
        if (this->mManageResources) {
            this->mManagedMemObjects.push_back(image);
        }
 
        return image;
    }
 
    std::shared_ptr<Algorithm> algorithm(
      const std::vector<std::shared_ptr<Memory>>& memObjects = {},
      const std::vector<uint32_t>& spirv = {},
      const Workgroup& workgroup = {},
      const std::vector<float>& specializationConstants = {},
      const std::vector<float>& pushConstants = {})
    {
        return this->algorithm<>(
          memObjects, spirv, workgroup, specializationConstants, pushConstants);
    }
 
    template<typename S = float, typename P = float>
    std::shared_ptr<Algorithm> algorithm(
      const std::vector<std::shared_ptr<Memory>>& memObjects,
      const std::vector<uint32_t>& spirv,
      const Workgroup& workgroup,
      const std::vector<S>& specializationConstants,
      const std::vector<P>& pushConstants)
    {
 
        KP_LOG_DEBUG("Kompute Manager algorithm creation triggered");
 
        std::shared_ptr<Algorithm> algorithm{ new kp::Algorithm(
          this->mDevice,
          memObjects,
          spirv,
          workgroup,
          specializationConstants,
          pushConstants) };
 
        if (this->mManageResources) {
            this->mManagedAlgorithms.push_back(algorithm);
        }
 
        return algorithm;
    }
 
    void destroy();
    void clear();
 
    vk::PhysicalDeviceProperties getDeviceProperties() const;
 
    std::vector<vk::PhysicalDevice> listDevices() const;
 
    std::shared_ptr<vk::Instance> getVkInstance() const;
 
  private:
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::Instance> mInstance = nullptr;
    bool mFreeInstance = false;
    std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice = nullptr;
    std::shared_ptr<vk::Device> mDevice = nullptr;
    bool mFreeDevice = false;
 
    // -------------- ALWAYS OWNED RESOURCES
    std::vector<std::weak_ptr<Memory>> mManagedMemObjects;
    std::vector<std::weak_ptr<Sequence>> mManagedSequences;
    std::vector<std::weak_ptr<Algorithm>> mManagedAlgorithms;
 
    std::vector<uint32_t> mComputeQueueFamilyIndices;
    std::vector<std::shared_ptr<vk::Queue>> mComputeQueues;
 
    bool mManageResources = false;
 
#ifndef KOMPUTE_DISABLE_VK_DEBUG_LAYERS
    vk::DebugReportCallbackEXT mDebugReportCallback;
#ifdef VK_VERSION_1_4
    vk::detail::DispatchLoaderDynamic mDebugDispatcher;
#else
    vk::DispatchLoaderDynamic mDebugDispatcher;
#endif // VK_VERSION_1_4
#endif // KOMPUTE_DISABLE_VK_DEBUG_LAYERS
 
    // Create functions
    void createInstance();
    void createDevice(const std::vector<uint32_t>& familyQueueIndices = {},
                      uint32_t hysicalDeviceIndex = 0,
                      const std::vector<std::string>& desiredExtensions = {});
};
 
} // End namespace kp