fg.h

cross platform rendering playground

src/passes/fg.h

11.86 KB
#pragma once

#include "backend/rhi.h"
#include "core/globals.h"
#include "core/math_rtm.h"
#include "shaders/renderer_types.h"
#include <cassert>
#include <functional>
#include <string>
#include <vector>

struct ImageHandle {
    u32 id = UINT32_MAX;
    bool is_valid() const {
        return id != UINT32_MAX;
    }
};

struct BufferHandle {
    u32 id = UINT32_MAX;
    bool is_valid() const {
        return id != UINT32_MAX;
    }
};

enum class Operation : u8 {
    Read,
    Write,
    ReadWrite
};

enum class Lifetime : u8 {
    Transient,  // graph-allocated & memory aliased
    Persistent, // graph-managed across frames
    History,    // cross-frame temporal history (2 versions, ping-ponged per frame)
    Imported    // externally-managed
};

// I haven't actually tested this path much, TODO: TAA
inline u32 get_version_count(Lifetime lifetime) {
    return (lifetime == Lifetime::History) ? 2u : 1u;
}

enum class SizeOp : u8 {
    Absolute,
    Relative,
    Swapchain
};

// read/write: Operation says read vs write vs both; state says in which
// layout/access the pass needs it. No usage/stage params (implied by state).
struct ImageResourceRef {
    ImageHandle img;
    Operation operation = Operation::Read;
    rhi::ResourceState state = rhi::ResourceState::Idle;
    i32 version_offset = 0; // 0 = current, -1 = previous (History resources)
};

// read/write
struct BufferResourceRef {
    BufferHandle buf;
    Operation operation = Operation::Read;
    rhi::ResourceState state = rhi::ResourceState::Idle;
    i32 version_offset = 0; // 0 = current, -1 = previous (History resources)
};

// add_image()
struct ImageResource {
    rhi::ImageDesc desc;
    std::string name;
    Lifetime lifetime = Lifetime::Transient;
    SizeOp size_class = SizeOp::Swapchain;
    vec2 scale{1.0f, 1.0f};
    rhi::ImageViewDesc view{
        .type = rhi::ImageViewType::None
    }; // None is an RHI sentinal just used here, not sure if it'l stick around
    rhi::ResourceState entry_state = rhi::ResourceState::Idle; // entry state to reset an imported resource to
    rhi::Image *external = nullptr;                            // set by import_image(); borrowed, never owned
};

// add_buffer()
struct BufferResource {
    rhi::BufferDesc desc;
    std::string name;
    u64 (*size_fn)(u32 width, u32 height) =
        nullptr; // really only need this for cluster bounds? Not sure of a better solution yet
    Lifetime lifetime = Lifetime::Transient;
    SizeOp size_class = SizeOp::Absolute;
    vec2 scale{1.0f, 1.0f}; // TODO: nuke
    rhi::ResourceState entry_state = rhi::ResourceState::Idle;
    rhi::Buffer *external = nullptr; // set by import_buffer(); borrowed, never owned
};

constexpr u32 INVALID_VIEW_INDEX = 0xFFFFFFFFu;

struct PhysicalStorage {
    std::vector<rhi::Image> owned_images;
    std::vector<rhi::Buffer> owned_buffers;
    std::vector<rhi::Image *> images;
    std::vector<rhi::Buffer *> buffers;

    std::vector<u32> image_base;           // per logical image id: base slot index
    std::vector<u32> buffer_base;          // per logical buffer id: base slot index
    std::vector<u32> image_version_count;  // per logical image id: version stride
    std::vector<u32> buffer_version_count; // per logical buffer id: version stride

    std::vector<u32> image_view_index;

    std::vector<rhi::ImageView> transient_views;

    rhi::Image &owned_image(u32 id, u32 version) {
        return owned_images[image_base[id] + version];
    }
    rhi::Buffer &owned_buffer(u32 id, u32 version) {
        return owned_buffers[buffer_base[id] + version];
    }
    rhi::Image *image(u32 id, u32 version) {
        return images[image_base[id] + version];
    }
    rhi::Buffer *buffer(u32 id, u32 version) {
        return buffers[buffer_base[id] + version];
    }
};

struct FrameData;
struct RenderData;
struct Scene;

struct ExtraImageView {
    ImageHandle img;
    // true: one storage slot per mip (e.g. HiZ downsample targets);
    // false: a single whole-image storage slot (e.g. DDGI atlas UAV).
    bool all_mips = false;
    // version_count * mips_per_version bindless slots.
    std::vector<u32> slots;
    u32 mips_per_version = 0;
};

