GpuResourceManager.cpp

Go to the documentation of this file.

#include "GpuResourceManager.hpp"
 
#include "MayaFlux/Core/Backends/Graphics/Vulkan/VKImage.hpp"
#include "MayaFlux/Portal/Graphics/TextureLoom.hpp"
 
#include "MayaFlux/Journal/Archivist.hpp"
 
namespace MayaFlux::Yantra {
 
//==============================================================================
// Buffer slot
//==============================================================================
 
struct VulkanBufferSlot {
    vk::Buffer buffer;
    vk::DeviceMemory memory;
    void* mapped_ptr { nullptr };
    size_t allocated_bytes {};
};
 
//==============================================================================
// PIMPL — one per pipeline unit
//==============================================================================
 
struct GpuResourceManagerImpl {
    std::vector<VulkanBufferSlot> buffers;
    std::unordered_map<std::string, VulkanBufferSlot> shared_buffers;
};
 
//==============================================================================
// File-local helpers
//==============================================================================
 
namespace {
 
    std::optional<uint32_t> try_memory_type(vk::PhysicalDevice phys,
        uint32_t type_filter,
        vk::MemoryPropertyFlags props)
    {
        auto mem_props = phys.getMemoryProperties();
        for (uint32_t i = 0; i < mem_props.memoryTypeCount; ++i) {
            if ((type_filter & (1U << i))
                && (mem_props.memoryTypes[i].propertyFlags & props) == props) {
                return i;
            }
        }
        return std::nullopt;
    }
 
    uint32_t find_memory_type(vk::PhysicalDevice phys,
        uint32_t type_filter,
        vk::MemoryPropertyFlags props,
        vk::MemoryPropertyFlags fallback_props = {})
    {
        if (const auto type = try_memory_type(phys, type_filter, props))
            return *type;
        if (fallback_props) {
            if (const auto type = try_memory_type(phys, type_filter, fallback_props))
                return *type;
        }
        error<std::runtime_error>(
            Journal::Component::Yantra,
            Journal::Context::BufferProcessing,
            std::source_location::current(),
            "GpuResourceManager: no suitable memory type found");
    }
 
    void free_slot(vk::Device device, VulkanBufferSlot& slot)
    {
        if (slot.mapped_ptr) {
            device.unmapMemory(slot.memory);
            slot.mapped_ptr = nullptr;
        }
        if (slot.buffer) {
            device.destroyBuffer(slot.buffer);
            slot.buffer = vk::Buffer {};
        }
        if (slot.memory) {
            device.freeMemory(slot.memory);
            slot.memory = vk::DeviceMemory {};
        }
        slot.allocated_bytes = 0;
    }
 
    /**
     * @brief Allocate and map a host visible buffer slot.
     *
     * By default the memory is host cached system memory, which suits buffers
     * the host reads back. With prefer_device_local the slot first tries
     * device local, host visible, host coherent memory (the BAR heap), so a
     * buffer that only shaders touch is read and written at device speed
     * instead of across the bus. When no such memory type exists, or the heap
     * is full, the default memory is used, so the request never fails where
     * the default would succeed. Host reads of device local memory are
     * uncached, so it is meant for small or write only host access.
     */
    void allocate_slot(vk::Device device, vk::PhysicalDevice phys,
        VulkanBufferSlot& slot, size_t byte_size,
        vk::BufferUsageFlags extra_usage = vk::BufferUsageFlagBits::eStorageBuffer,
        bool prefer_device_local = false)
    {
        free_slot(device, slot);
 
        vk::BufferCreateInfo bi;
        bi.size = byte_size;
        bi.usage = extra_usage;
        bi.sharingMode = vk::SharingMode::eExclusive;
        slot.buffer = device.createBuffer(bi);
 
        auto req = device.getBufferMemoryRequirements(slot.buffer);
 
        vk::MemoryAllocateInfo ai;
        ai.allocationSize = req.size;
 
        if (prefer_device_local) {
            const auto resident = try_memory_type(phys, req.memoryTypeBits,
                vk::MemoryPropertyFlagBits::eDeviceLocal
                    | vk::MemoryPropertyFlagBits::eHostVisible
                    | vk::MemoryPropertyFlagBits::eHostCoherent);
            if (resident) {
                ai.memoryTypeIndex = *resident;
                try {
                    slot.memory = device.allocateMemory(ai);
                } catch (const vk::SystemError&) {
                    slot.memory = vk::DeviceMemory {};
                }
            }
        }
 
        if (!slot.memory) {
            ai.memoryTypeIndex = find_memory_type(phys, req.memoryTypeBits,
                vk::MemoryPropertyFlagBits::eHostVisible
                    | vk::MemoryPropertyFlagBits::eHostCoherent
                    | vk::MemoryPropertyFlagBits::eHostCached,
                vk::MemoryPropertyFlagBits::eHostVisible
                    | vk::MemoryPropertyFlagBits::eHostCoherent);
            slot.memory = device.allocateMemory(ai);
        }
 
        device.bindBufferMemory(slot.buffer, slot.memory, 0);
        slot.mapped_ptr = device.mapMemory(slot.memory, 0, VK_WHOLE_SIZE);
        slot.allocated_bytes = byte_size;
    }
 
