bluenoise-raytracer/ray_tracing__advance/raytrace.cpp

439 lines
19 KiB
C++

/*
* Copyright (c) 2014-2021, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION
* SPDX-License-Identifier: Apache-2.0
*/
#include "raytrace.hpp"
#include "nvh/fileoperations.hpp"
#include "nvvk/descriptorsets_vk.hpp"
#include "nvh/alignment.hpp"
#include "nvvk/shaders_vk.hpp"
#include "obj_loader.h"
#include "nvvk/buffers_vk.hpp"
extern std::vector<std::string> defaultSearchPaths;
void Raytracer::setup(const VkDevice& device, const VkPhysicalDevice& physicalDevice, nvvk::ResourceAllocator* allocator, uint32_t queueFamily)
{
m_device = device;
m_physicalDevice = physicalDevice;
m_alloc = allocator;
m_graphicsQueueIndex = queueFamily;
// Requesting ray tracing properties
VkPhysicalDeviceProperties2 prop2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2};
prop2.pNext = &m_rtProperties;
vkGetPhysicalDeviceProperties2(m_physicalDevice, &prop2);
m_rtBuilder.setup(m_device, allocator, m_graphicsQueueIndex);
m_sbtWrapper.setup(device, queueFamily, allocator, m_rtProperties);
m_debug.setup(device);
}
void Raytracer::destroy()
{
m_sbtWrapper.destroy();
m_rtBuilder.destroy();
vkDestroyDescriptorPool(m_device, m_rtDescPool, nullptr);
vkDestroyDescriptorSetLayout(m_device, m_rtDescSetLayout, nullptr);
vkDestroyPipeline(m_device, m_rtPipeline, nullptr);
vkDestroyPipelineLayout(m_device, m_rtPipelineLayout, nullptr);
m_alloc->destroy(m_rtSBTBuffer);
}
//--------------------------------------------------------------------------------------------------
// Converting a OBJ primitive to the ray tracing geometry used for the BLAS
//
auto Raytracer::objectToVkGeometryKHR(const ObjModel& model)
{
// Building part
VkDeviceAddress vertexAddress = nvvk::getBufferDeviceAddress(m_device, model.vertexBuffer.buffer);
VkDeviceAddress indexAddress = nvvk::getBufferDeviceAddress(m_device, model.indexBuffer.buffer);
VkAccelerationStructureGeometryTrianglesDataKHR triangles{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR};
triangles.vertexFormat = VK_FORMAT_R32G32B32A32_SFLOAT;
triangles.vertexData.deviceAddress = vertexAddress;
triangles.vertexStride = sizeof(VertexObj);
triangles.indexType = VK_INDEX_TYPE_UINT32;
triangles.indexData.deviceAddress = indexAddress;
triangles.transformData = {};
triangles.maxVertex = model.nbVertices;
// Setting up the build info of the acceleration
VkAccelerationStructureGeometryKHR asGeom{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR};
asGeom.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
asGeom.flags = VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR; // For AnyHit
asGeom.geometry.triangles = triangles;
VkAccelerationStructureBuildRangeInfoKHR offset;
offset.firstVertex = 0;
offset.primitiveCount = model.nbIndices / 3; // Nb triangles
offset.primitiveOffset = 0;
offset.transformOffset = 0;
nvvk::RaytracingBuilderKHR::BlasInput input;
input.asGeometry.emplace_back(asGeom);
input.asBuildOffsetInfo.emplace_back(offset);
return input;
}
//--------------------------------------------------------------------------------------------------
// Returning the ray tracing geometry used for the BLAS, containing all spheres
//
auto Raytracer::implicitToVkGeometryKHR(const ImplInst& implicitObj)
{
VkDeviceAddress dataAddress = nvvk::getBufferDeviceAddress(m_device, implicitObj.implBuf.buffer);
VkAccelerationStructureGeometryAabbsDataKHR aabbs{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR};
aabbs.data.deviceAddress = dataAddress;
aabbs.stride = sizeof(ObjImplicit);
// Setting up the build info of the acceleration
VkAccelerationStructureGeometryKHR asGeom{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR};
asGeom.geometryType = VK_GEOMETRY_TYPE_AABBS_KHR;
asGeom.flags = VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR; // For AnyHit
asGeom.geometry.aabbs = aabbs;
VkAccelerationStructureBuildRangeInfoKHR offset;
offset.firstVertex = 0;
offset.primitiveCount = static_cast<uint32_t>(implicitObj.objImpl.size()); // Nb aabb
offset.primitiveOffset = 0;
offset.transformOffset = 0;
nvvk::RaytracingBuilderKHR::BlasInput input;
input.asGeometry.emplace_back(asGeom);
input.asBuildOffsetInfo.emplace_back(offset);
return input;
}
void Raytracer::createBottomLevelAS(std::vector<ObjModel>& models, ImplInst& implicitObj)
{
// BLAS - Storing each primitive in a geometry
std::vector<nvvk::RaytracingBuilderKHR::BlasInput> allBlas;
allBlas.reserve(models.size());
for(const auto& obj : models)
{
auto blas = objectToVkGeometryKHR(obj);
// We could add more geometry in each BLAS, but we add only one for now
allBlas.emplace_back(blas);
}
// Adding implicit
if(!implicitObj.objImpl.empty())
{
auto blas = implicitToVkGeometryKHR(implicitObj);
allBlas.emplace_back(blas);
implicitObj.blasId = static_cast<int>(allBlas.size() - 1); // remember blas ID for tlas
}
m_rtBuilder.buildBlas(allBlas, VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR
| VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR);
}
void Raytracer::createTopLevelAS(std::vector<ObjInstance>& instances, ImplInst& implicitObj)
{
std::vector<VkAccelerationStructureInstanceKHR> tlas;
auto nbObj = static_cast<uint32_t>(instances.size()) - 1; // minus the implicit (for material)
tlas.reserve(instances.size());
for(uint32_t i = 0; i < nbObj; i++)
{
VkAccelerationStructureInstanceKHR rayInst;
rayInst.transform = nvvk::toTransformMatrixKHR(instances[i].transform); // Position of the instance
rayInst.instanceCustomIndex = instances[i].objIndex; // gl_InstanceCustomIndexEXT
rayInst.accelerationStructureReference = m_rtBuilder.getBlasDeviceAddress(instances[i].objIndex);
rayInst.instanceShaderBindingTableRecordOffset = 0; // We will use the same hit group for all objects
rayInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR;
rayInst.mask = 0xFF;
tlas.emplace_back(rayInst);
}
// Add the blas containing all implicit
if(!implicitObj.objImpl.empty())
{
VkAccelerationStructureInstanceKHR rayInst;
rayInst.transform = nvvk::toTransformMatrixKHR(implicitObj.transform); // Position of the instance
rayInst.instanceCustomIndex = nbObj; // Same for material index
rayInst.accelerationStructureReference = m_rtBuilder.getBlasDeviceAddress(static_cast<uint32_t>(implicitObj.blasId));
rayInst.instanceShaderBindingTableRecordOffset = 1; // We will use the same hit group for all objects (the second one)
rayInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR;
rayInst.mask = 0xFF;
tlas.emplace_back(rayInst);
}
m_rtBuilder.buildTlas(tlas, VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR);
}
//--------------------------------------------------------------------------------------------------
// This descriptor set holds the Acceleration structure and the output image
//
void Raytracer::createRtDescriptorSet(const VkImageView& outputImage)
{
using vkDSLB = VkDescriptorSetLayoutBinding;
m_rtDescSetLayoutBind.addBinding(0, VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 1,
VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR); // TLAS
m_rtDescSetLayoutBind.addBinding(1, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, VK_SHADER_STAGE_RAYGEN_BIT_KHR); // Output image
m_rtDescPool = m_rtDescSetLayoutBind.createPool(m_device);
m_rtDescSetLayout = m_rtDescSetLayoutBind.createLayout(m_device);
VkDescriptorSetAllocateInfo allocateInfo{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO};
allocateInfo.descriptorPool = m_rtDescPool;
allocateInfo.descriptorSetCount = 1;
allocateInfo.pSetLayouts = &m_rtDescSetLayout;
vkAllocateDescriptorSets(m_device, &allocateInfo, &m_rtDescSet);
VkAccelerationStructureKHR tlas = m_rtBuilder.getAccelerationStructure();
VkWriteDescriptorSetAccelerationStructureKHR descASInfo{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR};
descASInfo.accelerationStructureCount = 1;
descASInfo.pAccelerationStructures = &tlas;
VkDescriptorImageInfo imageInfo{{}, outputImage, VK_IMAGE_LAYOUT_GENERAL};
std::vector<VkWriteDescriptorSet> writes;
writes.emplace_back(m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, 0, &descASInfo));
writes.emplace_back(m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, 1, &imageInfo));
vkUpdateDescriptorSets(m_device, static_cast<uint32_t>(writes.size()), writes.data(), 0, nullptr);
}
//--------------------------------------------------------------------------------------------------
// Writes the output image to the descriptor set
// - Required when changing resolution
//
void Raytracer::updateRtDescriptorSet(const VkImageView& outputImage)
{
VkDescriptorImageInfo imageInfo{{}, outputImage, VK_IMAGE_LAYOUT_GENERAL};
VkWriteDescriptorSet wds = m_rtDescSetLayoutBind.makeWrite(m_rtDescSet, 1, &imageInfo);
vkUpdateDescriptorSets(m_device, 1, &wds, 0, nullptr);
}
//--------------------------------------------------------------------------------------------------
// Pipeline for the ray tracer: all shaders, raygen, chit, miss
//
void Raytracer::createRtPipeline(VkDescriptorSetLayout& sceneDescLayout)
{
enum StageIndices
{
eRaygen,
eMiss,
eMiss2,
eClosestHit,
eAnyHit,
eClosestHit1,
eAnyHit1,
eIntersect,
eCall0,
eCall1,
eCall2,
eShaderGroupCount
};
// All stages
std::array<VkPipelineShaderStageCreateInfo, eShaderGroupCount> stages{};
VkPipelineShaderStageCreateInfo stage{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO};
stage.pName = "main"; // All the same entry point
// Raygen
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rgen.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_RAYGEN_BIT_KHR;
stages[eRaygen] = stage;
// Miss
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rmiss.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_MISS_BIT_KHR;
stages[eMiss] = stage;
// The second miss shader is invoked when a shadow ray misses the geometry. It simply indicates that no occlusion has been found
stage.module =
nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytraceShadow.rmiss.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_MISS_BIT_KHR;
stages[eMiss2] = stage;
// Hit Group - Closest Hit
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rchit.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR;
stages[eClosestHit] = stage;
// Hit Group - Closest Hit
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rahit.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR;
stages[eAnyHit] = stage;
// Hit Group - 1
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace2.rchit.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR;
stages[eClosestHit1] = stage;
// Hit
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace2.rahit.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_ANY_HIT_BIT_KHR;
stages[eAnyHit1] = stage;
// Hit
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/raytrace.rint.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_INTERSECTION_BIT_KHR;
stages[eIntersect] = stage;
// Call0
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_point.rcall.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR;
stages[eCall0] = stage;
// Call1
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_spot.rcall.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR;
stages[eCall1] = stage;
// Call2
stage.module = nvvk::createShaderModule(m_device, nvh::loadFile("spv/light_inf.rcall.spv", true, defaultSearchPaths, true));
stage.stage = VK_SHADER_STAGE_CALLABLE_BIT_KHR;
stages[eCall2] = stage;
// Shader groups
VkRayTracingShaderGroupCreateInfoKHR group{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR};
group.anyHitShader = VK_SHADER_UNUSED_KHR;
group.closestHitShader = VK_SHADER_UNUSED_KHR;
group.generalShader = VK_SHADER_UNUSED_KHR;
group.intersectionShader = VK_SHADER_UNUSED_KHR;
// Raygen
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
group.generalShader = eRaygen;
m_rtShaderGroups.push_back(group);
// Miss
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
group.generalShader = eMiss;
m_rtShaderGroups.push_back(group);
// Shadow Miss
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
group.generalShader = eMiss2;
m_rtShaderGroups.push_back(group);
// closest hit shader
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
group.generalShader = VK_SHADER_UNUSED_KHR;
group.closestHitShader = eClosestHit;
group.anyHitShader = eAnyHit;
m_rtShaderGroups.push_back(group);
// closest hit shader
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR;
group.generalShader = VK_SHADER_UNUSED_KHR;
group.closestHitShader = eClosestHit1;
group.anyHitShader = eAnyHit1;
group.intersectionShader = eIntersect;
m_rtShaderGroups.push_back(group);
// Callable shaders
group.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
group.closestHitShader = VK_SHADER_UNUSED_KHR;
group.anyHitShader = VK_SHADER_UNUSED_KHR;
group.intersectionShader = VK_SHADER_UNUSED_KHR;
group.generalShader = eCall0;
m_rtShaderGroups.push_back(group);
group.generalShader = eCall1;
m_rtShaderGroups.push_back(group);
group.generalShader = eCall2;
m_rtShaderGroups.push_back(group);
// Push constant: we want to be able to update constants used by the shaders
VkPushConstantRange pushConstant{VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR
| VK_SHADER_STAGE_MISS_BIT_KHR | VK_SHADER_STAGE_CALLABLE_BIT_KHR,
0, sizeof(RtPushConstants)};
VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO};
pipelineLayoutCreateInfo.pushConstantRangeCount = 1;
pipelineLayoutCreateInfo.pPushConstantRanges = &pushConstant;
// Descriptor sets: one specific to ray tracing, and one shared with the rasterization pipeline
std::vector<VkDescriptorSetLayout> rtDescSetLayouts = {m_rtDescSetLayout, sceneDescLayout};
pipelineLayoutCreateInfo.setLayoutCount = static_cast<uint32_t>(rtDescSetLayouts.size());
pipelineLayoutCreateInfo.pSetLayouts = rtDescSetLayouts.data();
vkCreatePipelineLayout(m_device, &pipelineLayoutCreateInfo, nullptr, &m_rtPipelineLayout);
// Assemble the shader stages and recursion depth info into the ray tracing pipeline
VkRayTracingPipelineCreateInfoKHR rayPipelineInfo{VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR};
rayPipelineInfo.stageCount = static_cast<uint32_t>(stages.size()); // Stages are shaders
rayPipelineInfo.pStages = stages.data();
rayPipelineInfo.groupCount = static_cast<uint32_t>(m_rtShaderGroups.size());
rayPipelineInfo.pGroups = m_rtShaderGroups.data();
// The ray tracing process can shoot rays from the camera, and a shadow ray can be shot from the
// hit points of the camera rays, hence a recursion level of 2. This number should be kept as low
// as possible for performance reasons. Even recursive ray tracing should be flattened into a loop
// in the ray generation to avoid deep recursion.
rayPipelineInfo.maxPipelineRayRecursionDepth = 2; // Ray depth
rayPipelineInfo.layout = m_rtPipelineLayout;
vkCreateRayTracingPipelinesKHR(m_device, {}, {}, 1, &rayPipelineInfo, nullptr, &m_rtPipeline);
m_sbtWrapper.create(m_rtPipeline, rayPipelineInfo);
for(auto& s : stages)
vkDestroyShaderModule(m_device, s.module, nullptr);
}
//--------------------------------------------------------------------------------------------------
// Ray Tracing the scene
//
void Raytracer::raytrace(const VkCommandBuffer& cmdBuf,
const nvmath::vec4f& clearColor,
VkDescriptorSet& sceneDescSet,
VkExtent2D& size,
ObjPushConstants& sceneConstants)
{
m_debug.beginLabel(cmdBuf, "Ray trace");
// Initializing push constant values
m_rtPushConstants.clearColor = clearColor;
m_rtPushConstants.lightPosition = sceneConstants.lightPosition;
m_rtPushConstants.lightIntensity = sceneConstants.lightIntensity;
m_rtPushConstants.lightDirection = sceneConstants.lightDirection;
m_rtPushConstants.lightSpotCutoff = sceneConstants.lightSpotCutoff;
m_rtPushConstants.lightSpotOuterCutoff = sceneConstants.lightSpotOuterCutoff;
m_rtPushConstants.lightType = sceneConstants.lightType;
m_rtPushConstants.frame = sceneConstants.frame;
std::vector<VkDescriptorSet> descSets{m_rtDescSet, sceneDescSet};
vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, m_rtPipeline);
vkCmdBindDescriptorSets(cmdBuf, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, m_rtPipelineLayout, 0,
(uint32_t)descSets.size(), descSets.data(), 0, nullptr);
vkCmdPushConstants(cmdBuf, m_rtPipelineLayout,
VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR | VK_SHADER_STAGE_MISS_BIT_KHR
| VK_SHADER_STAGE_CALLABLE_BIT_KHR,
0, sizeof(RtPushConstants), &m_rtPushConstants);
auto& regions = m_sbtWrapper.getRegions();
vkCmdTraceRaysKHR(cmdBuf, &regions[0], &regions[1], &regions[2], &regions[3], size.width, size.height, 1);
m_debug.endLabel(cmdBuf);
}