Manager.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <set>
#include <unordered_map>
 
#include "kompute/Core.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();
 
    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,
      Tensor::TensorTypes tensorType = Tensor::TensorTypes::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->mManagedTensors.push_back(tensor);
        }
 
        return tensor;
    }
 
    std::shared_ptr<TensorT<float>> tensor(
      const std::vector<float>& data,
      Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)
    {
        return this->tensorT<float>(data, tensorType);
    }
 
    std::shared_ptr<Tensor> tensor(
      void* data,
      uint32_t elementTotalCount,
      uint32_t elementMemorySize,
      const Tensor::TensorDataTypes& dataType,
      Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)
    {
        std::shared_ptr<Tensor> tensor{ new kp::Tensor(this->mPhysicalDevice,
                                                       this->mDevice,
                                                       data,
                                                       elementTotalCount,
                                                       elementMemorySize,
                                                       dataType,
                                                       tensorType) };
 
        if (this->mManageResources) {
            this->mManagedTensors.push_back(tensor);
        }
 
        return tensor;
    }
 
    std::shared_ptr<Algorithm> algorithm(
      const std::vector<std::shared_ptr<Tensor>>& tensors = {},
      const std::vector<uint32_t>& spirv = {},
      const Workgroup& workgroup = {},
      const std::vector<float>& specializationConstants = {},
      const std::vector<float>& pushConstants = {})
    {
        return this->algorithm<>(
          tensors, spirv, workgroup, specializationConstants, pushConstants);
    }
 
    template<typename S = float, typename P = float>
    std::shared_ptr<Algorithm> algorithm(
      const std::vector<std::shared_ptr<Tensor>>& tensors,
      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,
          tensors,
          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<Tensor>> mManagedTensors;
    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;
    vk::DispatchLoaderDynamic mDebugDispatcher;
#endif
 
    // 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