Algorithm.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
 
#include "fmt/format.h"
#include "kompute/Tensor.hpp"
#include "logger/Logger.hpp"
 
namespace kp {
 
class Algorithm
{
  public:
    template<typename S = float, typename P = float>
    Algorithm(std::shared_ptr<vk::Device> device,
              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 Algorithm Constructor with device");
 
        this->mDevice = device;
 
        if (tensors.size() && spirv.size()) {
            KP_LOG_INFO(
              "Kompute Algorithm initialising with tensor size: {} and "
              "spirv size: {}",
              tensors.size(),
              spirv.size());
            this->rebuild(tensors,
                          spirv,
                          workgroup,
                          specializationConstants,
                          pushConstants);
        } else {
            KP_LOG_INFO(
              "Kompute Algorithm constructor with empty tensors and or "
              "spirv so not rebuilding vulkan components");
        }
    }
 
    template<typename S = float, typename P = float>
    void rebuild(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 Algorithm rebuild started");
 
        this->mTensors = tensors;
        this->mSpirv = spirv;
 
        if (specializationConstants.size()) {
            if (this->mSpecializationConstantsData) {
                free(this->mSpecializationConstantsData);
            }
            uint32_t memorySize =
              sizeof(decltype(specializationConstants.back()));
            uint32_t size = specializationConstants.size();
            uint32_t totalSize = size * memorySize;
            this->mSpecializationConstantsData = malloc(totalSize);
            memcpy(this->mSpecializationConstantsData,
                   specializationConstants.data(),
                   totalSize);
            this->mSpecializationConstantsDataTypeMemorySize = memorySize;
            this->mSpecializationConstantsSize = size;
        }
 
        if (pushConstants.size()) {
            if (this->mPushConstantsData) {
                free(this->mPushConstantsData);
            }
            uint32_t memorySize = sizeof(decltype(pushConstants.back()));
            uint32_t size = pushConstants.size();
            uint32_t totalSize = size * memorySize;
            this->mPushConstantsData = malloc(totalSize);
            memcpy(this->mPushConstantsData, pushConstants.data(), totalSize);
            this->mPushConstantsDataTypeMemorySize = memorySize;
            this->mPushConstantsSize = size;
        }
 
        this->setWorkgroup(
          workgroup, this->mTensors.size() ? this->mTensors[0]->size() : 1);
 
        // Descriptor pool is created first so if available then destroy all
        // before rebuild
        if (this->isInit()) {
            this->destroy();
        }
 
        this->createParameters();
        this->createShaderModule();
        this->createPipeline();
    }
 
    ~Algorithm();
 
    void recordDispatch(const vk::CommandBuffer& commandBuffer);
 
    void recordBindCore(const vk::CommandBuffer& commandBuffer);
 
    void recordBindPush(const vk::CommandBuffer& commandBuffer);
 
    bool isInit();
 
    void setWorkgroup(const Workgroup& workgroup, uint32_t minSize = 1);
    template<typename T>
    void setPushConstants(const std::vector<T>& pushConstants)
    {
        uint32_t memorySize = sizeof(decltype(pushConstants.back()));
        uint32_t size = pushConstants.size();
 
        this->setPushConstants(pushConstants.data(), size, memorySize);
    }
 
    void setPushConstants(void* data, uint32_t size, uint32_t memorySize)
    {
 
        uint32_t totalSize = memorySize * size;
        uint32_t previousTotalSize =
          this->mPushConstantsDataTypeMemorySize * this->mPushConstantsSize;
 
        if (totalSize != previousTotalSize) {
            throw std::runtime_error(fmt::format(
              "Kompute Algorithm push "
              "constant total memory size provided is {} but expected {} bytes",
              totalSize,
              previousTotalSize));
        }
        if (this->mPushConstantsData) {
            free(this->mPushConstantsData);
        }
 
        this->mPushConstantsData = malloc(totalSize);
        memcpy(this->mPushConstantsData, data, totalSize);
        this->mPushConstantsDataTypeMemorySize = memorySize;
        this->mPushConstantsSize = size;
    }
 
    const Workgroup& getWorkgroup();
    template<typename T>
    const std::vector<T> getSpecializationConstants()
    {
        return { (T*)this->mSpecializationConstantsData,
                 ((T*)this->mSpecializationConstantsData) +
                   this->mSpecializationConstantsSize };
    }
    template<typename T>
    const std::vector<T> getPushConstants()
    {
        return { (T*)this->mPushConstantsData,
                 ((T*)this->mPushConstantsData) + this->mPushConstantsSize };
    }
    const std::vector<std::shared_ptr<Tensor>>& getTensors();
 
    void destroy();
 
  private:
    // -------------- NEVER OWNED RESOURCES
    std::shared_ptr<vk::Device> mDevice;
    std::vector<std::shared_ptr<Tensor>> mTensors;
 
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::DescriptorSetLayout> mDescriptorSetLayout;
    bool mFreeDescriptorSetLayout = false;
    std::shared_ptr<vk::DescriptorPool> mDescriptorPool;
    bool mFreeDescriptorPool = false;
    std::shared_ptr<vk::DescriptorSet> mDescriptorSet;
    bool mFreeDescriptorSet = false;
    std::shared_ptr<vk::ShaderModule> mShaderModule;
    bool mFreeShaderModule = false;
    std::shared_ptr<vk::PipelineLayout> mPipelineLayout;
    bool mFreePipelineLayout = false;
    std::shared_ptr<vk::PipelineCache> mPipelineCache;
    bool mFreePipelineCache = false;
    std::shared_ptr<vk::Pipeline> mPipeline;
    bool mFreePipeline = false;
 
    // -------------- ALWAYS OWNED RESOURCES
    std::vector<uint32_t> mSpirv;
    void* mSpecializationConstantsData = nullptr;
    uint32_t mSpecializationConstantsDataTypeMemorySize = 0;
    uint32_t mSpecializationConstantsSize = 0;
    void* mPushConstantsData = nullptr;
    uint32_t mPushConstantsDataTypeMemorySize = 0;
    uint32_t mPushConstantsSize = 0;
    Workgroup mWorkgroup;
 
    // Create util functions
    void createShaderModule();
    void createPipeline();
 
    // Parameters
    void createParameters();
};
 
} // End namespace kp