struct ImageBarrierBaked {
    u32 image_id;
    i32 version_offset; // 0 = current, -1 = previous (resolved per frame)
    rhi::ResourceState before;
    rhi::ResourceState after;
    rhi::ImageAspect aspect;
};

struct BufferBarrierBaked {
    u32 buffer_id;
    i32 version_offset; // 0 = current, -1 = previous (resolved per frame)
    rhi::ResourceState before;
    rhi::ResourceState after;
};

struct PassContext {
    rhi::CmdBuffer &cmd;
    rhi::Device &device;
    PhysicalStorage &storage;

    const std::vector<ExtraImageView> *extra_views = nullptr;
    FrameData *frame_data = nullptr;
    RenderData *render_data = nullptr;
    Scene *scene = nullptr;
    u32 frame_index = 0;
    u64 frame_count = 0;
    u32 render_width = 0;
    u32 render_height = 0;

    rhi::Image &get_image(ImageHandle handle, i32 version_offset = 0) const {
        u32 vc = storage.image_version_count[handle.id];
        i32 v = ((i32)frame_index + version_offset) % (i32)vc;
        if (v < 0) {
            v += (i32)vc;
        }
        return *storage.image(handle.id, (u32)v);
    }
    rhi::Buffer &get_buffer(BufferHandle handle, i32 version_offset = 0) const {
        u32 vc = storage.buffer_version_count[handle.id];
        i32 v = ((i32)frame_index + version_offset) % (i32)vc;
        if (v < 0) {
            v += (i32)vc;
        }
        return *storage.buffer(handle.id, (u32)v);
    }

    rhi::ImageView image_view(ImageHandle handle, const rhi::ImageViewDesc &desc) const;

    u32 view(ImageHandle handle, i32 version_offset = 0) const {
        u32 vc = storage.image_version_count[handle.id];
        i32 v = ((i32)frame_index + version_offset) % (i32)vc;
        if (v < 0) {
            v += (i32)vc;
        }
        u32 idx = storage.image_view_index[storage.image_base[handle.id] + (u32)v];
        assert(idx != INVALID_VIEW_INDEX && "view(): image declares no primary view (or arena missing at compile)");
        return idx;
    }

    rhi::BufferView buffer_view(BufferHandle handle, const rhi::BufferViewDesc &desc) const {
        rhi::Buffer &buf = get_buffer(handle);
        rhi::BufferViewDesc full = desc;
        full.buffer = &buf;
        if (full.size == 0) {
            full.size = buf.desc.size;
        }
        rhi::BufferView v{};
        v.buffer = &buf;
        v.desc = full;
        return v;
    }

    u32 view_extra(u32 extra_idx, i32 version_offset = 0, u32 mip = 0) const;
};

struct GraphExecInfo {
    FrameData *frame_data = nullptr;
    RenderData *render_data = nullptr;
    Scene *scene = nullptr;
    u32 frame_index = 0;
    u64 frame_count = 0;
};

using PassExecuteFn = std::function<void(PassContext &ctx)>;

struct FrameGraph;
struct PassTimings;
struct PassBuilder {
    const char *name;
    FrameGraph *graph_ = nullptr;

    // recording-time: slice ranges into image_refs / buffer_refs
    u32 image_ref_start = 0;
    u32 image_ref_count = 0;

    u32 buffer_ref_start = 0;
    u32 buffer_ref_count = 0;

    // compile-time: slice ranges into the baked barrier lists
    u32 image_barrier_start = 0;
    u32 image_barrier_count = 0;

    u32 buffer_barrier_start = 0;
    u32 buffer_barrier_count = 0;

    PassExecuteFn execute_fn = nullptr;
    // void *user_data = nullptr;

    PassBuilder &execute(PassExecuteFn fn) {
        execute_fn = fn;
        return *this;
    }

    // Explicit graph form. State implies the sync (no usage/stage params).
    PassBuilder &read(FrameGraph &g, ImageHandle img, rhi::ResourceState state);
    PassBuilder &write(FrameGraph &g, ImageHandle image, rhi::ResourceState state);
    PassBuilder &read_write(FrameGraph &g, ImageHandle image, rhi::ResourceState state);

