async_compute_task.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "async_compute.hpp"
#include "gpu_sync.hpp"
#include "shader_layout.hpp"
#include "vulkan_pipeline.hpp"
#include "vulkan_renderer.hpp"
 
#include <array>
#include <cstddef>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <string>
#include <toast/export.hpp>
#include <vector>
#include <vulkan-memory-allocator-hpp/vk_mem_alloc_raii.hpp>
 
namespace renderer {

class TOAST_API AsyncComputeTask {
public:
  using BufferId = uint32_t;

  enum class Copy : uint8_t {
    current,
    previous,
  };
 
  struct Config {
    std::string shader_uri;
    std::vector<std::string> entries {"computeMain"};
    std::string debug_name = "AsyncComputeTask";
  };

  class TOAST_API StepContext {
  public:
    template<typename Push>
    void dispatch(uint32_t groups_x, const Push& push, uint32_t entry = 0, uint32_t groups_y = 1, uint32_t groups_z = 1) {
      dispatchRaw(entry, {groups_x, groups_y, groups_z}, &push, sizeof(Push));
    }
 
    void dispatch(uint32_t groups_x, uint32_t entry = 0, uint32_t groups_y = 1, uint32_t groups_z = 1) {
      dispatchRaw(entry, {groups_x, groups_y, groups_z}, nullptr, 0);
    }

    [[nodiscard]]
    auto commandBuffer() const -> vk::CommandBuffer {
      return m_cmd;
    }
 
    [[nodiscard]]
    auto hostData(BufferId id) const -> std::span<std::byte>;
 
    template<typename T>
    [[nodiscard]]
    auto hostData(BufferId id) const -> std::span<T> {
      const std::span<std::byte> bytes = hostData(id);
      return {reinterpret_cast<T*>(bytes.data()), bytes.size() / sizeof(T)};
    }
 
  private:
    friend class AsyncComputeTask;
 
    StepContext(AsyncComputeTask& task, vk::CommandBuffer cmd, uint32_t slot) : m_task(task), m_cmd(cmd), m_slot(slot) { }
 
    void dispatchRaw(uint32_t entry, std::array<uint32_t, 3> groups, const void* push, uint32_t push_size);
 
    AsyncComputeTask& m_task;
    vk::CommandBuffer m_cmd;
    uint32_t m_slot;
    bool m_dispatched = false;
  };

  class FrameView {
  public:
    [[nodiscard]]
    explicit operator bool() const {
      return m_task != nullptr;
    }

    [[nodiscard]]
    auto buffer(BufferId id) const -> vk::Buffer;

    [[nodiscard]]
    auto step() const -> uint64_t {
      return m_step;
    }
 
  private:
    friend class AsyncComputeTask;
    const AsyncComputeTask* m_task = nullptr;
    uint32_t m_slot = 0;
    uint64_t m_step = 0;
  };
 
  virtual ~AsyncComputeTask();
 
  AsyncComputeTask(const AsyncComputeTask&) = delete;
  auto operator=(const AsyncComputeTask&) -> AsyncComputeTask& = delete;
  AsyncComputeTask(AsyncComputeTask&&) = delete;
  auto operator=(AsyncComputeTask&&) -> AsyncComputeTask& = delete;

  auto acquireForFrame(vk::CommandBuffer cmd, uint32_t frame_index, sync::Usage next) -> FrameView;

  auto read(BufferId id, std::span<std::byte> out) -> bool;
 
  template<typename T>
  auto read(BufferId id, std::span<T> out) -> bool {
    return read(id, std::as_writable_bytes(out));
  }

  [[nodiscard]]
  auto finishedStep() const -> uint64_t;
 
  [[nodiscard]]
  auto isReady() const -> bool {
    return m_ready;
  }
 
protected:
  explicit AsyncComputeTask(Config config);

  auto addRingBuffer(vk::DeviceSize size, bool cpu_readback = false) -> BufferId;
  auto addPersistentBuffer(vk::DeviceSize size) -> BufferId;
  auto addHostBuffer(vk::DeviceSize size) -> BufferId;
  void bind(uint32_t binding, BufferId id, Copy copy = Copy::current);
  void bindExternal(uint32_t set, uint32_t binding, std::function<vk::Buffer()> source);
  void setAcquireUsage(BufferId id, sync::Usage usage);

  void upload(BufferId id, std::span<const std::byte> data);
 
  template<typename T>
  void upload(BufferId id, std::span<const T> data) {
    upload(id, std::as_bytes(data));
  }

  auto step(
      const std::function<void(StepContext&)>& record, std::string_view debug_name = {},
      std::shared_ptr<const void> keep_alive = {}
  ) -> uint64_t;
 
private:
  static constexpr uint32_t k_slot_count = VulkanRenderer::k_frames_in_flight + 3;
 
  struct Buffer {
    bool ring = false;
    bool host = false;
    bool cpu_readback = false;
    vk::DeviceSize size = 0;
    std::vector<vma::raii::Buffer> copies;
    std::vector<vma::raii::Buffer> readbacks;
    std::optional<sync::Usage> acquire_usage;
  };
 
  struct Binding {
    uint32_t binding = 0;
    BufferId id = 0;
    Copy copy = Copy::current;
  };
 
  struct ExternalBinding {
    uint32_t set = 0;
    uint32_t binding = 0;
    std::function<vk::Buffer()> source;
  };
 
  struct Slot {
    std::vector<vk::raii::DescriptorSet> sets;
    AsyncCompute::Ticket ticket = AsyncCompute::k_no_ticket;
    uint32_t pins = 0;
    uint64_t step = 0;
    std::shared_ptr<const void> keep_alive;
  };
 
  void createDescriptors();
  void writeDescriptors(uint32_t slot);
  auto claimFreeSlot() -> uint32_t;
  void publish(uint32_t slot, AsyncCompute::Ticket ticket);
  auto newestFinishedLocked() const -> int32_t;
 
  Config m_config;
  const VulkanCore* m_core = nullptr;
  bool m_ready = false;
 
  ShaderLayout m_layout;
  std::vector<std::unique_ptr<VulkanPipeline>> m_pipelines;
  vk::raii::DescriptorPool m_pool = nullptr;
 
  std::vector<Buffer> m_buffers;
  std::vector<Binding> m_bindings;
  std::vector<ExternalBinding> m_externals;
  vma::raii::Buffer m_placeholder = nullptr;
 
  mutable std::mutex m_mutex;
  std::array<Slot, k_slot_count> m_slots;
  uint64_t m_step_counter = 0;
  int32_t m_last_written = -1;

  std::array<int32_t, VulkanRenderer::k_frames_in_flight> m_frame_slot {};
  std::array<uint64_t, k_slot_count> m_acquired_step {};
};
 
}