Line data    Source code

       1             : // SPDX-License-Identifier: Apache-2.0
       2             : #include <fstream>
       3             : 
       4             : #include "kompute/Algorithm.hpp"
       5             : 
       6             : namespace kp {
       7             : 
       8          37 : Algorithm::~Algorithm()
       9             : {
      10          74 :     KP_LOG_DEBUG("Kompute Algorithm Destructor started");
      11             : 
      12          37 :     this->destroy();
      13          37 : }
      14             : 
      15             : bool
      16          38 : Algorithm::isInit()
      17             : {
      18          40 :     return this->mPipeline && this->mPipelineCache && this->mPipelineLayout &&
      19           3 :            this->mDescriptorPool && this->mDescriptorSet &&
      20          40 :            this->mDescriptorSetLayout && this->mShaderModule;
      21             : }
      22             : 
      23             : void
      24          38 : Algorithm::destroy()
      25             : {
      26             :     // We don't have to free memory on destroy as it's freed by the
      27             :     // commandBuffer destructor if (this->mPushConstantsData) {
      28             :     //     free(this->mPushConstantsData);
      29             :     // }
      30             :     // if (this->mSpecializationConstantsData) {
      31             :     //     free(this->mSpecializationConstantsData);
      32             :     // }
      33             : 
      34          38 :     if (!this->mDevice) {
      35           0 :         KP_LOG_WARN("Kompute Algorithm destroy function reached with null "
      36             :                     "Device pointer");
      37           0 :         return;
      38             :     }
      39             : 
      40          38 :     if (this->mFreePipeline && this->mPipeline) {
      41          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying pipeline");
      42          38 :         if (!this->mPipeline) {
      43           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy "
      44             :                         "pipeline but it is null");
      45             :         }
      46          38 :         this->mDevice->destroy(
      47          38 :           *this->mPipeline,
      48             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
      49          38 :         this->mPipeline = nullptr;
      50             :     }
      51             : 
      52          38 :     if (this->mFreePipelineCache && this->mPipelineCache) {
      53          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying pipeline cache");
      54          38 :         if (!this->mPipelineCache) {
      55           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy "
      56             :                         "pipeline cache but it is null");
      57             :         }
      58          38 :         this->mDevice->destroy(
      59          38 :           *this->mPipelineCache,
      60             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
      61          38 :         this->mPipelineCache = nullptr;
      62             :     }
      63             : 
      64          38 :     if (this->mFreePipelineLayout && this->mPipelineLayout) {
      65          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying pipeline layout");
      66          38 :         if (!this->mPipelineLayout) {
      67           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy "
      68             :                         "pipeline layout but it is null");
      69             :         }
      70          38 :         this->mDevice->destroy(
      71          38 :           *this->mPipelineLayout,
      72             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
      73          38 :         this->mPipelineLayout = nullptr;
      74             :     }
      75             : 
      76          38 :     if (this->mFreeShaderModule && this->mShaderModule) {
      77          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying shader module");
      78          38 :         if (!this->mShaderModule) {
      79           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy shader "
      80             :                         "module but it is null");
      81             :         }
      82          38 :         this->mDevice->destroy(
      83          38 :           *this->mShaderModule,
      84             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
      85          38 :         this->mShaderModule = nullptr;
      86             :     }
      87             : 
      88             :     // We don't call freeDescriptorSet as the descriptor pool is not created
      89             :     // with VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT more at
      90             :     // (https://www.khronos.org/registry/vulkan/specs/1.0/html/vkspec.html#VUID-vkFreeDescriptorSets-descriptorPool-00312))
      91             :     // if (this->mFreeDescriptorSet && this->mDescriptorSet) {
      92             :     //    KP_LOG_DEBUG("Kompute Algorithm Freeing Descriptor Set");
      93             :     //    if (!this->mDescriptorSet) {
      94             :     //        KP_LOG_WARN(
      95             :     //          "Kompute Algorithm Error requested to free descriptor set");
      96             :     //    }
      97             :     //    this->mDevice->freeDescriptorSets(
      98             :     //      *this->mDescriptorPool, 1, this->mDescriptorSet.get());
      99             :     //    this->mDescriptorSet = nullptr;
     100             :     //}
     101             : 
     102          38 :     if (this->mFreeDescriptorSetLayout && this->mDescriptorSetLayout) {
     103          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying Descriptor Set Layout");
     104          38 :         if (!this->mDescriptorSetLayout) {
     105           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy "
     106             :                         "descriptor set layout but it is null");
     107             :         }
     108          38 :         this->mDevice->destroy(
     109          38 :           *this->mDescriptorSetLayout,
     110             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
     111          38 :         this->mDescriptorSetLayout = nullptr;
     112             :     }
     113             : 
     114          38 :     if (this->mFreeDescriptorPool && this->mDescriptorPool) {
     115          76 :         KP_LOG_DEBUG("Kompute Algorithm Destroying Descriptor Pool");
     116          38 :         if (!this->mDescriptorPool) {
     117           0 :             KP_LOG_WARN("Kompute Algorithm Error requested to destroy "
     118             :                         "descriptor pool but it is null");
     119             :         }
     120          38 :         this->mDevice->destroy(
     121          38 :           *this->mDescriptorPool,
     122             :           (vk::Optional<const vk::AllocationCallbacks>)nullptr);
     123          38 :         this->mDescriptorPool = nullptr;
     124             :     }
     125             : }
     126             : 
     127             : void
     128          38 : Algorithm::createParameters()
     129             : {
     130          76 :     KP_LOG_DEBUG("Kompute Algorithm createParameters started");
     131             : 
     132             :     std::vector<vk::DescriptorPoolSize> descriptorPoolSizes = {
     133             :         vk::DescriptorPoolSize(
     134             :           vk::DescriptorType::eStorageBuffer,
     135          38 :           static_cast<uint32_t>(this->mTensors.size()) // Descriptor count
     136             :           )
     137          38 :     };
     138             : 
     139             :     vk::DescriptorPoolCreateInfo descriptorPoolInfo(
     140             :       vk::DescriptorPoolCreateFlags(),
     141             :       1, // Max sets
     142          38 :       static_cast<uint32_t>(descriptorPoolSizes.size()),
     143          76 :       descriptorPoolSizes.data());
     144             : 
     145          76 :     KP_LOG_DEBUG("Kompute Algorithm creating descriptor pool");
     146          38 :     this->mDescriptorPool = std::make_shared<vk::DescriptorPool>();
     147          38 :     this->mDevice->createDescriptorPool(
     148             :       &descriptorPoolInfo, nullptr, this->mDescriptorPool.get());
     149          38 :     this->mFreeDescriptorPool = true;
     150             : 
     151          38 :     std::vector<vk::DescriptorSetLayoutBinding> descriptorSetBindings;
     152         122 :     for (size_t i = 0; i < this->mTensors.size(); i++) {
     153          84 :         descriptorSetBindings.push_back(
     154         168 :           vk::DescriptorSetLayoutBinding(i, // Binding index
     155             :                                          vk::DescriptorType::eStorageBuffer,
     156             :                                          1, // Descriptor count
     157             :                                          vk::ShaderStageFlagBits::eCompute));
     158             :     }
     159             : 
     160             :     // This is the component that is fed into the pipeline
     161             :     vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutInfo(
     162             :       vk::DescriptorSetLayoutCreateFlags(),
     163          38 :       static_cast<uint32_t>(descriptorSetBindings.size()),
     164          76 :       descriptorSetBindings.data());
     165             : 
     166          76 :     KP_LOG_DEBUG("Kompute Algorithm creating descriptor set layout");
     167          38 :     this->mDescriptorSetLayout = std::make_shared<vk::DescriptorSetLayout>();
     168          38 :     this->mDevice->createDescriptorSetLayout(
     169             :       &descriptorSetLayoutInfo, nullptr, this->mDescriptorSetLayout.get());
     170          38 :     this->mFreeDescriptorSetLayout = true;
     171             : 
     172             :     vk::DescriptorSetAllocateInfo descriptorSetAllocateInfo(
     173          38 :       *this->mDescriptorPool,
     174             :       1, // Descriptor set layout count
     175          76 :       this->mDescriptorSetLayout.get());
     176             : 
     177          76 :     KP_LOG_DEBUG("Kompute Algorithm allocating descriptor sets");
     178          38 :     this->mDescriptorSet = std::make_shared<vk::DescriptorSet>();
     179          38 :     this->mDevice->allocateDescriptorSets(&descriptorSetAllocateInfo,
     180             :                                           this->mDescriptorSet.get());
     181          38 :     this->mFreeDescriptorSet = true;
     182             : 
     183          76 :     KP_LOG_DEBUG("Kompute Algorithm updating descriptor sets");
     184         122 :     for (size_t i = 0; i < this->mTensors.size(); i++) {
     185          84 :         std::vector<vk::WriteDescriptorSet> computeWriteDescriptorSets;
     186             : 
     187             :         vk::DescriptorBufferInfo descriptorBufferInfo =
     188          84 :           this->mTensors[i]->constructDescriptorBufferInfo();
     189             : 
     190          84 :         computeWriteDescriptorSets.push_back(
     191          84 :           vk::WriteDescriptorSet(*this->mDescriptorSet,
     192             :                                  i, // Destination binding
     193             :                                  0, // Destination array element
     194             :                                  1, // Descriptor count
     195             :                                  vk::DescriptorType::eStorageBuffer,
     196             :                                  nullptr, // Descriptor image info
     197             :                                  &descriptorBufferInfo));
     198             : 
     199          84 :         this->mDevice->updateDescriptorSets(computeWriteDescriptorSets,
     200             :                                             nullptr);
     201          84 :     }
     202             : 
     203          76 :     KP_LOG_DEBUG("Kompute Algorithm successfully run init");
     204          38 : }
     205             : 
     206             : void
     207          38 : Algorithm::createShaderModule()
     208             : {
     209          76 :     KP_LOG_DEBUG("Kompute Algorithm createShaderModule started");
     210             : 
     211             :     vk::ShaderModuleCreateInfo shaderModuleInfo(vk::ShaderModuleCreateFlags(),
     212             :                                                 sizeof(uint32_t) *
     213          38 :                                                   this->mSpirv.size(),
     214          76 :                                                 this->mSpirv.data());
     215             : 
     216         114 :     KP_LOG_DEBUG("Kompute Algorithm Creating shader module. ShaderFileSize: {}",
     217             :                  this->mSpirv.size());
     218          38 :     this->mFreeShaderModule = true;
     219          38 :     this->mShaderModule = std::make_shared<vk::ShaderModule>();
     220          38 :     this->mDevice->createShaderModule(
     221             :       &shaderModuleInfo, nullptr, this->mShaderModule.get());
     222          38 :     this->mFreeShaderModule = true;
     223             : 
     224          76 :     KP_LOG_DEBUG("Kompute Algorithm create shader module success");
     225          38 : }
     226             : 
     227             : void
     228          38 : Algorithm::createPipeline()
     229             : {
     230          76 :     KP_LOG_DEBUG("Kompute Algorithm calling create Pipeline");
     231             : 
     232             :     vk::PipelineLayoutCreateInfo pipelineLayoutInfo(
     233             :       vk::PipelineLayoutCreateFlags(),
     234             :       1, // Set layout count
     235          38 :       this->mDescriptorSetLayout.get());
     236             : 
     237          38 :     vk::PushConstantRange pushConstantRange;
     238          38 :     if (this->mPushConstantsSize) {
     239           8 :         pushConstantRange.setStageFlags(vk::ShaderStageFlagBits::eCompute);
     240           8 :         pushConstantRange.setOffset(0);
     241           8 :         pushConstantRange.setSize(this->mPushConstantsDataTypeMemorySize *
     242           8 :                                   this->mPushConstantsSize);
     243             : 
     244           8 :         pipelineLayoutInfo.setPushConstantRangeCount(1);
     245           8 :         pipelineLayoutInfo.setPPushConstantRanges(&pushConstantRange);
     246             :     }
     247             : 
     248          38 :     this->mPipelineLayout = std::make_shared<vk::PipelineLayout>();
     249          38 :     this->mDevice->createPipelineLayout(
     250             :       &pipelineLayoutInfo, nullptr, this->mPipelineLayout.get());
     251          38 :     this->mFreePipelineLayout = true;
     252             : 
     253          38 :     std::vector<vk::SpecializationMapEntry> specializationEntries;
     254             : 
     255          46 :     for (uint32_t i = 0; i < this->mSpecializationConstantsSize; i++) {
     256             :         vk::SpecializationMapEntry specializationEntry(
     257             :           static_cast<uint32_t>(i),
     258             :           static_cast<uint32_t>(
     259           8 :             this->mSpecializationConstantsDataTypeMemorySize * i),
     260           8 :           this->mSpecializationConstantsDataTypeMemorySize);
     261             : 
     262           8 :         specializationEntries.push_back(specializationEntry);
     263             :     }
     264             : 
     265             :     // This passes ownership of the memory so we remove ownership from
     266             :     // specialization container by using "transferDataOwnership"
     267             :     vk::SpecializationInfo specializationInfo(
     268          38 :       static_cast<uint32_t>(specializationEntries.size()),
     269          38 :       specializationEntries.data(),
     270          38 :       this->mSpecializationConstantsDataTypeMemorySize *
     271          38 :         this->mSpecializationConstantsSize,
     272          76 :       this->mSpecializationConstantsData);
     273             : 
     274             :     vk::PipelineShaderStageCreateInfo shaderStage(
     275             :       vk::PipelineShaderStageCreateFlags(),
     276             :       vk::ShaderStageFlagBits::eCompute,
     277          38 :       *this->mShaderModule,
     278             :       "main",
     279          38 :       &specializationInfo);
     280             : 
     281             :     vk::ComputePipelineCreateInfo pipelineInfo(vk::PipelineCreateFlags(),
     282             :                                                shaderStage,
     283          38 :                                                *this->mPipelineLayout,
     284             :                                                vk::Pipeline(),
     285          76 :                                                0);
     286             : 
     287          38 :     vk::PipelineCacheCreateInfo pipelineCacheInfo =
     288             :       vk::PipelineCacheCreateInfo();
     289          38 :     this->mPipelineCache = std::make_shared<vk::PipelineCache>();
     290          38 :     this->mDevice->createPipelineCache(
     291             :       &pipelineCacheInfo, nullptr, this->mPipelineCache.get());
     292          38 :     this->mFreePipelineCache = true;
     293             : 
     294             : #ifdef KOMPUTE_CREATE_PIPELINE_RESULT_VALUE
     295             :     vk::ResultValue<vk::Pipeline> pipelineResult =
     296             :       this->mDevice->createComputePipeline(*this->mPipelineCache, pipelineInfo);
     297             : 
     298             :     if (pipelineResult.result != vk::Result::eSuccess) {
     299             :         throw std::runtime_error("Failed to create pipeline result: " +
     300             :                                  vk::to_string(pipelineResult.result));
     301             :     }
     302             : 
     303             :     vk::Pipeline& pipeline = pipelineResult.value;
     304             :     this->mPipeline = std::make_shared<vk::Pipeline>(pipeline);
     305             :     this->mFreePipeline = true;
     306             : #else
     307             :     vk::Pipeline pipeline =
     308          38 :       this->mDevice->createComputePipeline(*this->mPipelineCache, pipelineInfo)
     309          38 :         .value;
     310          38 :     this->mPipeline = std::make_shared<vk::Pipeline>(pipeline);
     311          38 :     this->mFreePipeline = true;
     312             : #endif
     313             : 
     314             :     // TODO: Update to consistent
     315             :     // this->mPipeline = std::make_shared<vk::Pipeline>();
     316             :     // this->mDevice->createComputePipelines(
     317             :     //         *this->mPipelineCache, 1, &pipelineInfo, nullptr,
     318             :     //         this->mPipeline.get());
     319             : 
     320          76 :     KP_LOG_DEBUG("Kompute Algorithm Create Pipeline Success");
     321          38 : }
     322             : 
     323             : void
     324          46 : Algorithm::recordBindCore(const vk::CommandBuffer& commandBuffer)
     325             : {
     326          92 :     KP_LOG_DEBUG("Kompute Algorithm binding pipeline");
     327             : 
     328          92 :     commandBuffer.bindPipeline(vk::PipelineBindPoint::eCompute,
     329          46 :                                *this->mPipeline);
     330             : 
     331          92 :     KP_LOG_DEBUG("Kompute Algorithm binding descriptor sets");
     332             : 
     333          92 :     commandBuffer.bindDescriptorSets(vk::PipelineBindPoint::eCompute,
     334          46 :                                      *this->mPipelineLayout,
     335             :                                      0, // First set
     336          46 :                                      *this->mDescriptorSet,
     337             :                                      nullptr // Dispatcher
     338             :     );
     339          46 : }
     340             : 
     341             : void
     342          46 : Algorithm::recordBindPush(const vk::CommandBuffer& commandBuffer)
     343             : {
     344          46 :     if (this->mPushConstantsSize) {
     345          36 :         KP_LOG_DEBUG("Kompute Algorithm binding push constants memory size: {}",
     346             :                      this->mPushConstantsSize *
     347             :                        this->mPushConstantsDataTypeMemorySize);
     348             : 
     349          36 :         commandBuffer.pushConstants(*this->mPipelineLayout,
     350             :                                     vk::ShaderStageFlagBits::eCompute,
     351             :                                     0,
     352          12 :                                     this->mPushConstantsSize *
     353          12 :                                       this->mPushConstantsDataTypeMemorySize,
     354          12 :                                     this->mPushConstantsData);
     355             :     }
     356          46 : }
     357             : 
     358             : void
     359          46 : Algorithm::recordDispatch(const vk::CommandBuffer& commandBuffer)
     360             : {
     361          92 :     KP_LOG_DEBUG("Kompute Algorithm recording dispatch");
     362             : 
     363          92 :     commandBuffer.dispatch(
     364          46 :       this->mWorkgroup[0], this->mWorkgroup[1], this->mWorkgroup[2]);
     365          46 : }
     366             : 
     367             : void
     368          38 : Algorithm::setWorkgroup(const Workgroup& workgroup, uint32_t minSize)
     369             : {
     370             : 
     371          76 :     KP_LOG_INFO("Kompute OpAlgoCreate setting dispatch size");
     372             : 
     373             :     // The dispatch size is set up based on either explicitly provided template
     374             :     // parameters or by default it would take the shape and size of the tensors
     375          38 :     if (workgroup[0] > 0) {
     376             :         // If at least the x value is provided we use mainly the parameters
     377             :         // provided
     378          20 :         this->mWorkgroup = { workgroup[0],
     379          10 :                              workgroup[1] > 0 ? workgroup[1] : 1,
     380          10 :                              workgroup[2] > 0 ? workgroup[2] : 1 };
     381             :     } else {
     382          28 :         this->mWorkgroup = { minSize, 1, 1 };
     383             :     }
     384             : 
     385         114 :     KP_LOG_INFO("Kompute OpAlgoCreate set dispatch size X: {}, Y: {}, Z: {}",
     386             :                 this->mWorkgroup[0],
     387             :                 this->mWorkgroup[1],
     388             :                 this->mWorkgroup[2]);
     389          38 : }
     390             : 
     391             : const Workgroup&
     392           1 : Algorithm::getWorkgroup()
     393             : {
     394           1 :     return this->mWorkgroup;
     395             : }
     396             : 
     397             : const std::vector<std::shared_ptr<Tensor>>&
     398          51 : Algorithm::getTensors()
     399             : {
     400          51 :     return this->mTensors;
     401             : }
     402             : 
     403             : }