OpAlgoDispatch.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Algorithm.hpp"
#include "kompute/Core.hpp"
#include "kompute/Tensor.hpp"
#include "kompute/operations/OpBase.hpp"
 
namespace kp {
 
class OpAlgoDispatch : public OpBase
{
  public:
    template<typename T = float>
    OpAlgoDispatch(const std::shared_ptr<kp::Algorithm>& algorithm,
                   const std::vector<T>& pushConstants = {})
    {
        KP_LOG_DEBUG("Kompute OpAlgoDispatch constructor");
 
        this->mAlgorithm = algorithm;
 
        if (pushConstants.size()) {
            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;
        }
    }
 
    virtual ~OpAlgoDispatch() override;
 
    virtual void record(const vk::CommandBuffer& commandBuffer) override;
 
    virtual void preEval(const vk::CommandBuffer& commandBuffer) override;
 
    virtual void postEval(const vk::CommandBuffer& commandBuffer) override;
 
  private:
    // -------------- ALWAYS OWNED RESOURCES
    std::shared_ptr<Algorithm> mAlgorithm;
    void* mPushConstantsData = nullptr;
    uint32_t mPushConstantsDataTypeMemorySize = 0;
    uint32_t mPushConstantsSize = 0;
};
 
} // End namespace kp