    [[nodiscard]] vk::DescriptorType element_type_to_vk(GpuBufferBinding::ElementType et)
    {
        switch (et) {
        case GpuBufferBinding::ElementType::IMAGE_STORAGE:
            return vk::DescriptorType::eStorageImage;
        case GpuBufferBinding::ElementType::IMAGE_SAMPLED:
            return vk::DescriptorType::eCombinedImageSampler;
        default:
            return vk::DescriptorType::eStorageBuffer;
        }
    }
 
} // anonymous namespace
 
struct GpuResourceManager::SharedBuffers {
    std::map<std::pair<uint32_t, size_t>, VulkanBufferSlot> slots;
};
 
//==============================================================================
// Lifecycle
//==============================================================================
 
GpuResourceManager::GpuResourceManager()
    : m_shared(std::make_unique<SharedBuffers>())
{
}
 
GpuResourceManager::~GpuResourceManager()
{
    cleanup();
}
 
GpuResourceManager::PipelineUnit& GpuResourceManager::unit_for(const std::string& key)
{
    auto it = m_units.find(key);
    if (it == m_units.end()) {
        error<std::runtime_error>(
            Journal::Component::Yantra,
            Journal::Context::BufferProcessing,
            std::source_location::current(),
            "GpuResourceManager: no unit for key '{}' — call initialise() first", key);
    }
    return *it->second;
}
 
const GpuResourceManager::PipelineUnit* GpuResourceManager::find_unit(const std::string& key) const
{
    auto it = m_units.find(key);
    return it == m_units.end() ? nullptr : it->second.get();
}
 
bool GpuResourceManager::is_ready(const std::string& key) const
{
    const auto* unit = find_unit(key);
    return unit && unit->ready;
}
 
bool GpuResourceManager::initialise(const std::string& key,
    const GpuComputeConfig& config,
    const std::vector<GpuBufferBinding>& bindings)
{
    if (const auto* existing = find_unit(key); existing && existing->ready)
        return true;
 
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    auto unit = std::make_unique<PipelineUnit>();
 
    if (config.shader_id != Portal::Graphics::INVALID_SHADER) {
        unit->shader_id = config.shader_id;
    } else {
        unit->shader_id = foundry.load_shader(config.shader_path);
    }
 
    if (unit->shader_id == Portal::Graphics::INVALID_SHADER) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::BufferProcessing,
            "GpuResourceManager: failed to load shader '{}' for key '{}'",
            config.shader_path.empty() ? "<generated>" : config.shader_path, key);
        return false;
    }
 
    std::map<uint32_t, std::vector<Portal::Graphics::DescriptorBindingInfo>> by_set;
    for (const auto& b : bindings) {
        const auto et = b.element_type;
        const bool is_image = et == GpuBufferBinding::ElementType::IMAGE_STORAGE
            || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
        if (is_image) {
            by_set[b.set].push_back({
                .set = b.set,
                .binding = b.binding,
                .type = element_type_to_vk(et),
            });
        }
    }
 
    for (const auto& b : bindings) {
        const auto et = b.element_type;
        const bool is_image = et == GpuBufferBinding::ElementType::IMAGE_STORAGE
            || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
        if (!is_image) {
            by_set[b.set].push_back({
                .set = b.set,
                .binding = b.binding,
                .type = vk::DescriptorType::eStorageBuffer,
            });
        }
    }
 
    std::vector<std::vector<Portal::Graphics::DescriptorBindingInfo>> descriptor_sets;
    descriptor_sets.reserve(by_set.size());
    for (auto& [set_idx, set_bindings] : by_set)
        descriptor_sets.push_back(std::move(set_bindings));
 
    unit->pipeline_id = compute_press.create_pipeline(
        unit->shader_id, descriptor_sets, config.push_constant_size);
 
    if (unit->pipeline_id == Portal::Graphics::INVALID_COMPUTE_PIPELINE) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::BufferProcessing,
            "GpuResourceManager: failed to create pipeline for key '{}'", key);
        foundry.destroy_shader(unit->shader_id);
        return false;
    }
 
    unit->descriptor_set_ids = compute_press.allocate_pipeline_descriptors(unit->pipeline_id);
    if (unit->descriptor_set_ids.empty()) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::BufferProcessing,
            "GpuResourceManager: failed to allocate descriptor sets for key '{}'", key);
        compute_press.destroy_pipeline(unit->pipeline_id);
        foundry.destroy_shader(unit->shader_id);
        return false;
    }
 
    unit->impl = std::make_unique<GpuResourceManagerImpl>();
    size_t max_binding = 0;
    for (const auto& b : bindings)
        max_binding = std::max(max_binding, static_cast<size_t>(b.binding));
    const size_t capacity = bindings.empty() ? 0 : max_binding + 1;
    unit->impl->buffers.resize(capacity);
    unit->buffer_slots.resize(capacity);
    unit->image_slots.resize(capacity);
    unit->ready = true;
 
    m_units[key] = std::move(unit);
    return true;
}
 