    PassBuilder &read(FrameGraph &g, BufferHandle buf, rhi::ResourceState state);
    PassBuilder &write(FrameGraph &g, BufferHandle buffer, rhi::ResourceState state);
    PassBuilder &read_write(FrameGraph &g, BufferHandle buffer, rhi::ResourceState state);

    // Owner-aware forms: forward to the graph this pass was recorded on.
    PassBuilder &read(ImageHandle img, rhi::ResourceState state);
    PassBuilder &write(ImageHandle image, rhi::ResourceState state);
    PassBuilder &read_write(ImageHandle image, rhi::ResourceState state);

    PassBuilder &read(BufferHandle buf, rhi::ResourceState state);
    PassBuilder &write(BufferHandle buffer, rhi::ResourceState state);
    PassBuilder &read_write(BufferHandle buffer, rhi::ResourceState state);

    // history resource versioning: write the current frame's version, read the
    // previous frame's version. Declared so the planner keeps prev-read and
    // curr-write on separate physical versions (no false dependency).
    PassBuilder &write_current(FrameGraph &g, ImageHandle img, rhi::ResourceState state);
    PassBuilder &read_previous(FrameGraph &g, ImageHandle img, rhi::ResourceState state);
    PassBuilder &write_current(FrameGraph &g, BufferHandle buf, rhi::ResourceState state);
    PassBuilder &read_previous(FrameGraph &g, BufferHandle buf, rhi::ResourceState state);

    PassBuilder &write_current(ImageHandle img, rhi::ResourceState state);
    PassBuilder &read_previous(ImageHandle img, rhi::ResourceState state);
    PassBuilder &write_current(BufferHandle buf, rhi::ResourceState state);
    PassBuilder &read_previous(BufferHandle buf, rhi::ResourceState state);

    // shared ref-recording
    PassBuilder &
    add_image_ref(FrameGraph &g, ImageHandle img, rhi::ResourceState state, Operation op, i32 version_offset);
    PassBuilder &
    add_buffer_ref(FrameGraph &g, BufferHandle buf, rhi::ResourceState state, Operation op, i32 version_offset);
};

struct FrameGraph {
    rhi::Device *device = nullptr;

    u32 request_extra_view(ImageHandle img, bool all_mips);
    std::vector<ExtraImageView> extra_views;

    u32 width;
    u32 height;

    bool image_readback = true; // TODO: nuke

    // logical resources
    std::vector<ImageResource> logical_images;
    std::vector<BufferResource> logical_buffers;

    // flat lists of typed pass dependencies
    std::vector<ImageResourceRef> image_refs;
    std::vector<BufferResourceRef> buffer_refs;

    PhysicalStorage physical;
    std::vector<PassBuilder> passes;

    // baked at compile(): flat barrier lists, sliced by PassBuilder
    std::vector<ImageBarrierBaked> image_barriers;
    std::vector<BufferBarrierBaked> buffer_barriers;

    // NOTE: execute() issues one state barrier per baked edge directly;
    // no staging vectors (batching is a future backend-internal optimization).

    void init();
    bool compile(rhi::Device &device, rhi::Extent2D swapchain_extent);
    void execute(rhi::CmdBuffer &cmd, const GraphExecInfo &info);
    void destroy(); // releases all owned GPU storage, then init()
    void reset();   // destroy() + init(): full teardown, ready to re-record

    ImageHandle add_image(ImageResource img);
    BufferHandle add_buffer(BufferResource buf);
    ImageHandle import_image(
        rhi::Image *img,
        rhi::ResourceState entry_state = rhi::ResourceState::Idle,
        const rhi::ImageViewDesc &view = {.type = rhi::ImageViewType::None}
    );
    BufferHandle import_buffer(rhi::Buffer *buf, rhi::ResourceState entry_state = rhi::ResourceState::Idle);

    void overwrite_imported_image(ImageHandle h, rhi::Image *img);

    PassBuilder &add_pass(const char *name);

    PassTimings *timings = nullptr;
    void set_timings(PassTimings *t) {
        timings = t;
    }

    // First execute() after compile patches the initial barrier per slot with live resource state
    bool pending_state_fixup = false;

    // execute() scratch for fixup patching: first-touch flags per physical
    // slot, sized in compile(), cleared on each fixup execute.
    std::vector<u8> touch_img;
    std::vector<u8> touch_buf;

    // true after a successful compile(); execute() refuses to run an un-compiled
    // graph. Reset by init()/destroy()/reset().
    bool is_compiled = false;
};