path: root/hk21/vulkan/sync/sync.cpp

// Copyright 2025 Steven Le Rouzic
//
// SPDX-License-Identifier: BSD-3-Clause

#include "hk21/vulkan/sync/sync.hpp"

#include <asl/types/option.hpp>

// All of this is largely inspired by nicegraf's synchronization utility.
// See https://github.com/nicebyte/nicegraf/blob/3a291433fdb4fd9cf38356f297ff1d851617f0f5/source/ngf-vk/impl.c#L2718

namespace
{

enum StageAccess : uint32_t
{
    kClearStageTransferWrite    = 0x0000'0001U,
    kFragmentStageShaderSampled = 0x0000'0002U,
    kVertexStageShaderSampled   = 0x0000'0004U,
    kColorAttachmentWrite       = 0x0000'0008U,
};

constexpr VkAccessFlags kWriteAccessMask =
      VK_ACCESS_2_SHADER_WRITE_BIT
    | VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT
    | VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
    | VK_ACCESS_2_TRANSFER_WRITE_BIT
    | VK_ACCESS_2_HOST_WRITE_BIT
    | VK_ACCESS_2_MEMORY_WRITE_BIT
    | VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT;

struct UsageInfo
{
    uint32_t              stage_access_mask{};
    VkAccessFlags2        access_flags{};
    VkPipelineStageFlags2 pipeline_stage_flags{};
    VkImageLayout         image_layout{};
};

const auto kUsageInfos = ([]() static {
    using namespace vulkan_sync;

    static UsageInfo info[asl::to_underlying(Usage::kCount_)]{};
    
    info[asl::to_underlying(Usage::kImageClear)] = UsageInfo{
        .stage_access_mask     = StageAccess::kClearStageTransferWrite,
        .access_flags          = VK_ACCESS_2_TRANSFER_WRITE_BIT,
        .pipeline_stage_flags  = VK_PIPELINE_STAGE_2_CLEAR_BIT,
        .image_layout          = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
    };
    
    info[asl::to_underlying(Usage::kImagePresent)] = UsageInfo{
        .stage_access_mask     = 0,
        .access_flags          = VK_ACCESS_2_NONE,
        .pipeline_stage_flags  = VK_PIPELINE_STAGE_2_NONE,
        .image_layout          = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
    };
    
    info[asl::to_underlying(Usage::kImageSampledInFragmentShader)] = UsageInfo{
        .stage_access_mask     = StageAccess::kFragmentStageShaderSampled,
        .access_flags          = VK_ACCESS_2_SHADER_SAMPLED_READ_BIT,
        .pipeline_stage_flags  = VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT,
        .image_layout          = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    };
    
    info[asl::to_underlying(Usage::kImageSampledInVertexShader)] = UsageInfo{
        .stage_access_mask     = StageAccess::kVertexStageShaderSampled,
        .access_flags          = VK_ACCESS_2_SHADER_SAMPLED_READ_BIT,
        .pipeline_stage_flags  = VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT,
        .image_layout          = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    };
    
    info[asl::to_underlying(Usage::kImageColorWriteAttachment)] = UsageInfo{
        .stage_access_mask     = StageAccess::kColorAttachmentWrite,
        .access_flags          = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
        .pipeline_stage_flags  = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
        .image_layout          = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
    };
    
    return asl::span(info);
})();



// We use an image barrier as a common structure for image and buffer barriers.
// The only differences are the resource, the subresource range, and the image layouts.
// We just discard whatever we don't need and fill the more specific fields outside.
asl::option<VkImageMemoryBarrier2> synchronize_resource_(
    vulkan_sync::ResourceState* state,
    VkImageLayout* state_layout,
    vulkan_sync::Usage new_usage)
{
    const UsageInfo& usage_info = kUsageInfos[asl::to_underlying(new_usage)];

    const bool is_read_only_access = (usage_info.access_flags & kWriteAccessMask) == 0U;

    const bool needs_layout_transition = *state_layout != usage_info.image_layout;
   
    const bool needs_write = needs_layout_transition || !is_read_only_access;

    VkImageMemoryBarrier2 barrier{
        .sType                    = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
        .pNext                    = nullptr,
        .srcStageMask             = 0,
        .srcAccessMask            = 0,
        .dstStageMask             = 0,
        .dstAccessMask            = 0,
        .oldLayout                = VK_IMAGE_LAYOUT_UNDEFINED,
        .newLayout                = VK_IMAGE_LAYOUT_UNDEFINED,
        .srcQueueFamilyIndex      = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex      = VK_QUEUE_FAMILY_IGNORED,
        .image                    = VK_NULL_HANDLE,
        .subresourceRange         = {},
    };

    if (needs_write)
    {
        barrier.srcStageMask  |= asl::exchange(state->active_readers_pipeline_stage_mask, VK_PIPELINE_STAGE_2_NONE);
        barrier.srcAccessMask |= asl::exchange(state->active_readers_access_mask, VK_ACCESS_2_NONE);
        state->has_seen_last_write = 0;
        
        // If there was to read since last write, but there was a write,
        // synchronize with last write instead.
        if (barrier.srcStageMask == VK_PIPELINE_STAGE_2_NONE &&
            state->last_writer_pipeline_stage_mask != VK_PIPELINE_STAGE_2_NONE)
        {
            barrier.srcStageMask  |= state->last_writer_pipeline_stage_mask;
            barrier.srcAccessMask |= state->last_writer_access_mask;
        }

        // Last write is now the new usage.
        state->last_writer_pipeline_stage_mask = usage_info.pipeline_stage_flags;
        state->last_writer_access_mask         = usage_info.access_flags;

        // If this is a read-only that is considered a write (layout transition)
        // we also record the reader info, because it acts as if this read has been
        // synchronized.
        if (is_read_only_access)
        {
            state->has_seen_last_write                |= usage_info.stage_access_mask;
            state->active_readers_access_mask         |= usage_info.access_flags;
            state->active_readers_pipeline_stage_mask |= usage_info.pipeline_stage_flags;
        }
    }
    else
    {
        // If there was a previous write we need to synchronize with, and this
        // access has not been synchronized with it yet, synchronize.
        if (state->last_writer_pipeline_stage_mask != VK_PIPELINE_STAGE_2_NONE
            && (state->has_seen_last_write & usage_info.stage_access_mask) != usage_info.stage_access_mask)
        {
            barrier.srcStageMask  |= state->last_writer_pipeline_stage_mask;
            barrier.srcAccessMask |= state->last_writer_access_mask;
        }

        // Record this reader info.
        state->has_seen_last_write                |= usage_info.stage_access_mask;
        state->active_readers_access_mask         |= usage_info.access_flags;
        state->active_readers_pipeline_stage_mask |= usage_info.pipeline_stage_flags;
    }

    // If the barrier has been filled or we need a layout transition, emit a barrier.
    if (barrier.srcStageMask != VK_PIPELINE_STAGE_2_NONE || needs_layout_transition)
    {
        barrier.dstStageMask  |= usage_info.pipeline_stage_flags;
        barrier.dstAccessMask |= usage_info.access_flags;

        if (barrier.dstStageMask == VK_PIPELINE_STAGE_2_NONE)
        {
            barrier.dstStageMask = VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT;
        }

        if (barrier.srcStageMask == VK_PIPELINE_STAGE_2_NONE)
        {
            barrier.srcStageMask = VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT;
        }

        if (needs_layout_transition)
        {
            barrier.oldLayout = asl::exchange(*state_layout, usage_info.image_layout);
            barrier.newLayout = usage_info.image_layout;
        }

        return barrier;
    }

    return asl::nullopt;
}

} // anonymous namespace