void GpuResourceManager::release(const std::string& key)
{
    auto it = m_units.find(key);
    if (it == m_units.end())
        return;
 
    auto& unit = *it->second;
    if (!unit.ready) {
        m_units.erase(it);
        return;
    }
 
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
    auto device = foundry.get_device();
 
    if (unit.impl) {
        for (auto& slot : unit.impl->buffers)
            free_slot(device, slot);
    }
    unit.impl.reset();
    unit.buffer_slots.clear();
    unit.image_slots.clear();
 
    if (unit.pipeline_id != Portal::Graphics::INVALID_COMPUTE_PIPELINE)
        compute_press.destroy_pipeline(unit.pipeline_id);
 
    if (unit.shader_id != Portal::Graphics::INVALID_SHADER)
        foundry.destroy_shader(unit.shader_id);
 
    m_units.erase(it);
}
 
void GpuResourceManager::cleanup()
{
    std::vector<std::string> keys;
    keys.reserve(m_units.size());
    for (const auto& [k, v] : m_units)
        keys.push_back(k);
    for (const auto& k : keys)
        release(k);
}
 
//==============================================================================
// Buffer operations
//==============================================================================
 
void GpuResourceManager::ensure_buffer(const std::string& key, size_t index,
    size_t required_bytes, Portal::Graphics::BufferUsageHint usage_hint)
{
    auto& unit = unit_for(key);
    auto& vk_slot = unit.impl->buffers[index];
    if (vk_slot.allocated_bytes >= required_bytes) {
        return;
    }
 
    auto& foundry = Portal::Graphics::get_shader_foundry();
    allocate_slot(foundry.get_device(), foundry.get_physical_device(),
        vk_slot, required_bytes, Portal::Graphics::to_buffer_usage_flags(usage_hint));
 
    unit.buffer_slots[index].allocated_bytes = required_bytes;
}
 
void GpuResourceManager::upload(const std::string& key, size_t index, const float* data, size_t byte_size)
{
    auto& vk_slot = unit_for(key).impl->buffers[index];
    std::memcpy(vk_slot.mapped_ptr, data, byte_size);
}
 
void GpuResourceManager::upload_raw(const std::string& key, size_t index, const uint8_t* data, size_t byte_size)
{
    auto& vk_slot = unit_for(key).impl->buffers[index];
    std::memcpy(vk_slot.mapped_ptr, data, byte_size);
}
 
