query_pool.cpp

#include <volk.h>

#include "backend/vulkan/debug.h"
#include "backend/vulkan/utils.h"
#include "backend/vulkan/vk_api.h"
#include "backend/vulkan/vk_conversion.h"

namespace rhi {
using namespace vk;

bool create_query_pool(Device &device, const QueryPoolDesc &desc, QueryPool &pool) {
    if (desc.queryType == QueryType::TIMESTAMP || desc.queryType == QueryType::TIMESTAMP_COPY_QUEUE) {
        pool._type = VK_QUERY_TYPE_TIMESTAMP;
    } else if (desc.queryType == QueryType::OCCLUSION) {
        pool._type = VK_QUERY_TYPE_OCCLUSION;
    } else if (desc.queryType == QueryType::PIPELINE_STATISTICS) {
        pool._type = VK_QUERY_TYPE_PIPELINE_STATISTICS;
    } else if (desc.queryType == QueryType::ACCELERATION_STRUCTURE_SIZE) {
        pool._type = VK_QUERY_TYPE_ACCELERATION_STRUCTURE_SIZE_KHR;
    } else if (desc.queryType == QueryType::ACCELERATION_STRUCTURE_COMPACTED_SIZE) {
        pool._type = VK_QUERY_TYPE_ACCELERATION_STRUCTURE_COMPACTED_SIZE_KHR;
    } else if (desc.queryType == QueryType::MICROMAP_COMPACTED_SIZE) {
        pool._type = VK_QUERY_TYPE_MICROMAP_COMPACTED_SIZE_EXT;
    } else {
        return false;
    }

    VkQueryPipelineStatisticFlags pipeline_statistics =
        VK_QUERY_PIPELINE_STATISTIC_INPUT_ASSEMBLY_VERTICES_BIT |
        VK_QUERY_PIPELINE_STATISTIC_INPUT_ASSEMBLY_PRIMITIVES_BIT |
        VK_QUERY_PIPELINE_STATISTIC_VERTEX_SHADER_INVOCATIONS_BIT |
        VK_QUERY_PIPELINE_STATISTIC_GEOMETRY_SHADER_INVOCATIONS_BIT |
        VK_QUERY_PIPELINE_STATISTIC_GEOMETRY_SHADER_PRIMITIVES_BIT |
        VK_QUERY_PIPELINE_STATISTIC_CLIPPING_INVOCATIONS_BIT | VK_QUERY_PIPELINE_STATISTIC_CLIPPING_PRIMITIVES_BIT |
        VK_QUERY_PIPELINE_STATISTIC_FRAGMENT_SHADER_INVOCATIONS_BIT |
        VK_QUERY_PIPELINE_STATISTIC_TESSELLATION_CONTROL_SHADER_PATCHES_BIT |
        VK_QUERY_PIPELINE_STATISTIC_TESSELLATION_EVALUATION_SHADER_INVOCATIONS_BIT |
        VK_QUERY_PIPELINE_STATISTIC_COMPUTE_SHADER_INVOCATIONS_BIT;

    // if (m_Device.GetDesc().features.meshShader)
    //     pipelineStatistics |= VK_QUERY_PIPELINE_STATISTIC_TASK_SHADER_INVOCATIONS_BIT_EXT |
    //     VK_QUERY_PIPELINE_STATISTIC_MESH_SHADER_INVOCATIONS_BIT_EXT;

    VkQueryPoolCreateInfo pool_info = {
        VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO,
        nullptr,
        (VkQueryPoolCreateFlags)0,
        (VkQueryType)pool._type,
        desc.capacity,
        pipeline_statistics
    };

    VkQueryPool vk_pool;
    VK_CHECK(vkCreateQueryPool(device.logical, &pool_info, nullptr, &vk_pool));
    pool._handle = vk_pool;
    pool.desc = desc;
    if (!desc.name.empty()) {
        VK_NAME_EX(device, vk_pool, desc.name.c_str());
    }
    return true;
}

void reset_query_pool(CmdBuffer &cmd, QueryPool &queryPool, u32 offset, u32 count) {
    vkCmdResetQueryPool(cmd.handle, queryPool._handle, offset, count);
}

void write_timestamp(CmdBuffer &cmd, QueryPool &queryPool, u32 queryIndex, rhi::PipelineStages stage) {
    vkCmdWriteTimestamp(cmd.handle, static_cast<VkPipelineStageFlagBits>(stage), queryPool._handle, queryIndex);
}

void destroy_query_pool(Device &device, QueryPool &queryPool) {
    vkDestroyQueryPool(device.logical, queryPool._handle, nullptr);
}

bool query_pool_results(
    Device &device,
    const QueryPool &pool,
    u32 first_query,
    u32 query_count,
    u64 *out,
    usize data_size_bytes,
    usize stride_bytes,
    bool with_availability
) {
    VkQueryResultFlags flags = VK_QUERY_RESULT_64_BIT;
    if (with_availability) {
        flags |= VK_QUERY_RESULT_WITH_AVAILABILITY_BIT;
    }
    VkResult res = vkGetQueryPoolResults(
        device.logical, pool._handle, first_query, query_count, data_size_bytes, out, stride_bytes, flags
    );
    return res == VK_SUCCESS;
}

} // namespace rhi