text: // SPDX-License-Identifier: Apache-2.0
  #pragma once
  
  #include "kompute/Core.hpp"
  
  #include "kompute/Tensor.hpp"
  
  namespace kp {
  
  /**
      Abstraction for compute shaders that are run on top of tensors grouped via
     ParameterGroups (which group descriptorsets)
  */
  class Algorithm
  {
    public:
      /**
       *  Main constructor for algorithm with configuration parameters to create
       *  the underlying resources.
       *
       *  @param device The Vulkan device to use for creating resources
       *  @param tensors (optional) The tensors to use to create the descriptor
       * resources
       *  @param spirv (optional) The spirv code to use to create the algorithm
       *  @param workgroup (optional) The kp::Workgroup to use for the dispatch
       * which defaults to kp::Workgroup(tensor[0].size(), 1, 1) if not set.
       *  @param specializationConstants (optional) The templatable param is to be used to
       * initialize the specialization constants which cannot be changed once set.
       *  @param pushConstants (optional) This templatable param is to be used when
       * initializing the pipeline, which set the size of the push constants -
       * these can be modified but all new values must have the same data type and length
       * as otherwise it will result in errors.
       */
      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");
          }
      }
  
      /**
       *  Rebuild function to reconstruct algorithm with configuration parameters
       * to create the underlying resources.
       *
       *  @param tensors The tensors to use to create the descriptor resources
       *  @param spirv The spirv code to use to create the algorithm
       *  @param workgroup (optional) The kp::Workgroup to use for the dispatch
       * which defaults to kp::Workgroup(tensor[0].size(), 1, 1) if not set.
       *  @param specializationConstants (optional) The std::vector<float> to use to
       * initialize the specialization constants which cannot be changed once set.
       *  @param pushConstants (optional) The std::vector<float> to use when
       * initializing the pipeline, which set the size of the push constants -
       * these can be modified but all new values must have the same vector size
       * as this initial value.
       */
      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();
      }
  
      /**
       * Destructor for Algorithm which is responsible for freeing and desroying
       * respective pipelines and owned parameter groups.
       */
      ~Algorithm();
  
      /**
       * Records the dispatch function with the provided template parameters or
       * alternatively using the size of the tensor by default.
       *
       * @param commandBuffer Command buffer to record the algorithm resources to
       */
      void recordDispatch(const vk::CommandBuffer& commandBuffer);
  
      /**
       * Records command that binds the "core" algorithm components which consist
       * of binding the pipeline and binding the descriptorsets.
       *
       * @param commandBuffer Command buffer to record the algorithm resources to
       */
      void recordBindCore(const vk::CommandBuffer& commandBuffer);
  
      /**
       * Records command that binds the push constants to the command buffer
       * provided
       * - it is required that the pushConstants provided are of the same size as
       * the ones provided during initialization.
       *
       * @param commandBuffer Command buffer to record the algorithm resources to
       */
      void recordBindPush(const vk::CommandBuffer& commandBuffer);
  
      /**
       * function that checks all the gpu resource components to verify if these
       * have been created and returns true if all are valid.
       *
       * @returns returns true if the algorithm is currently initialized.
       */
      bool isInit();
  
      /**
       * Sets the work group to use in the recordDispatch
       *
       * @param workgroup The kp::Workgroup value to use to update the algorithm.
       * It must have a value greater than 1 on the x value (index 1) otherwise it
       * will be initialized on the size of the first tensor (ie.
       * this->mTensor[0]->size())
       */
      void setWorkgroup(const Workgroup& workgroup, uint32_t minSize = 1);
      /**
       * Sets the push constants to the new value provided to use in the next
       * bindPush()
       *
       * @param pushConstants The templatable vector is to be used to set the push constants to use in the
       * next bindPush(...) calls. The constants provided must be of the same size
       * as the ones created during initialization.
       */
      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);
      }
  
      /**
       * Sets the push constants to the new value provided to use in the next
       * bindPush() with the raw memory block location and memory size to be used.
       *
       * @param data The raw data point to copy the data from, without modifying the pointer.
       * @param size The number of data elements provided in the data
       * @param memorySize The memory size of each of the data elements in bytes.
       */
      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;
      }
  
      /**
       * Gets the current workgroup from the algorithm.
       *
       * @param The kp::Constant to use to set the push constants to use in the
       * next bindPush(...) calls. The constants provided must be of the same size
       * as the ones created during initialization.
       */
      const Workgroup& getWorkgroup();
      /**
       * Gets the specialization constants of the current algorithm.
       *
       * @returns The std::vector<float> currently set for specialization constants
       */
      template<typename T>
      const std::vector<T> getSpecializationConstants()
      {
          return { (T*)this->mSpecializationConstantsData,
              ((T*)this->mSpecializationConstantsData) + this->mSpecializationConstantsSize };
      }
      /**
       * Gets the specialization constants of the current algorithm.
       *
       * @returns The std::vector<float> currently set for push constants
       */
      template<typename T>
      const std::vector<T> getPushConstants()
      {
          return { (T*)this->mPushConstantsData,
              ((T*)this->mPushConstantsData) + this->mPushConstantsSize };
      }
      /**
       * Gets the current tensors that are used in the algorithm.
       *
       * @returns The list of tensors used in the algorithm.
       */
      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
  