Sequence.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "kompute/Core.hpp"
 
#include "kompute/operations/OpAlgoDispatch.hpp"
#include "kompute/operations/OpBase.hpp"
 
namespace kp {
 
class Sequence : public std::enable_shared_from_this<Sequence>
{
  public:
    Sequence(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
             std::shared_ptr<vk::Device> device,
             std::shared_ptr<vk::Queue> computeQueue,
             uint32_t queueIndex,
             uint32_t totalTimestamps = 0);
    ~Sequence();
 
    std::shared_ptr<Sequence> record(std::shared_ptr<OpBase> op);
 
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> record(
      std::vector<std::shared_ptr<Tensor>> tensors,
      TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(tensors, std::forward<TArgs>(params)...) };
        return this->record(op);
    }
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> record(std::shared_ptr<Algorithm> algorithm,
                                     TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(algorithm,
                                     std::forward<TArgs>(params)...) };
        return this->record(op);
    }
 
    std::shared_ptr<Sequence> eval();
 
    std::shared_ptr<Sequence> eval(std::shared_ptr<OpBase> op);
 
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> eval(std::vector<std::shared_ptr<Tensor>> tensors,
                                   TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(tensors, std::forward<TArgs>(params)...) };
        return this->eval(op);
    }
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> eval(std::shared_ptr<Algorithm> algorithm,
                                   TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(algorithm,
                                     std::forward<TArgs>(params)...) };
        return this->eval(op);
    }
 
    std::shared_ptr<Sequence> evalAsync();
    std::shared_ptr<Sequence> evalAsync(std::shared_ptr<OpBase> op);
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> evalAsync(
      std::vector<std::shared_ptr<Tensor>> tensors,
      TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(tensors, std::forward<TArgs>(params)...) };
        return this->evalAsync(op);
    }
    template<typename T, typename... TArgs>
    std::shared_ptr<Sequence> evalAsync(std::shared_ptr<Algorithm> algorithm,
                                        TArgs&&... params)
    {
        std::shared_ptr<T> op{ new T(algorithm,
                                     std::forward<TArgs>(params)...) };
        return this->evalAsync(op);
    }
 
    std::shared_ptr<Sequence> evalAwait(uint64_t waitFor = UINT64_MAX);
 
    void clear();
 
    std::vector<std::uint64_t> getTimestamps();
 
    void begin();
 
    void end();
 
    bool isRecording() const;
 
    bool isInit() const;
 
    void rerecord();
 
    bool isRunning() const;
 
    void destroy();
 
  private:
    // -------------- NEVER OWNED RESOURCES
    std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice = nullptr;
    std::shared_ptr<vk::Device> mDevice = nullptr;
    std::shared_ptr<vk::Queue> mComputeQueue = nullptr;
    uint32_t mQueueIndex = -1;
 
    // -------------- OPTIONALLY OWNED RESOURCES
    std::shared_ptr<vk::CommandPool> mCommandPool = nullptr;
    bool mFreeCommandPool = false;
    std::shared_ptr<vk::CommandBuffer> mCommandBuffer = nullptr;
    bool mFreeCommandBuffer = false;
 
    // -------------- ALWAYS OWNED RESOURCES
    vk::Fence mFence;
    std::vector<std::shared_ptr<OpBase>> mOperations{};
    std::shared_ptr<vk::QueryPool> timestampQueryPool = nullptr;
 
    // State
    bool mRecording = false;
    bool mIsRunning = false;
 
    // Create functions
    void createCommandPool();
    void createCommandBuffer();
    void createTimestampQueryPool(uint32_t totalTimestamps);
};
 
} // End namespace kp