1 files changed, 104 insertions, 104 deletions

diff --git a/src/timestamp_pool.cc b/src/timestamp_pool.cc
index b4dc3c9..cf48873 100644
--- a/src/timestamp_pool.cc
+++ b/src/timestamp_pool.cc
@@ -3,142 +3,152 @@
 #include "queue_context.hh"
 
 #include <ranges>
+#include <vulkan/utility/vk_dispatch_table.h>
 #include <vulkan/vulkan_core.h>
 
 namespace low_latency {
 
-TimestampPool::Block TimestampPool::allocate() {
-    const auto& device_context = this->queue_context.device_context;
+TimestampPool::QueryChunk::QueryChunk(const QueueContext& queue_context) {
+    const auto& device_context = queue_context.device_context;
+    const auto& vtable = device_context.vtable;
 
-    const auto query_pool = [&]() -> VkQueryPool {
+    this->query_pool = [&]() {
         const auto qpci = VkQueryPoolCreateInfo{
             .sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO,
             .queryType = VK_QUERY_TYPE_TIMESTAMP,
-            .queryCount = this->TIMESTAMP_QUERY_POOL_SIZE};
-
-        auto query_pool = VkQueryPool{};
+            .queryCount = QueryChunk::CHUNK_SIZE};
 
-        device_context.vtable.CreateQueryPool(device_context.device, &qpci,
-                                              nullptr, &query_pool);
-        return query_pool;
+        auto qp = VkQueryPool{};
+        vtable.CreateQueryPool(device_context.device, &qpci, nullptr, &qp);
+        return qp;
     }();
 
-    const auto key_range =
-        std::views::iota(0u, this->TIMESTAMP_QUERY_POOL_SIZE / 2) |
-        std::views::transform([](const std::uint64_t& i) { return 2 * i; });
-
-    auto available_indices = std::make_unique<available_query_indicies_t>(
-        available_query_indicies_t{std::begin(key_range), std::end(key_range)});
-
-    auto command_buffers = [&, this]() -> auto {
-        auto command_buffers =
-            std::vector<VkCommandBuffer>(this->TIMESTAMP_QUERY_POOL_SIZE);
+    constexpr auto key_range = std::views::iota(0u, QueryChunk::CHUNK_SIZE);
+    this->free_indices = std::make_unique<free_indices_t>(std::begin(key_range),
+                                                          std::end(key_range));
 
+    this->command_buffers = [&, this]() -> auto {
+        auto cbs = std::make_unique<std::vector<VkCommandBuffer>>(CHUNK_SIZE);
         const auto cbai = VkCommandBufferAllocateInfo{
             .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
-            .commandPool = this->queue_context.command_pool,
+            .commandPool = queue_context.command_pool,
             .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
-            .commandBufferCount =
-                static_cast<std::uint32_t>(std::size(command_buffers)),
+            .commandBufferCount = static_cast<std::uint32_t>(std::size(*cbs)),
         };
-        device_context.vtable.AllocateCommandBuffers(
-            device_context.device, &cbai, std::data(command_buffers));
-        std::ranges::for_each(command_buffers, [&](const auto& cb) {
-            device_context.sdld(device_context.device, cb);
-        });
-        return std::make_unique<std::vector<VkCommandBuffer>>(command_buffers);
+        vtable.AllocateCommandBuffers(device_context.device, &cbai,
+                                      std::data(*cbs));
+        return cbs;
     }();
-
-    return Block{.query_pool = query_pool,
-                 .available_indicies = std::move(available_indices),
-                 .command_buffers = std::move(command_buffers)};
 }
 
+TimestampPool::QueryChunk::~QueryChunk() {}
+
 TimestampPool::TimestampPool(QueueContext& queue_context)
     : queue_context(queue_context) {
 
-    // Allocate one block on construction, it's likely more than enough!
-    this->blocks.emplace_back(this->allocate());
+    // Allocate one block on construction, it's likely more than enough.
+    auto query_chunk = std::make_shared<QueryChunk>(this->queue_context);
+    this->query_chunks.emplace(std::move(query_chunk));
 }
 
 std::shared_ptr<TimestampPool::Handle> TimestampPool::acquire() {
-    const auto vacant_iter = [this]() -> auto {
-        const auto it =
-            std::ranges::find_if(this->blocks, [](const auto& block) {
-                return std::size(*block.available_indicies);
+
+    // Gets the empty one, or inserts a new one and returns it.
+    const auto not_empty_iter = [this]() -> auto {
+        const auto not_empty_iter =
+            std::ranges::find_if(this->query_chunks, [](const auto& qc) {
+                assert(qc);
+                return std::size(*qc->free_indices);
             });
 
-        if (it != std::end(this->blocks)) {
-            return it;
+        if (not_empty_iter != std::end(this->query_chunks)) {
+            return not_empty_iter;
         }
-        this->blocks.emplace_back(this->allocate());
-        return std::prev(std::end(this->blocks));
-    }();
-
-    const auto query_pool = vacant_iter->query_pool;
-    auto& available_indices = *vacant_iter->available_indicies;
 
-    // Grab any element from our set and erase it immediately after.
-    const auto query_index = *std::begin(available_indices);
-    available_indices.erase(std::begin(available_indices));
-
-    const auto command_buffers = [&]() -> auto {
-        auto command_buffers = std::array<VkCommandBuffer, 2>{};
-        std::ranges::copy_n(
-            std::next(std::begin(*vacant_iter->command_buffers), query_index),
-            std::size(command_buffers), std::begin(command_buffers));
-        return command_buffers;
+        const auto insert = std::make_shared<QueryChunk>(this->queue_context);
+        const auto [iter, did_insert] = this->query_chunks.emplace(insert);
+        assert(did_insert);
+        return iter;
     }();
 
-    const auto block_index = static_cast<std::size_t>(
-        std::distance(std::begin(this->blocks), vacant_iter));
+    // Grab any element from our set and erase it immediately after.
+    auto& indices = *(*not_empty_iter)->free_indices;
+    const auto query_index = *std::begin(indices);
+    assert(indices.erase(query_index));
 
-    return std::make_shared<Handle>(available_indices, block_index, query_pool,
-                                    query_index, command_buffers);
+    return std::make_shared<Handle>(*not_empty_iter, query_index);
 }
 
-TimestampPool::Handle::Handle(
-    TimestampPool::available_query_indicies_t& index_origin,
-    const std::size_t block_index, const VkQueryPool& query_pool,
-    const std::uint64_t query_index,
-    const std::array<VkCommandBuffer, 2>& command_buffers)
-    : index_origin(index_origin), block_index(block_index),
-      query_pool(query_pool), query_index(query_index),
-      command_buffers(command_buffers) {}
+TimestampPool::Handle::Handle(const std::shared_ptr<QueryChunk>& origin_chunk,
+                              const std::uint64_t& query_index)
+    : query_pool(origin_chunk->query_pool), query_index(query_index),
+      origin_chunk(origin_chunk),
+      command_buffer((*origin_chunk->command_buffers)[query_index]) {}
 
 TimestampPool::Handle::~Handle() {
-    assert(this->index_origin.insert(this->query_index).second);
+    // Parent destructing shouldn't mean we should have a bunch of insertions
+    // for zero reason.
+    if (const auto ptr = this->origin_chunk.lock(); ptr) {
+        assert(ptr->free_indices->insert(this->query_index).second);
+    }
 }
 
 void TimestampPool::Handle::setup_command_buffers(
-    const VkuDeviceDispatchTable& vtable) const {
-
-    const auto& [head, tail] = this->command_buffers;
+    const Handle& tail, const QueueContext& queue_context) const {
 
     const auto cbbi = VkCommandBufferBeginInfo{
         .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
     };
-    // Heads
-    vtable.ResetCommandBuffer(head, 0);
-    vtable.BeginCommandBuffer(head, &cbbi);
-    // Reset the next two and make them unavailable when they are run!
-    vtable.CmdResetQueryPool(head, this->query_pool, this->query_index, 2);
-    vtable.CmdWriteTimestamp2KHR(head, VK_PIPELINE_STAGE_2_TOP_OF_PIPE_BIT,
+
+    const auto& device_context = queue_context.device_context;
+    const auto& vtable = device_context.vtable;
+
+    vtable.ResetQueryPoolEXT(device_context.device, this->query_pool,
+                             this->query_index, 1);
+
+    vtable.BeginCommandBuffer(this->command_buffer, &cbbi);
+    vtable.CmdWriteTimestamp2KHR(this->command_buffer,
+                                 VK_PIPELINE_STAGE_2_TOP_OF_PIPE_BIT,
                                  this->query_pool, this->query_index);
-    vtable.EndCommandBuffer(head);
-
-    // Tails
-    vtable.ResetCommandBuffer(tail, 0);
-    vtable.BeginCommandBuffer(tail, &cbbi);
-    vtable.CmdWriteTimestamp2KHR(tail, VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT,
-                                 this->query_pool, this->query_index + 1);
-    vtable.EndCommandBuffer(tail);
+    vtable.EndCommandBuffer(this->command_buffer);
+
+    vtable.ResetQueryPoolEXT(device_context.device, tail.query_pool,
+                             tail.query_index, 1);
+    vtable.ResetCommandBuffer(tail.command_buffer, 0);
+    vtable.BeginCommandBuffer(tail.command_buffer, &cbbi);
+    vtable.CmdWriteTimestamp2KHR(tail.command_buffer,
+                                 VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT,
+                                 tail.query_pool, tail.query_index);
+    vtable.EndCommandBuffer(tail.command_buffer);
 }
 
-void TimestampPool::poll() {
-    this->cached_timestamps.clear();
-    this->cached_timestamps.reserve(std::size(this->blocks));
+std::optional<std::uint64_t>
+TimestampPool::Handle::get_ticks(const TimestampPool& pool) {
+
+    const auto& device_context = pool.queue_context.device_context;
+    const auto& vtable = device_context.vtable;
+
+    struct QueryResult {
+        std::uint64_t value;
+        std::uint64_t available;
+    };
+    auto query_result = QueryResult{};
+
+    const auto r = vtable.GetQueryPoolResults(
+        device_context.device, query_pool, this->query_index, 1,
+        sizeof(query_result), &query_result, sizeof(query_result),
+        VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WITH_AVAILABILITY_BIT);
+
+    assert(r == VK_SUCCESS || r == VK_NOT_READY);
+
+    if (!query_result.available) {
+        return std::nullopt;
+    }
+    return query_result.value;
+}
 
+/*
+void TimestampPool::poll() {
     const auto& device_context = this->queue_context.device_context;
 
     std::ranges::transform(
@@ -163,26 +173,16 @@ void TimestampPool::poll() {
             return timestamps;
         });
 };
-
-std::uint64_t TimestampPool::get_polled(const Handle& handle, const bool hack) {
-
-    assert(handle.block_index < std::size(this->cached_timestamps));
-
-    const auto& cached_timestamp = this->cached_timestamps[handle.block_index];
-    assert(cached_timestamp != nullptr);
-    assert(handle.query_index < std::size(*cached_timestamp));
-
-    return (*cached_timestamp)[handle.query_index + hack];
-}
+*/
 
 TimestampPool::~TimestampPool() {
     const auto& device = this->queue_context.device_context.device;
     const auto& vtable = this->queue_context.device_context.vtable;
-    for (const auto& block : this->blocks) {
+    for (const auto& query_chunk : this->query_chunks) {
         vtable.FreeCommandBuffers(device, this->queue_context.command_pool,
-                                  std::size(*block.command_buffers),
-                                  std::data(*block.command_buffers));
-        vtable.DestroyQueryPool(device, block.query_pool, nullptr);
+                                  std::size(*query_chunk->command_buffers),
+                                  std::data(*query_chunk->command_buffers));
+        vtable.DestroyQueryPool(device, query_chunk->query_pool, nullptr);
     }
 }