void GpuResourceManager::download(const std::string& key, size_t index, float* dest, size_t byte_size)
{
    auto& vk_slot = unit_for(key).impl->buffers[index];
    std::memcpy(dest, vk_slot.mapped_ptr, byte_size);
}
 
void GpuResourceManager::bind_descriptor(const std::string& key, size_t index, const GpuBufferBinding& spec)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& vk_slot = unit.impl->buffers[index];
 
    foundry.update_descriptor_buffer(
        unit.descriptor_set_ids[spec.set],
        spec.binding,
        vk::DescriptorType::eStorageBuffer,
        vk_slot.buffer, 0, vk_slot.allocated_bytes);
}
 
void GpuResourceManager::bind_descriptors_batch(const std::string& key, const std::vector<GpuBufferBinding>& bindings)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
 
    std::vector<Portal::Graphics::DescriptorBufferWrite> writes;
    writes.reserve(bindings.size());
 
    for (const auto& b : bindings) {
        const auto et = b.element_type;
        if (et == GpuBufferBinding::ElementType::IMAGE_STORAGE
            || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED) {
            continue;
        }
 
        const auto shared_it = m_shared->slots.find({ b.set, static_cast<size_t>(b.binding) });
        vk::Buffer buffer;
        size_t allocated_bytes = 0;
        if (shared_it != m_shared->slots.end()) {
            buffer = shared_it->second.buffer;
            allocated_bytes = shared_it->second.allocated_bytes;
        } else {
            if (static_cast<size_t>(b.binding) >= unit.impl->buffers.size()) {
                continue;
            }
            const auto& vk_slot = unit.impl->buffers[b.binding];
            buffer = vk_slot.buffer;
            allocated_bytes = vk_slot.allocated_bytes;
        }
 
        if (!buffer)
            continue;
 
        writes.push_back({
            .descriptor_set_id = unit.descriptor_set_ids[b.set],
            .binding = b.binding,
            .type = vk::DescriptorType::eStorageBuffer,
            .buffer = buffer,
            .offset = 0,
            .size = allocated_bytes,
        });
    }
 
    foundry.update_descriptor_buffers(writes);
}
 
size_t GpuResourceManager::buffer_allocated_bytes(const std::string& key, size_t index) const
{
    return find_unit(key)->buffer_slots[index].allocated_bytes;
}
 
void GpuResourceManager::ensure_shared_buffer(uint32_t set, size_t binding_index, size_t element_count,
    GpuBufferBinding::ElementType element_type,
    Portal::Graphics::BufferUsageHint usage_hint)
{
    const size_t width = Portal::Graphics::element_type_bytes(element_type);
    if (width == 0) {
        error<std::runtime_error>(
            Journal::Component::Yantra,
            Journal::Context::BufferProcessing,
            std::source_location::current(),
            "GpuResourceManager: ensure_shared_buffer requires a sized element_type");
    }
 
    auto& slot = m_shared->slots[{ set, binding_index }];
    const size_t required_bytes = element_count * width;
    if (slot.allocated_bytes >= required_bytes)
        return;
 
    auto& foundry = Portal::Graphics::get_shader_foundry();
    const bool gpu_resident = usage_hint == Portal::Graphics::BufferUsageHint::DEVICE
        || usage_hint == Portal::Graphics::BufferUsageHint::COMPUTE;
    allocate_slot(foundry.get_device(), foundry.get_physical_device(),
        slot, required_bytes, Portal::Graphics::to_buffer_usage_flags(usage_hint), gpu_resident);
}
 
Portal::Graphics::GpuBufferHandle GpuResourceManager::shared_buffer_handle(uint32_t set, size_t binding_index) const
{
    const auto it = m_shared->slots.find({ set, binding_index });
    if (it == m_shared->slots.end())
        return {};
    return { .buffer = it->second.buffer, .mapped_ptr = it->second.mapped_ptr, .size_bytes = it->second.allocated_bytes };
}
 