namespace vulkan_sync
{

void synchronize_resource(
    VkImage image, VkImageAspectFlags aspects,
    ImageState* state, Usage new_usage, DependencyInfoBuilder* builder)
{
    const UsageInfo& usage_info = kUsageInfos[asl::to_underlying(new_usage)];
    ASL_ASSERT(usage_info.image_layout != VK_IMAGE_LAYOUT_UNDEFINED);

    auto barrier_opt = synchronize_resource_(state, &state->current_layout, new_usage);
    if (barrier_opt.has_value())
    {
        auto& barrier = barrier_opt.value();

        barrier.image            = image;
        barrier.subresourceRange = {
            .aspectMask          = aspects,
            .baseMipLevel        = 0,
            .levelCount          = VK_REMAINING_MIP_LEVELS,
            .baseArrayLayer      = 0,
            .layerCount          = VK_REMAINING_ARRAY_LAYERS,
        };
        
        builder->add_image_barrier(barrier);
    }
}

void synchronize_resource(
    VkBuffer buffer, BufferState* state,
    Usage new_usage, DependencyInfoBuilder* builder)
{
    const UsageInfo& usage_info = kUsageInfos[asl::to_underlying(new_usage)];
    ASL_ASSERT(usage_info.image_layout == VK_IMAGE_LAYOUT_UNDEFINED);

    VkImageLayout dummy_layout = VK_IMAGE_LAYOUT_UNDEFINED;
    auto barrier_opt = synchronize_resource_(state, &dummy_layout, new_usage);
    if (barrier_opt.has_value())
    {
        const auto& image_barrier = barrier_opt.value();

        VkBufferMemoryBarrier2 barrier{
            .sType                    = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
            .pNext                    = nullptr,
            .srcStageMask             = image_barrier.srcStageMask,
            .srcAccessMask            = image_barrier.srcAccessMask,
            .dstStageMask             = image_barrier.dstStageMask,
            .dstAccessMask            = image_barrier.dstAccessMask,
            .srcQueueFamilyIndex      = image_barrier.srcQueueFamilyIndex,
            .dstQueueFamilyIndex      = image_barrier.dstQueueFamilyIndex,
            .buffer                   = buffer,
            .offset                   = 0,
            .size                     = VK_WHOLE_SIZE,
        };
        
        builder->add_buffer_barrier(barrier);
    }
}

} // namespace vulkan_sync