Line data    Source code

       1             : // SPDX-License-Identifier: Apache-2.0
       2             : #pragma once
       3             : 
       4             : #include "kompute/Core.hpp"
       5             : 
       6             : #include "fmt/format.h"
       7             : #include "kompute/Tensor.hpp"
       8             : #include "logger/Logger.hpp"
       9             : 
      10             : namespace kp {
      11             : 
      12             : /**
      13             :     Abstraction for compute shaders that are run on top of tensors grouped via
      14             :    ParameterGroups (which group descriptorsets)
      15             : */
      16             : class Algorithm
      17             : {
      18             :   public:
      19             :     /**
      20             :      *  Main constructor for algorithm with configuration parameters to create
      21             :      *  the underlying resources.
      22             :      *
      23             :      *  @param device The Vulkan device to use for creating resources
      24             :      *  @param tensors (optional) The tensors to use to create the descriptor
      25             :      * resources
      26             :      *  @param spirv (optional) The spirv code to use to create the algorithm
      27             :      *  @param workgroup (optional) The kp::Workgroup to use for the dispatch
      28             :      * which defaults to kp::Workgroup(tensor[0].size(), 1, 1) if not set.
      29             :      *  @param specializationConstants (optional) The templatable param is to be
      30             :      * used to initialize the specialization constants which cannot be changed
      31             :      * once set.
      32             :      *  @param pushConstants (optional) This templatable param is to be used
      33             :      * when initializing the pipeline, which set the size of the push constants
      34             :      * - these can be modified but all new values must have the same data type
      35             :      * and length as otherwise it will result in errors.
      36             :      */
      37             :     template<typename S = float, typename P = float>
      38          37 :     Algorithm(std::shared_ptr<vk::Device> device,
      39             :               const std::vector<std::shared_ptr<Tensor>>& tensors = {},
      40             :               const std::vector<uint32_t>& spirv = {},
      41             :               const Workgroup& workgroup = {},
      42             :               const std::vector<S>& specializationConstants = {},
      43             :               const std::vector<P>& pushConstants = {})
      44          37 :     {
      45          74 :         KP_LOG_DEBUG("Kompute Algorithm Constructor with device");
      46             : 
      47          37 :         this->mDevice = device;
      48             : 
      49          37 :         if (tensors.size() && spirv.size()) {
      50          99 :             KP_LOG_INFO(
      51             :               "Kompute Algorithm initialising with tensor size: {} and "
      52             :               "spirv size: {}",
      53             :               tensors.size(),
      54             :               spirv.size());
      55          33 :             this->rebuild(tensors,
      56             :                           spirv,
      57             :                           workgroup,
      58             :                           specializationConstants,
      59             :                           pushConstants);
      60             :         } else {
      61           8 :             KP_LOG_INFO(
      62             :               "Kompute Algorithm constructor with empty tensors and or "
      63             :               "spirv so not rebuilding vulkan components");
      64             :         }
      65          37 :     }
      66             : 
      67             :     /**
      68             :      *  Rebuild function to reconstruct algorithm with configuration parameters
      69             :      * to create the underlying resources.
      70             :      *
      71             :      *  @param tensors The tensors to use to create the descriptor resources
      72             :      *  @param spirv The spirv code to use to create the algorithm
      73             :      *  @param workgroup (optional) The kp::Workgroup to use for the dispatch
      74             :      * which defaults to kp::Workgroup(tensor[0].size(), 1, 1) if not set.
      75             :      *  @param specializationConstants (optional) The std::vector<float> to use
      76             :      * to initialize the specialization constants which cannot be changed once
      77             :      * set.
      78             :      *  @param pushConstants (optional) The std::vector<float> to use when
      79             :      * initializing the pipeline, which set the size of the push constants -
      80             :      * these can be modified but all new values must have the same vector size
      81             :      * as this initial value.
      82             :      */
      83             :     template<typename S = float, typename P = float>
      84          38 :     void rebuild(const std::vector<std::shared_ptr<Tensor>>& tensors,
      85             :                  const std::vector<uint32_t>& spirv,
      86             :                  const Workgroup& workgroup = {},
      87             :                  const std::vector<S>& specializationConstants = {},
      88             :                  const std::vector<P>& pushConstants = {})
      89             :     {
      90          76 :         KP_LOG_DEBUG("Kompute Algorithm rebuild started");
      91             : 
      92          38 :         this->mTensors = tensors;
      93          38 :         this->mSpirv = spirv;
      94             : 
      95          38 :         if (specializationConstants.size()) {
      96           6 :             if (this->mSpecializationConstantsData) {
      97           0 :                 free(this->mSpecializationConstantsData);
      98             :             }
      99           6 :             uint32_t memorySize =
     100             :               sizeof(decltype(specializationConstants.back()));
     101           6 :             uint32_t size = specializationConstants.size();
     102           6 :             uint32_t totalSize = size * memorySize;
     103           6 :             this->mSpecializationConstantsData = malloc(totalSize);
     104          12 :             memcpy(this->mSpecializationConstantsData,
     105           6 :                    specializationConstants.data(),
     106             :                    totalSize);
     107           6 :             this->mSpecializationConstantsDataTypeMemorySize = memorySize;
     108           6 :             this->mSpecializationConstantsSize = size;
     109             :         }
     110             : 
     111          38 :         if (pushConstants.size()) {
     112           8 :             if (this->mPushConstantsData) {
     113           0 :                 free(this->mPushConstantsData);
     114             :             }
     115           8 :             uint32_t memorySize = sizeof(decltype(pushConstants.back()));
     116           8 :             uint32_t size = pushConstants.size();
     117           8 :             uint32_t totalSize = size * memorySize;
     118           8 :             this->mPushConstantsData = malloc(totalSize);
     119           8 :             memcpy(this->mPushConstantsData, pushConstants.data(), totalSize);
     120           8 :             this->mPushConstantsDataTypeMemorySize = memorySize;
     121           8 :             this->mPushConstantsSize = size;
     122             :         }
     123             : 
     124          76 :         this->setWorkgroup(
     125          76 :           workgroup, this->mTensors.size() ? this->mTensors[0]->size() : 1);
     126             : 
     127             :         // Descriptor pool is created first so if available then destroy all
     128             :         // before rebuild
     129          38 :         if (this->isInit()) {
     130           1 :             this->destroy();
     131             :         }
     132             : 
     133          38 :         this->createParameters();
     134          38 :         this->createShaderModule();
     135          38 :         this->createPipeline();
     136          38 :     }
     137             : 
     138             :     /**
     139             :      * Destructor for Algorithm which is responsible for freeing and desroying
     140             :      * respective pipelines and owned parameter groups.
     141             :      */
     142             :     ~Algorithm();
     143             : 
     144             :     /**
     145             :      * Records the dispatch function with the provided template parameters or
     146             :      * alternatively using the size of the tensor by default.
     147             :      *
     148             :      * @param commandBuffer Command buffer to record the algorithm resources to
     149             :      */
     150             :     void recordDispatch(const vk::CommandBuffer& commandBuffer);
     151             : 
     152             :     /**
     153             :      * Records command that binds the "core" algorithm components which consist
     154             :      * of binding the pipeline and binding the descriptorsets.
     155             :      *
     156             :      * @param commandBuffer Command buffer to record the algorithm resources to
     157             :      */
     158             :     void recordBindCore(const vk::CommandBuffer& commandBuffer);
     159             : 
     160             :     /**
     161             :      * Records command that binds the push constants to the command buffer
     162             :      * provided
     163             :      * - it is required that the pushConstants provided are of the same size as
     164             :      * the ones provided during initialization.
     165             :      *
     166             :      * @param commandBuffer Command buffer to record the algorithm resources to
     167             :      */
     168             :     void recordBindPush(const vk::CommandBuffer& commandBuffer);
     169             : 
     170             :     /**
     171             :      * function that checks all the gpu resource components to verify if these
     172             :      * have been created and returns true if all are valid.
     173             :      *
     174             :      * @returns returns true if the algorithm is currently initialized.
     175             :      */
     176             :     bool isInit();
     177             : 
     178             :     /**
     179             :      * Sets the work group to use in the recordDispatch
     180             :      *
     181             :      * @param workgroup The kp::Workgroup value to use to update the algorithm.
     182             :      * It must have a value greater than 1 on the x value (index 1) otherwise it
     183             :      * will be initialized on the size of the first tensor (ie.
     184             :      * this->mTensor[0]->size())
     185             :      */
     186             :     void setWorkgroup(const Workgroup& workgroup, uint32_t minSize = 1);
     187             :     /**
     188             :      * Sets the push constants to the new value provided to use in the next
     189             :      * bindPush()
     190             :      *
     191             :      * @param pushConstants The templatable vector is to be used to set the push
     192             :      * constants to use in the next bindPush(...) calls. The constants provided
     193             :      * must be of the same size as the ones created during initialization.
     194             :      */
     195             :     template<typename T>
     196             :     void setPushConstants(const std::vector<T>& pushConstants)
     197             :     {
     198             :         uint32_t memorySize = sizeof(decltype(pushConstants.back()));
     199             :         uint32_t size = pushConstants.size();
     200             : 
     201             :         this->setPushConstants(pushConstants.data(), size, memorySize);
     202             :     }
     203             : 
     204             :     /**
     205             :      * Sets the push constants to the new value provided to use in the next
     206             :      * bindPush() with the raw memory block location and memory size to be used.
     207             :      *
     208             :      * @param data The raw data point to copy the data from, without modifying
     209             :      * the pointer.
     210             :      * @param size The number of data elements provided in the data
     211             :      * @param memorySize The memory size of each of the data elements in bytes.
     212             :      */
     213          11 :     void setPushConstants(void* data, uint32_t size, uint32_t memorySize)
     214             :     {
     215             : 
     216          11 :         uint32_t totalSize = memorySize * size;
     217          11 :         uint32_t previousTotalSize =
     218          11 :           this->mPushConstantsDataTypeMemorySize * this->mPushConstantsSize;
     219             : 
     220          11 :         if (totalSize != previousTotalSize) {
     221           2 :             throw std::runtime_error(fmt::format(
     222             :               "Kompute Algorithm push "
     223             :               "constant total memory size provided is {} but expected {} bytes",
     224             :               totalSize,
     225           2 :               previousTotalSize));
     226             :         }
     227          10 :         if (this->mPushConstantsData) {
     228          10 :             free(this->mPushConstantsData);
     229             :         }
     230             : 
     231          10 :         this->mPushConstantsData = malloc(totalSize);
     232          10 :         memcpy(this->mPushConstantsData, data, totalSize);
     233          10 :         this->mPushConstantsDataTypeMemorySize = memorySize;
     234          10 :         this->mPushConstantsSize = size;
     235          10 :     }
     236             : 
     237             :     /**
     238             :      * Gets the current workgroup from the algorithm.
     239             :      *
     240             :      * @param The kp::Constant to use to set the push constants to use in the
     241             :      * next bindPush(...) calls. The constants provided must be of the same size
     242             :      * as the ones created during initialization.
     243             :      */
     244             :     const Workgroup& getWorkgroup();
     245             :     /**
     246             :      * Gets the specialization constants of the current algorithm.
     247             :      *
     248             :      * @returns The std::vector<float> currently set for specialization
     249             :      * constants
     250             :      */
     251             :     template<typename T>
     252           1 :     const std::vector<T> getSpecializationConstants()
     253             :     {
     254           1 :         return { (T*)this->mSpecializationConstantsData,
     255           1 :                  ((T*)this->mSpecializationConstantsData) +
     256           1 :                    this->mSpecializationConstantsSize };
     257             :     }
     258             :     /**
     259             :      * Gets the specialization constants of the current algorithm.
     260             :      *
     261             :      * @returns The std::vector<float> currently set for push constants
     262             :      */
     263             :     template<typename T>
     264           1 :     const std::vector<T> getPushConstants()
     265             :     {
     266           1 :         return { (T*)this->mPushConstantsData,
     267           1 :                  ((T*)this->mPushConstantsData) + this->mPushConstantsSize };
     268             :     }
     269             :     /**
     270             :      * Gets the current tensors that are used in the algorithm.
     271             :      *
     272             :      * @returns The list of tensors used in the algorithm.
     273             :      */
     274             :     const std::vector<std::shared_ptr<Tensor>>& getTensors();
     275             : 
     276             :     void destroy();
     277             : 
     278             :   private:
     279             :     // -------------- NEVER OWNED RESOURCES
     280             :     std::shared_ptr<vk::Device> mDevice;
     281             :     std::vector<std::shared_ptr<Tensor>> mTensors;
     282             : 
     283             :     // -------------- OPTIONALLY OWNED RESOURCES
     284             :     std::shared_ptr<vk::DescriptorSetLayout> mDescriptorSetLayout;
     285             :     bool mFreeDescriptorSetLayout = false;
     286             :     std::shared_ptr<vk::DescriptorPool> mDescriptorPool;
     287             :     bool mFreeDescriptorPool = false;
     288             :     std::shared_ptr<vk::DescriptorSet> mDescriptorSet;
     289             :     bool mFreeDescriptorSet = false;
     290             :     std::shared_ptr<vk::ShaderModule> mShaderModule;
     291             :     bool mFreeShaderModule = false;
     292             :     std::shared_ptr<vk::PipelineLayout> mPipelineLayout;
     293             :     bool mFreePipelineLayout = false;
     294             :     std::shared_ptr<vk::PipelineCache> mPipelineCache;
     295             :     bool mFreePipelineCache = false;
     296             :     std::shared_ptr<vk::Pipeline> mPipeline;
     297             :     bool mFreePipeline = false;
     298             : 
     299             :     // -------------- ALWAYS OWNED RESOURCES
     300             :     std::vector<uint32_t> mSpirv;
     301             :     void* mSpecializationConstantsData = nullptr;
     302             :     uint32_t mSpecializationConstantsDataTypeMemorySize = 0;
     303             :     uint32_t mSpecializationConstantsSize = 0;
     304             :     void* mPushConstantsData = nullptr;
     305             :     uint32_t mPushConstantsDataTypeMemorySize = 0;
     306             :     uint32_t mPushConstantsSize = 0;
     307             :     Workgroup mWorkgroup;
     308             : 
     309             :     // Create util functions
     310             :     void createShaderModule();
     311             :     void createPipeline();
     312             : 
     313             :     // Parameters
     314             :     void createParameters();
     315             : };
     316             : 
     317             : } // End namespace kp