void GpuResourceManager::bind_shared_descriptor(const std::string& key, uint32_t set, size_t binding_index, const GpuBufferBinding& spec)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& slot = m_shared->slots.at({ set, binding_index });
 
    foundry.update_descriptor_buffer(
        unit.descriptor_set_ids[spec.set],
        spec.binding,
        vk::DescriptorType::eStorageBuffer,
        slot.buffer, 0, slot.allocated_bytes);
}
 
void GpuResourceManager::download_shared(uint32_t set, size_t binding_index, void* dest, size_t byte_size)
{
    auto& slot = m_shared->slots.at({ set, binding_index });
    std::memcpy(dest, slot.mapped_ptr, byte_size);
}
 
void GpuResourceManager::upload_shared_raw(uint32_t set, size_t binding_index, const uint8_t* data, size_t byte_size)
{
    auto& slot = m_shared->slots.at({ set, binding_index });
    std::memcpy(slot.mapped_ptr, data, byte_size);
}
 
Portal::Graphics::HazardResource GpuResourceManager::make_shared_buffer_hazard(
    const GpuBufferBinding& spec) const
{
    const auto it = m_shared->slots.find({ spec.set, static_cast<size_t>(spec.binding) });
    vk::Buffer handle = it != m_shared->slots.end() ? it->second.buffer : vk::Buffer {};
    return Portal::Graphics::HazardResource {
        .binding = spec,
        .image = nullptr,
        .buffer = handle,
    };
}
 
void GpuResourceManager::bind_image_storage(
    const std::string& key, size_t index,
    const std::shared_ptr<Core::VKImage>& image,
    const GpuBufferBinding& spec)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
 
    if (index >= unit.image_slots.size())
        unit.image_slots.resize(index + 1);
    unit.image_slots[index] = image;
 
    foundry.update_descriptor_storage_image(
        unit.descriptor_set_ids[spec.set],
        spec.binding,
        image->get_image_view(),
        vk::ImageLayout::eGeneral);
}
 
void GpuResourceManager::bind_image_sampled(
    const std::string& key, size_t index,
    const std::shared_ptr<Core::VKImage>& image,
    vk::Sampler sampler,
    const GpuBufferBinding& spec)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
 
    if (index >= unit.image_slots.size())
        unit.image_slots.resize(index + 1);
    unit.image_slots[index] = image;
 
    foundry.update_descriptor_image(
        unit.descriptor_set_ids[spec.set],
        spec.binding,
        image->get_image_view(),
        sampler,
        vk::ImageLayout::eShaderReadOnlyOptimal);
}
 
void GpuResourceManager::bind_images_batch(const std::string& key, const std::vector<ImageBind>& images)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
 
    std::vector<Portal::Graphics::DescriptorImageWrite> writes;
    writes.reserve(images.size());
 
    for (const auto& entry : images) {
        if (entry.index >= unit.image_slots.size())
            unit.image_slots.resize(entry.index + 1);
        unit.image_slots[entry.index] = entry.image;
 
        const bool is_storage = entry.spec.element_type == GpuBufferBinding::ElementType::IMAGE_STORAGE;
        writes.push_back({
            .descriptor_set_id = unit.descriptor_set_ids[entry.spec.set],
            .binding = entry.spec.binding,
            .type = is_storage ? vk::DescriptorType::eStorageImage : vk::DescriptorType::eCombinedImageSampler,
            .image_view = entry.image->get_image_view(),
            .sampler = entry.sampler,
            .layout = is_storage ? vk::ImageLayout::eGeneral : vk::ImageLayout::eShaderReadOnlyOptimal,
        });
    }
 
    foundry.update_descriptor_images(writes);
}
 
void GpuResourceManager::transition_image(
    const std::shared_ptr<Core::VKImage>& image,
    vk::ImageLayout old_layout,
    vk::ImageLayout new_layout)
{
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& backend = Portal::Graphics::get_texture_manager(); // TextureLoom -> backend ref
 
    backend.transition_layout(
        image,
        old_layout,
        new_layout,
        1, 1, vk::ImageAspectFlagBits::eColor);
}
 
//==============================================================================
// Dispatch
//==============================================================================
 
