ibl.cpp

cross platform rendering playground

src/scene/ibl.cpp

15.79 KB
#include "ibl.h"

#include <array>
#include <cstring>

#include "core/globals.h"
#include "passes/fg.h"
#include "quick_submit.h"
#include "shaders/renderer_types.h"

static void transition_whole(rhi::CmdBuffer &cmd, rhi::Image &img, rhi::ResourceState after) {
    if (img.state == after) {
        return;
    }
    rhi::ImageRange range{};
    rhi::barrier(cmd, img, range, img.state, after);
}

static void barrier_mip_layers(
    rhi::CmdBuffer &cmd, rhi::Image &img, u32 mip, u32 layers, rhi::ResourceState before, rhi::ResourceState after
) {
    rhi::ImageRange range{};
    range.mip = mip;
    range.mip_count = 1;
    range.layer_count = layers;
    rhi::barrier(cmd, img, range, before, after);
}

static void push_ibl_constants(rhi::CmdBuffer &cmd, rhi::RenderPipeline &pipeline, const IblConstants &pc) {
    static_assert(sizeof(pc) <= 128, "IblConstants exceeds global push-constant range");
    char buf[128] = {};
    std::memcpy(buf, &pc, sizeof(pc));
    rhi::set_constants(cmd, pipeline, rhi::ShaderStage::ALL, sizeof(buf), buf);
}

constexpr u32 ibl_floor_log2(u32 v) {
    return v <= 1 ? 0 : 1 + ibl_floor_log2(v >> 1);
}
constexpr u32 ibl_mip_count(u32 size) {
    return 1 + ibl_floor_log2(size);
}

constexpr u32 ENV_SIZE = 2048;
constexpr u32 ENV_MIP_COUNT = ibl_mip_count(ENV_SIZE);
constexpr u32 IBL_IRRADIANCE_SIZE = 16;
constexpr u32 PREFILTER_SIZE = 256;
constexpr u32 PREFILTER_MIPS = ibl_mip_count(PREFILTER_SIZE);
constexpr u32 BRDF_LUT_SIZE = 512;

