project-aster/samples/03_model_render/ibl_helpers.cpp

415 lines
16 KiB
C++

// =============================================
// Aster: ibl_helpers.cpp
// Copyright (c) 2020-2024 Anish Bhobe
// =============================================
#include "ibl_helpers.h"
#include "aster/core/device.h"
#include "aster/core/image.h"
#include "asset_loader.h"
#include "helpers.h"
#include "pipeline_utils.h"
#include "aster/systems/commit_manager.h"
#include "aster/systems/resource_manager.h"
#include <EASTL/fixed_vector.h>
#include <EASTL/tuple.h>
constexpr cstr EQUIRECT_TO_CUBE_SHADER_FILE = "shader/eqrect_to_cube.cs.hlsl.spv";
constexpr cstr DIFFUSE_IRRADIANCE_SHADER_FILE = "shader/diffuse_irradiance.cs.hlsl.spv";
constexpr cstr PREFILTER_SHADER_FILE = "shader/prefilter.cs.hlsl.spv";
constexpr cstr BRDF_LUT_SHADER_FILE = "shader/brdf_lut.cs.hlsl.spv";
Environment
CreateCubeFromHdrEnv(AssetLoader *assetLoader, vk::Queue computeQueue, const u32 cubeSide,
systems::ResId<Texture> hdrEnv, const cstr name)
{
systems::ResourceManager *resourceManager = assetLoader->m_ResourceManager;
systems::CommitManager *commitManager = assetLoader->m_CommitManager;
const Device *pDevice = commitManager->m_Device;
auto skybox = resourceManager->Images().CreateTextureCube<StorageTexture>({
.m_Format = vk::Format::eR16G16B16A16Sfloat,
.m_Side = cubeSide,
.m_Name = "Skybox",
.m_IsSampled = true,
.m_IsMipMapped = true,
.m_IsStorage = true,
});
auto skyboxHandle = commitManager->CommitTexture(skybox);
auto skyboxStorageHandle = commitManager->CommitStorageImage(skybox);
auto diffuseIrradiance = resourceManager->Images().CreateTextureCube<StorageTexture>({
.m_Format = vk::Format::eR16G16B16A16Sfloat,
.m_Side = 64,
.m_Name = "Diffuse Irradiance",
.m_IsSampled = true,
.m_IsMipMapped = false,
.m_IsStorage = true,
});
auto diffuseIrradianceHandle = commitManager->CommitTexture(diffuseIrradiance);
auto diffuseIrradianceStorageHandle = commitManager->CommitStorageImage(diffuseIrradiance);
auto prefilterCube = resourceManager->Images().CreateTextureCube<StorageTextureCube>({
.m_Format = vk::Format::eR16G16B16A16Sfloat,
.m_Side = cubeSide,
.m_Name = "Prefilter",
.m_IsSampled = true,
.m_IsMipMapped = true,
.m_IsStorage = true,
});
auto prefilterHandle = commitManager->CommitTexture(prefilterCube); // This stores the original view for us.
constexpr u32 prefilterMipCountMax = 6;
eastl::fixed_vector<systems::ResId<StorageImage>, prefilterMipCountMax> prefilterStorageHandles;
// All non-owning copies.
for (u32 mipLevel = 0; mipLevel < prefilterMipCountMax; ++mipLevel)
{
vk::ImageViewCreateInfo imageViewCreateInfo = {
.image = prefilterCube->m_Image,
.viewType = vk::ImageViewType::eCube,
.format = vk::Format::eR16G16B16A16Sfloat,
.components = vk::ComponentMapping{},
.subresourceRange =
{
.aspectMask = vk::ImageAspectFlagBits::eColor,
.baseMipLevel = mipLevel,
.levelCount = 1,
.baseArrayLayer = 0,
.layerCount = 6,
},
};
AbortIfFailed(pDevice->m_Device.createImageView(&imageViewCreateInfo, nullptr, &prefilterCube->m_View));
// TODO: FIXME: This is an issue. This needs copying but we don't do that anymore.
// The views need to be separated from the images.
TODO("This is bad");
prefilterStorageHandles.push_back(commitManager->CommitStorageImage(prefilterCube));
}
auto brdfLut = resourceManager->Images().CreateTexture2D<StorageTexture>({.m_Format = vk::Format::eR16G16Sfloat,
.m_Extent = {512, 512},
.m_Name = "BRDF LUT",
.m_IsSampled = true,
.m_IsMipMapped = true,
.m_IsStorage = true});
auto brdfLutSampler = resourceManager->Samplers().CreateSampler({
.m_AddressModeU = vk::SamplerAddressMode::eClampToEdge,
.m_AddressModeV = vk::SamplerAddressMode::eClampToEdge,
.m_AddressModeW = vk::SamplerAddressMode::eClampToEdge,
});
auto brdfLutHandle = commitManager->CommitTexture(brdfLut, brdfLutSampler);
auto brdfLutStorageHandle = commitManager->CommitStorageImage(brdfLut);
#pragma region Dependencies and Copies
vk::ImageSubresourceRange cubeSubresRange = {
.aspectMask = vk::ImageAspectFlagBits::eColor,
.baseMipLevel = 0,
.levelCount = 1,
.baseArrayLayer = 0,
.layerCount = 6,
};
vk::ImageSubresourceRange lutSubresRange = {
.aspectMask = vk::ImageAspectFlagBits::eColor,
.baseMipLevel = 0,
.levelCount = 1,
.baseArrayLayer = 0,
.layerCount = 1,
};
vk::ImageMemoryBarrier2 readyToWriteBarrierTemplate = {
.srcStageMask = vk::PipelineStageFlagBits2::eTopOfPipe,
.srcAccessMask = vk::AccessFlagBits2::eNone,
.dstStageMask = vk::PipelineStageFlagBits2::eComputeShader,
.dstAccessMask = vk::AccessFlagBits2::eShaderStorageWrite,
.oldLayout = vk::ImageLayout::eUndefined,
.newLayout = vk::ImageLayout::eGeneral,
.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
.subresourceRange = cubeSubresRange,
};
eastl::fixed_vector<vk::ImageMemoryBarrier2, 4> readyToWriteBarriers(4, readyToWriteBarrierTemplate);
readyToWriteBarriers[0].image = skybox->m_Image;
readyToWriteBarriers[1].image = diffuseIrradiance->m_Image;
readyToWriteBarriers[2].image = prefilterCube->m_Image;
readyToWriteBarriers[3].image = brdfLut->m_Image;
readyToWriteBarriers[3].subresourceRange = lutSubresRange;
vk::DependencyInfo readyToWriteDependency = {
.imageMemoryBarrierCount = Cast<u32>(readyToWriteBarriers.size()),
.pImageMemoryBarriers = readyToWriteBarriers.data(),
};
vk::ImageMemoryBarrier2 readyToSampleBarrierTemplate = {
.srcStageMask = vk::PipelineStageFlagBits2::eComputeShader,
.srcAccessMask = vk::AccessFlagBits2::eShaderStorageWrite | vk::AccessFlagBits2::eShaderStorageRead,
.dstStageMask = vk::PipelineStageFlagBits2::eBottomOfPipe,
.oldLayout = vk::ImageLayout::eGeneral,
.newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
.subresourceRange = cubeSubresRange,
};
auto skyboxToSampleBarrier = readyToSampleBarrierTemplate;
skyboxToSampleBarrier.image = skybox->m_Image;
auto diffIrrToSampleBarrier = readyToSampleBarrierTemplate;
diffIrrToSampleBarrier.image = diffuseIrradiance->m_Image;
auto prefilterToSampleBarrier = readyToSampleBarrierTemplate;
prefilterToSampleBarrier.image = prefilterCube->m_Image;
auto brdfToSampleBarrier = readyToSampleBarrierTemplate;
prefilterToSampleBarrier.image = brdfLut->m_Image;
prefilterToSampleBarrier.subresourceRange = lutSubresRange;
vk::DependencyInfo skyboxToSampleDependency = {
.imageMemoryBarrierCount = 1,
.pImageMemoryBarriers = &skyboxToSampleBarrier,
};
vk::DependencyInfo diffIrrToSampleDependency = {
.imageMemoryBarrierCount = 1,
.pImageMemoryBarriers = &diffIrrToSampleBarrier,
};
vk::DependencyInfo prefilterToSampleDependency = {
.imageMemoryBarrierCount = 1,
.pImageMemoryBarriers = &prefilterToSampleBarrier,
};
vk::DependencyInfo brdfToSampleDependency = {
.imageMemoryBarrierCount = 1,
.pImageMemoryBarriers = &brdfToSampleBarrier,
};
#pragma endregion
struct SkyboxPushConstants
{
systems::ResId<Texture> m_HdrEnvHandle;
systems::ResId<StorageImage> m_OutputTexture;
u32 m_CubeSide;
};
struct DiffuseIrradiancePushConstants
{
systems::ResId<Texture> m_SkyboxHandle;
systems::ResId<StorageImage> m_OutputTexture;
u32 m_CubeSide;
};
struct PrefilterPushConstants
{
systems::ResId<Texture> m_SkyboxHandle;
systems::ResId<StorageImage> m_OutputTexture;
u32 m_CubeSide;
f32 m_Roughness;
u32 m_EnvSide;
};
struct BrdfLutPushConstants
{
systems::ResId<StorageImage> m_OutputTexture;
};
#pragma region Pipeline Creation etc
vk::PushConstantRange pcr = {
.stageFlags = vk::ShaderStageFlagBits::eCompute,
.offset = 0,
.size =
Cast<u32>(eastl::max(eastl::max(sizeof(SkyboxPushConstants), sizeof(BrdfLutPushConstants)),
eastl::max(sizeof(DiffuseIrradiancePushConstants), sizeof(PrefilterPushConstants)))),
};
vk::PipelineLayout pipelineLayout;
const vk::PipelineLayoutCreateInfo layoutCreateInfo = {
.setLayoutCount = 1,
.pSetLayouts = &commitManager->GetDescriptorSetLayout(),
.pushConstantRangeCount = 1,
.pPushConstantRanges = &pcr,
};
AbortIfFailed(pDevice->m_Device.createPipelineLayout(&layoutCreateInfo, nullptr, &pipelineLayout));
const auto eqRectToCubeShader = CreateShader(pDevice, EQUIRECT_TO_CUBE_SHADER_FILE);
const auto diffuseRadianceShader = CreateShader(pDevice, DIFFUSE_IRRADIANCE_SHADER_FILE);
const auto prefilterShader = CreateShader(pDevice, PREFILTER_SHADER_FILE);
const auto brdfLutShader = CreateShader(pDevice, BRDF_LUT_SHADER_FILE);
eastl::array computePipelineCreateInfo = {
vk::ComputePipelineCreateInfo{
.stage =
{
.stage = vk::ShaderStageFlagBits::eCompute,
.module = eqRectToCubeShader,
.pName = "main",
},
.layout = pipelineLayout,
},
vk::ComputePipelineCreateInfo{
.stage =
{
.stage = vk::ShaderStageFlagBits::eCompute,
.module = diffuseRadianceShader,
.pName = "main",
},
.layout = pipelineLayout,
},
vk::ComputePipelineCreateInfo{
.stage =
{
.stage = vk::ShaderStageFlagBits::eCompute,
.module = prefilterShader,
.pName = "main",
},
.layout = pipelineLayout,
},
vk::ComputePipelineCreateInfo{
.stage =
{
.stage = vk::ShaderStageFlagBits::eCompute,
.module = brdfLutShader,
.pName = "main",
},
.layout = pipelineLayout,
},
};
eastl::array<vk::Pipeline, computePipelineCreateInfo.size()> pipelines;
AbortIfFailed(
pDevice->m_Device.createComputePipelines(pDevice->m_PipelineCache, Cast<u32>(computePipelineCreateInfo.size()),
computePipelineCreateInfo.data(), nullptr, pipelines.data()));
vk::Pipeline eqRectToCubePipeline = pipelines[0];
vk::Pipeline diffuseIrradiancePipeline = pipelines[1];
vk::Pipeline prefilterPipeline = pipelines[2];
vk::Pipeline brdfLutPipeline = pipelines[3];
for (auto &createInfos : computePipelineCreateInfo)
{
pDevice->m_Device.destroy(createInfos.stage.module, nullptr);
}
#pragma endregion
SkyboxPushConstants skyboxPushConstant = {
.m_HdrEnvHandle = hdrEnv,
.m_OutputTexture = skyboxStorageHandle,
.m_CubeSide = cubeSide,
};
DiffuseIrradiancePushConstants diffuseIrradiancePushConstants = {
.m_SkyboxHandle = skyboxHandle,
.m_OutputTexture = diffuseIrradianceStorageHandle,
.m_CubeSide = diffuseIrradiance->m_Extent.width,
};
PrefilterPushConstants prefilterPushConstants = {
.m_SkyboxHandle = skyboxHandle,
.m_OutputTexture = systems::NullId{},
.m_EnvSide = cubeSide,
};
BrdfLutPushConstants brdfLutPushConstants = {
.m_OutputTexture = brdfLutStorageHandle,
};
commitManager->Update();
auto cmd = assetLoader->m_CommandBuffer;
constexpr vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
AbortIfFailed(cmd.begin(&beginInfo));
#if !defined(ASTER_NDEBUG)
StackString<128> labelName = "Eqrect -> Cubemap: ";
labelName += name ? name : "<unknown env>";
vk::DebugUtilsLabelEXT label = {
.pLabelName = labelName.c_str(),
.color = std::array{1.0f, 1.0f, 1.0f, 1.0f},
};
cmd.beginDebugUtilsLabelEXT(&label);
#endif
cmd.pipelineBarrier2(&readyToWriteDependency);
cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, pipelineLayout, 0, 1, &commitManager->GetDescriptorSet(), 0,
nullptr);
cmd.bindPipeline(vk::PipelineBindPoint::eCompute, eqRectToCubePipeline);
cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof skyboxPushConstant,
&skyboxPushConstant);
assert(skybox->m_Extent.width % 16 == 0 && skybox->m_Extent.height % 16 == 0);
cmd.dispatch(skybox->m_Extent.width / 16, skybox->m_Extent.height / 16, 6);
GenerateMipMaps(cmd, skybox, vk::ImageLayout::eGeneral, vk::ImageLayout::eGeneral,
vk::PipelineStageFlagBits2::eComputeShader, vk::PipelineStageFlagBits2::eComputeShader);
cmd.bindPipeline(vk::PipelineBindPoint::eCompute, diffuseIrradiancePipeline);
cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof skyboxPushConstant,
&diffuseIrradiancePushConstants);
assert(diffuseIrradiance->m_Extent.width % 16 == 0 && diffuseIrradiance->m_Extent.height % 16 == 0);
cmd.dispatch(diffuseIrradiance->m_Extent.width / 16, diffuseIrradiance->m_Extent.width / 16, 6);
cmd.pipelineBarrier2(&diffIrrToSampleDependency);
cmd.bindPipeline(vk::PipelineBindPoint::eCompute, prefilterPipeline);
u32 mipSize = prefilterCube->m_Extent.width;
assert(mipSize % 16 == 0);
for (u32 mipCount = 0; auto &tex : prefilterStorageHandles)
{
prefilterPushConstants.m_OutputTexture = tex;
prefilterPushConstants.m_CubeSide = mipSize;
prefilterPushConstants.m_Roughness = Cast<f32>(mipCount) / Cast<f32>(prefilterMipCountMax);
cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof prefilterPushConstants,
&prefilterPushConstants);
u32 groupCount = eastl::max(mipSize / 16u, 1u);
cmd.dispatch(groupCount, groupCount, 6);
++mipCount;
mipSize = mipSize >> 1;
}
cmd.pipelineBarrier2(&skyboxToSampleDependency);
cmd.pipelineBarrier2(&prefilterToSampleDependency);
cmd.bindPipeline(vk::PipelineBindPoint::eCompute, brdfLutPipeline);
cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof brdfLutPushConstants,
&brdfLutPushConstants);
assert(brdfLut->m_Extent.width % 16 == 0 && brdfLut->m_Extent.height % 16 == 0);
cmd.dispatch(brdfLut->m_Extent.width / 16, brdfLut->m_Extent.height / 16, 1);
#if !defined(ASTER_NDEBUG)
cmd.endDebugUtilsLabelEXT();
#endif
AbortIfFailed(cmd.end());
vk::SubmitInfo submitInfo = {
.waitSemaphoreCount = 0,
.pWaitDstStageMask = nullptr,
.commandBufferCount = 1,
.pCommandBuffers = &cmd,
};
vk::Fence fence;
vk::FenceCreateInfo fenceCreateInfo = {};
AbortIfFailed(pDevice->m_Device.createFence(&fenceCreateInfo, nullptr, &fence));
AbortIfFailed(computeQueue.submit(1, &submitInfo, fence));
AbortIfFailed(pDevice->m_Device.waitForFences(1, &fence, true, MaxValue<u32>));
pDevice->m_Device.destroy(fence, nullptr);
AbortIfFailed(pDevice->m_Device.resetCommandPool(assetLoader->m_CommandPool, {}));
skybox = {};
for (auto &_ : prefilterStorageHandles)
{
StorageTextureCube st;
// TODO: This needs fixing
pDevice->m_Device.destroy(st.m_View, nullptr);
}
for (auto &pipeline : pipelines)
{
pDevice->m_Device.destroy(pipeline, nullptr);
}
pDevice->m_Device.destroy(pipelineLayout, nullptr);
return {
.m_Skybox = skyboxHandle,
.m_Diffuse = diffuseIrradianceHandle,
.m_Prefilter = prefilterHandle,
.m_BrdfLut = brdfLutHandle,
};
}