void GpuResourceManager::dispatch(const std::string& key,
    const std::array<uint32_t, 3>& groups,
    const std::vector<GpuBufferBinding>& bindings,
    const uint8_t* push_constant_data,
    size_t push_constant_size)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    auto cmd_id = foundry.begin_commands(
        Portal::Graphics::ShaderFoundry::CommandBufferType::COMPUTE);
 
    compute_press.bind_all(
        cmd_id, unit.pipeline_id, unit.descriptor_set_ids,
        push_constant_data, push_constant_size);
 
    compute_press.dispatch(cmd_id, groups[0], groups[1], groups[2]);
 
    for (const auto& b : bindings) {
        const auto et = b.element_type;
        const bool is_image = et == GpuBufferBinding::ElementType::IMAGE_STORAGE
            || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
        const bool is_output = b.direction == GpuBufferBinding::Direction::OUTPUT
            || b.direction == GpuBufferBinding::Direction::INPUT_OUTPUT;
 
        const bool is_shared = m_shared->slots.contains({ b.set, static_cast<size_t>(b.binding) });
        if (is_output && !is_image && !is_shared
            && static_cast<size_t>(b.binding) < unit.impl->buffers.size()
            && unit.impl->buffers[b.binding].buffer) {
            foundry.buffer_barrier(
                cmd_id,
                unit.impl->buffers[b.binding].buffer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eHostRead,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eHost);
        }
    }
 
    foundry.submit_and_wait(cmd_id);
}
 
void GpuResourceManager::dispatch_batched(const std::string& key,
    const std::array<uint32_t, 3>& groups,
    const std::vector<GpuBufferBinding>& bindings,
    size_t push_constant_size,
    const ExecutionContext& ctx)
{
    const auto& params = safe_variant_get_or_throw<ChainedParams>(ctx.parameters,
        "GpuResourceManager: dispatch_batched requires ChainedParams");
 
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    const uint32_t workgroups_per_pass = groups[0] * groups[1] * groups[2];
 
    const uint32_t default_passes = std::max(1U, 65536U / std::max(1U, workgroups_per_pass));
 
    const uint32_t effective_passes_per_batch = params.passes_per_batch.value_or(default_passes);
    for (uint32_t base = 0; base < params.pass_count; base += effective_passes_per_batch) {
        const uint32_t batch_end = std::min(base + effective_passes_per_batch, params.pass_count);
        auto cmd_id = foundry.begin_commands(
            Portal::Graphics::ShaderFoundry::CommandBufferType::COMPUTE);
 
        for (uint32_t pass = base; pass < batch_end; ++pass) {
            std::vector<uint8_t> pc_data(push_constant_size);
            params.pc_updater(pass, pc_data.data());
            compute_press.bind_all(
                cmd_id, unit.pipeline_id, unit.descriptor_set_ids,
                pc_data.data(), push_constant_size);
            compute_press.dispatch(cmd_id, groups[0], groups[1], groups[2]);
 
            for (const auto& b : bindings) {
                if (b.direction != GpuBufferBinding::Direction::INPUT_OUTPUT)
                    continue;
 
                const bool is_image = b.element_type == GpuBufferBinding::ElementType::IMAGE_STORAGE
                    || b.element_type == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
                const bool is_shared = m_shared->slots.contains({ b.set, static_cast<size_t>(b.binding) });
 
                if (is_image) {
                    if (static_cast<size_t>(b.binding) < unit.image_slots.size() && unit.image_slots[b.binding]) {
                        foundry.image_barrier(
                            cmd_id,
                            unit.image_slots[b.binding]->get_image(),
                            vk::ImageLayout::eGeneral,
                            vk::ImageLayout::eGeneral,
                            vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead,
                            vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead,
                            vk::PipelineStageFlagBits::eComputeShader,
                            vk::PipelineStageFlagBits::eComputeShader);
                    }
                } else if (!is_shared && static_cast<size_t>(b.binding) < unit.impl->buffers.size()
                    && unit.impl->buffers[b.binding].buffer) {
                    foundry.buffer_barrier(
                        cmd_id,
                        unit.impl->buffers[b.binding].buffer,
                        vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead,
                        vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead,
                        vk::PipelineStageFlagBits::eComputeShader,
                        vk::PipelineStageFlagBits::eComputeShader);
                }
            }
        }
 
        for (const auto& b : bindings) {
            if (b.direction == GpuBufferBinding::Direction::OUTPUT
                || b.direction == GpuBufferBinding::Direction::INPUT_OUTPUT) {
                const bool is_image = b.element_type == GpuBufferBinding::ElementType::IMAGE_STORAGE
                    || b.element_type == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
                const bool is_shared = m_shared->slots.contains({ b.set, static_cast<size_t>(b.binding) });
                if (!is_image && !is_shared && static_cast<size_t>(b.binding) < unit.impl->buffers.size()
                    && unit.impl->buffers[b.binding].buffer) {
                    foundry.buffer_barrier(
                        cmd_id,
                        unit.impl->buffers[b.binding].buffer,
                        vk::AccessFlagBits::eShaderWrite,
                        vk::AccessFlagBits::eHostRead,
                        vk::PipelineStageFlagBits::eComputeShader,
                        vk::PipelineStageFlagBits::eHost);
                }
            }
        }
 
        foundry.submit_and_wait(cmd_id);
    }
}
 
