RelaxationStepProcessor.cpp

Go to the documentation of this file.

#include "RelaxationStepProcessor.hpp"
#include "RelaxationGridBuffer.hpp"
 
#include "MayaFlux/Kriya/Awaiters/BroadcastAwaiter.hpp"
 
#include "MayaFlux/Buffers/Staging/StagingUtils.hpp"
#include "MayaFlux/Registry/BackendRegistry.hpp"
#include "MayaFlux/Registry/Service/BufferService.hpp"
 
namespace MayaFlux::Buffers {
 
RelaxationStepProcessor::RelaxationStepProcessor(const std::string& shader_path, uint32_t workgroup_x)
    : ComputeProcessor(shader_path, workgroup_x)
{
    m_config.bindings["state_in"] = ShaderBinding(0, 0, vk::DescriptorType::eStorageBuffer);
    m_config.bindings["state_out"] = ShaderBinding(0, 1, vk::DescriptorType::eStorageBuffer);
}
 
RelaxationStepProcessor::RelaxationStepProcessor(const Portal::Graphics::ShaderSpec& spec)
    : ComputeProcessor(spec)
{
    m_config.bindings["state_in"] = ShaderBinding(0, 0, vk::DescriptorType::eStorageBuffer);
    m_config.bindings["state_out"] = ShaderBinding(0, 1, vk::DescriptorType::eStorageBuffer);
}
 
void RelaxationStepProcessor::write_state_descriptors(const std::shared_ptr<RelaxationGridBuffer>& grid)
{
    auto& resources = grid->get_buffer_resources();
    if (resources.back_buffers.size() < 2) {
        return;
    }
 
    if (m_descriptor_set_ids.empty()) {
        return;
    }
 
    auto& foundry = Portal::Graphics::get_shader_foundry();
    const size_t bytes = grid->get_state_bytes();
    foundry.update_descriptor_buffer(
        m_descriptor_set_ids[0], 0, vk::DescriptorType::eStorageBuffer,
        resources.back_buffers[grid->front_index()].buffer, 0, bytes);
 
    foundry.update_descriptor_buffer(
        m_descriptor_set_ids[0], 1, vk::DescriptorType::eStorageBuffer,
        resources.back_buffers[grid->back_index()].buffer, 0, bytes);
}
 
void RelaxationStepProcessor::processing_function(const std::shared_ptr<Buffer>& buffer)
{
    if (m_step_predicate && !m_step_predicate()) {
        return;
    }
 
    m_dispatched = false;
 
    if (m_grid && are_descriptors_ready()) {
        write_state_descriptors(m_grid);
    }
 
    ComputeProcessor::processing_function(buffer);
 
    if (!m_dispatched || !m_grid) {
        return;
    }
 
    m_grid->swap_generation();
 
    if (!is_deferred_submission()) {
        serve_snapshot();
    }
}
 
void RelaxationStepProcessor::write_grid_extent_constants()
{
    if (m_config.push_constant_size < sizeof(GridExtent)) {
        set_push_constant_size(sizeof(GridExtent));
    }
 
    auto& data = get_push_constant_data();
    if (data.size() < m_config.push_constant_size) {
        data.resize(m_config.push_constant_size);
    }
 
    const GridExtent extent {
        .width = m_grid->get_grid_width(),
        .height = m_grid->get_grid_height()
    };
    std::memcpy(data.data(), &extent, sizeof(GridExtent));
}
 
void RelaxationStepProcessor::on_attach(const std::shared_ptr<Buffer>& buffer)
{
    ComputeProcessor::on_attach(buffer);
    m_grid = std::dynamic_pointer_cast<RelaxationGridBuffer>(buffer);
 
    if (m_grid) {
        set_workgroup_size(16, 16, 1);
        set_dispatch_mode(ShaderDispatchConfig::DispatchMode::MANUAL);
        uint32_t gx = (m_grid->get_grid_width() + 15) / 16;
        uint32_t gy = (m_grid->get_grid_height() + 15) / 16;
        set_manual_dispatch(gx, gy, 1);
        write_grid_extent_constants();
    }
}
 
void RelaxationStepProcessor::on_descriptors_created()
{
    if (m_grid) {
        write_state_descriptors(m_grid);
    }
}
 
bool RelaxationStepProcessor::on_before_execute(
    Portal::Graphics::CommandBufferID /*cmd_id*/,
    const std::shared_ptr<VKBuffer>& buffer)
{
    return std::dynamic_pointer_cast<RelaxationGridBuffer>(buffer) != nullptr;
}
 
bool RelaxationStepProcessor::on_iteration(
    Portal::Graphics::CommandBufferID /*cmd_id*/,
    const std::shared_ptr<VKBuffer>& /*buffer*/,
    uint32_t /*index*/)
{
    m_dispatched = true;
    return true;
}
 
void RelaxationStepProcessor::on_dispatch_complete(const std::shared_ptr<VKBuffer>& /*buffer*/)
{
    if (m_grid) {
        serve_snapshot();
    }
}
 
void RelaxationStepProcessor::serve_snapshot()
{
    if (!m_grid->consume_snapshot_request()) {
        return;
    }
 
    const auto& front = m_grid->get_buffer_resources().back_buffers[m_grid->front_index()];
 
    std::vector<uint8_t> bytes(m_grid->get_state_bytes());
    download_back_buffer(front, bytes.data(), bytes.size(), m_snapshot_staging);
 
    m_grid->snapshot_source()->signal(bytes);
}
 
} // namespace MayaFlux::Buffers