id_blit.cpp

#include "id_blit.h"

#include "passes/fg.h"
#include "pipelines.h"
#include "renderer.h"
#include "shaders/renderer_types.h"

namespace id_blit {

ImageHandle record(
    rhi::Device &device,
    Pipelines &pipelines,
    rhi::ShaderCompiler &sc,
    FrameGraph &fg,
    BufferHandle indirect,
    BufferHandle draw_count,
    BufferHandle visible,
    ImageHandle depth
) {

    ImageHandle out = fg.add_image(
        {.desc = {.format = rhi::ImageFormat::R32_UINT}, .name = "fg/idblit/id", .size_class = SizeOp::Swapchain}
    );

    rhi::RenderPipelineDesc id_blit_p;
    id_blit_p.device = &device;
    id_blit_p.set_shaders(sc, "src/passes/id_blit/id_blit.slang", "vtx_main", "fs_main");
    id_blit_p.set_cull_mode(rhi::PipelineCullMode::NONE, rhi::PipelineTriangleWindingOrder::CW);
    id_blit_p.set_color_format(rhi::ImageFormat::R32_UINT);
    rhi::RenderPipeline &pipeline = pipelines.add_render("IDBlitPipeline", id_blit_p);

    auto &builder = fg.add_pass("ID_Blit");
    builder.read(depth, rhi::ResourceState::TextureSample);
    builder.read(indirect, rhi::ResourceState::IndirectFetch);
    builder.read(draw_count, rhi::ResourceState::IndirectFetch);
    builder.read(visible, rhi::ResourceState::StorageRead);
    builder.write(out, rhi::ResourceState::ColorDraw);

    builder.execute([=, &pipeline](PassContext &ctx) {
        FrameData &frame = ctx.frame_data[ctx.frame_index];

        auto display_view = ctx.image_view(out, {.aspect = rhi::ImageAspect::Color});

        rhi::AttachmentDesc color_att{};
        color_att.img = &display_view;
        color_att.loadOp = rhi::LoadOp::CLEAR;
        color_att.storeOp = rhi::StoreOp::STORE;
        color_att.clearValue.color.uint32[0] = 0u; // 0 = "no hit"; editor subtracts 1 from stored ids

        rhi::RenderingInfo ri{};
        ri.colorAttachments[0] = color_att;
        ri.colorAttachmentCount = 1;
        ri.width = ctx.render_width;
        ri.height = ctx.render_height;

        rhi::begin_rendering(ctx.cmd, ri);
        rhi::set_pipeline(ctx.cmd, pipeline);

        rhi::Viewport vp{};
        vp.width = (f32)ctx.render_width;
        vp.height = (f32)ctx.render_height;
        rhi::set_viewports(ctx.cmd, &vp);

        rhi::Rect scissor{};
        scissor.width = (i32)ctx.render_width;
        scissor.height = (i32)ctx.render_height;
        rhi::set_scissors(ctx.cmd, &scissor);

        IdBlitPushConstants pc{};
        pc.frame = (u64)ctx.render_data->frame_ubo_view[ctx.frame_index].slot;
        pc.vertex = (u64)ctx.render_data->vertex_view.slot;
        pc.visible_items = rhi::buffer_device_address(ctx.device, ctx.get_buffer(visible));
        pc.depth = (u32)ctx.view(depth);

        {
            static_assert(sizeof(pc) <= 128, "push constants exceed global 128B range");
            char pc128[128] = {};
            std::memcpy(pc128, &pc, sizeof(pc));
            rhi::set_constants(ctx.cmd, pipeline, rhi::ShaderStage::ALL, sizeof(pc128), pc128);
        }

        rhi::set_index_buffer(ctx.cmd, ctx.render_data->index_buffer);
        rhi::draw_indexed_indirect(
            ctx.cmd,
            ctx.get_buffer(indirect),
            &ctx.get_buffer(draw_count),
            ctx.render_data->render_item_count,
            rhi::DRAW_INDEXED_INDIRECT_STRIDE,
            0,
            0
        );

        rhi::end_rendering(ctx.cmd);
    });

    return out;
};

} // namespace id_blit