void GpuResourceManager::dispatch_batched_indirect(const std::string& key,
    uint32_t indirect_set, size_t indirect_binding,
    const std::array<uint32_t, 3>& groups,
    const std::vector<GpuBufferBinding>& bindings,
    size_t push_constant_size,
    const ExecutionContext& ctx)
{
    const auto& params = safe_variant_get_or_throw<ChainedIndirectParams>(ctx.parameters,
        "GpuResourceManager: dispatch_batched_indirect requires ChainedIndirectParams");
 
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    auto& indirect_slot = m_shared->slots.at({ indirect_set, indirect_binding });
    const vk::Buffer indirect_buffer = indirect_slot.buffer;
    const uint32_t init_cmd[3] = { groups[0], groups[1], groups[2] };
    std::memcpy(indirect_slot.mapped_ptr, init_cmd, sizeof(init_cmd));
 
    const uint32_t default_passes = std::max(1U, 65536U / std::max(1U, groups[0] * groups[1] * groups[2]));
    const uint32_t effective_passes_per_batch = params.passes_per_batch.value_or(default_passes);
 
    for (uint32_t base = 0; base < params.pass_count; base += effective_passes_per_batch) {
        const uint32_t batch_end = std::min(base + effective_passes_per_batch, params.pass_count);
        auto cmd_id = foundry.begin_commands(Portal::Graphics::ShaderFoundry::CommandBufferType::COMPUTE);
 
        for (uint32_t pass = base; pass < batch_end; ++pass) {
            std::vector<uint8_t> pc_data(push_constant_size);
            params.pc_updater(pass, 1, pc_data.data());
            compute_press.bind_all(cmd_id, unit.pipeline_id, unit.descriptor_set_ids, pc_data.data(), push_constant_size);
            compute_press.dispatch_indirect(cmd_id, indirect_buffer);
 
            for (const auto& b : bindings) {
                const auto et = b.element_type;
                const bool is_image = et == GpuBufferBinding::ElementType::IMAGE_STORAGE
                    || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
                const bool is_output = b.direction == GpuBufferBinding::Direction::OUTPUT
                    || b.direction == GpuBufferBinding::Direction::INPUT_OUTPUT;
                const bool is_shared = m_shared->slots.contains({ b.set, static_cast<size_t>(b.binding) });
                if (is_output && !is_image && !is_shared
                    && static_cast<size_t>(b.binding) < unit.impl->buffers.size()
                    && unit.impl->buffers[b.binding].buffer) {
                    foundry.buffer_barrier(
                        cmd_id,
                        unit.impl->buffers[b.binding].buffer,
                        vk::AccessFlagBits::eShaderWrite,
                        vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite,
                        vk::PipelineStageFlagBits::eComputeShader,
                        vk::PipelineStageFlagBits::eComputeShader);
                }
            }
        }
        foundry.submit_and_wait(cmd_id);
    }
}
 