void GenerateIBL(
    rhi::Device &device,
    QuickSubmit &qs,
    rhi::ShaderCompiler &sc,
    EnvMap &map,
    FrameGraph & /*fg*/,
    rhi::ImageView &hdr_view,
    u64 samp_linear_clamp
) {
    // Environment cubemap
    {
        rhi::ImageDesc env_img_desc = {
            .width = ENV_SIZE,
            .height = ENV_SIZE,
            .depth = 1,
            .mip_levels = ENV_MIP_COUNT,
            .layers = 6,
            .is_cubemap = true,
            .format = rhi::ImageFormat::RGBA16_FLOAT,
            .usage = rhi::ImageUsage::ColorAttachment | rhi::ImageUsage::Sampled | rhi::ImageUsage::TransferSrc |
                     rhi::ImageUsage::TransferDst,
            .name = "ibl/environment",
        };
        rhi::create_image(device, env_img_desc, map.environment_img);

        rhi::ImageViewDesc env_img_view_desc = {
            .image = &map.environment_img,
            .dimension = rhi::TextureViewDimension::TEXTURE_CUBE,
            .mip_start = 0,
            .mip_count = ENV_MIP_COUNT,
            .layer_start = 0,
            .layer_count = 6
        };
        rhi::create_image_view(device, env_img_view_desc, map.environment_view);
    }

    // Irradiance cubemap
    {
        rhi::ImageDesc irradiance_img_desc = {
            .width = IBL_IRRADIANCE_SIZE,
            .height = IBL_IRRADIANCE_SIZE,
            .depth = 1,
            .mip_levels = 1,
            .layers = 6,
            .is_cubemap = true,
            .format = rhi::ImageFormat::RGBA16_FLOAT,
            .usage = rhi::ImageUsage::ColorAttachment | rhi::ImageUsage::Sampled,
            .name = "ibl/irradiance",
        };
        rhi::create_image(device, irradiance_img_desc, map.irradiance_img);

        rhi::ImageViewDesc irradiance_img_view_desc = {
            .image = &map.irradiance_img,
            .dimension = rhi::TextureViewDimension::TEXTURE_CUBE,
            .mip_start = 0,
            .mip_count = 1,
            .layer_start = 0,
            .layer_count = 6,
        };
        rhi::create_image_view(device, irradiance_img_view_desc, map.irradiance_view);
    }

    // Prefilter cubemap
    {
        rhi::ImageDesc prefiltered_img_desc = {
            .width = PREFILTER_SIZE,
            .height = PREFILTER_SIZE,
            .depth = 1,
            .mip_levels = PREFILTER_MIPS,
            .layers = 6,
            .is_cubemap = true,
            .format = rhi::ImageFormat::RGBA16_FLOAT,
            .usage = rhi::ImageUsage::ColorAttachment | rhi::ImageUsage::Sampled,
            .name = "ibl/prefiltered",
        };
        rhi::create_image(device, prefiltered_img_desc, map.prefiltered_img);
        map.prefilter_mip_count = PREFILTER_MIPS;

        rhi::ImageViewDesc prefiltered_img_view_desc = {
            .image = &map.prefiltered_img,
            .dimension = rhi::TextureViewDimension::TEXTURE_CUBE,
            .mip_start = 0,
            .mip_count = PREFILTER_MIPS,
            .layer_start = 0,
            .layer_count = 6
        };
        rhi::create_image_view(device, prefiltered_img_view_desc, map.prefiltered_view);
    }

    // BRDF LUT
    {
        rhi::ImageDesc brdf_lut_desc = {
            .width = BRDF_LUT_SIZE,
            .height = BRDF_LUT_SIZE,
            .depth = 1,
            .mip_levels = 1,
            .layers = 1,
            .format = rhi::ImageFormat::RGBA16_FLOAT,
            .usage = rhi::ImageUsage::ColorAttachment | rhi::ImageUsage::Sampled,
            .name = "ibl/brdf-lut",
        };
        rhi::create_image(device, brdf_lut_desc, map.brdf_lut_img);

        rhi::ImageViewDesc brdf_lut_view_desc = {
            .image = &map.brdf_lut_img,
            .dimension = rhi::TextureViewDimension::TEXTURE_2D,
            .mip_start = 0,
            .mip_count = 1,
            .layer_start = 0,
            .layer_count = 1,
        };
        rhi::create_image_view(device, brdf_lut_view_desc, map.brdf_lut_view);
    }

    map.environmentIdx = rhi::handle_id(device, map.environment_view);
    map.irradianceIdx = rhi::handle_id(device, map.irradiance_view);
    map.prefilteredIdx = rhi::handle_id(device, map.prefiltered_view);
    map.brdfLutIdx = rhi::handle_id(device, map.brdf_lut_view);

    rhi::RenderPipelineDesc equirect_pd;
    equirect_pd.device = &device;
    equirect_pd.set_shaders(sc, "src/shaders/ibl/equirect_to_cube.slang", "vsMain", "fsMain");
    equirect_pd.set_depth_testing(false, false, rhi::PipelineCompareOp::NEVER);
    equirect_pd.set_cull_mode(rhi::PipelineCullMode::NONE, rhi::PipelineTriangleWindingOrder::CW);
    equirect_pd.set_color_format(rhi::ImageFormat::RGBA16_FLOAT);
    rhi::RenderPipeline equirect_pipeline = equirect_pd.build();

    rhi::RenderPipelineDesc irradiance_pd;
    irradiance_pd.device = &device;
    irradiance_pd.set_shaders(sc, "src/shaders/ibl/irradiance_convolution.slang", "vsMain", "fsMain");
    irradiance_pd.set_depth_testing(false, false, rhi::PipelineCompareOp::NEVER);
    irradiance_pd.set_cull_mode(rhi::PipelineCullMode::NONE, rhi::PipelineTriangleWindingOrder::CW);
    irradiance_pd.set_color_format(rhi::ImageFormat::RGBA16_FLOAT);
    rhi::RenderPipeline irradiance_pipeline = irradiance_pd.build();

    rhi::RenderPipelineDesc prefilter_pd;
    prefilter_pd.device = &device;
    prefilter_pd.set_shaders(sc, "src/shaders/ibl/prefilter_env.slang", "vsMain", "fsMain");
    prefilter_pd.set_depth_testing(false, false, rhi::PipelineCompareOp::NEVER);
    prefilter_pd.set_cull_mode(rhi::PipelineCullMode::NONE, rhi::PipelineTriangleWindingOrder::CW);
    prefilter_pd.set_color_format(rhi::ImageFormat::RGBA16_FLOAT);
    rhi::RenderPipeline prefilter_pipeline = prefilter_pd.build();

    rhi::RenderPipelineDesc brdf_pd;
    brdf_pd.device = &device;
    brdf_pd.set_shaders(sc, "src/shaders/ibl/brdf_lut.slang", "vsMain", "fsMain");
    brdf_pd.set_depth_testing(false, false, rhi::PipelineCompareOp::NEVER);
    brdf_pd.set_cull_mode(rhi::PipelineCullMode::NONE, rhi::PipelineTriangleWindingOrder::CW);
    brdf_pd.set_color_format(rhi::ImageFormat::RGBA16_FLOAT);
    rhi::RenderPipeline brdf_lut_pipeline = brdf_pd.build();

    QsCmd *cmd = qs_acquire(qs);

    const u32 hdr_handle = (u32)rhi::handle_id(device, hdr_view);

    transition_whole(cmd->cmd, map.environment_img, rhi::ResourceState::ColorDraw);
    transition_whole(cmd->cmd, map.irradiance_img, rhi::ResourceState::ColorDraw);
    transition_whole(cmd->cmd, map.prefiltered_img, rhi::ResourceState::ColorDraw);
    transition_whole(cmd->cmd, map.brdf_lut_img, rhi::ResourceState::ColorDraw);

    // 1. Equirectangular -> Cubemap
    for (u32 face = 0; face < 6; ++face) {
        rhi::ImageViewDesc face_view_desc{
            .image = &map.environment_img, .mip_start = 0, .mip_count = 1, .layer_start = face, .layer_count = 1
        };
        rhi::ImageView face_view{};
        rhi::create_image_view(device, face_view_desc, face_view);

        rhi::AttachmentDesc attachment{};
        attachment.img = &face_view;
        attachment.loadOp = rhi::LoadOp::CLEAR;
        attachment.storeOp = rhi::StoreOp::STORE;

        rhi::RenderingInfo ri{};
        ri.colorAttachments[0] = attachment;
        ri.colorAttachmentCount = 1;
        ri.width = ENV_SIZE;
        ri.height = ENV_SIZE;

        rhi::begin_rendering(cmd->cmd, ri);
        rhi::set_pipeline(cmd->cmd, equirect_pipeline);

        IblConstants pc{};
        pc.hdrTexture = hdr_handle;
        pc.faceidx = face;
        pc.samp = samp_linear_clamp;

        push_ibl_constants(cmd->cmd, equirect_pipeline, pc);

        rhi::Viewport viewport{};
        viewport.height = (f32)ENV_SIZE;
        viewport.width = (f32)ENV_SIZE;
        rhi::set_viewports(cmd->cmd, &viewport);

        rhi::Rect scissor{};
        scissor.width = ENV_SIZE;
        scissor.height = ENV_SIZE;
        rhi::set_scissors(cmd->cmd, &scissor);

        rhi::draw(cmd->cmd, 3, 1, 0, 0);

        rhi::end_rendering(cmd->cmd);

        rhi::destroy_image_view(device, face_view);
    }

    // Generate environment cubemap mips
    for (u32 mip = 0; mip < ENV_MIP_COUNT - 1; ++mip) {
        u32 src_width = std::max(ENV_SIZE >> mip, 1u);
        u32 src_height = std::max(ENV_SIZE >> mip, 1u);

        u32 out_width = std::max(src_width >> 1, 1u);
        u32 out_height = std::max(src_height >> 1, 1u);

        barrier_mip_layers(
            cmd->cmd,
            map.environment_img,
            mip,
            6,
            mip == 0 ? rhi::ResourceState::ColorDraw : rhi::ResourceState::TransferTo,
            rhi::ResourceState::TransferFrom
        );
        barrier_mip_layers(
            cmd->cmd, map.environment_img, mip + 1, 6, rhi::ResourceState::ColorDraw, rhi::ResourceState::TransferTo
        );

        u32 constexpr blit_count = 6;
        rhi::ImageBlit blits[blit_count]{};
        for (u32 face = 0; face < 6; ++face) {
            blits[face].src_mip = mip;
            blits[face].dst_mip = mip + 1;
            blits[face].src_layer = face;
            blits[face].dst_layer = face;
            blits[face].src_x0 = 0;
            blits[face].src_y0 = 0;
            blits[face].src_x1 = (i32)src_width;
            blits[face].src_y1 = (i32)src_height;
            blits[face].dst_x0 = 0;
            blits[face].dst_y0 = 0;
            blits[face].dst_x1 = (i32)out_width;
            blits[face].dst_y1 = (i32)out_height;
        }

        rhi::blit_image(cmd->cmd, map.environment_img, blits, blit_count);
    }

    for (u32 mip = 0; mip < ENV_MIP_COUNT; ++mip) {
        barrier_mip_layers(
            cmd->cmd,
            map.environment_img,
            mip,
            6,
            mip == ENV_MIP_COUNT - 1 ? rhi::ResourceState::TransferTo : rhi::ResourceState::TransferFrom,
            rhi::ResourceState::TextureSample
        );
    }

    map.environment_img.state = rhi::ResourceState::TextureSample;

    // 2. Generate irradiance cubemap
    for (u32 face = 0; face < 6; ++face) {
        rhi::ImageViewDesc face_view_desc{
            .image = &map.irradiance_img, .mip_start = 0, .mip_count = 1, .layer_start = face, .layer_count = 1
        };

        rhi::ImageView face_view{};
        rhi::create_image_view(device, face_view_desc, face_view);
        // Irradiance image is already ColorDraw (transitioned above).

        rhi::AttachmentDesc attachment{};
        attachment.img = &face_view;
        attachment.loadOp = rhi::LoadOp::CLEAR;
        attachment.storeOp = rhi::StoreOp::STORE;

        rhi::RenderingInfo ri{};
        ri.colorAttachments[0] = attachment;
        ri.colorAttachmentCount = 1;
        ri.width = IBL_IRRADIANCE_SIZE;
        ri.height = IBL_IRRADIANCE_SIZE;

        rhi::begin_rendering(cmd->cmd, ri);
        rhi::set_pipeline(cmd->cmd, irradiance_pipeline);

        IblConstants pc{};
        pc.environment = map.environmentIdx;
        pc.faceidx = face;
        pc.samp = samp_linear_clamp;

        push_ibl_constants(cmd->cmd, irradiance_pipeline, pc);

        rhi::Viewport viewport{};
        viewport.width = (f32)IBL_IRRADIANCE_SIZE;
        viewport.height = (f32)IBL_IRRADIANCE_SIZE;
        rhi::set_viewports(cmd->cmd, &viewport);

        rhi::Rect scissor{};
        scissor.width = IBL_IRRADIANCE_SIZE;
        scissor.height = IBL_IRRADIANCE_SIZE;
        rhi::set_scissors(cmd->cmd, &scissor);

        rhi::draw(cmd->cmd, 3, 1, 0, 0);

        rhi::end_rendering(cmd->cmd);

        rhi::destroy_image_view(device, face_view);
    }

    // 3. Generate prefiltered cubemap
    for (u32 mip = 0; mip < PREFILTER_MIPS; ++mip) {
        u32 mip_width = PREFILTER_SIZE >> mip;
        u32 mip_height = PREFILTER_SIZE >> mip;

        f32 roughness = (f32)mip / (f32)(PREFILTER_MIPS - 1);

        for (u32 face = 0; face < 6; ++face) {
            rhi::ImageViewDesc face_view_desc{
                .image = &map.prefiltered_img, .mip_start = mip, .mip_count = 1, .layer_start = face, .layer_count = 1
            };

            rhi::ImageView face_view{};
            rhi::create_image_view(device, face_view_desc, face_view);
            // Prefiltered image is already ColorDraw (transitioned above).

            rhi::AttachmentDesc attachment{};
            attachment.img = &face_view;
            attachment.loadOp = rhi::LoadOp::CLEAR;
            attachment.storeOp = rhi::StoreOp::STORE;

            rhi::RenderingInfo ri{};
            ri.colorAttachments[0] = attachment;
            ri.colorAttachmentCount = 1;
            ri.width = mip_width;
            ri.height = mip_height;

            rhi::begin_rendering(cmd->cmd, ri);
            rhi::set_pipeline(cmd->cmd, prefilter_pipeline);

            IblConstants pc{};
            pc.faceidx = face;
            pc.environment = map.environmentIdx;
            pc.roughness = roughness;
            pc.environment_resolution = (f32)ENV_SIZE;
            pc.samp = samp_linear_clamp;

            push_ibl_constants(cmd->cmd, prefilter_pipeline, pc);

            rhi::Viewport viewport{};
            viewport.width = (f32)mip_width;
            viewport.height = (f32)mip_height;
            rhi::set_viewports(cmd->cmd, &viewport);

            rhi::Rect scissor{};
            scissor.width = mip_width;
            scissor.height = mip_height;
            rhi::set_scissors(cmd->cmd, &scissor);

            rhi::draw(cmd->cmd, 3, 1, 0, 0);

            rhi::end_rendering(cmd->cmd);

            rhi::destroy_image_view(device, face_view);
        }
    }

    // 4. Generate brdf integration texture2d
    {
        rhi::AttachmentDesc attachment{};
        attachment.img = &map.brdf_lut_view;
        attachment.loadOp = rhi::LoadOp::CLEAR;
        attachment.storeOp = rhi::StoreOp::STORE;

        rhi::RenderingInfo ri{};
        ri.colorAttachments[0] = attachment;
        ri.colorAttachmentCount = 1;
        ri.width = BRDF_LUT_SIZE;
        ri.height = BRDF_LUT_SIZE;

        rhi::begin_rendering(cmd->cmd, ri);
        rhi::set_pipeline(cmd->cmd, brdf_lut_pipeline);

        rhi::Viewport viewport{};
        viewport.width = (f32)BRDF_LUT_SIZE;
        viewport.height = (f32)BRDF_LUT_SIZE;
        rhi::set_viewports(cmd->cmd, &viewport);

        rhi::Rect scissor{};
        scissor.width = BRDF_LUT_SIZE;
        scissor.height = BRDF_LUT_SIZE;
        rhi::set_scissors(cmd->cmd, &scissor);

        rhi::draw(cmd->cmd, 3, 1, 0, 0);

        rhi::end_rendering(cmd->cmd);
    }

    transition_whole(cmd->cmd, map.irradiance_img, rhi::ResourceState::TextureSample);
    transition_whole(cmd->cmd, map.prefiltered_img, rhi::ResourceState::TextureSample);
    transition_whole(cmd->cmd, map.brdf_lut_img, rhi::ResourceState::TextureSample);

    qs_submit(qs, *cmd);
}