Portal::Graphics::FenceID GpuResourceManager::dispatch_async(const std::string& key,
    const std::array<uint32_t, 3>& groups,
    const std::vector<GpuBufferBinding>& bindings,
    const uint8_t* push_constant_data,
    size_t push_constant_size)
{
    auto& unit = unit_for(key);
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    auto cmd_id = foundry.begin_commands(
        Portal::Graphics::ShaderFoundry::CommandBufferType::COMPUTE);
 
    compute_press.bind_all(
        cmd_id, unit.pipeline_id, unit.descriptor_set_ids,
        push_constant_data, push_constant_size);
 
    compute_press.dispatch(cmd_id, groups[0], groups[1], groups[2]);
 
    for (const auto& b : bindings) {
        const auto et = b.element_type;
        const bool is_image = et == GpuBufferBinding::ElementType::IMAGE_STORAGE
            || et == GpuBufferBinding::ElementType::IMAGE_SAMPLED;
        const bool is_output = b.direction == GpuBufferBinding::Direction::OUTPUT
            || b.direction == GpuBufferBinding::Direction::INPUT_OUTPUT;
        const bool is_shared = m_shared->slots.contains({ b.set, static_cast<size_t>(b.binding) });
 
        if (is_output && !is_image && !is_shared
            && static_cast<size_t>(b.binding) < unit.impl->buffers.size()
            && unit.impl->buffers[b.binding].buffer) {
            foundry.buffer_barrier(
                cmd_id,
                unit.impl->buffers[b.binding].buffer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eHostRead,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eHost);
        }
    }
 
    return foundry.submit_async(cmd_id);
}
 
Portal::Graphics::CommandBufferID GpuResourceManager::record_sequence_commands(
    const std::vector<std::string>& keys,
    const std::vector<std::array<uint32_t, 3>>& groups_per_key,
    const std::vector<std::vector<uint8_t>>& push_constants_per_key,
    const std::vector<std::vector<Portal::Graphics::HazardResource>>& hazards_per_key,
    const std::vector<Portal::Graphics::IndirectDispatch>& indirect_per_key)
{
    auto& foundry = Portal::Graphics::get_shader_foundry();
    auto& compute_press = Portal::Graphics::get_compute_press();
 
    std::vector<Portal::Graphics::ComputeStage> stages;
    stages.reserve(keys.size());
 
    for (size_t i = 0; i < keys.size(); ++i) {
        auto& unit = unit_for(keys[i]);
        stages.push_back(Portal::Graphics::ComputeStage {
            .pipeline_id = unit.pipeline_id,
            .descriptor_set_ids = unit.descriptor_set_ids,
            .groups = groups_per_key[i],
            .push_constant_data = push_constants_per_key[i],
            .hazard_resources = hazards_per_key[i],
            .indirect = i < indirect_per_key.size() ? indirect_per_key[i] : Portal::Graphics::IndirectDispatch {},
        });
    }
 
    auto cmd_id = foundry.begin_commands(
        Portal::Graphics::ShaderFoundry::CommandBufferType::COMPUTE);
    compute_press.record_sequence(cmd_id, stages);
    return cmd_id;
}
 
void GpuResourceManager::dispatch_sequence(
    const std::vector<std::string>& keys,
    const std::vector<std::array<uint32_t, 3>>& groups_per_key,
    const std::vector<std::vector<uint8_t>>& push_constants_per_key,
    const std::vector<std::vector<Portal::Graphics::HazardResource>>& hazards_per_key,
    const std::vector<Portal::Graphics::IndirectDispatch>& indirect_per_key)
{
    auto& foundry = Portal::Graphics::get_shader_foundry();
    foundry.submit_and_wait(
        record_sequence_commands(keys, groups_per_key, push_constants_per_key, hazards_per_key, indirect_per_key));
}
 
Portal::Graphics::FenceID GpuResourceManager::dispatch_sequence_async(
    const std::vector<std::string>& keys,
    const std::vector<std::array<uint32_t, 3>>& groups_per_key,
    const std::vector<std::vector<uint8_t>>& push_constants_per_key,
    const std::vector<std::vector<Portal::Graphics::HazardResource>>& hazards_per_key,
    const std::vector<Portal::Graphics::IndirectDispatch>& indirect_per_key)
{
    auto& foundry = Portal::Graphics::get_shader_foundry();
    return foundry.submit_async(
        record_sequence_commands(keys, groups_per_key, push_constants_per_key, hazards_per_key, indirect_per_key));
}
 
} // namespace MayaFlux::Yantra